import math
import torch
import torch.nn as nn
import torch.nn.functional as F


def get_alibi_slopes(num_heads: int) -> torch.Tensor:
    n = 2 ** math.ceil(math.log2(num_heads))
    m = torch.tensor([2 ** (-(i + 1)) for i in range(n)])
    m = m[:num_heads]
    return m


def build_alibi_bias(num_heads: int, seq_len: int, device: torch.device, dtype: torch.dtype) -> torch.Tensor:
    slopes = get_alibi_slopes(num_heads).to(device=device, dtype=dtype)
    pos = torch.arange(seq_len, device=device, dtype=dtype)
    mask = pos.view(1, seq_len) - pos.view(seq_len, 1)
    mask = mask.abs().neg().unsqueeze(0).unsqueeze(0)
    bias = mask * slopes.view(-1, 1, 1)
    return bias


class ALiBiAttention(nn.Module):
    def __init__(self, hidden: int, num_heads: int, num_kv_heads: int):
        super().__init__()
        self.hidden = hidden
        self.num_heads = num_heads
        self.num_kv_heads = num_kv_heads
        self.head_dim = hidden // num_heads
        self.num_groups = num_heads // num_kv_heads

        self.q_proj = nn.Linear(hidden, hidden, bias=False)
        self.k_proj = nn.Linear(hidden, num_kv_heads * self.head_dim, bias=False)
        self.v_proj = nn.Linear(hidden, num_kv_heads * self.head_dim, bias=False)
        self.o_proj = nn.Linear(hidden, hidden, bias=False)
        self._alibi_bias = None
        self._alibi_seq_len = 0

    def _get_alibi(self, seq_len: int, device: torch.device, dtype: torch.dtype) -> torch.Tensor:
        if self._alibi_bias is None or seq_len > self._alibi_seq_len:
            bias = build_alibi_bias(self.num_heads, seq_len, device, dtype)
            causal = torch.triu(torch.full((seq_len, seq_len), float('-inf'), device=device, dtype=dtype), diagonal=1)
            self._alibi_bias = bias + causal
            self._alibi_seq_len = seq_len
        return self._alibi_bias[:, :, :seq_len, :seq_len]

    def forward(self, x):
        B, T, _ = x.shape
        q = self.q_proj(x).view(B, T, self.num_heads, self.head_dim).transpose(1, 2)
        k = self.k_proj(x).view(B, T, self.num_kv_heads, self.head_dim).transpose(1, 2)
        v = self.v_proj(x).view(B, T, self.num_kv_heads, self.head_dim).transpose(1, 2)

        if self.num_groups > 1:
            k = k.repeat_interleave(self.num_groups, dim=1)
            v = v.repeat_interleave(self.num_groups, dim=1)

        scores = (q @ k.transpose(-2, -1)) / math.sqrt(self.head_dim)
        scores = scores + self._get_alibi(T, x.device, scores.dtype)

        attn = F.softmax(scores, dim=-1, dtype=torch.float32).to(x.dtype)
        out = (attn @ v).transpose(1, 2).contiguous().view(B, T, self.hidden)
        return self.o_proj(out)


class MLP(nn.Module):
    def __init__(self, hidden: int, intermediate: int):
        super().__init__()
        self.gate = nn.Linear(hidden, intermediate, bias=False)
        self.up = nn.Linear(hidden, intermediate, bias=False)
        self.down = nn.Linear(intermediate, hidden, bias=False)

    def forward(self, x):
        return self.down(F.gelu(self.gate(x)) * self.up(x))


class TransformerBlock(nn.Module):
    def __init__(self, hidden: int, intermediate: int, num_heads: int, num_kv_heads: int):
        super().__init__()
        self.ln1 = nn.LayerNorm(hidden)
        self.attn = ALiBiAttention(hidden, num_heads, num_kv_heads)
        self.ln2 = nn.LayerNorm(hidden)
        self.mlp = MLP(hidden, intermediate)

    def forward(self, x):
        x = x + self.attn(self.ln1(x))
        x = x + self.mlp(self.ln2(x))
        return x


