"""Rose X1 model implementation for Hugging Face transformers.

Architecture (T-X4 family with XSA refresh gate):
  RoPE (half-split) + RMSNorm + SwiGLU + grouped-query attention with
  per-head QK-norm, plus an XSA *refresh gate* that re-injects the original
  token embedding through a gated depthwise-causal-conv path on a subset of
  layers (``config.refresh_gate_inject_layers``).

Cache design (after the T-X4 reference implementation):
  KV cache uses HF's ``DynamicCache``.  The refresh gate's conv history
  (last ``kernel-1`` timesteps of the normalised attention output) is stored
  in a plain dict monkey-patched onto the same ``DynamicCache`` object as
  ``_refresh_conv_state``, so both share one lifetime and no custom Cache
  subclass is needed.
"""
from typing import Optional

import torch
import torch.nn as nn
from torch.nn import functional as F
from transformers import PreTrainedModel
from transformers.cache_utils import DynamicCache
from transformers.generation.utils import GenerationMixin
from transformers.modeling_outputs import CausalLMOutputWithPast

try:
    from .configuration_rose_x1 import RoseX1Config
except ImportError:
    from configuration_rose_x1 import RoseX1Config


# ═══════════════════════════════════════════════════════════════════════════
#  Primitives
# ═══════════════════════════════════════════════════════════════════════════

class RMSNorm(nn.Module):
    """RMSNorm with fp32 internal computation, cast back to input dtype."""
    def __init__(self, dim: int, eps: float = 1e-5):
        super().__init__()
        self.eps = eps
        self.weight = nn.Parameter(torch.ones(dim))

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        in_dtype = x.dtype
        xf = x.float()
        out = xf * torch.rsqrt(xf.pow(2).mean(-1, keepdim=True) + self.eps)
        return (out * self.weight.float()).to(in_dtype)


def precompute_rope_cos_sin(head_dim: int, seq_len: int, theta: float = 100000.0):
    """Precompute RoPE cos/sin tables.  Returns (cos, sin) each (seq_len, head_dim//2)."""
    freqs = 1.0 / (theta ** (torch.arange(0, head_dim, 2, dtype=torch.float32) / head_dim))
    t = torch.arange(seq_len, dtype=torch.float32)
    angles = torch.outer(t, freqs)                       # (seq_len, head_dim//2)
    return angles.cos(), angles.sin()


def apply_rotary_emb(q: torch.Tensor, k: torch.Tensor,
                     cos: torch.Tensor, sin: torch.Tensor):
    """Half-split RoPE (matches the trainer's ``apply_rope``).

    cos / sin: (T, head_dim//2)  — already sliced to the right positions.
    q, k:      (B, H, T, head_dim)
    """
    cos = cos.unsqueeze(0).unsqueeze(0).to(q.dtype)      # (1,1,T,d//2)
    sin = sin.unsqueeze(0).unsqueeze(0).to(q.dtype)
    d2 = q.shape[-1] // 2

    q1, q2 = q[..., :d2], q[..., d2:]
    k1, k2 = k[..., :d2], k[..., d2:]

    q_out = torch.cat([q1 * cos - q2 * sin, q2 * cos + q1 * sin], dim=-1)
    k_out = torch.cat([k1 * cos - k2 * sin, k2 * cos + k1 * sin], dim=-1)
    return q_out, k_out


# ═══════════════════════════════════════════════════════════════════════════
#  Attention  (GQA + QK-norm + RoPE)
# ═══════════════════════════════════════════════════════════════════════════

