"""
GOAT-GPT Nano v5 — Hugging Face Spaces Demo
─────────────────────────────────────────────
Loads the trained checkpoint (goat_gpt_nano_gen.pt) and serves a simple
text-generation UI with Gradio.

Expected files in the Space repo root:
  - app.py                  (this file)
  - requirements.txt
  - goat_gpt_nano_gen.pt    (your trained checkpoint — upload this too)

The checkpoint is expected to be a dict with:
  {'model_state_dict': ..., 'model_config': {...}}
as produced by `save_model_for_hf` / the generation-checkpoint save in the
training script.
"""

import math
import os
from dataclasses import dataclass
from typing import Optional, Tuple

import gradio as gr
import torch
import torch.nn as nn
from torch.nn import functional as F
from transformers import AutoTokenizer

CHECKPOINT_PATH = os.environ.get("CHECKPOINT_PATH", "goat_gpt_nano_gen.pt")
TOKENIZER_NAME = os.environ.get("TOKENIZER_NAME", "NousResearch/Llama-2-7b-hf")
DEVICE = "cuda" if torch.cuda.is_available() else "cpu"


# ═══════════════════════════════════════════════════════════════════════════
# MODEL DEFINITION (must match the architecture used for training — v5)
# ═══════════════════════════════════════════════════════════════════════════

@dataclass
class ModelConfig:
    block_size: int = 512
    vocab_size: int = 32000
    n_layer: int = 10
    n_head: int = 10
    n_embd: int = 640
    intermediate_size: int = 1707
    dropout: float = 0.0
    bias: bool = False
    rope_theta: float = 10000.0
    use_qk_norm: bool = True
    tie_weights: bool = True


class RotaryEmbedding(nn.Module):
    def __init__(self, dim: int, max_seq_len: int = 2048, theta: float = 10000.0):
        super().__init__()
        inv_freq = 1.0 / (theta ** (torch.arange(0, dim, 2).float() / dim))
        self.register_buffer("inv_freq", inv_freq)
        t = torch.arange(max_seq_len, device=inv_freq.device, dtype=inv_freq.dtype)
        freqs = torch.outer(t, inv_freq)
        emb = torch.cat((freqs, freqs), dim=-1)
        self.register_buffer("cos_cached", emb.cos()[None, None, :, :])
        self.register_buffer("sin_cached", emb.sin()[None, None, :, :])

    def forward(self, x, seq_len: int) -> Tuple[torch.Tensor, torch.Tensor]:
        return self.cos_cached[:, :, :seq_len, :], self.sin_cached[:, :, :seq_len, :]


def rotate_half(x):
    x1, x2 = x.chunk(2, dim=-1)
    return torch.cat((-x2, x1), dim=-1)


def apply_rotary_pos_emb(q, k, cos, sin):
    return (q * cos) + (rotate_half(q) * sin), (k * cos) + (rotate_half(k) * sin)


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

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


class MultiHeadAttention(nn.Module):
    def __init__(self, config: ModelConfig):
        super().__init__()
        assert config.n_embd % config.n_head == 0
        self.n_head = config.n_head
        self.head_dim = config.n_embd // config.n_head
        self.use_qk_norm = config.use_qk_norm

        self.qkv_proj = nn.Linear(config.n_embd, 3 * config.n_embd, bias=config.bias)
        self.o_proj = nn.Linear(config.n_embd, config.n_embd, bias=config.bias)
        self.rotary = RotaryEmbedding(self.head_dim, max_seq_len=config.block_size, theta=config.rope_theta)

        if self.use_qk_norm:
            self.q_norm = RMSNorm(self.head_dim, eps=1e-6)
            self.k_norm = RMSNorm(self.head_dim, eps=1e-6)

    def forward(self, x):
        B, T, C = x.size()
        qkv = self.qkv_proj(x)
        q, k, v = qkv.split(C, dim=2)
        q = q.view(B, T, self.n_head, self.head_dim).transpose(1, 2)
        k = k.view(B, T, self.n_head, self.head_dim).transpose(1, 2)
        v = v.view(B, T, self.n_head, self.head_dim).transpose(1, 2)

        cos, sin = self.rotary(q, seq_len=T)
        q, k = apply_rotary_pos_emb(q, k, cos, sin)

        if self.use_qk_norm:
            q, k = self.q_norm(q), self.k_norm(k)

        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, config: ModelConfig):
        super().__init__()
        self.gate_proj = nn.Linear(config.n_embd, config.intermediate_size, bias=config.bias)
        self.up_proj = nn.Linear(config.n_embd, config.intermediate_size, bias=config.bias)
        self.down_proj = nn.Linear(config.intermediate_size, config.n_embd, bias=config.bias)

    def forward(self, x):
        return self.down_proj(F.silu(self.gate_proj(x)) * self.up_proj(x))


class TransformerBlock(nn.Module):
    def __init__(self, config: ModelConfig, layer_idx: int):
        super().__init__()
        self.layer_idx = layer_idx
        self.attn = MultiHeadAttention(config)
        self.input_norm = RMSNorm(config.n_embd)
        self.post_attn_norm = RMSNorm(config.n_embd)
        self.mlp = SwiGLU(config)

    def forward(self, x):
        h = x + self.attn(self.input_norm(x))
        h = h + self.mlp(self.post_attn_norm(h))
        return h


