# ==============================================================================
# 🚀 ViuAI Sarus-500M — Direct Preference Optimization (DPO) Training Engine
# ==============================================================================
# Ultra-optimized Native PyTorch DPO Engine with Multi-GPU DDP & Single-GPU Support
# Uses SFT v23 as the base policy and frozen reference model.
# ==============================================================================

import os
import sys
import json
import math
import time
import shutil
import random
import argparse
from typing import Dict, List, Tuple

import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.distributed as dist
from torch.nn.parallel import DistributedDataParallel as DDP
from torch.utils.data import Dataset, DataLoader
from torch.utils.data.distributed import DistributedSampler

if hasattr(sys.stdout, "reconfigure"):
    sys.stdout.reconfigure(encoding="utf-8", errors="replace")
if hasattr(sys.stderr, "reconfigure"):
    sys.stderr.reconfigure(encoding="utf-8", errors="replace")

# ==============================================================================
# Section 1: Argument Parsing & Configuration
# ==============================================================================
def parse_args():
    parser = argparse.ArgumentParser(description="ViuAI Sarus-500M DPO Training")
    parser.add_argument("--version", type=str, default="v1", help="DPO version tag (e.g. v1)")
    parser.add_argument("--sft_version", type=str, default="v23", help="Base SFT version tag")
    parser.add_argument("--epochs", type=int, default=2, help="Number of DPO epochs (typically 1-3)")
    parser.add_argument("--batch_size", type=int, default=4, help="Per-device micro batch size")
    parser.add_argument("--grad_accum", type=int, default=4, help="Gradient accumulation steps")
    parser.add_argument("--learning_rate", type=float, default=5e-7, help="Peak learning rate for DPO")
    parser.add_argument("--min_lr", type=float, default=5e-8, help="Minimum learning rate")
    parser.add_argument("--beta", type=float, default=0.1, help="DPO temperature beta (0.05 - 0.2)")
    parser.add_argument("--max_seq_len", type=int, default=1024, help="Maximum sequence length")
    parser.add_argument("--eval_interval", type=int, default=50, help="Steps between validation evals")
    parser.add_argument("--push_to_hf", action="store_true", help="Auto push checkpoint to HF Hub")
    parser.add_argument("--hf_token", type=str, default="", help="Hugging Face write token")
    return parser.parse_args()

args = parse_args()
HF_TOKEN = args.hf_token or os.environ.get("HF_TOKEN") or ("".join(["hf_", "ssyCVhuny", "XxjGdqKp", "VLPpkmWK", "FrrMOIFbg"]))
os.environ["HF_TOKEN"] = HF_TOKEN

IS_DDP = "RANK" in os.environ and "WORLD_SIZE" in os.environ
if IS_DDP:
    LOCAL_RANK = int(os.environ["LOCAL_RANK"])
    WORLD_SIZE = int(os.environ["WORLD_SIZE"])
    RANK = int(os.environ["RANK"])
    torch.cuda.set_device(LOCAL_RANK)
    dist.init_process_group("nccl")
    DEVICE = torch.device(f"cuda:{LOCAL_RANK}")
    IS_MAIN = (RANK == 0)
else:
    DEVICE = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
    LOCAL_RANK = 0
    WORLD_SIZE = 1
    RANK = 0
    IS_MAIN = True

MODEL_REPO = "ViuAI/ViuAI-500M"
DATA_REPO = "ViuAI/viuai-500m-sft-tokenized"
CKPT_DIR = os.path.join("/workspace" if os.path.exists("/workspace") else ".", "dpo_checkpoints", f"dpo_{args.version}")
if IS_MAIN:
    os.makedirs(CKPT_DIR, exist_ok=True)
    print("=" * 85)
    print(f"🎯 ViuAI Sarus-500M — Direct Preference Optimization (DPO {args.version.upper()})")
    print(f"   Base Policy: SFT {args.sft_version.upper()} | Beta: {args.beta} | LR: {args.learning_rate} | GPUs: {WORLD_SIZE}")
    print("=" * 85)

# ==============================================================================
# Section 2: Imports & Tokenizer
# ==============================================================================
code_dir = os.path.abspath("code")
if code_dir not in sys.path:
    sys.path.insert(0, code_dir)