class RoseX1Attention(nn.Module):
    def __init__(self, config: RoseX1Config, layer_idx: int):
        super().__init__()
        self.layer_idx = layer_idx
        self.n_head = config.num_attention_heads
        self.n_kv_heads = config.num_key_value_heads
        self.head_dim = config.head_dim
        self.n_rep = self.n_head // self.n_kv_heads

        self.q_proj = nn.Linear(config.hidden_size, self.n_head * self.head_dim, bias=False)
        self.k_proj = nn.Linear(config.hidden_size, self.n_kv_heads * self.head_dim, bias=False)
        self.v_proj = nn.Linear(config.hidden_size, self.n_kv_heads * self.head_dim, bias=False)
        self.o_proj = nn.Linear(self.n_head * self.head_dim, config.hidden_size, bias=False)

        # QK-norm: per-head RMSNorm on Q & K, applied BEFORE RoPE
        self.use_qk_norm = bool(getattr(config, "use_qk_norm", False))
        if self.use_qk_norm:
            self.q_norm = RMSNorm(self.head_dim, eps=config.rms_norm_eps)
            self.k_norm = RMSNorm(self.head_dim, eps=config.rms_norm_eps)

    def forward(self, x, rope_cos, rope_sin,
                past_key_value: Optional[DynamicCache] = None,
                use_cache: bool = False,
                attention_mask: Optional[torch.Tensor] = None):
        B, T, _ = x.size()

        q = self.q_proj(x).view(B, T, self.n_head, self.head_dim).transpose(1, 2)
        k = self.k_proj(x).view(B, T, self.n_kv_heads, self.head_dim).transpose(1, 2)
        v = self.v_proj(x).view(B, T, self.n_kv_heads, self.head_dim).transpose(1, 2)

        # ── QK-norm BEFORE RoPE (== trainer) ──────────────────────────────
        if self.use_qk_norm:
            q = self.q_norm(q)
            k = self.k_norm(k)

        # ── RoPE (half-split, cos/sin already sliced to current positions) ─
        q, k = apply_rotary_emb(q, k, rope_cos, rope_sin)

        # ── KV cache (DynamicCache.update handles concat internally) ──────
        if past_key_value is not None:
            k, v = past_key_value.update(k, v, self.layer_idx)

        S = k.size(2)

        # ── GQA expansion ─────────────────────────────────────────────────
        k = k.unsqueeze(2).expand(B, self.n_kv_heads, self.n_rep, S, self.head_dim) \
             .reshape(B, self.n_head, S, self.head_dim)
        v = v.unsqueeze(2).expand(B, self.n_kv_heads, self.n_rep, S, self.head_dim) \
             .reshape(B, self.n_head, S, self.head_dim)

        # ── Attention mask ────────────────────────────────────────────────
        # is_causal=True only for prefill (no cache) with T>1 and no padding
        # mask.  For decode (T=1) causal is trivially satisfied.
        is_causal = (past_key_value is None
                     or past_key_value.get_seq_length(self.layer_idx) == T)
        attn_mask = None
        if attention_mask is not None:
            key_pad = attention_mask.to(torch.bool)[:, None, None, :]   # (B,1,1,S)
            if is_causal and T > 1:
                causal = torch.ones(T, S, dtype=torch.bool, device=x.device) \
                             .tril(diagonal=S - T)
                attn_mask = key_pad & causal[None, None, :, :]
            else:
                attn_mask = key_pad.expand(B, 1, T, S)
            is_causal = False

        y = F.scaled_dot_product_attention(q, k, v,
                                           attn_mask=attn_mask,
                                           is_causal=is_causal)
        y = y.transpose(1, 2).contiguous().view(B, T, self.n_head * self.head_dim)
        return self.o_proj(y)


# ═══════════════════════════════════════════════════════════════════════════
#  MLP  (SwiGLU)
# ═══════════════════════════════════════════════════════════════════════════

class RoseX1MLP(nn.Module):
    def __init__(self, config: RoseX1Config):
        super().__init__()
        self.gate_proj = nn.Linear(config.hidden_size, config.intermediate_size, bias=False)
        self.up_proj   = nn.Linear(config.hidden_size, config.intermediate_size, bias=False)
        self.down_proj = nn.Linear(config.intermediate_size, config.hidden_size, bias=False)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        return self.down_proj(F.silu(self.gate_proj(x)) * self.up_proj(x))


# ═══════════════════════════════════════════════════════════════════════════
#  XSA Refresh Gate
# ═══════════════════════════════════════════════════════════════════════════

