import torch
import sentencepiece as spm
from gpt_model import MiniGPT
import math
import time
import os

# ======================
# CONFIG
# ======================
SEQ_LEN = 128
BATCH_SIZE = 32
EPOCHS = 4
LR = 3e-4
MAX_TOKENS_LIMIT = 3000000
CHECKPOINT_FILE = "/content/drive/MyDrive/xr_ai_checkpoint.pt"

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

print("Using device:", device)

# ======================
# TOKENIZER
# ======================
print("📦 Loading tokenizer...")

sp = spm.SentencePieceProcessor()
sp.load("hindi_tokenizer.model")

vocab_size = sp.get_piece_size()

# ======================
# DATASET
# ======================
print("📂 Loading dataset...")

text = open("sangraha_hindi_final_dataset.txt", encoding="utf-8").read()

tokens = sp.encode(text)

if MAX_TOKENS_LIMIT:
    tokens = tokens[:MAX_TOKENS_LIMIT]

data = torch.tensor(tokens, dtype=torch.long)

split = int(0.9 * len(data))

train_data = data[:split]
val_data = data[split:]

# ======================
# MODEL
# ======================
model = MiniGPT(vocab_size).to(device)

optimizer = torch.optim.AdamW(model.parameters(), lr=LR)

loss_fn = torch.nn.CrossEntropyLoss()

resume_step = 0

# ======================
# RESUME CHECKPOINT
# ======================
if os.path.exists(CHECKPOINT_FILE):

    print("🔄 Loading checkpoint...")

    checkpoint = torch.load(CHECKPOINT_FILE)

    model.load_state_dict(checkpoint["model"])
    optimizer.load_state_dict(checkpoint["optimizer"])

    resume_step = checkpoint["step"]

    print("Resuming from step:", resume_step)

# ======================
# TRAINING
# ======================
print("🚀 XR_AI Training Started...")

step = resume_step

for epoch in range(EPOCHS):

    model.train()

    total_train_loss = 0
    train_steps = 0

    for i in range(0, len(train_data) - SEQ_LEN - 1, SEQ_LEN * BATCH_SIZE):

        batch_inputs = []
        batch_targets = []

        for j in range(BATCH_SIZE):

            start = i + j * SEQ_LEN

            if start + SEQ_LEN + 1 >= len(train_data):
                continue

            x = train_data[start:start+SEQ_LEN]
            y = train_data[start+1:start+SEQ_LEN+1]

            batch_inputs.append(x)
            batch_targets.append(y)

        if len(batch_inputs) == 0:
            continue

        inputs = torch.stack(batch_inputs).to(device)
        targets = torch.stack(batch_targets).to(device)

        optimizer.zero_grad()

        outputs = model(inputs)

        loss = loss_fn(outputs.view(-1, vocab_size), targets.view(-1))

        loss.backward()

        optimizer.step()

        step += 1
        train_steps += 1
        total_train_loss += loss.item()

        # checkpoint
        if step % 200 == 0:

            torch.save({
                "model": model.state_dict(),
                "optimizer": optimizer.state_dict(),
                "step": step
            }, CHECKPOINT_FILE)

            print("💾 checkpoint saved at step", step)

    avg_train_loss = total_train_loss / train_steps

    # ======================
    # VALIDATION
    # ======================
    model.eval()

    total_val_loss = 0
    val_steps = 0

    with torch.no_grad():

        for i in range(0, len(val_data) - SEQ_LEN - 1, SEQ_LEN):

            inputs = val_data[i:i+SEQ_LEN].unsqueeze(0).to(device)
            targets = val_data[i+1:i+SEQ_LEN+1].unsqueeze(0).to(device)

            outputs = model(inputs)

            loss = loss_fn(outputs.view(-1, vocab_size), targets.view(-1))

            total_val_loss += loss.item()
            val_steps += 1

    avg_val_loss = total_val_loss / val_steps

    perplexity = math.exp(avg_val_loss)

    print(
        f"Epoch {epoch+1}/{EPOCHS} | "
        f"Train: {avg_train_loss:.4f} | "
        f"Val: {avg_val_loss:.4f} | "
        f"PPL: {perplexity:.2f}"
    )

    torch.save(model.state_dict(), "best_model.pt")

print("🏁 Training complete")