import math

import torch
import torch.nn as nn
import torch.nn.functional as F
# CMA is text-only. Disable Transformers' optional vision backend before it
# imports PreTrainedModel: some cloud images expose a system torchvision built
# against a different PyTorch, which otherwise crashes on torchvision::nms.
from transformers.utils import import_utils as _transformers_import_utils
_transformers_import_utils._torchvision_available = False
_transformers_import_utils.is_torchvision_available = lambda: False
try:
    from transformers import GenerationMixin
except ImportError:
    from transformers.generation import GenerationMixin
from transformers import PretrainedConfig, PreTrainedModel
from transformers.modeling_outputs import CausalLMOutputWithPast


class CMAConfig(PretrainedConfig):
    model_type = "cma"

    def __init__(
        self,
        vocab_size=256,
        seq_len=4096,
        d_model=96,
        n_layers=12,
        n_heads=6,
        n_kv_heads=2,
        chunk=8,
        cma_heads=2,
        expand=2,
        cma_identity_prob=0.90,
        max_position_embeddings=None,
        n_positions=None,
        n_ctx=None,
        **kwargs,
    ):
        super().__init__(**kwargs)
        self.vocab_size = vocab_size
        self.seq_len = seq_len
        self.max_position_embeddings = max_position_embeddings or seq_len
        self.n_positions = n_positions or self.max_position_embeddings
        self.n_ctx = n_ctx or self.max_position_embeddings
        self.d_model = d_model
        self.n_layers = n_layers
        self.n_heads = n_heads
        self.n_kv_heads = n_kv_heads
        self.chunk = chunk
        self.cma_heads = cma_heads
        self.expand = expand
        self.cma_identity_prob = cma_identity_prob


class RMSNorm(nn.Module):
    def __init__(self, d):
        super().__init__()
        self.w = nn.Parameter(torch.ones(d))

    def forward(self, x):
        return F.rms_norm(
            x, (x.shape[-1],), self.w.to(dtype=x.dtype), eps=1e-6
        )


def rope_cache(seq_len, head_dim, device, base=10000.0):
    inv = 1.0 / (
        base ** (torch.arange(0, head_dim, 2, device=device).float() / head_dim)
    )
    t = torch.arange(seq_len, device=device).float()
    freqs = torch.outer(t, inv)
    return torch.cos(freqs), torch.sin(freqs)


def apply_rope(x, cos, sin):
    x1, x2 = x.chunk(2, dim=-1)
    return torch.cat([x1 * cos - x2 * sin, x2 * cos + x1 * sin], dim=-1)


class TokenAttention(nn.Module):
    def __init__(self, d, n_heads, n_kv_heads):
        super().__init__()
        if d % n_heads or n_heads % n_kv_heads:
            raise ValueError("d_model % n_heads and n_heads % n_kv_heads must be zero.")
        self.h, self.kv_h, self.hd = n_heads, n_kv_heads, d // n_heads
        self.q = nn.Linear(d, d, bias=False)
        self.k = nn.Linear(d, n_kv_heads * self.hd, bias=False)
        self.v = nn.Linear(d, n_kv_heads * self.hd, bias=False)
        self.o = nn.Linear(d, d, bias=False)
        self.qn, self.kn = RMSNorm(self.hd), RMSNorm(self.hd)

    def forward(self, x, cos, sin):
        B, T, d = x.shape
        q = self.q(x).view(B, T, self.h, self.hd).transpose(1, 2)
        k = self.k(x).view(B, T, self.kv_h, self.hd).transpose(1, 2)
        v = self.v(x).view(B, T, self.kv_h, self.hd).transpose(1, 2)
        q, k = self.qn(q), self.kn(k)
        q, k = apply_rope(q, cos, sin), apply_rope(k, cos, sin)
        # Native GQA keeps compact K/V heads and lets SDPA select a fused CUDA
        # implementation without materializing repeated K/V activations.
        y = F.scaled_dot_product_attention(
            q,
            k,
            v,
            is_causal=True,
            enable_gqa=self.h != self.kv_h,
        )
        return self.o(y.transpose(1, 2).reshape(B, T, d))