class RoseX1RefreshGate(nn.Module):
    """Re-injects the original token embedding (e0) into the residual stream,
    gated by a causal depthwise conv over the (detached) attention output.

    Conv history for cached generation is read/written via ``conv_state``
    (a plain dict living on the DynamicCache object).
    """
    def __init__(self, config: RoseX1Config):
        super().__init__()
        H = config.hidden_size
        self.kernel_size = int(getattr(config, "refresh_gate_kernel_size", 9))

        self.attn_norm  = RMSNorm(H, eps=config.rms_norm_eps)
        self.emb_norm   = RMSNorm(H, eps=config.rms_norm_eps)
        self.gate_proj  = nn.Linear(H, H, bias=False)
        self.value_proj = nn.Linear(H, H, bias=False)
        self.out_proj   = nn.Linear(H, H, bias=False)
        self.out_norm   = RMSNorm(H, eps=config.rms_norm_eps)

        # padding attribute documents intent; forward uses F.conv1d(padding=0)
        # with manual left-pad so the same weight works for cached & non-cached.
        self.causal_conv = nn.Conv1d(H, H, self.kernel_size,
                                     groups=H, bias=False,
                                     padding=self.kernel_size - 1)
        self.alpha = nn.Parameter(torch.tensor(0.1))

    def forward(self, h, attn_out, e0, conv_state=None, layer_idx=None):
        a = self.attn_norm(attn_out.detach())
        e = self.emb_norm(e0)

        k = self.kernel_size
        B, T, D = a.shape

        if conv_state is not None:
            # ── cached generation: prepend stored history ─────────────────
            prev = conv_state.get(layer_idx)
            if prev is None or prev.size(0) != B:
                prev = a.new_zeros(B, k - 1, D)
            a_ext = torch.cat([prev, a], dim=1)                    # (B, k-1+T, D)
            conv_state[layer_idx] = a_ext[:, -(k - 1):, :].detach()
        else:
            # ── no cache (training / full recompute): left-pad zeros ──────
            a_ext = F.pad(a, (0, 0, k - 1, 0))                    # (B, k-1+T, D)

        # Manual left-pad + padding=0 conv  (== trainer's CausalDepthwiseConv1d)
        c = F.conv1d(a_ext.transpose(1, 2),
                     self.causal_conv.weight,
                     bias=None, padding=0, groups=D)
        c = c.transpose(1, 2)                                      # (B, T, D)

        gate  = self.gate_proj(a) + c
        value = self.value_proj(e)
        z = self.out_norm(self.out_proj(F.silu(gate) * value))
        return h + self.alpha * z


# ═══════════════════════════════════════════════════════════════════════════
#  Decoder layer
# ═══════════════════════════════════════════════════════════════════════════

class RoseX1DecoderLayer(nn.Module):
    def __init__(self, config: RoseX1Config, layer_idx: int):
        super().__init__()
        self.layer_idx = layer_idx
        self.input_layernorm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
        self.self_attn = RoseX1Attention(config, layer_idx)
        self.post_attention_layernorm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
        self.mlp = RoseX1MLP(config)

        inject = list(getattr(config, "refresh_gate_inject_layers", []) or [])
        self.has_refresh = (bool(getattr(config, "refresh_gate_enabled", False))
                            and layer_idx in inject)
        if self.has_refresh:
            self.refresh_gate = RoseX1RefreshGate(config)

    def forward(self, x, e0, rope_cos, rope_sin,
                past_key_value=None, use_cache=False,
                attention_mask=None, conv_state=None):
        attn_out = self.self_attn(self.input_layernorm(x), rope_cos, rope_sin,
                                  past_key_value, use_cache, attention_mask)
        x = x + attn_out
        # Refresh gate fires AFTER attention residual, BEFORE FFN (== trainer)
        if self.has_refresh:
            x = self.refresh_gate(x, attn_out, e0,
                                  conv_state=conv_state,
                                  layer_idx=self.layer_idx)
        x = x + self.mlp(self.post_attention_layernorm(x))
        return x


# ═══════════════════════════════════════════════════════════════════════════
#  Base / backbone / head
# ═══════════════════════════════════════════════════════════════════════════

class RoseX1PreTrainedModel(PreTrainedModel):
    config_class = RoseX1Config
    base_model_prefix = "model"
    supports_gradient_checkpointing = False
    _no_split_modules = ["RoseX1DecoderLayer"]
    _supports_sdpa = True

    def _init_weights(self, module):
        std = self.config.initializer_range
        if isinstance(module, nn.Linear):
            nn.init.normal_(module.weight, mean=0.0, std=std)
            if module.bias is not None:
                nn.init.zeros_(module.bias)
        elif isinstance(module, nn.Embedding):
            nn.init.normal_(module.weight, mean=0.0, std=std)
        elif isinstance(module, nn.Conv1d):
            nn.init.normal_(module.weight, mean=0.0, std=std)
        elif isinstance(module, RMSNorm):
            nn.init.ones_(module.weight)


class RoseX1Model(nn.Module):
    """Backbone: embed → N × decoder layer → final norm."""
    def __init__(self, config: RoseX1Config):
        super().__init__()
        self.embed_tokens = nn.Embedding(config.vocab_size, config.hidden_size)
        self.dropout = nn.Dropout(config.attention_dropout)
        self.layers = nn.ModuleList(
            [RoseX1DecoderLayer(config, i) for i in range(config.num_hidden_layers)]
        )
        self.norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)

    def forward(self, input_ids, e0, rope_cos, rope_sin,
                past_key_value=None, use_cache=False,
                attention_mask=None, conv_state=None):
        x = self.dropout(self.embed_tokens(input_ids))
        for layer in self.layers:
            x = layer(x, e0, rope_cos, rope_sin,
                      past_key_value, use_cache, attention_mask, conv_state)
        return self.norm(x)


