#!/usr/bin/env python3
"""Treina um pequeno GPT causal do zero com tokenizer BPE local."""

from __future__ import annotations

import argparse
import json
import math
import random
from pathlib import Path

import torch
from tokenizers import Tokenizer
from torch import nn
from torch.nn import functional as F


class GPTZero(nn.Module):
    def __init__(self, vocab_size: int, block_size: int, d_model: int, nhead: int, num_layers: int, dropout: float):
        super().__init__()
        self.block_size = block_size
        self.token_embedding = nn.Embedding(vocab_size, d_model)
        self.position_embedding = nn.Embedding(block_size, d_model)
        camada = nn.TransformerEncoderLayer(
            d_model=d_model,
            nhead=nhead,
            dim_feedforward=4 * d_model,
            dropout=dropout,
            activation="gelu",
            batch_first=True,
            norm_first=True,
        )
        self.transformer = nn.TransformerEncoder(camada, num_layers=num_layers)
        self.norm = nn.LayerNorm(d_model)
        self.lm_head = nn.Linear(d_model, vocab_size, bias=False)
        self.lm_head.weight = self.token_embedding.weight
        self.apply(self._init_weights)

    @staticmethod
    def _init_weights(modulo):
        if isinstance(modulo, (nn.Linear, nn.Embedding)):
            nn.init.normal_(modulo.weight, mean=0.0, std=0.02)
            if isinstance(modulo, nn.Linear) and modulo.bias is not None:
                nn.init.zeros_(modulo.bias)
        elif isinstance(modulo, nn.LayerNorm):
            nn.init.ones_(modulo.weight)
            nn.init.zeros_(modulo.bias)

    def forward(self, indices: torch.Tensor, alvos: torch.Tensor | None = None):
        _, comprimento = indices.shape
        if comprimento > self.block_size:
            raise ValueError("A sequência excede o contexto do modelo.")
        posicoes = torch.arange(comprimento, device=indices.device)
        x = self.token_embedding(indices) + self.position_embedding(posicoes)[None, :, :]
        mascara = torch.triu(
            torch.ones(comprimento, comprimento, device=indices.device, dtype=torch.bool),
            diagonal=1,
        )
        x = self.transformer(x, mask=mascara)
        logits = self.lm_head(self.norm(x))
        perda = None
        if alvos is not None:
            perda = F.cross_entropy(logits.reshape(-1, logits.size(-1)), alvos.reshape(-1))
        return logits, perda


def carregar_textos(path: Path) -> list[str]:
    textos = []
    with path.open(encoding="utf-8") as arquivo:
        for linha in arquivo:
            if linha.strip():
                registro = json.loads(linha)
                if isinstance(registro.get("text"), str):
                    textos.append(registro["text"])
    return textos


def ids_textos(textos: list[str], tokenizer: Tokenizer, bos_id: int, eos_id: int) -> list[torch.Tensor]:
    return [torch.tensor([bos_id] + tokenizer.encode(texto).ids + [eos_id], dtype=torch.long) for texto in textos]


def amostrar_lote(streams: list[torch.Tensor], block_size: int, batch_size: int, device: torch.device):
    possiveis = [stream for stream in streams if len(stream) > block_size]
    if not possiveis:
        raise ValueError("Nenhum exemplo é maior que o contexto escolhido.")
    xs, ys = [], []
    for _ in range(batch_size):
        stream = random.choice(possiveis)
        inicio = random.randint(0, len(stream) - block_size - 1)
        janela = stream[inicio : inicio + block_size + 1]
        xs.append(janela[:-1])
        ys.append(janela[1:])
    return torch.stack(xs).to(device), torch.stack(ys).to(device)


@torch.no_grad()
def avaliar(modelo, streams, block_size, batch_size, batches, device):
    modelo.eval()
    perdas = []
    for _ in range(batches):
        xb, yb = amostrar_lote(streams, block_size, batch_size, device)
        _, perda = modelo(xb, yb)
        perdas.append(float(perda))
    modelo.train()
    return sum(perdas) / len(perdas)


