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

def apply_rope(x, cos, sin):
    S = x.shape[2]
    c, s = cos[:, :, :S, :].to(x.dtype), sin[:, :, :S, :].to(x.dtype)
    x1, x2 = x[..., :x.shape[-1]//2], x[..., x.shape[-1]//2:]
    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, k, v = [t.view(B, S, self.n_heads, self.head_dim).transpose(1, 2) for t in (q, k, v)]
        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.res = config.residual_scale
        self.norm1 = nn.RMSNorm(config.dim)
        self.attn = TRMAttention(config)
        self.norm2 = nn.RMSNorm(config.dim)
        self.mlp = 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, c, s):
        x = x + self.res * torch.sigmoid(self.attn_gate).view(1,1,-1) * self.attn(self.norm1(x), mask, c, s)
        return x + self.res * torch.sigmoid(self.mlp_gate).view(1,1,-1) * self.mlp(self.norm2(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 = TRMBlock(config)
        self.norm = nn.RMSNorm(config.dim)
        self.lm_head = nn.Linear(config.dim, config.vocab_size, bias=False)
        pos = torch.arange(config.max_seq_len).float()
        theta = 1.0 / (10000.0 ** (torch.arange(0, config.head_dim//2).float() / (config.head_dim//2)))
        f = torch.outer(pos, theta)
        self.register_buffer("rope_cos", f.cos().view(1, 1, config.max_seq_len, -1))
        self.register_buffer("rope_sin", f.sin().view(1, 1, config.max_seq_len, -1))
        self.post_init()

    def get_input_embeddings(self):
        return self.token_emb

    def set_input_embeddings(self, value):
        self.token_emb = value

    def get_output_embeddings(self):
        return self.lm_head

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

    def tie_weights(self, *args, **kwargs):
        if hasattr(self, 'lm_head'):
            self.lm_head.weight = self.token_emb.weight

    def prepare_inputs_for_generation(self, input_ids, attention_mask=None, **kwargs):
        return {"input_ids": input_ids, "attention_mask": attention_mask, "use_cache": False}

    def forward(self, input_ids, attention_mask=None, **kwargs):
        B, S = input_ids.shape
        x = self.token_emb(input_ids)
        m = torch.tril(torch.ones(S, S, device=input_ids.device)).bool().unsqueeze(0).expand(B, -1, -1)
        if attention_mask is not None:
            m = m & attention_mask[:, None, :].bool()
        c, s = self.rope_cos, self.rope_sin
        for _ in range(self.config.recurrence_steps):
            x = self.block(x, m, c, s)
        logits = self.lm_head(self.norm(x))
        return CausalLMOutputWithPast(logits=logits)