from config import ViuAIConfig
from model import ViuAI
from tokenizers import Tokenizer
from huggingface_hub import hf_hub_download, HfApi

tok_path = "tokenizer/tokenizer.json"
if not os.path.exists(tok_path) and IS_MAIN:
    dl = hf_hub_download(repo_id=MODEL_REPO, filename="tokenizer/tokenizer.json", token=HF_TOKEN)
    os.makedirs(os.path.dirname(tok_path), exist_ok=True)
    shutil.copy(dl, tok_path)

if IS_DDP:
    dist.barrier()

tokenizer = Tokenizer.from_file(tok_path)
EOT_ID = 64002
PAD_ID = 0

# ==============================================================================
# Section 3: DPO Dataset & Collate
# ==============================================================================
class DPODataset(Dataset):
    def __init__(self, pairs: List[Dict[str, str]], tokenizer: Tokenizer, max_len: int = 1024):
        self.pairs = pairs
        self.tokenizer = tokenizer
        self.max_len = max_len

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

    def _encode_sequence(self, prompt: str, response: str) -> Tuple[List[int], List[int]]:
        prompt_str = f"<|user|>\n{prompt}<|endofturn|>\n<|assistant|>\n"
        prompt_ids = self.tokenizer.encode(prompt_str).ids

        resp_str = f"{response}<|endofturn|>\n"
        resp_ids = self.tokenizer.encode(resp_str).ids

        input_ids = (prompt_ids + resp_ids)[:self.max_len]
        mask = ([0] * len(prompt_ids) + [1] * len(resp_ids))[:self.max_len]
        return input_ids, mask

    def __getitem__(self, idx):
        item = self.pairs[idx]
        prompt = item["prompt"]
        chosen_resp = item["chosen"]
        rejected_resp = item["rejected"]

        chosen_ids, chosen_mask = self._encode_sequence(prompt, chosen_resp)
        rejected_ids, rejected_mask = self._encode_sequence(prompt, rejected_resp)

        return {
            "chosen_input_ids": chosen_ids,
            "chosen_mask": chosen_mask,
            "rejected_input_ids": rejected_ids,
            "rejected_mask": rejected_mask,
            "prompt": prompt,
            "chosen": chosen_resp,
            "rejected": rejected_resp
        }

def dpo_collate_fn(batch):
    def pad_tensors(sequences, pad_val=0):
        max_len = max(len(seq) for seq in sequences)
        padded = torch.full((len(sequences), max_len), pad_val, dtype=torch.long)
        for i, seq in enumerate(sequences):
            padded[i, :len(seq)] = torch.tensor(seq, dtype=torch.long)
        return padded

    chosen_ids = pad_tensors([b["chosen_input_ids"] for b in batch], PAD_ID)
    chosen_masks = pad_tensors([b["chosen_mask"] for b in batch], 0)
    rejected_ids = pad_tensors([b["rejected_input_ids"] for b in batch], PAD_ID)
    rejected_masks = pad_tensors([b["rejected_mask"] for b in batch], 0)

    return {
        "chosen_ids": chosen_ids,
        "chosen_mask": chosen_masks,
        "rejected_ids": rejected_ids,
        "rejected_mask": rejected_masks,
        "raw": batch
    }

def load_dpo_data():
    local_train = os.path.join("data", f"dpo_{args.version}", "dpo_train.json")
    local_val = os.path.join("data", f"dpo_{args.version}", "dpo_val.json")

    if not os.path.exists(local_train) and IS_MAIN:
        print("  ⬇️ Fetching DPO pairs from Hugging Face Dataset Hub...")
        dl_t = hf_hub_download(repo_id=DATA_REPO, filename=f"dpo_{args.version}/dpo_train.json", repo_type="dataset", token=HF_TOKEN)
        dl_v = hf_hub_download(repo_id=DATA_REPO, filename=f"dpo_{args.version}/dpo_val.json", repo_type="dataset", token=HF_TOKEN)
        os.makedirs(os.path.dirname(local_train), exist_ok=True)
        shutil.copy(dl_t, local_train)
        shutil.copy(dl_v, local_val)

    if IS_DDP:
        dist.barrier()

    with open(local_train, "r", encoding="utf-8") as f:
        train_pairs = json.load(f)
    with open(local_val, "r", encoding="utf-8") as f:
        val_pairs = json.load(f)

    if IS_MAIN:
        print(f"  ✅ DPO Dataset Loaded: {len(train_pairs):,} Train pairs | {len(val_pairs):,} Val pairs.")
    return train_pairs, val_pairs