def selecionar(logits: torch.Tensor, temperature: float, top_k: int, top_p: float) -> torch.Tensor:
    logits = logits / max(temperature, 1e-5)
    if top_k > 0:
        valores, indices = torch.topk(logits, min(top_k, logits.size(-1)))
        filtrados = torch.full_like(logits, -float("inf"))
        filtrados.scatter_(1, indices, valores)
        logits = filtrados
    if 0 < top_p < 1:
        valores, indices = torch.sort(logits, descending=True)
        probabilidades = F.softmax(valores, dim=-1)
        acumuladas = torch.cumsum(probabilidades, dim=-1)
        remover = acumuladas > top_p
        remover[:, 1:] = remover[:, :-1].clone()
        remover[:, 0] = False
        valores[remover] = -float("inf")
        filtrados = torch.full_like(logits, -float("inf"))
        filtrados.scatter_(1, indices, valores)
        logits = filtrados
    return torch.multinomial(F.softmax(logits, dim=-1), num_samples=1)


@torch.no_grad()
def gerar(modelo, tokenizer, prompt: str, device: torch.device, max_new_tokens: int, temperature: float, top_k: int, top_p: float, bos_id: int, eos_id: int, repetition_penalty: float = 1.0) -> str:
    ids = [bos_id] + tokenizer.encode(prompt).ids
    contexto = torch.tensor([ids], dtype=torch.long, device=device)
    fim_id = tokenizer.token_to_id("<|end|>")
    modelo.eval()
    for _ in range(max_new_tokens):
        entrada = contexto[:, -modelo.block_size :]
        logits, _ = modelo(entrada)
        logits = logits[:, -1, :]
        if repetition_penalty > 1.0:
            vistos = torch.unique(contexto[0])
            logits[:, vistos] = torch.where(logits[:, vistos] < 0, logits[:, vistos] * repetition_penalty, logits[:, vistos] / repetition_penalty)
        proximo = selecionar(logits, temperature, top_k, top_p)
        contexto = torch.cat([contexto, proximo], dim=1)
        token_id = int(proximo.item())
        if token_id == eos_id or (fim_id is not None and token_id == fim_id):
            break
    return tokenizer.decode(contexto[0].tolist()[len(ids):], skip_special_tokens=True)


