import torch
import torch.nn as nn
import torch.nn.functional as F
from transformers import PreTrainedModel
from transformers.modeling_outputs import CausalLMOutputWithPast
from transformers.generation import GenerationMixin
from .configuration_trm_text_ism_v6 import TRMTextISMConfig

def apply_rope(x, cos, sin):
    B, H, S, D = x.shape
    half = D // 2
    # Slice to current seq length and ensure dtypes match x
    c = cos[:, :, :S, :].to(x.dtype)
    s = sin[:, :, :S, :].to(x.dtype)
    x1, x2 = x[..., :half], x[..., half:]
    return torch.cat([x1 * c - x2 * s, x2 * c + x1 * s], dim=-1)

class SwiGLUMLP(nn.Module):
    def __init__(self, config):
        super().__init__()
        h = config.mlp_hidden_size or int(config.dim * config.mlp_ratio)
        self.gate_proj = nn.Linear(config.dim, h, bias=False)
        self.up_proj = nn.Linear(config.dim, h, bias=False)
        self.down_proj = nn.Linear(h, config.dim, bias=False)
    def forward(self, x): return self.down_proj(F.silu(self.gate_proj(x)) * self.up_proj(x))

class TRMAttention(nn.Module):
    def __init__(self, config):
        super().__init__()
        self.n_heads, self.head_dim = config.n_heads, config.head_dim
        self.qkv = nn.Linear(config.dim, 3 * config.dim, bias=False)
        self.out = nn.Linear(config.dim, config.dim, bias=False)
    def forward(self, x, mask, cos, sin):
        B, S, _ = x.shape
        q, k, v = self.qkv(x).chunk(3, dim=-1)
        q = q.view(B, S, self.n_heads, self.head_dim).transpose(1, 2)
        k = k.view(B, S, self.n_heads, self.head_dim).transpose(1, 2)
        v = v.view(B, S, self.n_heads, self.head_dim).transpose(1, 2)
        q, k = apply_rope(q, cos, sin), apply_rope(k, cos, sin)
        y = F.scaled_dot_product_attention(q, k, v, attn_mask=mask[:, None, :, :])
        return self.out(y.transpose(1, 2).reshape(B, S, -1))

class TRMBlock(nn.Module):
    def __init__(self, config):
        super().__init__()
        self.residual_scale = config.residual_scale
        self.norm1, self.norm2 = nn.RMSNorm(config.dim), nn.RMSNorm(config.dim)
        self.attn, self.mlp = TRMAttention(config), SwiGLUMLP(config)
        self.attn_gate = nn.Parameter(torch.ones(config.dim))
        self.mlp_gate = nn.Parameter(torch.ones(config.dim))
    def forward(self, x, mask, cos, sin):
        x = x + self.residual_scale * torch.sigmoid(self.attn_gate) * self.attn(self.norm1(x), mask, cos, sin)
        x = x + self.residual_scale * torch.sigmoid(self.mlp_gate) * self.mlp(self.norm2(x))
        return x

class TRMTextISMForCausalLM(PreTrainedModel, GenerationMixin):
    config_class = TRMTextISMConfig
    def __init__(self, config):
        super().__init__(config)
        self.token_emb = nn.Embedding(config.vocab_size, config.dim)
        self.block, self.norm = TRMBlock(config), nn.RMSNorm(config.dim)
        self.lm_head = nn.Linear(config.dim, config.vocab_size, bias=False)
        half = config.head_dim // 2
        theta = 1.0 / (10000.0 ** (torch.arange(0, half).float() / half))
        pos = torch.arange(config.max_seq_len).float()
        freqs = torch.outer(pos, theta)
        self.register_buffer("rope_cos", freqs.cos().unsqueeze(0).unsqueeze(0), persistent=False)
        self.register_buffer("rope_sin", freqs.sin().unsqueeze(0).unsqueeze(0), persistent=False)
        self.post_init()
    def tie_weights(self, **kwargs):
        if self.config.tie_word_embeddings: self.lm_head.weight = self.token_emb.weight
    def prepare_inputs_for_generation(self, input_ids, **kwargs):
        return {"input_ids": input_ids, "attention_mask": kwargs.get("attention_mask", None)}
    def forward(self, input_ids, attention_mask=None, response_starts=None, **kwargs):
        B, S = input_ids.shape
        x = self.token_emb(input_ids)
        mask = torch.tril(torch.ones(S, S, device=x.device)).bool().unsqueeze(0).expand(B, -1, -1)
        if response_starts is not None:
            for b in range(B):
                rs = response_starts[b]
                mask[b, :rs, :rs] = True
        for _ in range(self.config.recurrence_steps):
            x = self.block(x, mask, self.rope_cos, self.rope_sin)
        return CausalLMOutputWithPast(logits=self.lm_head(self.norm(x)))