import gradio as gr
import torch
import torch.nn.functional as F
from torch.utils.data import Dataset, DataLoader
import time
import json
from tokenizer import encode
from model import ResumeEncoder

context_length = 64

class ResumeDataset(Dataset):
    def __init__(self, data, context_length):
        self.data = data
        self.context_length = context_length

    def __len__(self):
        return len(self.data) - self.context_length

    def __getitem__(self, idx):
        x = torch.tensor(self.data[idx : idx + self.context_length])
        y = torch.tensor(self.data[idx + 1 : idx + self.context_length + 1])
        return x, y

def train_model(epochs):
    with open("resume_text.txt", "r", encoding="utf-8") as f:
        text = f.read()

    data       = encode(text)
    dataset    = ResumeDataset(data, context_length)
    dataloader = DataLoader(dataset, batch_size=32, shuffle=True)

    device    = "cuda" if torch.cuda.is_available() else "cpu"
    model     = ResumeEncoder().to(device)
    optimizer = torch.optim.AdamW(model.parameters(), lr=3e-4)

    log = f"Device: {device}\nDataset: {len(dataset)} pairs\n\n"

    for epoch in range(1, epochs + 1):
        start      = time.time()
        total_loss = 0
        for x, y in dataloader:
            x, y   = x.to(device), y.to(device)
            logits = model(x)
            B, T, C = logits.shape
            loss   = F.cross_entropy(logits.view(B*T, C), y.view(B*T))
            optimizer.zero_grad()
            loss.backward()
            optimizer.step()
            total_loss += loss.item()

        avg_loss = total_loss / len(dataloader)
        elapsed  = time.time() - start

        if epoch % 10 == 0:
            log += f"Epoch {epoch}/{epochs}  |  Loss: {avg_loss:.4f}  |  Time: {elapsed:.1f}s\n"

    torch.save(model.state_dict(), "resume_encoder.pth")
    log += "\nTraining complete! Model saved as resume_encoder.pth"
    return log

with gr.Blocks() as app:
    gr.Markdown("## Resume Encoder Training")
    epochs_input = gr.Slider(50, 500, value=100, step=50, label="Epochs")
    run_btn      = gr.Button("Start Training")
    output_log   = gr.Textbox(label="Training Log", lines=20)
    run_btn.click(fn=train_model, inputs=epochs_input, outputs=output_log)

app.launch()