import torch
from dataset import MiniBPETokenizr, ChatDataset, train, SimpleTokenizr # SimpleTokenizr might be unused now
from model import MiniGPT
import json
from tokenizers import Tokenizer, models, trainers, pre_tokenizers, normalizers
from tokenizers.trainers import BpeTrainer
from tokenizers.normalizers import Lowercase, NFD, StripAccents
from tokenizers.pre_tokenizers import Whitespace

# For debugging purposes, turn on anomaly detection for gradients
torch.autograd.set_detect_anomaly(True)

# Load training data
# NOTE: For underfitting on "10 examples", ensure this file *only* contains those 10 examples,
# and they are long enough (as you confirmed).
with open("./data/overfit_data.jsonl", "r", encoding="utf-8") as f:
    texts = [(json.loads(line)["input"] + ' ' + json.loads(line)["output"]) for line in f if line.strip()]

def main():
    # 🧠 Initialize HuggingFace BPE tokenizer
    tokenizer = Tokenizer(models.BPE(unk_token="<UNK>"))
    tokenizer.normalizer = normalizers.Sequence([Lowercase(), NFD(), StripAccents()])
    tokenizer.pre_tokenizer = Whitespace()

    # 🛠️ BPE Training
    trainer = BpeTrainer(
        vocab_size=28517,
        special_tokens=["<PAD>", "<UNK>", "<END>", "^User:", "MiniGPT:"]
    )
    tokenizer.train_from_iterator(texts, trainer)

    # 💾 Save tokenizer
    tokenizer.save("./trained-mini-gpt/tokenizer.json")
    hf_tokenizer = Tokenizer.from_file("./trained-mini-gpt/tokenizer.json")

    # 🧾 Dataset & Model Init
    dataset = ChatDataset(
        data="./data/overfit_data.jsonl", # Ensure this path points to your 10-example dataset for testing
        tokenizer=hf_tokenizer
    )
    model = MiniGPT(vocab_size=hf_tokenizer.get_vocab_size())
    model.reset_params()

    # 🚂 Train
    # 🎯 CHANGE 2: Pass an increased learning rate (e.g., 1e-4) to the train function.
    # Set epochs to a high number for clear overfitting.
    train(model, dataset, hf_tokenizer, epochs=200, filepathh="./data/merged_data.jsonl", learning_rate=1e-4)

if __name__ == "__main__":
    main()