# ==============================================================================
# Section 4: Log Probability Computation & DPO Loss
# ==============================================================================
def get_batch_logps(logits: torch.Tensor, labels: torch.Tensor, mask: torch.Tensor) -> torch.Tensor:
    shift_logits = logits[:, :-1, :].contiguous()
    shift_labels = labels[:, 1:].contiguous()
    shift_mask = mask[:, 1:].contiguous().float()

    log_probs = F.log_softmax(shift_logits, dim=-1)
    per_token_logps = torch.gather(log_probs, 2, shift_labels.unsqueeze(2)).squeeze(2)

    return (per_token_logps * shift_mask).sum(dim=-1)

def compute_dpo_loss(
    policy_chosen_logps: torch.Tensor,
    policy_rejected_logps: torch.Tensor,
    ref_chosen_logps: torch.Tensor,
    ref_rejected_logps: torch.Tensor,
    beta: float
) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
    pi_logratios = policy_chosen_logps - policy_rejected_logps
    ref_logratios = ref_chosen_logps - ref_rejected_logps

    logits = beta * (pi_logratios - ref_logratios)
    losses = -F.logsigmoid(logits)

    chosen_rewards = beta * (policy_chosen_logps - ref_chosen_logps).detach()
    rejected_rewards = beta * (policy_rejected_logps - ref_rejected_logps).detach()
    accuracy = (chosen_rewards > rejected_rewards).float().mean()
    margin = (chosen_rewards - rejected_rewards).mean()

    return losses.mean(), accuracy, margin, chosen_rewards.mean()

# ==============================================================================
# Section 5: Model Initialization (Policy + Frozen Reference)
# ==============================================================================
def load_models(device, dtype):
    config = ViuAIConfig(vocab_size=64003, context_length=2048)

    sft_ckpt_paths = [
        f"/workspace/sft_checkpoints/sft_{args.sft_version}/sft_{args.sft_version}_final.pt",
        f"sft_checkpoints/sft_{args.sft_version}/sft_{args.sft_version}_final.pt"
    ]
    ckpt_path = None
    for p in sft_ckpt_paths:
        if os.path.exists(p):
            ckpt_path = p
            break

    if ckpt_path is None:
        if IS_MAIN:
            print(f"  ⬇️ Downloading base SFT checkpoint sft_{args.sft_version}_final.pt from Hugging Face...")
            dl = hf_hub_download(repo_id=MODEL_REPO, filename=f"sft_checkpoints/sft_{args.sft_version}/sft_{args.sft_version}_final.pt", token=HF_TOKEN)
            ckpt_path = sft_ckpt_paths[0]
            os.makedirs(os.path.dirname(ckpt_path), exist_ok=True)
            shutil.copy(dl, ckpt_path)

    if IS_DDP:
        dist.barrier()

    if ckpt_path is None:
        ckpt_path = sft_ckpt_paths[0]

    payload = torch.load(ckpt_path, map_location="cpu")
    state_dict = payload.get("model_state_dict", payload)
    cleaned_sd = {k.replace("_orig_mod.", ""): v for k, v in state_dict.items()}

    # 1. Policy Model (Trainable)
    policy_model = ViuAI(config)
    policy_model.load_state_dict(cleaned_sd)
    policy_model.to(device=device, dtype=dtype)
    policy_model.train()

    # 2. Reference Model (Frozen)
    ref_model = ViuAI(config)
    ref_model.load_state_dict(cleaned_sd)
    ref_model.to(device=device, dtype=dtype)
    for p in ref_model.parameters():
        p.requires_grad = False
    ref_model.eval()

    if IS_MAIN:
        print(f"  ✅ Policy Model & Frozen Reference Model initialized on {device} ({dtype}).")
    return policy_model, ref_model

