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


# ============================================================
# RoPE: Rotary Position Embedding (tự code từ đầu)
# Dùng trong LLaMA, Mistral, Gemma - thay thế Position Embedding tuyệt đối
# Nguyên lý: Xoay các vector Query/Key trong không gian phức theo vị trí
# ============================================================

class RotaryEmbedding(nn.Module):
    def __init__(self, dim, max_len=512, base=10000):
        super().__init__()
        # Tính tần số xoay: theta_i = 1 / (base ^ (2i / dim))
        inv_freq = 1.0 / (base ** (torch.arange(0, dim, 2).float() / dim))
        self.register_buffer('inv_freq', inv_freq)
        self.max_len = max_len
        # Cache cos/sin để tránh tính lại mỗi lần
        self._build_cache(max_len)

    def _build_cache(self, seq_len):
        t = torch.arange(seq_len, device=self.inv_freq.device).float()
        freqs = torch.einsum('i,j->ij', t, self.inv_freq)  # (seq_len, dim/2)
        emb = torch.cat([freqs, freqs], dim=-1)             # (seq_len, dim)
        self.register_buffer('cos_cached', emb.cos().unsqueeze(0).unsqueeze(0))  # (1,1,seq,dim)
        self.register_buffer('sin_cached', emb.sin().unsqueeze(0).unsqueeze(0))  # (1,1,seq,dim)

    def _rotate_half(self, x):
        """Xoay nửa sau của vector lên nửa trước (phép xoay phức)."""
        half = x.shape[-1] // 2
        x1, x2 = x[..., :half], x[..., half:]
        return torch.cat([-x2, x1], dim=-1)

    def forward(self, q, k):
        """
        Áp dụng RoPE lên Query và Key.
        q, k: (batch, n_head, seq_len, head_dim)
        """
        seq_len = q.shape[2]
        if seq_len > self.max_len:
            self._build_cache(seq_len)

        cos = self.cos_cached[:, :, :seq_len, :q.shape[-1]]
        sin = self.sin_cached[:, :, :seq_len, :q.shape[-1]]

        q_rot = (q * cos) + (self._rotate_half(q) * sin)
        k_rot = (k * cos) + (self._rotate_half(k) * sin)
        return q_rot, k_rot


# ============================================================
# SwiGLU: Swish Gated Linear Unit (tự code từ đầu)
# Dùng trong LLaMA, PaLM - thay thế khối MLP + GELU truyền thống
# Nguyên lý: SwiGLU(x) = (xW) ⊗ swish(xV) => học phi tuyến tính tốt hơn GELU
# ============================================================

class SwiGLU(nn.Module):
    def __init__(self, d_model, dropout=0.1):
        super().__init__()
        # Dùng tỉ lệ 8/3 thay vì 4 để bù đắp cho cơ chế gate
        hidden_dim = int(d_model * 8 / 3)
        # Làm tròn thành bội số của 64 để tối ưu phần cứng GPU
        hidden_dim = (hidden_dim + 63) // 64 * 64

        self.w1 = nn.Linear(d_model, hidden_dim, bias=False)  # Gate branch
        self.w2 = nn.Linear(hidden_dim, d_model, bias=False)  # Output projection
        self.w3 = nn.Linear(d_model, hidden_dim, bias=False)  # Value branch
        self.dropout = nn.Dropout(dropout)

    def forward(self, x):
        # SwiGLU(x) = w2(swish(w1(x)) ⊗ w3(x))
        gate   = F.silu(self.w1(x))   # swish(xW1)
        value  = self.w3(x)           # xW3
        return self.dropout(self.w2(gate * value))


# ============================================================
# CausalSelfAttention: Tích hợp RoPE
# ============================================================

class CausalSelfAttention(nn.Module):
    def __init__(self, d_model, n_head, dropout, rope):
        super().__init__()
        assert d_model % n_head == 0
        self.c_attn        = nn.Linear(d_model, 3 * d_model, bias=False)
        self.c_proj        = nn.Linear(d_model, d_model, bias=False)
        self.attn_dropout  = nn.Dropout(dropout)
        self.resid_dropout = nn.Dropout(dropout)
        self.n_head  = n_head
        self.d_model = d_model
        self.rope    = rope  # RotaryEmbedding được chia sẻ giữa các lớp

    def forward(self, x, return_attention=False):
        B, T, C = x.size()
        head_dim = C // self.n_head

        q, k, v = self.c_attn(x).split(self.d_model, dim=2)
        q = q.view(B, T, self.n_head, head_dim).transpose(1, 2)  # (B, nh, T, hd)
        k = k.view(B, T, self.n_head, head_dim).transpose(1, 2)
        v = v.view(B, T, self.n_head, head_dim).transpose(1, 2)

        # Áp dụng RoPE lên Q và K
        q, k = self.rope(q, k)

        att  = (q @ k.transpose(-2, -1)) * (1.0 / math.sqrt(head_dim))
        mask = torch.tril(torch.ones(T, T, device=x.device)).view(1, 1, T, T)
        att  = att.masked_fill(mask == 0, float('-inf'))
        att  = F.softmax(att, dim=-1)

        attn_weights = att.clone() if return_attention else None
        att  = self.attn_dropout(att)
        y    = att @ v
        y    = y.transpose(1, 2).contiguous().view(B, T, C)
        y    = self.resid_dropout(self.c_proj(y))

        if return_attention:
            return y, attn_weights
        return y


