import torch
import torch.nn as nn
import torch.nn.functional as F
import math
import os
import json
from safetensors.torch import load_file


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

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


def precompute_rope_freqs(dim, max_len=24, theta=10000.0):
    freqs = 1.0 / (theta ** (torch.arange(0, dim, 2).float() / dim))
    t = torch.arange(max_len)
    freqs = torch.outer(t, freqs)
    cos = torch.cos(freqs).repeat_interleave(2, dim=-1)
    sin = torch.sin(freqs).repeat_interleave(2, dim=-1)
    return cos, sin


def apply_rope(x, cos, sin):
    T = x.shape[2]
    cos = cos[:T].unsqueeze(0).unsqueeze(0)
    sin = sin[:T].unsqueeze(0).unsqueeze(0)
    x_real = x.float()
    x_rot = torch.stack([-x_real[..., 1::2], x_real[..., ::2]], dim=-1).flatten(-2)
    return (x_real * cos + x_rot * sin).to(x.dtype)


class CausalSelfAttn(nn.Module):
    def __init__(self, n_embd, n_head, max_seq_len=24):
        super().__init__()
        self.n_head = n_head
        self.n_embd = n_embd
        self.head_dim = n_embd // n_head

        self.q_proj = nn.Linear(n_embd, n_embd, bias=False)
        self.k_proj = nn.Linear(n_embd, n_embd, bias=False)
        self.v_proj = nn.Linear(n_embd, n_embd, bias=False)
        self.o_proj = nn.Linear(n_embd, n_embd, bias=False)

        cos, sin = precompute_rope_freqs(self.head_dim, max_seq_len)
        self.register_buffer('rope_cos', cos, persistent=False)
        self.register_buffer('rope_sin', sin, persistent=False)

    def forward(self, x):
        B, T, C = x.shape
        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_head, self.head_dim).transpose(1, 2)
        v = self.v_proj(x).view(B, T, self.n_head, self.head_dim).transpose(1, 2)

        q = apply_rope(q, self.rope_cos, self.rope_sin)
        k = apply_rope(k, self.rope_cos, self.rope_sin)

        y = F.scaled_dot_product_attention(q, k, v, is_causal=True)
        y = y.transpose(1, 2).contiguous().view(B, T, C)
        return self.o_proj(y)


class SwiGLU(nn.Module):
    def __init__(self, n_embd, hidden_mult=8/3):
        super().__init__()
        hidden = int(n_embd * hidden_mult)
        self.w1 = nn.Linear(n_embd, hidden, bias=False)
        self.w2 = nn.Linear(hidden, n_embd, bias=False)
        self.w3 = nn.Linear(n_embd, hidden, bias=False)

    def forward(self, x):
        return self.w2(F.silu(self.w1(x)) * self.w3(x))


class Block(nn.Module):
    def __init__(self, n_embd, n_head, max_seq_len=24):
        super().__init__()
        self.ln1 = RMSNorm(n_embd)
        self.attn = CausalSelfAttn(n_embd, n_head, max_seq_len)
        self.ln2 = RMSNorm(n_embd)
        self.mlp = SwiGLU(n_embd)

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


class LLaMAModel(nn.Module):
    def __init__(self, vocab_size, n_layer=6, n_embd=384, n_head=6, max_seq_len=24):
        super().__init__()
        self.wte = nn.Embedding(vocab_size, n_embd)
        self.blocks = nn.ModuleList([
            Block(n_embd, n_head, max_seq_len) for _ in range(n_layer)
        ])
        self.ln_f = RMSNorm(n_embd)
        self.lm_head = nn.Linear(n_embd, vocab_size, bias=False)
        self.max_seq_len = max_seq_len

    def forward(self, x):
        x = self.wte(x)
        for block in self.blocks:
            x = block(x)
        x = self.ln_f(x)
        return self.lm_head(x)

    @classmethod
    def from_pretrained(cls, pretrained_path, device='cpu', **kwargs):
        if os.path.isdir(pretrained_path):
            config_path = os.path.join(pretrained_path, 'config.json')
            weights_path = os.path.join(pretrained_path, 'model.safetensors')
        else:
            config_path = os.path.join(pretrained_path, 'config.json')
            weights_path = os.path.join(pretrained_path, 'model.safetensors')

        with open(config_path) as f:
            config = json.load(f)

        model = cls(
            vocab_size=config['vocab_size'],
            n_layer=config['num_hidden_layers'],
            n_embd=config['hidden_size'],
            n_head=config['num_attention_heads'],
            max_seq_len=config['max_position_embeddings'],
        ).to(device)

        if os.path.exists(weights_path):
            state_dict = load_file(weights_path, device=device)
            model.load_state_dict(state_dict, strict=False)

        return model

    @torch.no_grad()
    def generate(self, tokenizer, temperature=1.0, top_k=50,
                 max_len=24, min_len=4, prefix_ids=None, device='cuda'):
        bos_id = tokenizer.token_to_id('<BOS>')
        eos_id = tokenizer.token_to_id('<EOS>')

        self.eval()

        if prefix_ids is not None:
            ids = torch.tensor([prefix_ids], dtype=torch.long, device=device)
        else:
            ids = torch.tensor([[bos_id]], dtype=torch.long, device=device)

        for _ in range(max_len - ids.shape[1]):
            logits = self(ids)[0, -1, :] / temperature

            if top_k > 0:
                top_k_vals, _ = torch.topk(logits, min(top_k, logits.size(-1)))
                logits[logits < top_k_vals[-1]] = float('-inf')

            probs = F.softmax(logits, dim=-1)
            next_id = torch.multinomial(probs, 1)

            if next_id.item() == eos_id:
                break

            ids = torch.cat([ids, next_id.unsqueeze(0)], dim=1)

            if ids.shape[1] >= max_len:
                break

        pw = tokenizer.decode(ids[0].tolist())
        pw = pw.replace('<BOS>', '').replace('<EOS>', '').replace('<PAD>', '').replace('<UNK>', '').strip()

        if prefix_ids is not None:
            prefix_str = tokenizer.decode(prefix_ids)
            pw = pw.replace(prefix_str, '').strip()

        if len(pw) < min_len:
            return None

        return pw