from __future__ import annotations

from math import isqrt

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


FACTOR_NAMES = ("line", "color", "texture", "layout")


class ValueOnlyResampler(nn.Module):
    """Position-free, bias-free resampling whose output is zero for zero values."""

    def __init__(self, width: int, tokens: int = 144, heads: int = 8) -> None:
        super().__init__()
        if width % heads:
            raise ValueError("width must be divisible by heads")
        self.heads = heads
        self.head_dim = width // heads
        self.queries = nn.Parameter(torch.empty(tokens, width))
        self.q_proj = nn.Linear(width, width, bias=False)
        self.k_proj = nn.Linear(width, width, bias=False)
        self.v_proj = nn.Linear(width, width, bias=False)
        self.out_proj = nn.Linear(width, width, bias=False)
        nn.init.normal_(self.queries, std=width**-0.5)

    def forward(self, values: torch.Tensor) -> torch.Tensor:
        batch, source_tokens, width = values.shape
        target_tokens = self.queries.shape[0]
        queries = self.queries.unsqueeze(0).expand(batch, -1, -1)
        q = self.q_proj(queries).view(batch, target_tokens, self.heads, self.head_dim).transpose(1, 2)
        k = self.k_proj(values).view(batch, source_tokens, self.heads, self.head_dim).transpose(1, 2)
        v = self.v_proj(values).view(batch, source_tokens, self.heads, self.head_dim).transpose(1, 2)
        result = F.scaled_dot_product_attention(q, k, v)
        return self.out_proj(result.transpose(1, 2).reshape(batch, target_tokens, width))


