import torch
import torch.nn as nn
import torch.nn.functional as F
from transformers import PreTrainedModel
from .configuration_pebble import PebbleConfig

try:
    from mamba_ssm import Mamba2
except ImportError:
    Mamba2 = None
    print("Warning: mamba-ssm not installed. Please install it to use PebbleLM.")

class RMSNorm(nn.Module):
    def __init__(self, dim, eps=1e-6):
        super().__init__()
        self.eps = eps
        self.weight = nn.Parameter(torch.ones(dim))

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

class AttentionBlock(nn.Module):
    def __init__(self, dim, n_heads, hidden, rope_theta=10000.0):
        super().__init__()
        assert dim % n_heads == 0
        self.nh, self.hd = n_heads, dim // n_heads
        self.rope_theta = rope_theta
        self.wqkv = nn.Linear(dim, 3 * dim, bias=False)
        self.wo   = nn.Linear(dim, dim, bias=False)
        self.fc1  = nn.Linear(dim, hidden, bias=False)
        self.fc2  = nn.Linear(hidden, dim, bias=False)
        self.ln1  = RMSNorm(dim)
        self.ln2  = RMSNorm(dim)

    def forward(self, x):
        B, T, C = x.shape
        h = self.ln1(x)
        qkv = self.wqkv(h).view(B, T, 3, self.nh, self.hd).permute(2, 0, 3, 1, 4)
        q, k, v = qkv[0], qkv[1], qkv[2]

        half = self.hd // 2
        invf = 1.0 / (self.rope_theta ** (
            torch.arange(0, half, device=x.device, dtype=torch.float32) * 2.0 / self.hd))
        ang = torch.outer(torch.arange(T, device=x.device, dtype=torch.float32), invf)
        cos, sin = ang.cos()[None, None], ang.sin()[None, None]

        q1, q2 = q.float()[..., :half], q.float()[..., half:]
        k1, k2 = k.float()[..., :half], k.float()[..., half:]
        
        q = torch.cat([q1 * cos - q2 * sin, q1 * sin + q2 * cos], dim=-1).to(v.dtype)
        k = torch.cat([k1 * cos - k2 * sin, k1 * sin + k2 * cos], dim=-1).to(v.dtype)

        y = F.scaled_dot_product_attention(q, k, v, is_causal=True)
        y = y.transpose(1, 2).reshape(B, T, C)
        x = x + self.wo(y)
        x = x + self.fc2(F.gelu(self.fc1(self.ln2(x))))
        return x

class MambaBlock(nn.Module):
    def __init__(self, dim, d_state=128, d_conv=4, expand=2, headdim=64):
        super().__init__()
        if Mamba2 is None:
            raise ImportError("mamba-ssm is not installed. Please install via `pip install mamba-ssm`")
        self.ln = RMSNorm(dim)
        self.mixer = Mamba2(
            d_model=dim,
            d_state=d_state,
            d_conv=d_conv,
            expand=expand,
            headdim=headdim,
            use_mem_eff_path=True,
        )

    def forward(self, x):
        return x + self.mixer(self.ln(x))

class PebbleLM(PreTrainedModel):
    config_class = PebbleConfig
    base_model_prefix = "model"
    supports_gradient_checkpointing = True

    def __init__(self, config):
        super().__init__(config)
        self.wte = nn.Embedding(config.vocab_size, config.d_model)
        self.blocks = nn.ModuleList([
            MambaBlock(
                config.d_model, 
                config.mamba_d_state, 
                config.mamba_d_conv, 
                config.mamba_expand, 
                config.mamba_headdim
            ) if i % 4 < 3 else AttentionBlock(
                config.d_model, 
                config.n_heads, 
                config.att_hidden, 
                config.rope_theta
            )
            for i in range(config.n_blocks)
        ])
        self.lnf = RMSNorm(config.d_model)
        self.lm_head = nn.Linear(config.d_model, config.vocab_size, bias=False)
        self.lm_head.weight = self.wte.weight  # weight sharing

    def forward(self, input_ids=None, labels=None, targets=None, **kwargs):
        x = self.wte(input_ids)
        for blk in self.blocks:
            x = blk(x)
        logits = self.lm_head(self.lnf(x))
        
        loss = None
        if labels is not None:
            loss = F.cross_entropy(logits.view(-1, logits.size(-1)), labels.reshape(-1))
        elif targets is not None:
            loss = F.cross_entropy(logits.view(-1, logits.size(-1)), targets.reshape(-1))
            
        return {"logits": logits, "loss": loss}

# Register the model for AutoModel
from transformers import AutoModelForCausalLM
AutoModelForCausalLM.register(PebbleConfig, PebbleLM)