import os, json, base64, re, time
import numpy as np
import gradio as gr
import torch
import torch.nn as nn
import torch.nn.functional as F

CONFIG = {
    "corpus_file": "Recognizement_Corpus.txt",
    "context_length": 12,
    "embed_dim": 16,
    "num_heads": 2,
    "num_layers": 1,
    "ffn_hidden": 48,
    "batch_size": 128,
    "max_vocab": 700,
    "learning_rate": 0.015,
    "focal_gamma": 0.5,
    "loss_norm": True,
    "max_epochs": 600,
    "target_loss": 0.10,
}
PAD, UNK, EOS = "<pad>", "<unk>", "<eos>"
PAD_ID, UNK_ID, EOS_ID = 0, 1, 2

def tokenize(text):
    text = text.replace("<pad>", " __PAD__ ").replace("<unk>", " __UNK__ ").replace("<eos>", " __EOS__ ")
    text = text.lower()
    text = re.sub(r"([.,!?;:])", r" \1 ", text)
    return text.replace("__PAD__", "<pad>").replace("__UNK__", "<unk>").replace("__EOS__", "<eos>").split()

def build_vocab(tokens, max_vocab=None):
    from collections import Counter
    freq = Counter(tokens)
    vocab = {PAD: 0, UNK: 1, EOS: 2}
    for tok, _ in freq.most_common():
        if tok in vocab: continue
        if max_vocab and len(vocab) >= max_vocab: break
        vocab[tok] = len(vocab)
    return vocab

def build_dataset(lines, vocab, ctx):
    X_all, Y_all, M_all = [], [], []
    for line in lines:
        toks = tokenize(line) + [EOS]
        ids = [vocab.get(t, UNK_ID) for t in toks][:ctx+1]
        real_len = len(ids)
        padded = ids + [PAD_ID] * (ctx + 1 - len(ids))
        X_all.append(padded[:-1]); Y_all.append(padded[1:])
        M_all.append([1.0 if i < real_len - 1 else 0.0 for i in range(ctx)])
    return (torch.tensor(X_all, dtype=torch.long),
            torch.tensor(Y_all, dtype=torch.long),
            torch.tensor(M_all, dtype=torch.float32))

class MHSA(nn.Module):
    def __init__(self, D, H, ctx):
        super().__init__()
        self.D, self.H, self.hd = D, H, D // H
        self.Wq = nn.Linear(D, D); self.Wk = nn.Linear(D, D)
        self.Wv = nn.Linear(D, D); self.Wo = nn.Linear(D, D)
        self.register_buffer("mask", torch.triu(torch.full((ctx, ctx), -1e9), diagonal=1))
    def forward(self, x):
        B, T, D = x.shape
        Q = self.Wq(x).view(B, T, self.H, self.hd).transpose(1, 2)
        K = self.Wk(x).view(B, T, self.H, self.hd).transpose(1, 2)
        V = self.Wv(x).view(B, T, self.H, self.hd).transpose(1, 2)
        S = (Q @ K.transpose(-2, -1)) / (self.hd ** 0.5) + self.mask[:T, :T]
        A = F.softmax(S, dim=-1)
        O = (A @ V).transpose(1, 2).reshape(B, T, D)
        return self.Wo(O)

class Block(nn.Module):
    def __init__(self, D, H, Ff, ctx):
        super().__init__()
        self.ln1 = nn.LayerNorm(D); self.attn = MHSA(D, H, ctx)
        self.ln2 = nn.LayerNorm(D)
        self.W1 = nn.Linear(D, Ff); self.W2 = nn.Linear(Ff, D)
    def forward(self, x):
        x = x + self.attn(self.ln1(x))
        x = x + self.W2(F.relu(self.W1(self.ln2(x))))
        return x

class Transformer(nn.Module):
    def __init__(self, V, ctx, D, H, L, Ff):
        super().__init__()
        self.V, self.ctx, self.D = V, ctx, D
        self.H, self.L, self.Ff = H, L, Ff
        self.tok_emb = nn.Embedding(V, D)
        self.pos_emb = nn.Embedding(ctx, D)
        self.blocks = nn.ModuleList([Block(D, H, Ff, ctx) for _ in range(L)])
        self.ln_f = nn.LayerNorm(D)
        self.Wout = nn.Linear(D, V, bias=True)
    def forward(self, X):
        B, T = X.shape
        pos = torch.arange(T, device=X.device).unsqueeze(0)
        x = self.tok_emb(X) + self.pos_emb(pos)
        for b in self.blocks: x = b(x)
        return self.Wout(self.ln_f(x))

def focal_loss(logits, Y, M, gamma, log_V):
    B, T, V = logits.shape
    lf = logits.reshape(B*T, V)
    Yf = Y.reshape(B*T); Mf = M.reshape(B*T)
    log_probs = F.log_softmax(lf, dim=-1)
    probs = log_probs.exp()
    idx = torch.arange(B*T, device=logits.device)
    p_t = probs[idx, Yf].clamp(1e-9, 1.0 - 1e-9)
    log_pt = -torch.log(p_t)
    focal_val = ((1.0 - p_t) ** gamma) * log_pt
    Msum = Mf.sum() + 1e-9
    loss = (focal_val * Mf).sum() / Msum
    loss_norm = loss / log_V
    # grad: s * (probs - onehot) * mask / Msum
    s = ((1.0 - p_t) ** gamma) + gamma * p_t * ((1.0 - p_t) ** (gamma - 1.0)) * log_pt
    s = torch.nan_to_num(s, nan=0.0, posinf=0.0, neginf=0.0)
    grad_logits = probs.clone()
    grad_logits[idx, Yf] -= 1.0
    grad_logits = grad_logits * Mf.unsqueeze(1) * s.unsqueeze(1) / Msum
    return loss_norm, grad_logits.reshape(B, T, V)