class SharedAxisStyleEncoder(nn.Module):
    """One VAE field and one shared axis space for embedding, routing, and transfer."""

    def __init__(
        self,
        dino_dim: int,
        style_dim: int = 1024,
        axis_count: int = 256,
        embedding_dim: int = 1024,
        anima_dim: int = 0,
        resolution: str = "s12",
        set_layers: int = 2,
        set_heads: int = 8,
        direct_context_axes: bool = False,
        axis_aligned: bool = False,
        disable_token_film: bool = False,
        translation_invariant_stats: bool = False,
    ) -> None:
        super().__init__()
        if resolution not in {"s12", "s24r"}:
            raise ValueError(f"unknown spatial resolution: {resolution}")
        if axis_count % set_heads:
            raise ValueError("axis_count must be divisible by set_heads")
        self.resolution = resolution
        self.style_dim = style_dim
        self.axis_count = axis_count
        self.embedding_dim = embedding_dim
        self.set_layers = set_layers
        self.set_heads = set_heads
        self.direct_context_axes = direct_context_axes
        self.axis_aligned = axis_aligned
        self.disable_token_film = disable_token_film
        self.translation_invariant_stats = translation_invariant_stats
        if disable_token_film and not axis_aligned:
            raise ValueError("disable_token_film requires axis_aligned=True")
        if translation_invariant_stats and not axis_aligned:
            raise ValueError("translation-invariant statistics require axis_aligned=True")
        last_stride = 2 if resolution == "s12" else 1
        self.stem = nn.Sequential(
            nn.Conv2d(16, 256, 3, stride=2, padding=1),
            nn.GELU(),
            nn.Conv2d(256, 512, 3, stride=2, padding=1),
            nn.GELU(),
            nn.Conv2d(512, style_dim, 3, stride=last_stride, padding=1),
        )
        self.token_norm = nn.LayerNorm(style_dim, elementwise_affine=False)
        self.dino_norm = nn.LayerNorm(dino_dim)
        self.dino_film = nn.Linear(dino_dim, 2 * style_dim, bias=False)
        self.dino_relation = nn.Linear(dino_dim, axis_count, bias=False)
        nn.init.zeros_(self.dino_film.weight)

        self.axis_dictionary = nn.Parameter(torch.empty(style_dim, axis_count))
        nn.init.orthogonal_(self.axis_dictionary)
        self.membership_logits = nn.Parameter(torch.zeros(len(FACTOR_NAMES), axis_count))
        self.moment_projection = nn.Linear(2 * axis_count, axis_count, bias=False)
        relation_layer = nn.TransformerEncoderLayer(
            axis_count,
            nhead=set_heads,
            dim_feedforward=2 * axis_count,
            dropout=0.0,
            activation="gelu",
            batch_first=True,
            norm_first=True,
        )
        self.relation_encoder = nn.TransformerEncoder(
            relation_layer,
            set_layers,
            enable_nested_tensor=False,
        )
        self.axis_reliability = nn.Linear(axis_count, axis_count)
        self.face_reliability = nn.Linear(axis_count, 1)
        self.embedding_head = nn.Sequential(
            nn.LayerNorm(axis_count),
            nn.Linear(axis_count, embedding_dim),
        )
        self.resampler = (
            ValueOnlyResampler(style_dim, tokens=144, heads=set_heads)
            if resolution == "s24r"
            else None
        )
        # Construct optional modality modules last so matched runs share identical common initialization.
        self.anima_norm = nn.LayerNorm(anima_dim) if anima_dim else None
        self.anima_film = nn.Linear(anima_dim, 2 * style_dim, bias=False) if anima_dim else None
        self.anima_relation = nn.Linear(anima_dim, axis_count, bias=False) if anima_dim else None
        if self.anima_film is not None:
            nn.init.zeros_(self.anima_film.weight)
        if axis_aligned:
            self.axis_moment_weights = nn.Parameter(torch.empty(axis_count, 2))
            nn.init.normal_(self.axis_moment_weights, std=axis_count**-0.5)
        if translation_invariant_stats:
            self.axis_power_weights = nn.Parameter(torch.zeros(axis_count, 4))

    def model_config(self) -> dict[str, int | str | bool]:
        return {
            "style_dim": self.style_dim,
            "axis_count": self.axis_count,
            "embedding_dim": self.embedding_dim,
            "anima_dim": self.anima_norm.normalized_shape[0] if self.anima_norm is not None else 0,
            "resolution": self.resolution,
            "set_layers": self.set_layers,
            "set_heads": self.set_heads,
            "direct_context_axes": self.direct_context_axes,
            "axis_aligned": self.axis_aligned,
            "disable_token_film": self.disable_token_film,
            "translation_invariant_stats": self.translation_invariant_stats,
        }

    def _field(
        self,
        latent: torch.Tensor,
        scale: torch.Tensor,
        shift: torch.Tensor,
        axis_shift: torch.Tensor | None = None,
    ) -> tuple[torch.Tensor, torch.Tensor]:
        batch, references = latent.shape[:2]
        field = self.stem(latent.flatten(0, 1).float())
        field = field.flatten(2).transpose(1, 2).reshape(batch, references, -1, field.shape[1])
        field = self.token_norm(field)
        field = field * (1 + scale.unsqueeze(2)) + shift.unsqueeze(2)
        axes = torch.einsum("brnd,dk->brnk", field, self.axis_dictionary.to(field.dtype))
        if axis_shift is not None:
            axes = axes + axis_shift.unsqueeze(2)
        return field, axes

    def _moments(self, axes: torch.Tensor) -> torch.Tensor:
        mean = axes.mean(dim=2)
        log_std = torch.log(axes.std(dim=2, unbiased=False).clamp_min(1e-6))
        if self.axis_aligned:
            coordinates = (
                torch.stack((mean, log_std), dim=-1)
                * self.axis_moment_weights.to(mean.dtype)
            ).sum(-1)
            if self.translation_invariant_stats:
                power = self._power_statistics(axes).to(mean.dtype)
                coordinates = coordinates + (
                    power * self.axis_power_weights.to(mean.dtype)
                ).sum(-1)
            return coordinates
        moments = torch.cat((mean, log_std), dim=-1)
        return self.moment_projection(moments)

    def _power_statistics(self, axes: torch.Tensor) -> torch.Tensor:
        """Phase-free radial and directional power for each aligned spatial axis."""
        side = isqrt(axes.shape[2])
        if side * side != axes.shape[2]:
            raise ValueError(f"unexpected spatial token count: {axes.shape[2]}")
        fy = torch.fft.fftfreq(side, device=axes.device)[:, None]
        fx = torch.fft.rfftfreq(side, device=axes.device)[None, :]
        radius2 = fy.square() + fx.square()
        radius = radius2.sqrt()
        basis = torch.stack((
            ((radius > 0) & (radius <= 0.25)).float(),
            ((radius > 0.25) & (radius <= 0.5)).float(),
            (fx.square() - fy.square()) / radius2.clamp_min(1e-12),
            4 * fx.square() * fy.square() / radius2.square().clamp_min(1e-12),
        ))
        rfft_weight = torch.ones(side // 2 + 1, device=axes.device)
        rfft_weight[1:-1] = 2
        maps = axes.float().permute(0, 1, 3, 2).reshape(
            *axes.shape[:2],
            axes.shape[-1],
            side,
            side,
        )
        maps = maps - maps.mean(dim=(-2, -1), keepdim=True)
        power = torch.fft.rfft2(maps, norm="ortho").abs().square()
        power = power * rfft_weight[None, None, None, None, :]
        power = power / power.sum(dim=(-2, -1), keepdim=True).clamp_min(1e-12)
        return torch.einsum("brkhw,phw->brkp", power, basis)

    def factor_coordinates(self, axis_coordinates: torch.Tensor) -> torch.Tensor:
        """Apply overlapping factor memberships without creating factor branches."""
        memberships = torch.sigmoid(self.membership_logits).to(axis_coordinates.dtype)
        return axis_coordinates.unsqueeze(-2) * memberships

    def encode_vae_axes(self, latents: torch.Tensor) -> torch.Tensor:
        """Encode VAE latents without DINO/Anima context for the frozen transfer teacher."""
        if latents.ndim != 5 or latents.shape[2] != 16:
            raise ValueError(f"expected latents [B,R,16,H,W], got {tuple(latents.shape)}")
        batch, references = latents.shape[:2]
        zeros = torch.zeros(
            batch,
            references,
            self.axis_dictionary.shape[0],
            device=latents.device,
            dtype=latents.dtype,
        )
        _, axes = self._field(latents, zeros, zeros)
        return self._moments(axes)

    def _user_axis_gate(self, user_weights: torch.Tensor) -> torch.Tensor:
        memberships = torch.sigmoid(self.membership_logits)
        factors = user_weights[..., : len(FACTOR_NAMES)].unsqueeze(-1)
        return 1 - torch.prod(1 - factors * memberships[None, None], dim=-2)

    def _controlled_tokens(
        self,
        axes: torch.Tensor,
        beta: torch.Tensor,
        view_gate: torch.Tensor,
    ) -> torch.Tensor:
        gated = axes * beta.unsqueeze(2) * view_gate[:, :, None, None]
        tokens = torch.einsum(
            "brnk,dk->brnd",
            gated,
            self.axis_dictionary.to(gated.dtype),
        )
        if self.resampler is not None:
            batch, references, token_count, width = tokens.shape
            tokens = self.resampler(tokens.reshape(batch * references, token_count, width))
            tokens = tokens.reshape(batch, references, -1, width)
        return tokens

    def soft_orthogonality_loss(self) -> torch.Tensor:
        dictionary = F.normalize(self.axis_dictionary, dim=0)
        gram = dictionary.T @ dictionary
        return (gram - torch.eye(gram.shape[0], device=gram.device, dtype=gram.dtype)).square().mean()

    def forward(
        self,
        full_latents: torch.Tensor,
        dino: torch.Tensor,
        reference_valid: torch.Tensor,
        face_latents: torch.Tensor | None = None,
        face_valid: torch.Tensor | None = None,
        user_weights: torch.Tensor | None = None,
        mode: str = "auto",
        overall_style_gain: float | torch.Tensor = 1.0,
        build_memory: bool = True,
        anima: torch.Tensor | None = None,
    ) -> dict[str, torch.Tensor]:
        if full_latents.ndim != 5 or full_latents.shape[2] != 16:
            raise ValueError(f"expected full latents [B,R,16,H,W], got {tuple(full_latents.shape)}")
        if not reference_valid.any(dim=1).all():
            raise ValueError("every sample needs at least one valid reference")
        if mode not in {"auto", "assisted", "manual"}:
            raise ValueError(f"unknown routing mode: {mode}")

        batch, references = full_latents.shape[:2]
        if user_weights is None:
            user_weights = torch.ones(
                batch,
                references,
                len(FACTOR_NAMES) + 1,
                device=full_latents.device,
                dtype=full_latents.dtype,
            )
        context = self.dino_norm(dino.float())
        scale, shift = self.dino_film(context).chunk(2, dim=-1)
        relation_context = self.dino_relation(context)
        if self.anima_norm is not None:
            if anima is None:
                raise ValueError("Anima features are required when anima_dim is configured")
            anima_context = self.anima_norm(anima.float().flatten(start_dim=2))
            anima_scale, anima_shift = self.anima_film(anima_context).chunk(2, dim=-1)
            scale = scale + anima_scale
            shift = shift + anima_shift
            relation_context = relation_context + self.anima_relation(anima_context)
        elif anima is not None:
            raise ValueError("Anima features were provided but anima_dim=0")
        if self.disable_token_film:
            scale = torch.zeros_like(scale)
            shift = torch.zeros_like(shift)
        axis_shift = (
            relation_context
            if self.axis_aligned and self.direct_context_axes
            else None
        )
        full_field, full_axes = self._field(
            full_latents,
            scale,
            shift,
            axis_shift,
        )
        full_coordinates = self._moments(full_axes)
        if self.direct_context_axes and not self.axis_aligned:
            full_coordinates = full_coordinates + relation_context

        if face_latents is None:
            face_valid = torch.zeros_like(reference_valid)
            face_axes = None
            face_coordinates = torch.zeros_like(full_coordinates)
        else:
            if face_valid is None:
                raise ValueError("face_valid is required with face_latents")
            _, face_axes = self._field(
                face_latents,
                scale,
                shift,
                axis_shift,
            )
            face_coordinates = self._moments(face_axes)

        preliminary = (
            full_coordinates
            + face_coordinates * face_valid.unsqueeze(-1)
            + (0 if self.direct_context_axes else relation_context)
        )
        contextual = self.relation_encoder(preliminary, src_key_padding_mask=~reference_valid)
        reliability = torch.sigmoid(self.axis_reliability(contextual))
        predicted_face = torch.sigmoid(self.face_reliability(contextual).squeeze(-1))

        if mode == "auto":
            face_weight = predicted_face
        elif mode == "assisted":
            face_weight = predicted_face * user_weights[..., -1]
        else:
            face_weight = user_weights[..., -1]
        face_weight = face_weight * face_valid * reference_valid
        coordinates = (
            full_coordinates + face_weight.unsqueeze(-1) * face_coordinates
        ) / (1 + face_weight.unsqueeze(-1))

        axis_gate = self._user_axis_gate(user_weights.float())
        valid = reference_valid.unsqueeze(-1).to(axis_gate.dtype)
        if mode == "auto":
            weights = reliability * valid
            gamma = torch.ones(batch, self.axis_dictionary.shape[1], device=weights.device, dtype=weights.dtype)
        elif mode == "assisted":
            weights = reliability * axis_gate * valid
            gamma = axis_gate.amax(dim=1)
        else:
            weights = axis_gate * valid
            gamma = axis_gate.amax(dim=1)

        denominator = weights.sum(dim=1, keepdim=True)
        alpha = torch.where(
            denominator > 0,
            weights / denominator.clamp_min(torch.finfo(weights.dtype).tiny),
            torch.zeros_like(weights),
        )
        beta = alpha * gamma.unsqueeze(1)
        mixed = gamma * torch.einsum("brk,brk->bk", alpha, coordinates)

        gain = torch.as_tensor(overall_style_gain, device=mixed.device, dtype=mixed.dtype)
        while gain.ndim < mixed.ndim:
            gain = gain.unsqueeze(-1)
        mixed = mixed * gain
        beta = beta * gain.unsqueeze(1) if gain.ndim == 2 else beta * gain

        pre_embedding = self.embedding_head(mixed)
        output = {
            "embedding": F.normalize(pre_embedding, dim=-1),
            "pre_embedding": pre_embedding,
            "style_present": mixed.ne(0).any(dim=-1),
            "reference_axis_coordinates": coordinates,
            "axis_coordinates": mixed,
            "axis_reliability": reliability,
            "axis_weights": weights,
            "axis_alpha": alpha,
            "axis_gamma": gamma,
            "factor_memberships": torch.sigmoid(self.membership_logits),
            "reference_factor_coordinates": self.factor_coordinates(coordinates),
            "factor_coordinates": self.factor_coordinates(mixed),
            "face_weight": face_weight,
        }
        if build_memory:
            full_tokens = self._controlled_tokens(
                full_axes,
                beta,
                reference_valid.to(beta.dtype),
            )
            memories = [full_tokens.flatten(1, 2)]
            memory_masks = [
                reference_valid[:, :, None].expand(-1, -1, full_tokens.shape[2]).flatten(1)
            ]
            if face_axes is not None:
                face_tokens = self._controlled_tokens(face_axes, beta, face_weight)
                memories.append(face_tokens.flatten(1, 2))
                memory_masks.append(
                    (reference_valid & face_valid)[:, :, None]
                    .expand(-1, -1, face_tokens.shape[2])
                    .flatten(1)
                )
            output.update({
                "style_memory": torch.cat(memories, dim=1),
                "style_memory_valid": torch.cat(memory_masks, dim=1),
                "style_condition": torch.einsum(
                    "bk,dk->bd",
                    mixed,
                    self.axis_dictionary.to(mixed.dtype),
                ),
            })
        return output


def load_shared_axis_state(
    model: SharedAxisStyleEncoder,
    state: dict[str, torch.Tensor],
) -> None:
    """Load a checkpoint, deriving aligned per-axis moments from a legacy dense head."""
    state = dict(state)
    if model.axis_aligned and "axis_moment_weights" not in state:
        dense = state["moment_projection.weight"]
        axes = model.axis_count
        state["axis_moment_weights"] = torch.stack((
            dense.diagonal(),
            dense[:, axes:].diagonal(),
        ), dim=-1)
    if model.translation_invariant_stats and "axis_power_weights" not in state:
        state["axis_power_weights"] = torch.zeros_like(model.axis_power_weights)
    model.load_state_dict(state)