import torch
import torch.nn.functional as F
from torch.utils.data import DataLoader
from datasets import load_dataset
import sentencepiece as spm
import os
from model import Translation, TranslationConfig

dataset = load_dataset(
    "cfilt/iitb-english-hindi",
    split="train"
)

if not os.path.exists("spm.model"):
    with open("corpus.txt", "w", encoding="utf-8") as f:
        for ex in dataset:
            f.write(ex["translation"]["en"] + "\n")
            f.write(ex["translation"]["hi"] + "\n")

    spm.SentencePieceTrainer.train(
        input="corpus.txt",
        model_prefix="spm",
        vocab_size=8000
    )
    
sp = spm.SentencePieceProcessor()
sp.load("spm.model")


BOS = sp.bos_id()
EOS = sp.eos_id()
PAD = 0
MAX_LEN = 32

def encode(example):

    en = example["translation"]["en"]
    hi = example["translation"]["hi"]

    src = sp.encode(en)[:MAX_LEN - 1] + [EOS]

    tgt = [BOS] + sp.encode(hi)[:MAX_LEN - 2] + [EOS]

    tgt_input = tgt[:-1]

    targets = tgt[1:]

    src += [PAD] * (MAX_LEN - len(src))

    tgt_input += [PAD] * (MAX_LEN - len(tgt_input))

    targets += [PAD] * (MAX_LEN - len(targets))

    return {
        "src": torch.tensor(src),
        "tgt": torch.tensor(tgt_input),
        "targets": torch.tensor(targets)
    }

tokenized_dataset = dataset.map(encode)

tokenized_dataset.set_format(
    type="torch",
    columns=["src", "tgt", "targets"]
)

loader = DataLoader(
    tokenized_dataset,
    batch_size=32,
    shuffle=True
)

config = TranslationConfig()

config.vocab_size = 8000

model = Translation(config)


optimizer = torch.optim.AdamW(
    model.parameters(),
    lr=3e-4
)


epochs = 10

for epoch in range(epochs):
    model.train()
    total_loss = 0
    for batch in loader:
        src = batch["src"]
        tgt = batch["tgt"]
        targets = batch["targets"]
        logits, loss = model(
            src,
            tgt,
            targets
        )
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()
        total_loss += loss.item()
    avg_loss = total_loss / len(loader)
    print(f"Epoch {epoch+1} Loss: {avg_loss:.4f}")
    
    
checkpoint = {
    "model_state_dict": model.state_dict(),
    "config": vars(config)
}

torch.save(
    checkpoint,
    "translator_checkpoint.pth"
)

print("Model Saved")