class CMA(nn.Module):
    def __init__(self, d, chunk=32, heads=4, expand=2, identity_prob=0.90):
        super().__init__()
        if min(d, chunk, heads, expand) <= 0:
            raise ValueError("CMA dimensions and expansion must be positive.")
        if d % chunk or chunk % heads:
            raise ValueError("CMA requires d % chunk == 0 and chunk % heads == 0.")
        if d // chunk < 2:
            raise ValueError("CMA requires at least two channel chunks.")
        self.n, self.c, self.h = d // chunk, chunk, heads
        self.hd = chunk // heads
        self.chunk_emb = nn.Parameter(torch.randn(self.n, chunk) * 0.02)
        self.wqk = nn.Parameter(torch.randn(self.n, chunk, 2 * chunk) * 0.02)
        self.global_proj = nn.Linear(d, chunk, bias=False)
        self.wv = nn.Linear(d, d * expand, bias=False)
        self.bias = nn.Parameter(torch.zeros(heads, self.n, self.n))
        self.qn, self.kn = RMSNorm(self.hd), RMSNorm(self.hd)
        self.logit_scale = nn.Parameter(torch.zeros(heads))
        self.layer_gain = nn.Parameter(torch.zeros(heads))
        self.route_gate_weight = nn.Parameter(
            torch.randn(heads, self.hd) * 0.02
        )
        self.route_gate_bias = nn.Parameter(torch.zeros(heads))
        self.gate = nn.Linear(d, d * expand, bias=False)
        self.o = nn.Linear(d * expand, d, bias=False)
        nn.init.zeros_(self.o.weight)
        if not 0.0 < identity_prob < 1.0:
            raise ValueError("cma_identity_prob must be between zero and one.")
        diagonal_bias = math.log(
            (self.n - 1) * identity_prob / (1.0 - identity_prob)
        )
        with torch.no_grad():
            self.bias.add_(torch.eye(self.n) * diagonal_bias)
        self.expand = expand
        self.record_diagnostics = False
        self.last_diagnostics = None

    def forward(self, x):
        B, T, d = x.shape
        xc = x.view(B * T, self.n, self.c)
        global_state = self.global_proj(x).reshape(B * T, 1, self.c)
        qk_in = xc + self.chunk_emb + global_state
        value_slots = self.wv(x).reshape(
            B * T, self.n, self.h, self.hd, self.expand
        )
        # Key j summarizes value slot j, restoring the key/value contract while
        # retaining the full dense value projection and exact parameter count.
        key_in = value_slots.mean(dim=-1).reshape(B * T, self.n, self.c)
        key_in = key_in + self.chunk_emb
        q = torch.einsum("bnc,nco->bno", qk_in, self.wqk[..., : self.c])
        k = torch.einsum("bnc,nco->bno", key_in, self.wqk[..., self.c :])
        v = value_slots.reshape(
            B * T, self.n, self.h, self.hd * self.expand
        )
        q = self.qn(q.view(B * T, self.n, self.h, self.hd)).transpose(1, 2)
        k = self.kn(k.view(B * T, self.n, self.h, self.hd)).transpose(1, 2)
        v = v.transpose(1, 2)
        q = F.normalize(q.float(), dim=-1).to(v.dtype)
        k = F.normalize(k.float(), dim=-1).to(v.dtype)
        scale = self.logit_scale.clamp(max=math.log(100.0)).exp()
        scale = scale.view(1, self.h, 1, 1)
        attn_logits = (q @ k.transpose(-2, -1)) * scale.to(q.dtype)
        attn_logits = attn_logits + self.bias.unsqueeze(0).to(q.dtype)
        attn = F.softmax(attn_logits, dim=-1, dtype=torch.float32).to(v.dtype)
        routed = attn @ v
        route_signal = (
            torch.einsum(
                "bhnd,hd->bhn", q.float(), self.route_gate_weight.float()
            )
            + self.route_gate_bias.float().view(1, self.h, 1)
            + self.layer_gain.float().view(1, self.h, 1)
        )
        route_coeff = torch.tanh(route_signal).to(v.dtype)
        contribution = route_coeff.unsqueeze(-1) * (routed - v)
        y = v + contribution
        if self.record_diagnostics:
            probs = attn.float().clamp_min(1e-9)
            entropy = -(probs * probs.log()).sum(dim=-1) / math.log(self.n)
            diagonal_mass = probs.diagonal(dim1=-2, dim2=-1).mean()
            if self.h > 1:
                flattened = probs.transpose(0, 1).reshape(self.h, -1)
                normalized = F.normalize(flattened, dim=-1)
                similarity = normalized @ normalized.mT
                off_diagonal = (
                    similarity.sum() - similarity.diagonal().sum()
                ) / (self.h * (self.h - 1))
            else:
                off_diagonal = probs.new_zeros(())
            self.last_diagnostics = {
                "entropy": entropy.mean().item(),
                "diagonal_mass": diagonal_mass.item(),
                "head_similarity": off_diagonal.item(),
                "gate_mean": route_coeff.float().mean().item(),
                "gate_std": route_coeff.float().std(unbiased=False).item(),
                "gate_abs_mean": route_coeff.float().abs().mean().item(),
                "gate_saturation": (
                    route_coeff.float().abs() > 0.95
                ).float().mean().item(),
                "contribution_ratio": (
                    contribution.float().norm() / (v.float().norm() + 1e-9)
                ).item(),
            }
        y = y.transpose(1, 2).reshape(B, T, d * self.expand)
        return self.o(y * F.silu(self.gate(x)))


