"""Independent implementation of LiNO, arXiv:2606.03262, Eqs. (5)-(44).

Only the efficient scattering variant for one-dimensional fields is implemented.
Unspecified numerical choices are documented in results/reproduction_spec.md.
"""

import torch
from torch import nn
from torch.nn import functional as F


class AdaptiveMirror(nn.Module):
    def __init__(self, width, eps):
        super().__init__()
        self.direction = nn.Linear(width, width)
        self.eps = eps

    def forward(self, field):
        normal = F.normalize(self.direction(field), dim=-1, eps=self.eps)
        component = (field * normal).sum(-1, keepdim=True)
        return field - 2 * component * normal


class AdaptiveRefraction(nn.Module):
    def __init__(self, width, rate_range, eps):
        super().__init__()
        self.direction = nn.Linear(width, width)
        self.contrast = nn.Parameter(torch.zeros(()))
        self.rate_range, self.eps = rate_range, eps

    def forward(self, field):
        normal = F.normalize(self.direction(field), dim=-1, eps=self.eps)
        normal_component = (field * normal).sum(-1, keepdim=True) * normal
        return field + self.rate_range * self.contrast.tanh() * normal_component


class EfficientScattering(nn.Module):
    def __init__(self, width, eps):
        super().__init__()
        self.query = nn.Linear(width, width, bias=False)
        self.key = nn.Linear(width, width, bias=False)
        self.value = nn.Linear(width, width, bias=False)
        self.coordinates = nn.Sequential(nn.Linear(1, width), nn.GELU(), nn.Linear(width, width))
        self.diffusion = nn.Conv1d(width, width, 3, padding=1, groups=width, bias=False)
        self.local_projection = nn.Linear(width, width, bias=False)
        self.local_logit = nn.Parameter(torch.tensor(-1.0))
        self.log_strength = nn.Parameter(torch.tensor(-2.0))
        self.scale, self.eps = width ** -0.5, eps

    def forward(self, field, grid):
        position = self.coordinates(grid)
        query = F.elu(self.query(field) * self.scale + position) + 1
        key = F.elu(self.key(field) * self.scale + position) + 1
        value = self.value(field)
        # Eq. (33): aggregate over space before contracting each query.
        summary = key.transpose(1, 2) @ value
        numerator = query @ summary
        denominator = (query * key.sum(1, keepdim=True)).sum(-1, keepdim=True)
        global_average = numerator / (denominator + self.eps)
        local_field = self.diffusion(field.transpose(1, 2)).transpose(1, 2)
        local_average = self.local_projection(local_field)
        beta = self.local_logit.sigmoid()
        average = (1 - beta) * global_average + beta * local_average
        return self.log_strength.exp() * (average - field)


class LightBlock(nn.Module):
    def __init__(self, width, rate_range, eps):
        super().__init__()
        self.before_light = nn.LayerNorm(width)
        self.before_ffn = nn.LayerNorm(width)
        self.mirror = AdaptiveMirror(width, eps)
        self.refraction = AdaptiveRefraction(width, rate_range, eps)
        self.scattering = EfficientScattering(width, eps)
        self.gate = nn.Sequential(nn.Linear(width, 2 * width), nn.GELU(), nn.Linear(2 * width, 3))
        self.mix_projection = nn.Linear(width, width)
        self.ffn = nn.Sequential(nn.Linear(width, 2 * width), nn.GELU(), nn.Linear(2 * width, width))

    def forward(self, field, grid):
        normalized = self.before_light(field)
        weights = self.gate(normalized.mean(1)).softmax(-1)
        branches = (self.mirror(normalized), self.refraction(normalized), self.scattering(normalized, grid))
        mixture = sum(weights[:, i, None, None] * branch for i, branch in enumerate(branches))
        field = field + self.mix_projection(mixture)
        return field + self.ffn(self.before_ffn(field))


class LiNO(nn.Module):
    """[batch, space, 2] -> [batch, space, 1], normalized field and x input."""

    def __init__(self, in_channels=2, out_channels=1, width=128, depth=8,
                 scattering="efficient", rate_range=0.25, eps=1e-6):
        super().__init__()
        if scattering != "efficient":
            raise ValueError("This experiment requires efficient scattering")
        self.architecture = dict(in_channels=in_channels, out_channels=out_channels,
                                 width=width, depth=depth, scattering=scattering,
                                 rate_range=rate_range, eps=eps)
        self.lift = nn.Linear(in_channels, width)
        self.blocks = nn.ModuleList([LightBlock(width, rate_range, eps) for _ in range(depth)])
        self.project = nn.Linear(width, out_channels)

    def forward(self, inputs):
        if inputs.ndim != 3 or inputs.shape[-1] != self.architecture["in_channels"]:
            raise ValueError(f"Expected [B,N,{self.architecture['in_channels']}], got {inputs.shape}")
        grid = torch.linspace(0, 1, inputs.shape[1], device=inputs.device, dtype=inputs.dtype).view(1, -1, 1)
        field = self.lift(inputs)
        for block in self.blocks:
            field = block(field, grid)
        return self.project(field)