# ============================================================
# Block: Khối Transformer tích hợp SwiGLU + RoPE
# ============================================================

class Block(nn.Module):
    def __init__(self, d_model, n_head, dropout, rope):
        super().__init__()
        self.ln_1 = nn.LayerNorm(d_model)
        self.attn = CausalSelfAttention(d_model, n_head, dropout, rope)
        self.ln_2 = nn.LayerNorm(d_model)
        self.mlp  = SwiGLU(d_model, dropout)  # SwiGLU thay thế MLP + GELU

    def forward(self, x, return_attention=False):
        if return_attention:
            attn_out, attn_weights = self.attn(self.ln_1(x), return_attention=True)
            x = x + attn_out
            x = x + self.mlp(self.ln_2(x))
            return x, attn_weights
        else:
            x = x + self.attn(self.ln_1(x))
            x = x + self.mlp(self.ln_2(x))
            return x


# ============================================================
# Model: Kiến trúc Transformer Decoder thuần
# RoPE thay thế wpe (Absolute Position Embedding)
# SwiGLU thay thế MLP + GELU
# Weight Tying: wte.weight = lm_head.weight
# ============================================================

class Model(nn.Module):
    def __init__(self, vocab_size, n_layer=4, d_model=256, n_head=8, dropout=0.1, max_len=512):
        super().__init__()
        # RoPE: Chia sẻ 1 đối tượng RoPE dùng chung cho tất cả các lớp
        rope = RotaryEmbedding(dim=d_model // n_head, max_len=max_len)

        self.transformer = nn.ModuleDict(dict(
            wte  = nn.Embedding(vocab_size, d_model),   # Token Embedding
            drop = nn.Dropout(dropout),
            h    = nn.ModuleList([Block(d_model, n_head, dropout, rope) for _ in range(n_layer)]),
            ln_f = nn.LayerNorm(d_model),
        ))
        # Không còn wpe vì RoPE đã xử lý vị trí tương đối
        self.rope    = rope
        self.lm_head = nn.Linear(d_model, vocab_size, bias=False)
        self.max_len = max_len
        self.d_model = d_model

        # Weight Tying: Chia sẻ trọng số Embedding đầu vào và đầu ra
        self.transformer.wte.weight = self.lm_head.weight

        # Khởi tạo trọng số theo chuẩn GPT-2
        self.apply(self._init_weights)

    def _init_weights(self, module):
        if isinstance(module, nn.Linear):
            torch.nn.init.normal_(module.weight, mean=0.0, std=0.02)
            if module.bias is not None:
                torch.nn.init.zeros_(module.bias)
        elif isinstance(module, nn.Embedding):
            torch.nn.init.normal_(module.weight, mean=0.0, std=0.02)

    def forward(self, idx, targets=None):
        device = idx.device
        b, t   = idx.size()
        assert t <= self.max_len, f"Chuỗi dài {t} vượt quá max_len={self.max_len}"

        # Chỉ còn Token Embedding, RoPE xử lý vị trí bên trong Attention
        x = self.transformer.drop(self.transformer.wte(idx))
        for block in self.transformer.h:
            x = block(x)
        x      = self.transformer.ln_f(x)
        logits = self.lm_head(x)

        loss = None
        if targets is not None:
            loss = F.cross_entropy(
                logits.view(-1, logits.size(-1)),
                targets.view(-1),
                ignore_index=-1,
                label_smoothing=0.1
            )
        return logits, loss

    # ============================================================
    # HÀM: Trích xuất Vector Ngữ Nghĩa (Sentence Embedding)
    # Dùng cho bài toán So khớp Độ tương đồng câu hỏi (NLP)
    # ============================================================
    @torch.no_grad()
    def get_embedding(self, idx):
        """
        Trích xuất vector ngữ nghĩa (embedding) từ lớp Hidden States cuối cùng.
        Sử dụng Mean Pooling trên toàn bộ chuỗi token.

        Args:
            idx: tensor (batch, seq_len) chứa token IDs
        Returns:
            embedding: tensor (batch, d_model) - Vector ngữ nghĩa của câu
        """
        device = idx.device
        x = self.transformer.drop(self.transformer.wte(idx))
        for block in self.transformer.h:
            x = block(x)
        x = self.transformer.ln_f(x)  # (batch, seq_len, d_model)

        # Mean Pooling: lấy trung bình cộng tất cả các token (bỏ qua padding)
        mask = (idx != 0).unsqueeze(-1).float()  # (batch, seq_len, 1)
        x_masked = x * mask
        embedding = x_masked.sum(dim=1) / mask.sum(dim=1).clamp(min=1)  # (batch, d_model)

        return embedding

    # ============================================================
    # HÀM: Trích xuất Ma trận Attention (để vẽ Heatmap)
    # ============================================================
    @torch.no_grad()
    def get_attention_maps(self, idx):
        """
        Trích xuất attention weights từ tất cả các lớp.

        Args:
            idx: tensor (1, seq_len) chứa token IDs
        Returns:
            all_attn: list of tensors, mỗi tensor shape (1, n_head, seq_len, seq_len)
        """
        x = self.transformer.drop(self.transformer.wte(idx))
        all_attn = []
        for block in self.transformer.h:
            x, attn_weights = block(x, return_attention=True)
            all_attn.append(attn_weights)
        return all_attn

    # ============================================================
    # HÀM: Generate với Top-P (Nucleus Sampling) + Repetition Penalty
    # ============================================================
    @torch.no_grad()
    def generate(self, idx, max_new_tokens, temperature=1.0, top_k=None,
                 top_p=None, repetition_penalty=1.3):
        """
        Sinh text mới với hỗ trợ cả Top-K và Top-P (Nucleus Sampling).

        Args:
            idx: tensor (1, seq_len) chứa prompt token IDs
            max_new_tokens: số token tối đa cần sinh
            temperature: độ "sáng tạo" (cao = sáng tạo, thấp = chắc chắn)
            top_k: chỉ lấy K từ xác suất cao nhất
            top_p: chỉ lấy các từ có tổng xác suất cộng dồn <= p
            repetition_penalty: phạt các từ đã xuất hiện (tránh lặp)
        """
        for _ in range(max_new_tokens):
            idx_cond = idx if idx.size(1) <= self.max_len else idx[:, -self.max_len:]
            logits, _ = self(idx_cond)
            logits = logits[:, -1, :] / temperature

            # Repetition Penalty: phạt các token đã xuất hiện
            for row in range(idx.size(0)):
                for token_id in set(idx[row].tolist()):
                    if logits[row, token_id] > 0:
                        logits[row, token_id] /= repetition_penalty
                    else:
                        logits[row, token_id] *= repetition_penalty

            # Top-K Filtering
            if top_k is not None:
                v, _ = torch.topk(logits, min(top_k, logits.size(-1)))
                logits[logits < v[:, [-1]]] = -float('Inf')

            # Top-P (Nucleus) Filtering
            if top_p is not None:
                sorted_logits, sorted_indices = torch.sort(logits, descending=True)
                cumulative_probs = torch.cumsum(F.softmax(sorted_logits, dim=-1), dim=-1)

                sorted_indices_to_remove = cumulative_probs > top_p
                sorted_indices_to_remove[..., 1:] = sorted_indices_to_remove[..., :-1].clone()
                sorted_indices_to_remove[..., 0] = 0

                indices_to_remove = sorted_indices_to_remove.scatter(
                    1, sorted_indices, sorted_indices_to_remove
                )
                logits[indices_to_remove] = -float('Inf')

            probs    = F.softmax(logits, dim=-1)
            idx_next = torch.multinomial(probs, num_samples=1)
            idx      = torch.cat((idx, idx_next), dim=1)
            if idx_next.item() == 2:  # <eos> token
                break
        return idx

    # ============================================================
    # HÀM: Generate Stream (Streaming Token)
    # ============================================================
    @torch.no_grad()
    def generate_stream(self, idx, max_new_tokens, temperature=1.0, top_k=None,
                        top_p=None, repetition_penalty=1.3):
        for _ in range(max_new_tokens):
            idx_cond = idx if idx.size(1) <= self.max_len else idx[:, -self.max_len:]
            logits, _ = self(idx_cond)
            logits = logits[:, -1, :] / temperature

            for row in range(idx.size(0)):
                for token_id in set(idx[row].tolist()):
                    if logits[row, token_id] > 0:
                        logits[row, token_id] /= repetition_penalty
                    else:
                        logits[row, token_id] *= repetition_penalty

            if top_k is not None:
                v, _ = torch.topk(logits, min(top_k, logits.size(-1)))
                logits[logits < v[:, [-1]]] = -float('Inf')

            if top_p is not None:
                sorted_logits, sorted_indices = torch.sort(logits, descending=True)
                cumulative_probs = torch.cumsum(F.softmax(sorted_logits, dim=-1), dim=-1)
                sorted_indices_to_remove = cumulative_probs > top_p
                sorted_indices_to_remove[..., 1:] = sorted_indices_to_remove[..., :-1].clone()
                sorted_indices_to_remove[..., 0] = 0
                indices_to_remove = sorted_indices_to_remove.scatter(
                    1, sorted_indices, sorted_indices_to_remove
                )
                logits[indices_to_remove] = -float('Inf')

            probs    = F.softmax(logits, dim=-1)
            idx_next = torch.multinomial(probs, num_samples=1)
            idx      = torch.cat((idx, idx_next), dim=1)
            yield idx_next.item()
            if idx_next.item() == 2:  # <eos> token
                break


def make_model(vocab_size, n_layer=4, d_model=256, n_head=8, dropout=0.1):
    return Model(vocab_size, n_layer, d_model, n_head, dropout)