import torch
import torch.nn as nn
from torch.utils.data import Dataset, DataLoader
import time
import sys
import pickle
import os
import json
import shutil
from sklearn.model_selection import train_test_split
from model import make_model

sys.stdout.reconfigure(encoding='utf-8')

DATA_PATH  = 'data/processed/data.pkl'
VOCAB_PATH = 'data/processed/vocab.pkl'

MODEL_CONFIGS = {
    "small": {
        "D_MODEL":     256,
        "N_LAYER":     4,
        "N_HEAD":      8,
        "DROPOUT":     0.3,
        "BATCH_SIZE":  64,
        "EPOCHS":      200,
        "LR":          5e-4,
    },
    "medium": {
        "D_MODEL":     512,
        "N_LAYER":     6,
        "N_HEAD":      8,
        "DROPOUT":     0.35,
        "BATCH_SIZE":  32,
        "EPOCHS":      300,
        "LR":          3e-4,
    },
    "large": {
        "D_MODEL":     768,
        "N_LAYER":     12,
        "N_HEAD":      12,
        "DROPOUT":     0.35,
        "BATCH_SIZE":  128,
        "EPOCHS":      500,
        "LR":          3e-4,
    },
}


class EarlyStopping:
    def __init__(self, patience=30, delta=0.0):
        self.patience   = patience
        self.delta      = delta
        self.best_loss  = float('inf')
        self.counter    = 0
        self.early_stop = False

    def __call__(self, val_loss):
        if val_loss < self.best_loss - self.delta:
            self.best_loss = val_loss
            self.counter   = 0
        else:
            self.counter += 1
            if self.counter >= self.patience:
                self.early_stop = True


class SeqDataset(Dataset):
    def __init__(self, sequences):
        self.sequences = sequences

    def __len__(self):
        return len(self.sequences)

    def __getitem__(self, idx):
        return torch.tensor(self.sequences[idx])


def collate_fn(batch, pad_idx):
    seqs = nn.utils.rnn.pad_sequence(batch, batch_first=True, padding_value=pad_idx)
    x = seqs[:, :-1]
    y = seqs[:, 1:].clone()
    y[y == pad_idx] = -1
    return x, y


def count_params(model):
    return sum(p.numel() for p in model.parameters() if p.requires_grad)