class RoseX1ForCausalLM(RoseX1PreTrainedModel, GenerationMixin):
    # Dict format required by modern transformers' get_expanded_tied_weights_keys.
    # Tells HF: "lm_head.weight is tied to model.embed_tokens.weight — if it's
    # missing from the checkpoint, fill it from the embedding, don't warn."
    _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"}

    def __init__(self, config: RoseX1Config):
        super().__init__(config)
        self.model = RoseX1Model(config)

        # Always create lm_head; tie it when configured (standard HF pattern).
        self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False)
        if config.tie_word_embeddings:
            self.lm_head.weight = self.model.embed_tokens.weight

        self._rope_cache = None          # (cos, sin) cached on device
        self.post_init()

    # ── Embedding accessors (used by tie_weights / resize) ────────────────
    def get_input_embeddings(self):
        return self.model.embed_tokens

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

    def get_output_embeddings(self):
        return self.lm_head

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

    # ── RoPE cache ────────────────────────────────────────────────────────
    def _get_rope(self, seq_len: int, device: torch.device):
        cache = self._rope_cache
        if (cache is None
                or cache[0].device != device
                or cache[0].size(0) < seq_len):
            cos, sin = precompute_rope_cos_sin(
                self.config.head_dim, seq_len, self.config.rope_theta)
            cache = (cos.to(device), sin.to(device))
            self._rope_cache = cache
        return cache[0][:seq_len], cache[1][:seq_len]

    # ── Generation plumbing ───────────────────────────────────────────────
    def prepare_inputs_for_generation(self, input_ids,
                                      past_key_values=None,
                                      attention_mask=None, **kwargs):
        # When a cache with content exists, feed only the newest token.
        if past_key_values is not None and past_key_values.get_seq_length() > 0:
            input_ids = input_ids[:, -1:]
        return {
            "input_ids": input_ids,
            "attention_mask": attention_mask,
            "past_key_values": past_key_values,
            "use_cache": True,
        }

    # ── Forward ───────────────────────────────────────────────────────────
    def forward(
        self,
        input_ids: torch.Tensor,
        attention_mask: Optional[torch.Tensor] = None,
        labels: Optional[torch.Tensor] = None,
        past_key_values: Optional[DynamicCache] = None,
        use_cache: bool = False,
        **kwargs,
    ) -> CausalLMOutputWithPast:
        B, T = input_ids.size()

        # ── Conv-state cache for the refresh gate ─────────────────────────
        # Monkey-patched onto the DynamicCache so it shares the cache's
        # lifetime.  No custom Cache subclass needed.
        conv_state = None
        if use_cache:
            if past_key_values is None:
                past_key_values = DynamicCache()
            if not hasattr(past_key_values, "_refresh_conv_state"):
                past_key_values._refresh_conv_state = {}
            conv_state = past_key_values._refresh_conv_state

        # ── Position from cache length (no explicit position_ids needed) ──
        past_len = (past_key_values.get_seq_length()
                    if past_key_values is not None else 0)

        # ── Embeddings ────────────────────────────────────────────────────
        e0 = self.model.embed_tokens(input_ids)     # original embedding for refresh gate

        # ── RoPE: precompute up to past_len+T, slice to current positions ─
        cos, sin = self._get_rope(past_len + T, input_ids.device)
        cos, sin = cos[past_len:], sin[past_len:]   # (T, head_dim//2)

        # ── Backbone ──────────────────────────────────────────────────────
        hidden = self.model(
            input_ids, e0, cos, sin,
            past_key_values if use_cache else None,
            use_cache, attention_mask, conv_state,
        )

        # ── Head ──────────────────────────────────────────────────────────
        logits = self.lm_head(hidden).float()       # fp32 for stable logprobs

        loss = None
        if labels is not None:
            loss = F.cross_entropy(
                logits[..., :-1, :].contiguous().view(-1, self.config.vocab_size),
                labels[..., 1:].contiguous().view(-1),
                ignore_index=-100,
            )

        return CausalLMOutputWithPast(
            loss=loss,
            logits=logits,
            past_key_values=past_key_values if use_cache else None,
        )


# ── Optional registration (lets model_type="rose_x1" resolve without auto_map) ──
try:
    from transformers import AutoConfig, AutoModelForCausalLM
    try:
        AutoConfig.register("rose_x1", RoseX1Config)
    except Exception:
        pass
    try:
        AutoModelForCausalLM.register(RoseX1Config, RoseX1ForCausalLM)
    except Exception:
        pass
except Exception:
    pass