def main() -> None:
    parser = argparse.ArgumentParser()
    parser.add_argument("--data-dir", type=Path, required=True)
    parser.add_argument("--tokenizer", type=Path, required=True)
    parser.add_argument("--steps", type=int, default=3000)
    parser.add_argument("--batch-size", type=int, default=8)
    parser.add_argument("--block-size", type=int, default=256)
    parser.add_argument("--d-model", type=int, default=192)
    parser.add_argument("--nhead", type=int, default=6)
    parser.add_argument("--num-layers", type=int, default=4)
    parser.add_argument("--dropout", type=float, default=0.1)
    parser.add_argument("--lr", type=float, default=3e-4)
    parser.add_argument("--eval-interval", type=int, default=250)
    parser.add_argument("--eval-batches", type=int, default=10)
    parser.add_argument("--seed", type=int, default=42)
    parser.add_argument("--checkpoint", type=Path, default=Path("modelo_gpt_zero.pt"))
    parser.add_argument("--resume", type=Path, default=None, help="checkpoint do próprio GPT do zero para continuar")
    parser.add_argument("--device", choices=["auto", "cpu", "cuda"], default="auto")
    args = parser.parse_args()

    random.seed(args.seed)
    torch.manual_seed(args.seed)
    if torch.cuda.is_available():
        torch.cuda.manual_seed_all(args.seed)
    if args.device == "auto":
        device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    else:
        device = torch.device(args.device)
    if device.type == "cuda" and not torch.cuda.is_available():
        raise RuntimeError("CUDA foi solicitado, mas não está disponível.")

    tokenizer = Tokenizer.from_file(str(args.tokenizer))
    bos_id = tokenizer.token_to_id("<bos>")
    eos_id = tokenizer.token_to_id("<eos>")
    if bos_id is None or eos_id is None:
        raise ValueError("Tokenizer sem <bos> ou <eos>.")
    train_textos = carregar_textos(args.data_dir / "train_conversacional.jsonl")
    val_textos = carregar_textos(args.data_dir / "validation_conversacional.jsonl")
    test_textos = carregar_textos(args.data_dir / "test_conversacional.jsonl")
    train_streams = ids_textos(train_textos, tokenizer, bos_id, eos_id)
    val_streams = ids_textos(val_textos, tokenizer, bos_id, eos_id)
    test_streams = ids_textos(test_textos, tokenizer, bos_id, eos_id)
    # O corpus é concatenado para permitir janelas maiores que um exemplo isolado.
    # Os marcadores <bos>/<eos> preservam as fronteiras entre diálogos.
    train_streams = [torch.cat(train_streams)]
    val_streams = [torch.cat(val_streams)]
    test_streams = [torch.cat(test_streams)]

    modelo = GPTZero(tokenizer.get_vocab_size(), args.block_size, args.d_model, args.nhead, args.num_layers, args.dropout).to(device)
    if args.resume is not None:
        estado = torch.load(args.resume, map_location=device, weights_only=False)
        if estado.get("pretrained", True):
            raise ValueError("O checkpoint indicado não está marcado como treinado do zero.")
        modelo.load_state_dict(estado["model_state"])
        print(f"Checkpoint do zero retomado: {args.resume}")
    otimizador = torch.optim.AdamW(modelo.parameters(), lr=args.lr, betas=(0.9, 0.95), weight_decay=0.1)
    total_parametros = sum(param.numel() for param in modelo.parameters())
    melhor_validacao = math.inf

    print(f"Dispositivo: {device}")
    print(f"Pesos inicializados aleatoriamente; nenhum modelo pré-treinado é carregado.")
    print(f"Vocabulário: {tokenizer.get_vocab_size()} tokens")
    print(f"Parâmetros: {total_parametros:,}")
    print(f"Tokens: treino={sum(map(len, train_streams))}, validação={sum(map(len, val_streams))}, teste={sum(map(len, test_streams))}")

    for passo in range(1, args.steps + 1):
        xb, yb = amostrar_lote(train_streams, args.block_size, args.batch_size, device)
        _, perda = modelo(xb, yb)
        otimizador.zero_grad(set_to_none=True)
        perda.backward()
        torch.nn.utils.clip_grad_norm_(modelo.parameters(), 1.0)
        otimizador.step()
        if passo == 1 or passo % args.eval_interval == 0 or passo == args.steps:
            treino = avaliar(modelo, train_streams, args.block_size, args.batch_size, args.eval_batches, device)
            validacao = avaliar(modelo, val_streams, args.block_size, args.batch_size, args.eval_batches, device)
            print(f"passo={passo:06d} treino={treino:.4f} validação={validacao:.4f}", flush=True)
            if validacao < melhor_validacao:
                melhor_validacao = validacao
                torch.save(
                    {
                        "model_state": modelo.state_dict(),
                        "vocab_size": tokenizer.get_vocab_size(),
                        "block_size": args.block_size,
                        "d_model": args.d_model,
                        "nhead": args.nhead,
                        "num_layers": args.num_layers,
                        "dropout": args.dropout,
                        "tokenizer": str(args.tokenizer),
                        "validacao": validacao,
                        "seed": args.seed,
                        "pretrained": False,
                    },
                    args.checkpoint,
                )

    print(f"Checkpoint do zero salvo em: {args.checkpoint}")
    prompt = "<|system|>Você é um assistente útil.<|end|><|user|>Olá! Quem é você?<|end|><|assistant|>"
    print("Amostra:")
    print(gerar(modelo, tokenizer, prompt, device, 120, 0.7, 40, 0.9, bos_id, eos_id))


if __name__ == "__main__":
    main()