def train_one_model(name, cfg, device, train_seq, val_seq, pad_idx, vocab_size):
    D_MODEL    = cfg["D_MODEL"]
    N_LAYER    = cfg["N_LAYER"]
    N_HEAD     = cfg["N_HEAD"]
    DROPOUT    = cfg["DROPOUT"]
    BATCH_SIZE = cfg["BATCH_SIZE"]
    EPOCHS     = cfg["EPOCHS"]
    LR         = cfg["LR"]

    save_path_all = f'models/all/model_{name}.pt'

    print(f"\n{'='*60}")
    print(f"  TRAIN: {name.upper()}  |  D={D_MODEL} L={N_LAYER} H={N_HEAD}")
    print(f"{'='*60}")

    sorted_train = sorted(train_seq, key=len)
    curriculum_epochs = max(1, EPOCHS // 5)

    curriculum_dl = DataLoader(
        SeqDataset(sorted_train), batch_size=BATCH_SIZE, shuffle=False,
        collate_fn=lambda b: collate_fn(b, pad_idx), pin_memory=True
    )
    normal_dl = DataLoader(
        SeqDataset(train_seq), batch_size=BATCH_SIZE, shuffle=True,
        collate_fn=lambda b: collate_fn(b, pad_idx), pin_memory=True
    )
    val_dl = DataLoader(
        SeqDataset(val_seq), batch_size=BATCH_SIZE, shuffle=False,
        collate_fn=lambda b: collate_fn(b, pad_idx), pin_memory=True
    )

    log_file = f'logs/progress_{name}.txt'
    checkpoint_path = f'models/all/checkpoint_{name}.pt'

    # Phục hồi checkpoint (nếu có)
    start_epoch = 0
    best_val_loss = float('inf')
    best_val_acc  = 0.0
    best_train_loss = float('inf')
    best_epoch    = 0
    history = {"train_loss": [], "val_loss": [], "train_acc": [], "val_acc": [], "lr": []}
    early_stop    = EarlyStopping(patience=30, delta=0.001)

    model = make_model(vocab_size, n_layer=N_LAYER, d_model=D_MODEL, n_head=N_HEAD, dropout=DROPOUT)
    model.to(device)
    print(f"  Params: {count_params(model):,}")

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

    warmup_epochs = min(10, EPOCHS // 10)
    def lr_lambda(epoch):
        import math
        if epoch < warmup_epochs:
            return (epoch + 1) / warmup_epochs
        progress = (epoch - warmup_epochs) / max(1, EPOCHS - warmup_epochs)
        return 0.5 * (1 + math.cos(math.pi * progress))

    scheduler     = torch.optim.lr_scheduler.LambdaLR(optimizer, lr_lambda)
    scaler        = torch.amp.GradScaler('cuda') if torch.cuda.is_available() else None

    is_resumed = False

    if os.path.exists(checkpoint_path):
        print(f"  --> Tìm thấy checkpoint tại {checkpoint_path}. Đang tải lại...")
        try:
            checkpoint = torch.load(checkpoint_path, map_location=device)
            model.load_state_dict(checkpoint['model_state_dict'])
            optimizer.load_state_dict(checkpoint['optimizer_state_dict'])
            scheduler.load_state_dict(checkpoint['scheduler_state_dict'])
            if scaler and 'scaler_state_dict' in checkpoint and checkpoint['scaler_state_dict'] is not None:
                scaler.load_state_dict(checkpoint['scaler_state_dict'])
            
            start_epoch = checkpoint['epoch'] + 1
            best_val_loss = checkpoint['best_val_loss']
            best_val_acc = checkpoint['best_val_acc']
            best_train_loss = checkpoint['best_train_loss']
            best_epoch = checkpoint['best_epoch']
            history = checkpoint['history']
            
            early_stop.best_loss = checkpoint['early_stop_best_loss']
            early_stop.counter = checkpoint['early_stop_counter']
            early_stop.early_stop = checkpoint['early_stop_early_stop']
            is_resumed = True
            print(f"  --> Phục hồi thành công! Tiếp tục huấn luyện từ Epoch {start_epoch + 1}")
        except Exception as e:
            print(f"  [CẢNH BÁO] Lỗi khi tải checkpoint: {e}. Sẽ huấn luyện lại từ đầu.")

    os.makedirs('logs', exist_ok=True)
    mode = 'a' if is_resumed else 'w'
    with open(log_file, mode, encoding='utf-8') as f:
        if is_resumed:
            f.write(f"\n--- TIẾP TỤC HUẤN LUYỆN TỪ EPOCH {start_epoch + 1} ---\n")
        else:
            f.write(f"=== {name.upper()} | D={D_MODEL} L={N_LAYER} H={N_HEAD} ===\n\n")

    total_start = time.time()

    for epoch in range(start_epoch, EPOCHS):
        model.train()
        train_dl = curriculum_dl if epoch < curriculum_epochs else normal_dl
        total_loss = total_acc = 0

        for x, y in train_dl:
            x, y = x.to(device, non_blocking=True), y.to(device, non_blocking=True)
            optimizer.zero_grad()

            if scaler:
                with torch.amp.autocast('cuda'):
                    logits, loss = model(x, y)
                scaler.scale(loss).backward()
                scaler.unscale_(optimizer)
                torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
                scaler.step(optimizer)
                scaler.update()
            else:
                logits, loss = model(x, y)
                loss.backward()
                torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
                optimizer.step()

            total_loss += loss.item()
            mask = y != -1
            total_acc += (logits.argmax(-1) == y)[mask].float().mean().item()

        scheduler.step()
        current_lr      = optimizer.param_groups[0]['lr']
        avg_train_loss  = total_loss / len(train_dl)
        avg_train_acc   = total_acc  / len(train_dl) * 100

        model.eval()
        v_loss = v_acc = 0
        with torch.no_grad():
            for x, y in val_dl:
                x, y = x.to(device, non_blocking=True), y.to(device, non_blocking=True)
                if scaler:
                    with torch.amp.autocast('cuda'):
                        logits, loss = model(x, y)
                else:
                    logits, loss = model(x, y)
                v_loss += loss.item()
                mask = y != -1
                v_acc += (logits.argmax(-1) == y)[mask].float().mean().item()

        avg_val_loss = v_loss / len(val_dl)
        avg_val_acc  = v_acc  / len(val_dl) * 100

        history["train_loss"].append(avg_train_loss)
        history["val_loss"].append(avg_val_loss)
        history["train_acc"].append(avg_train_acc)
        history["val_acc"].append(avg_val_acc)
        history["lr"].append(current_lr)

        gap  = avg_val_loss - avg_train_loss
        warn = " [OVERFIT]" if gap > 1.0 else (" [~overfit]" if gap > 0.5 else "")
        msg  = (f"Ep {epoch+1:3d}/{EPOCHS} | "
                f"Train {avg_train_loss:.4f} ({avg_train_acc:.1f}%) | "
                f"Val {avg_val_loss:.4f} ({avg_val_acc:.1f}%) | "
                f"LR {current_lr:.2e}{warn}")
        print(msg)
        with open(log_file, 'a', encoding='utf-8') as f:
            f.write(msg + "\n")

        if avg_val_loss < best_val_loss:
            best_val_loss   = avg_val_loss
            best_val_acc    = avg_val_acc
            best_train_loss = avg_train_loss
            best_epoch      = epoch + 1
            torch.save({
                'model_state_dict': model.state_dict(),
                'v_size': vocab_size,
                'd_m': D_MODEL, 'n_l': N_LAYER, 'n_h': N_HEAD, 'dr': DROPOUT,
            }, save_path_all)

        # Lưu checkpoint sau mỗi epoch để phòng ngừa Colab sập
        try:
            torch.save({
                'epoch': epoch,
                'model_state_dict': model.state_dict(),
                'optimizer_state_dict': optimizer.state_dict(),
                'scheduler_state_dict': scheduler.state_dict(),
                'scaler_state_dict': scaler.state_dict() if scaler else None,
                'best_val_loss': best_val_loss,
                'best_val_acc': best_val_acc,
                'best_train_loss': best_train_loss,
                'best_epoch': best_epoch,
                'history': history,
                'early_stop_best_loss': early_stop.best_loss,
                'early_stop_counter': early_stop.counter,
                'early_stop_early_stop': early_stop.early_stop
            }, checkpoint_path)
            
            # --- TỰ ĐỘNG BACKUP SANG GOOGLE DRIVE (CHO COLAB) ---
            drive_base = '/content/drive/MyDrive'
            if os.path.exists(drive_base):
                drive_path = os.path.join(drive_base, 'AI_Chatbot_Backup')
                os.makedirs(drive_path, exist_ok=True)
                
                # Copy đè checkpoint, vocab, tokenizer và file log sang Drive
                shutil.copy(checkpoint_path, drive_path)
                if os.path.exists('data/processed/vocab.pkl'):
                    shutil.copy('data/processed/vocab.pkl', drive_path)
                if os.path.exists('data/processed/bpe_tokenizer.json'):
                    shutil.copy('data/processed/bpe_tokenizer.json', drive_path)
                if os.path.exists(log_file):
                    shutil.copy(log_file, drive_path)
                    
        except Exception as e:
            print(f"  [CẢNH BÁO] Không thể lưu checkpoint tại Epoch {epoch+1}: {e}")

        if epoch >= EPOCHS // 3:
            early_stop(avg_val_loss)
            if early_stop.early_stop:
                print(f"  Early stop tại epoch {epoch+1}.")
                break

    # Xóa checkpoint sau khi hoàn tất huấn luyện thành công
    if os.path.exists(checkpoint_path):
        try:
            os.remove(checkpoint_path)
        except Exception:
            pass

    total_time = time.time() - total_start
    with open(f'logs/history_{name}.json', 'w', encoding='utf-8') as f:
        json.dump(history, f)

    result = {
        "name":            name,
        "num_params":      count_params(model),
        "best_epoch":      best_epoch,
        "best_train_loss": best_train_loss,
        "best_val_loss":   best_val_loss,
        "best_val_acc":    best_val_acc,
        "total_time":      total_time,
        "epochs_ran":      len(history["train_loss"]),
        "overfit_gap":     best_val_loss - best_train_loss,
        "save_path":       save_path_all,
    }

    print(f"\n  [{name.upper()}] Best epoch {best_epoch} | Val Loss {best_val_loss:.4f} | Val Acc {best_val_acc:.1f}%")
    print(f"  Thời gian: {total_time/60:.1f} phút | Lưu: {save_path_all}")
    return result


def train_gpu():
    if not torch.cuda.is_available():
        device = torch.device("cpu")
        print("Không tìm thấy GPU, dùng CPU.")
    else:
        device = torch.device("cuda")
        print(f"GPU: {torch.cuda.get_device_name(0)} | VRAM: {torch.cuda.get_device_properties(0).total_memory/1e9:.1f} GB")
        torch.backends.cudnn.benchmark = True

    if not os.path.exists(DATA_PATH):
        print(f"Không tìm thấy {DATA_PATH}. Hãy chạy preprocess.py trước.")
        return

    with open(VOCAB_PATH, 'rb') as f: vocab = pickle.load(f)
    with open(DATA_PATH,  'rb') as f: data  = pickle.load(f)

    word2idx   = vocab['word2idx']
    pad_idx    = word2idx['<pad>']
    vocab_size = len(word2idx)

    train_seq, val_seq = data['train']['sequences'], data['val']['sequences']
    print(f"\nTrain: {len(train_seq)} | Val: {len(val_seq)} | Vocab: {vocab_size}")

    os.makedirs('models/all',  exist_ok=True)
    os.makedirs('models/best', exist_ok=True)

    results = []
    for name in ["large"]:
        try:
            r = train_one_model(
                name, MODEL_CONFIGS[name],
                device, train_seq, val_seq, pad_idx, vocab_size
            )
            results.append(r)
        except RuntimeError as e:
            if "out of memory" in str(e).lower():
                print(f"  [{name.upper()}] VRAM không đủ! Giảm BATCH_SIZE rồi thử lại.")
                if torch.cuda.is_available():
                    torch.cuda.empty_cache()
            else:
                raise

    if not results:
        return

    # Bảng so sánh
    print(f"\n{'='*70}")
    print(f"  BẢNG SO SÁNH")
    print(f"{'='*70}")
    print(f"{'Model':<10} | {'Params':>12} | {'Val Loss':>10} | {'Val Acc':>9} | {'Gap':>8} | {'Time':>8}")
    print("-" * 70)
    for r in results:
        print(f"{r['name']:<10} | {r['num_params']:>12,} | {r['best_val_loss']:>10.4f} | "
              f"{r['best_val_acc']:>8.1f}% | {r['overfit_gap']:>8.4f} | {r['total_time']/60:>6.1f}m")

    best = min(results, key=lambda x: x['best_val_loss'])

    best_dst = 'models/best/model.pt'
    shutil.copy(best['save_path'], best_dst)

    print(f"\n  BEST: {best['name'].upper()} | Val Loss {best['best_val_loss']:.4f}")
    print(f"  Đã copy -> {best_dst}  (dùng cho Web)")

    with open('logs/model_comparison.json', 'w', encoding='utf-8') as f:
        json.dump(results, f, indent=2, ensure_ascii=False)

    print(f"\n  Tất cả model: models/all/")
    print(f"  Model tốt nhất: models/best/model.pt")


if __name__ == "__main__":
    train_gpu()