def to_b64(t):
    return base64.b64encode(t.detach().cpu().numpy().astype("<f4").tobytes()).decode("ascii")

def export_weights(model, vocab):
    tokens_por_id = [None] * len(vocab)
    for tok, idx in vocab.items(): tokens_por_id[idx] = tok
    blocks_data = []
    for b in model.blocks:
        blocks_data.append({
            "ln1_gamma": to_b64(b.ln1.weight), "ln1_beta": to_b64(b.ln1.bias),
            "ln2_gamma": to_b64(b.ln2.weight), "ln2_beta": to_b64(b.ln2.bias),
            "attn": {
                "Wq": to_b64(b.attn.Wq.weight.T), "Wk": to_b64(b.attn.Wk.weight.T),
                "Wv": to_b64(b.attn.Wv.weight.T), "Wo": to_b64(b.attn.Wo.weight.T),
                "bq": to_b64(b.attn.Wq.bias),     "bk": to_b64(b.attn.Wk.bias),
                "bv": to_b64(b.attn.Wv.bias),     "bo": to_b64(b.attn.Wo.bias),
            },
            "W1": to_b64(b.W1.weight.T), "b1": to_b64(b.W1.bias),
            "W2": to_b64(b.W2.weight.T), "b2": to_b64(b.W2.bias),
        })
    return {
        "config": {
            "arch": "transformer", "vocab_size": model.V,
            "context_length": model.ctx, "embed_dim": model.D,
            "num_heads": model.H, "num_layers": model.L,
            "ffn_hidden": model.Ff, "activation": "relu",
        },
        "generation": {"temperature": 0.6, "top_k": 5, "repetition_penalty": 1.2},
        "vocab": tokens_por_id,
        "tok_emb": to_b64(model.tok_emb.weight),
        "pos_emb": to_b64(model.pos_emb.weight),
        "ln_f_gamma": to_b64(model.ln_f.weight),
        "ln_f_beta":  to_b64(model.ln_f.bias),
        "blocks": blocks_data,
        "output_layer": {"W": to_b64(model.Wout.weight), "b": to_b64(model.Wout.bias)},
    }

def train_fn(max_epochs, target_loss, progress=gr.Progress()):
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    with open(CONFIG["corpus_file"], "r", encoding="utf-8") as f:
        text = f.read()
    lines = [l.strip() for l in text.splitlines() if l.strip()]
    all_toks = []
    for line in lines: all_toks.extend(tokenize(line))
    vocab = build_vocab(all_toks, max_vocab=CONFIG["max_vocab"])
    ctx = CONFIG["context_length"]
    X, Y, M = build_dataset(lines, vocab, ctx)
    V = len(vocab); log_V = float(np.log(V))
    model = Transformer(V, ctx, CONFIG["embed_dim"], CONFIG["num_heads"],
                        CONFIG["num_layers"], CONFIG["ffn_hidden"]).to(device)
    opt = torch.optim.Adam(model.parameters(), lr=CONFIG["learning_rate"],
                           betas=(0.9, 0.999), eps=1e-8)
    n = len(X); bs = CONFIG["batch_size"]
    n_batches = (n + bs - 1) // bs
    X, Y, M = X.to(device), Y.to(device), M.to(device)
    t0 = time.time()
    log_lines = []
    for ep in range(1, max_epochs + 1):
        perm = torch.randperm(n, device=device)
        Xs, Ys, Ms = X[perm], Y[perm], M[perm]
        ep_loss, nb = 0.0, 0
        for i in range(0, n, bs):
            Xb, Yb, Mb = Xs[i:i+bs], Ys[i:i+bs], Ms[i:i+bs]
            opt.zero_grad()
            logits = model(Xb)
            loss, grad_logits = focal_loss(logits, Yb, Mb, CONFIG["focal_gamma"], log_V)
            logits.backward(grad_logits)
            opt.step()
            ep_loss += loss.item(); nb += 1
        avg = ep_loss / nb
        elapsed = time.time() - t0
        log_lines.append(f"Ep {ep:4d} | loss {avg:.4f} | {elapsed:6.1f}s | device={device.type}")
        progress(ep / max_epochs, desc=f"Ep {ep} | loss {avg:.4f}")
        if avg <= target_loss: break
    ckpt = export_weights(model, vocab)
    out_path = "/tmp/AI_Weights.json"
    with open(out_path, "w", encoding="utf-8") as f:
        json.dump(ckpt, f)
    return "\n".join(log_lines[-30:]), out_path

def launch_train(max_epochs, target_loss):
    log, path = train_fn(int(max_epochs), float(target_loss))
    return log, path

with gr.Blocks() as demo:
    gr.Markdown("# Transformer Trainer (ZeroGPU)")
    with gr.Row():
        ep_in = gr.Number(value=600, label="max_epochs", precision=0)
        tl_in = gr.Number(value=0.10, label="target_loss")
    btn = gr.Button("Treinar", variant="primary")
    log_out = gr.Textbox(label="Log", lines=20)
    file_out = gr.File(label="AI_Weights.json")
    btn.click(launch_train, inputs=[ep_in, tl_in], outputs=[log_out, file_out])

demo.launch()