class GOATGPT(nn.Module):
    def __init__(self, config: ModelConfig):
        super().__init__()
        self.config = config
        self.transformer = nn.ModuleDict({
            'wte': nn.Embedding(config.vocab_size, config.n_embd),
            'h': nn.ModuleList([TransformerBlock(config, i) for i in range(config.n_layer)]),
            'ln_f': RMSNorm(config.n_embd),
        })
        self.lm_head = nn.Linear(config.n_embd, config.vocab_size, bias=config.bias)
        if config.tie_weights:
            self.transformer['wte'].weight = self.lm_head.weight

    def forward(self, idx, targets=None):
        b, t = idx.size()
        x = 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)
        return logits, loss

    @torch.no_grad()
    def generate(self, idx, max_new_tokens, temperature=1.0, top_k=None, top_p=None):
        for _ in range(max_new_tokens):
            idx_cond = idx if idx.size(1) <= self.config.block_size else idx[:, -self.config.block_size:]
            logits, _ = self(idx_cond)
            logits = logits[:, -1, :] / max(temperature, 1e-5)
            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)
        return idx


# ═══════════════════════════════════════════════════════════════════════════
# LOAD MODEL + TOKENIZER (once, at startup)
# ═══════════════════════════════════════════════════════════════════════════

print(f"[Startup] Device: {DEVICE}")
print(f"[Startup] Loading checkpoint from {CHECKPOINT_PATH} ...")

if not os.path.exists(CHECKPOINT_PATH):
    raise FileNotFoundError(
        f"Checkpoint not found at '{CHECKPOINT_PATH}'. "
        f"Upload your trained .pt file to the Space and set CHECKPOINT_PATH "
        f"if it's not named goat_gpt_nano_gen.pt."
    )

checkpoint = torch.load(CHECKPOINT_PATH, map_location=DEVICE)
cfg_dict = checkpoint.get("model_config", {})
model_config = ModelConfig(**{k: v for k, v in cfg_dict.items() if k in ModelConfig.__dataclass_fields__})

model = GOATGPT(model_config)
state_dict = checkpoint.get("model_state_dict", checkpoint.get("model"))
model.load_state_dict(state_dict)
model.to(DEVICE)
model.eval()

print(f"[Startup] Model loaded: {model_config.n_layer}L/{model_config.n_head}H/{model_config.n_embd}D")

print(f"[Startup] Loading tokenizer: {TOKENIZER_NAME}")
tokenizer = AutoTokenizer.from_pretrained(TOKENIZER_NAME)
tokenizer.pad_token = tokenizer.eos_token


# ═══════════════════════════════════════════════════════════════════════════
# GENERATION FN
# ═══════════════════════════════════════════════════════════════════════════

@torch.no_grad()
def generate_text(prompt, max_new_tokens, temperature, top_k, top_p):
    if not prompt or not prompt.strip():
        return "Please enter a prompt."

    input_ids = tokenizer.encode(prompt, return_tensors="pt").to(DEVICE)
    top_k = int(top_k) if top_k and top_k > 0 else None
    top_p = float(top_p) if top_p and top_p < 1.0 else None

    output_ids = model.generate(
        input_ids,
        max_new_tokens=int(max_new_tokens),
        temperature=float(temperature),
        top_k=top_k,
        top_p=top_p,
    )
    return tokenizer.decode(output_ids[0], skip_special_tokens=True)


# ═══════════════════════════════════════════════════════════════════════════
# GRADIO UI
# ═══════════════════════════════════════════════════════════════════════════

with gr.Blocks(title="AM-64M") as demo:
    gr.Markdown(
        "# AM-64M\n"
        "A small from-scratch language model (RoPE, QK-Norm, RMSNorm, SwiGLU) "
        "trained on TinyTextbooks. Enter a prompt and generate a continuation."
    )

    with gr.Row():
        with gr.Column(scale=2):
            prompt = gr.Textbox(
                label="Prompt",
                value="The theory of relativity states that",
                lines=4,
            )
            generate_btn = gr.Button("Generate", variant="primary")
            output = gr.Textbox(label="Output", lines=10)

        with gr.Column(scale=1):
            max_new_tokens = gr.Slider(8, 512, value=128, step=8, label="Max new tokens")
            temperature = gr.Slider(0.1, 2.0, value=0.8, step=0.05, label="Temperature")
            top_k = gr.Slider(0, 100, value=40, step=1, label="Top-k (0 = disabled)")
            top_p = gr.Slider(0.0, 1.0, value=1.0, step=0.01, label="Top-p (1.0 = disabled)")

    generate_btn.click(
        fn=generate_text,
        inputs=[prompt, max_new_tokens, temperature, top_k, top_p],
        outputs=output,
    )

    gr.Examples(
        examples=[
            "The theory of relativity states that",
            "In the early history of computing,",
            "Photosynthesis is the process by which",
        ],
        inputs=prompt,
    )

if __name__ == "__main__":
    demo.queue().launch()