class Block(nn.Module):
    def __init__(self, config):
        super().__init__()
        d = config.d_model
        self.n1, self.n2 = RMSNorm(d), RMSNorm(d)
        self.attn = TokenAttention(d, config.n_heads, config.n_kv_heads)
        self.mix = CMA(
            d,
            config.chunk,
            config.cma_heads,
            config.expand,
            config.cma_identity_prob,
        )

    def forward(self, x, cos, sin):
        x = x + self.attn(self.n1(x), cos, sin)
        x = x + self.mix(self.n2(x))
        return x


class CMAModel(nn.Module):
    def __init__(self, config):
        super().__init__()
        d = config.d_model
        self.config = config
        self.emb = nn.Embedding(config.vocab_size, d)
        self.blocks = nn.ModuleList(Block(config) for _ in range(config.n_layers))
        self.norm = RMSNorm(d)
        hd = d // config.n_heads
        cos, sin = rope_cache(config.seq_len, hd, "cpu")
        self.register_buffer("cos", cos)
        self.register_buffer("sin", sin)

    def forward_hidden(self, idx):
        if idx.size(1) > self.config.seq_len:
            idx = idx[:, -self.config.seq_len :]
        x = self.emb(idx)
        device_type = x.device.type
        rope_dtype = (
            torch.get_autocast_dtype(device_type)
            if torch.is_autocast_enabled(device_type)
            else x.dtype
        )
        cos = self.cos[: idx.size(1)].to(device=idx.device, dtype=rope_dtype)
        sin = self.sin[: idx.size(1)].to(device=idx.device, dtype=rope_dtype)
        for b in self.blocks:
            x = b(x, cos, sin)
        return self.norm(x)


class CMAForCausalLM(PreTrainedModel, GenerationMixin):
    config_class = CMAConfig
    base_model_prefix = "model"
    _tied_weights_keys = {"head.weight": "model.emb.weight"}
    all_tied_weights_keys = {"head.weight": "model.emb.weight"}

    def __init__(self, config):
        super().__init__(config)
        self.model = CMAModel(config)
        self.head = nn.Linear(config.d_model, config.vocab_size, bias=False)
        # Modern Transformers creates loader metadata and performs configured
        # tying in post_init(); omitting it leaves all_tied_weights_keys absent.
        self.post_init()
        self.head.weight = self.model.emb.weight

    def get_input_embeddings(self):
        return self.model.emb

    def set_input_embeddings(self, value):
        self.model.emb = value

    def get_output_embeddings(self):
        return self.head

    def set_output_embeddings(self, new_embeddings):
        self.head = new_embeddings

    def raw_logits(self, idx):
        return self.head(self.model.forward_hidden(idx))

    def logits(self, idx):
        return self.raw_logits(idx)

    def _masked_logits(self, input_ids, attention_mask):
        if attention_mask is None or bool(attention_mask.all()):
            return self.logits(input_ids)

        B, T = input_ids.shape
        out = None
        for i in range(B):
            keep = attention_mask[i].bool().nonzero(as_tuple=False).flatten()
            if keep.numel() == 0:
                keep = torch.tensor([T - 1], device=input_ids.device)
            trimmed = input_ids[i, keep].unsqueeze(0)
            logits_i = self.logits(trimmed)
            if out is None:
                out = logits_i.new_zeros(B, T, logits_i.size(-1))
            out[i, keep, :] = logits_i[0, -keep.numel() :, :]
        return out

    def forward(
        self,
        input_ids=None,
        attention_mask=None,
        labels=None,
        use_cache=False,
        past_key_values=None,
        **kwargs,
    ):
        if input_ids.size(1) > self.config.seq_len:
            input_ids = input_ids[:, -self.config.seq_len :]
            if attention_mask is not None:
                attention_mask = attention_mask[:, -self.config.seq_len :]
            if labels is not None:
                labels = labels[:, -self.config.seq_len :]
        logits = self._masked_logits(input_ids, attention_mask)
        loss = None
        if labels is not None:
            shift_logits = logits[:, :-1, :].contiguous()
            shift_labels = labels[:, 1:].contiguous()
            loss = F.cross_entropy(
                shift_logits.view(-1, shift_logits.size(-1)).float(),
                shift_labels.view(-1),
                ignore_index=-100,
            )
        return CausalLMOutputWithPast(loss=loss, logits=logits, past_key_values=None)

    def prepare_inputs_for_generation(self, input_ids, **kwargs):
        attention_mask = kwargs.get("attention_mask")
        result = {"input_ids": input_ids[:, -self.config.seq_len :]}
        if attention_mask is not None:
            result["attention_mask"] = attention_mask[:, -self.config.seq_len :]
        return result