import torch
import torch.nn as nn

from .text_encoder import ConvNeXtWrapper


class ReferenceEncoder(nn.Module):
    def __init__(
        self,
        in_channels: int = 144,
        d_model: int = 256,
        hidden_dim: int = 1024,
        num_blocks: int = 6,
        num_tokens: int = 50,
        num_heads: int = 2,
        kernel_size: int = 5,
        dilation_lst: list = None,
        prototype_dim: int = 256,
        n_units: int = 256,
        style_value_dim: int = 256,
    ):
        super().__init__()
        self.d_model = d_model
        self.num_tokens = num_tokens

        if hidden_dim % d_model != 0:
            raise ValueError(f"hidden_dim ({hidden_dim}) must be divisible by d_model ({d_model})")
        mlp_ratio = hidden_dim // d_model

        self.input_proj = nn.Conv1d(in_channels, d_model, kernel_size=1)
        self.convnext = ConvNeXtWrapper(
            d_model,
            n_layers=num_blocks,
            expansion_factor=mlp_ratio,
            kernel_size=kernel_size,
            dilation_lst=dilation_lst,
        )

        self.ref_keys = nn.Parameter(torch.randn(num_tokens, prototype_dim) * 0.02)
        self.q_proj = nn.Linear(prototype_dim, n_units) if prototype_dim != n_units else nn.Identity()
        self.out_proj = nn.Linear(n_units, style_value_dim) if n_units != style_value_dim else nn.Identity()

        self.attn1 = nn.MultiheadAttention(
            embed_dim=n_units, num_heads=num_heads, kdim=d_model, vdim=d_model, batch_first=True
        )
        self.attn2 = nn.MultiheadAttention(
            embed_dim=n_units, num_heads=num_heads, kdim=d_model, vdim=d_model, batch_first=True
        )

    def forward(self, z_ref: torch.Tensor, mask: torch.Tensor = None):
        B = z_ref.shape[0]
        x = self.input_proj(z_ref)
        x = self.convnext(x, mask=mask)
        kv = x.transpose(1, 2)

        key_padding_mask = None
        if mask is not None:
            key_padding_mask = (mask.squeeze(1) == 0)

        q0 = self.ref_keys.unsqueeze(0).expand(B, -1, -1)
        q0 = self.q_proj(q0)

        q1, _ = self.attn1(query=q0, key=kv, value=kv, key_padding_mask=key_padding_mask, need_weights=False)
        q2 = q0 + q1
        out, _ = self.attn2(query=q2, key=kv, value=kv, key_padding_mask=key_padding_mask, need_weights=False)
        return self.out_proj(out)

    @staticmethod
    def remap_legacy_state_dict(state_dict: dict) -> dict:
        """Remap pre-refactor checkpoints (per-layer pre-norm + FFN) onto current layout."""
        remapped = {}
        legacy_prefix_map = {
            "attn_layers.0.attn.": "attn1.",
            "attn_layers.1.attn.": "attn2.",
        }
        drop_substrings = (".norm_q.", ".norm_kv.", ".ffn.", "pos_emb.")
        for k, v in state_dict.items():
            if any(s in k for s in drop_substrings):
                continue
            new_key = k
            for old, new in legacy_prefix_map.items():
                if new_key.startswith(old):
                    new_key = new + new_key[len(old):]
                    break
            remapped[new_key] = v
        return remapped