class TinyModel(nn.Module):
    def __init__(self, vocab_size=4096, hidden=128, intermediate=512,
                 num_layers=2, num_heads=8, num_kv_heads=4, max_seq_len=2048,
                 tie_weights=True):
        super().__init__()
        self.hidden = hidden
        self.max_seq_len = max_seq_len
        self.token_embed = nn.Embedding(vocab_size, hidden)
        self.blocks = nn.ModuleList([
            TransformerBlock(hidden, intermediate, num_heads, num_kv_heads)
            for _ in range(num_layers)
        ])
        self.ln_f = nn.LayerNorm(hidden)
        self.lm_head = nn.Linear(hidden, vocab_size, bias=False)
        if tie_weights:
            self.lm_head.weight = self.token_embed.weight

    def forward(self, input_ids, labels=None):
        x = self.token_embed(input_ids)
        for block in self.blocks:
            x = block(x)
        x = self.ln_f(x)
        logits = self.lm_head(x)

        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)),
                shift_labels.view(-1),
                ignore_index=-100,
            )
        return logits, loss

    def generate(self, input_ids, max_new_tokens=128, temperature=0.7, top_p=0.9):
        self.eval()
        for _ in range(max_new_tokens):
            seq_len = input_ids.size(1)
            if seq_len > self.max_seq_len:
                input_ids = input_ids[:, -self.max_seq_len:]
            with torch.no_grad():
                logits, _ = self.forward(input_ids)
            logits = logits[:, -1, :] / temperature

            if top_p < 1.0:
                sorted_logits, sorted_idx = logits.sort(dim=-1, descending=True)
                cum_probs = sorted_logits.softmax(dim=-1).cumsum(dim=-1)
                cutoff = cum_probs > top_p
                cutoff[..., 1:] = cutoff[..., :-1].clone()
                cutoff[..., 0] = False
                logits[~cutoff] = float('-inf')

            probs = F.softmax(logits, dim=-1)
            next_token = torch.multinomial(probs, num_samples=1)
            input_ids = torch.cat([input_ids, next_token], dim=1)

            if next_token.item() == 2:
                break
        return input_ids


class LoRALinear(nn.Module):
    def __init__(self, original: nn.Linear, r: int = 8, alpha: float = 16, dropout: float = 0.0):
        super().__init__()
        self.linear = original
        self.r = r
        self.alpha = alpha
        self.scaling = alpha / r
        self.dropout = nn.Dropout(dropout) if dropout > 0 else nn.Identity()
        in_dim, out_dim = original.in_features, original.out_features
        self.lora_A = nn.Parameter(torch.zeros(r, in_dim))
        self.lora_B = nn.Parameter(torch.zeros(out_dim, r))
        nn.init.kaiming_uniform_(self.lora_A, a=math.sqrt(5))
        self.linear.weight.requires_grad = False
        if self.linear.bias is not None:
            self.linear.bias.requires_grad = False

    def forward(self, x):
        return self.linear(x) + self.dropout(x) @ self.lora_A.T @ self.lora_B.T * self.scaling


def apply_lora(model, target_modules=None, r=8, alpha=16, dropout=0.0):
    if target_modules is None:
        target_modules = ["q_proj", "k_proj", "v_proj", "o_proj", "gate", "up", "down"]
    lora_params = 0
    for name, module in model.named_modules():
        if not isinstance(module, nn.Linear):
            continue
        key = name.split(".")[-1]
        if key not in target_modules:
            continue
        parent = model
        parts = name.split(".")
        for p in parts[:-1]:
            parent = getattr(parent, p)
        lora = LoRALinear(module, r=r, alpha=alpha, dropout=dropout)
        setattr(parent, parts[-1], lora)
        lora_params += 2 * r * module.in_features + module.out_features * r
    n = sum(p.numel() for p in model.parameters() if p.requires_grad)
    print(f"LoRA applied: {lora_params:,} trainable params (total trainable: {n:,})")
    return model


def count_params(model):
    return sum(p.numel() for p in model.parameters())


def create_model():
    model = TinyModel(
        vocab_size=4096, hidden=128, intermediate=512,
        num_layers=2, num_heads=8, num_kv_heads=4,
        max_seq_len=2048, tie_weights=True,
    )
    n = count_params(model)
    print(f"TinyModel: {n:,} params ({n/1e6:.2f}M)")
    return model


if __name__ == "__main__":
    m = create_model()
    x = torch.randint(0, 100, (1, 16))
    logits, loss = m(x, labels=x)
    print(f"Forward OK: logits {logits.shape}, loss {loss.item():.4f}")
    gen = m.generate(x, max_new_tokens=10)
    print(f"Generate OK: {gen.shape}")
    m2 = create_model()
    apply_lora(m2)
    x2 = torch.randint(0, 100, (1, 16))
    logits2, _ = m2(x2)
    print(f"LoRA forward OK: {logits2.shape}")