import torch
from torch import nn


class TwoChannelFractalBlock(nn.Module):
    """Ternary two-channel renormalization block.

    Inputs have shape [batch, 3, width], representing three child
    representations with separate neutral (a) and defect (b) channels.
    """

    def __init__(self, width, hidden=None, eps=1e-6):
        super().__init__()
        hidden = hidden or 2 * width
        self.eps = eps
        self.x_raw = nn.Parameter(torch.tensor(-1.0))
        self.lambda_raw = nn.Parameter(torch.tensor(0.0))
        self.align_a = nn.Linear(width, width, bias=False)
        self.align_b = nn.Linear(width, width, bias=False)
        self.defect = nn.Sequential(
            nn.Linear(6 * width, hidden),
            nn.GELU(),
            nn.Linear(hidden, width),
            nn.Tanh(),
        )
        self.gate = nn.Sequential(
            nn.Linear(6 * width, width),
            nn.Sigmoid(),
        )

    def forward(self, a, b):
        if a.ndim != 3 or b.shape != a.shape or a.shape[1] != 3:
            raise ValueError("a and b must both have shape [batch, 3, width]")
        aa, bb = self.align_a(a), self.align_b(b)
        lam = torch.nn.functional.softplus(self.lambda_raw)
        x = torch.nn.functional.softplus(self.x_raw)
        neutral = lam**3 * aa[:, 0] * aa[:, 1] * aa[:, 2]
        neutral = neutral + 2.0 * x**3 * bb[:, 0] * bb[:, 1] * bb[:, 2]
        neutral = neutral / (neutral.square().mean(-1, keepdim=True).add(self.eps).sqrt())
        flat = torch.cat((aa.reshape(a.shape[0], -1), bb.reshape(b.shape[0], -1)), dim=-1)
        proposed = self.defect(flat)
        gate = self.gate(flat)
        defect = gate * proposed + (1.0 - gate) * bb.mean(dim=1)
        return neutral, defect

    def coefficients(self):
        return (
            torch.nn.functional.softplus(self.x_raw).item(),
            torch.nn.functional.softplus(self.lambda_raw).item(),
        )