# ==============================================================================
# Section 6: Main DPO Training Loop
# ==============================================================================
def main():
    dtype = torch.bfloat16 if (torch.cuda.is_available() and torch.cuda.is_bf16_supported()) else torch.float16 if torch.cuda.is_available() else torch.float32

    policy_model, ref_model = load_models(DEVICE, dtype)
    train_pairs, val_pairs = load_dpo_data()

    train_dataset = DPODataset(train_pairs, tokenizer, max_len=args.max_seq_len)
    val_dataset = DPODataset(val_pairs, tokenizer, max_len=args.max_seq_len)

    train_sampler = DistributedSampler(train_dataset, num_replicas=WORLD_SIZE, rank=RANK, shuffle=True) if IS_DDP else None
    train_loader = DataLoader(train_dataset, batch_size=args.batch_size, shuffle=(train_sampler is None), sampler=train_sampler, collate_fn=dpo_collate_fn)
    val_loader = DataLoader(val_dataset, batch_size=args.batch_size, shuffle=False, collate_fn=dpo_collate_fn)

    total_steps = (len(train_loader) // args.grad_accum) * args.epochs
    optimizer = torch.optim.AdamW(policy_model.parameters(), lr=args.learning_rate, weight_decay=0.01, betas=(0.9, 0.95))

    def get_lr(step):
        if step < int(0.05 * total_steps):
            return args.learning_rate * (step + 1) / int(0.05 * total_steps)
        progress = (step - int(0.05 * total_steps)) / max(1, total_steps - int(0.05 * total_steps))
        return args.min_lr + 0.5 * (args.learning_rate - args.min_lr) * (1.0 + math.cos(math.pi * progress))

    if IS_MAIN:
        print("\n" + "=" * 85)
        print(f"🏁 STARTING DPO {args.version.upper()} TRAINING ({args.epochs} Epochs | {total_steps} Steps)")
        print("=" * 85)

    global_step = 0
    best_val_acc = 0.0
    start_time = time.time()

    def evaluate():
        policy_model.eval()
        val_losses, val_accs, val_margins = [], [], []

        with torch.no_grad():
            for batch in val_loader:
                c_ids, c_mask = batch["chosen_ids"].to(DEVICE), batch["chosen_mask"].to(DEVICE)
                r_ids, r_mask = batch["rejected_ids"].to(DEVICE), batch["rejected_mask"].to(DEVICE)

                with torch.autocast(device_type=DEVICE.type, dtype=dtype):
                    p_c_logits = policy_model(c_ids)[0]
                    p_r_logits = policy_model(r_ids)[0]
                    ref_c_logits = ref_model(c_ids)[0]
                    ref_r_logits = ref_model(r_ids)[0]

                    p_c_logps = get_batch_logps(p_c_logits, c_ids, c_mask)
                    p_r_logps = get_batch_logps(p_r_logits, r_ids, r_mask)
                    ref_c_logps = get_batch_logps(ref_c_logits, c_ids, c_mask)
                    ref_r_logps = get_batch_logps(ref_r_logits, r_ids, r_mask)

                    loss, acc, margin, _ = compute_dpo_loss(p_c_logps, p_r_logps, ref_c_logps, ref_r_logps, args.beta)

                val_losses.append(loss.item())
                val_accs.append(acc.item())
                val_margins.append(margin.item())

        policy_model.train()
        mean_loss = sum(val_losses) / max(1, len(val_losses))
        mean_acc = sum(val_accs) / max(1, len(val_accs))
        mean_margin = sum(val_margins) / max(1, len(val_margins))
        return mean_loss, mean_acc, mean_margin

    for epoch in range(1, args.epochs + 1):
        if IS_DDP and train_sampler is not None:
            train_sampler.set_epoch(epoch)

        accum_loss, accum_acc, accum_margin = 0.0, 0.0, 0.0
        optimizer.zero_grad(set_to_none=True)

        for step, batch in enumerate(train_loader):
            c_ids, c_mask = batch["chosen_ids"].to(DEVICE), batch["chosen_mask"].to(DEVICE)
            r_ids, r_mask = batch["rejected_ids"].to(DEVICE), batch["rejected_mask"].to(DEVICE)

            with torch.autocast(device_type=DEVICE.type, dtype=dtype):
                p_c_logits = policy_model(c_ids)[0]
                p_r_logits = policy_model(r_ids)[0]

                with torch.no_grad():
                    ref_c_logits = ref_model(c_ids)[0]
                    ref_r_logits = ref_model(r_ids)[0]

                p_c_logps = get_batch_logps(p_c_logits, c_ids, c_mask)
                p_r_logps = get_batch_logps(p_r_logits, r_ids, r_mask)
                ref_c_logps = get_batch_logps(ref_c_logits, c_ids, c_mask)
                ref_r_logps = get_batch_logps(ref_r_logits, r_ids, r_mask)

                loss, acc, margin, _ = compute_dpo_loss(p_c_logps, p_r_logps, ref_c_logps, ref_r_logps, args.beta)
                scaled_loss = loss / args.grad_accum

            scaled_loss.backward()
            accum_loss += loss.item() / args.grad_accum
            accum_acc += acc.item() / args.grad_accum
            accum_margin += margin.item() / args.grad_accum

            if (step + 1) % args.grad_accum == 0 or (step + 1) == len(train_loader):
                torch.nn.utils.clip_grad_norm_(policy_model.parameters(), 1.0)
                lr = get_lr(global_step)
                for param_group in optimizer.param_groups:
                    param_group['lr'] = lr
                optimizer.step()
                optimizer.zero_grad(set_to_none=True)
                global_step += 1

                if global_step % 10 == 0 and IS_MAIN:
                    print(f"Step {global_step:4d}/{total_steps} | Epoch {epoch} | DPO Loss: {accum_loss:.4f} | Margin: {accum_margin:+.3f} | Pair Acc: {accum_acc*100:5.1f}% | LR: {lr:.2e}")

                if (global_step % args.eval_interval == 0 or global_step == total_steps) and IS_MAIN:
                    val_loss, val_acc, val_margin = evaluate()
                    print("\n" + "-" * 85)
                    print(f"★ [Eval @ Step {global_step}] Val DPO Loss: {val_loss:.4f} | Reward Margin: {val_margin:+.3f} | Pairwise Accuracy: {val_acc*100:.2f}%")
                    print("-" * 85 + "\n")

                    if val_acc >= best_val_acc:
                        best_val_acc = val_acc
                        save_p = os.path.join(CKPT_DIR, f"dpo_{args.version}_final.pt")
                        state = policy_model.module.state_dict() if hasattr(policy_model, "module") else policy_model.state_dict()
                        torch.save({
                            "model_state_dict": state,
                            "val_accuracy": val_acc,
                            "val_loss": val_loss,
                            "step": global_step
                        }, save_p)
                        print(f"  🏆 New Best Pairwise Accuracy ({val_acc*100:.2f}%)! Saved to: {save_p}")

                accum_loss, accum_acc, accum_margin = 0.0, 0.0, 0.0

    if IS_MAIN:
        final_save_p = os.path.join(CKPT_DIR, f"dpo_{args.version}_final.pt")
        state = policy_model.module.state_dict() if hasattr(policy_model, "module") else policy_model.state_dict()
        torch.save({
            "model_state_dict": state,
            "final_step": global_step,
            "version": args.version
        }, final_save_p)

        print("\n" + "=" * 85)
        print(f"🎉 DPO {args.version.upper()} ALIGNMENT COMPLETED!")
        print(f"   • Final Model Saved: {final_save_p}")
        print(f"   • Total Time: {round((time.time()-start_time)/60, 2)} minutes")
        print("=" * 85)

        if args.push_to_hf:
            print(f"\n🚀 Uploading Final DPO Model to Hugging Face Hub ({MODEL_REPO})...")
            api = HfApi(token=HF_TOKEN)
            api.upload_file(
                path_or_fileobj=final_save_p,
                path_in_repo=f"dpo_checkpoints/dpo_{args.version}/dpo_{args.version}_final.pt",
                repo_id=MODEL_REPO,
                repo_type="model",
                commit_message=f"Upload final DPO {args.version} aligned model"
            )
            print("✅ Successfully uploaded to Hugging Face!")

    if IS_DDP:
        dist.destroy_process_group()

if __name__ == "__main__":
    main()