#!/usr/bin/env python3
"""
Stage 7: UniGRPO RL Fine-tuning — reduces hallucinations in dreamer + calibrates HyperNetwork.

Applies GRPO-style reinforcement learning to improve:
  1. Dreamer prediction accuracy (reward = cosine_sim(dream, actual))
  2. HyperNetwork calibration (confidence should correlate with actual error)

Only trains Dreamer + HyperNetwork calibration heads — base model is frozen.

Reference: UniGRPO (arXiv: 2603.17892)

Usage:
    torchrun --nproc_per_node=8 scripts/train_stage7_rl.py \
        --config configs/scale_1.3b.yaml \
        --dreamer_config configs/dreamer.yaml \
        --hn_config configs/hypernetwork.yaml \
        --stage3_ckpt checkpoints/v5_stage3_final.pt \
        --dreamer_ckpt checkpoints/v5_dreamer.pt \
        --hypernet_ckpt checkpoints/v5_hypernet.pt
"""

import argparse
import math
import os
import sys
import time

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

sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))

from model.vlm import VLJEPAModel
from model.latent_dreamer import LatentDreamer
from model.hypernetwork import HyperNetwork


# ---------------------------------------------------------------------------
# Dataset: video clip sequences for RL
# ---------------------------------------------------------------------------

class DummyVideoClipDataset(Dataset):
    """Dummy video clips: sequences of related embeddings simulating camera motion."""

    def __init__(self, num_samples=5000, seq_len=12, embed_dim=2048):
        self.num_samples = num_samples
        self.seq_len = seq_len
        self.embed_dim = embed_dim

    def __len__(self):
        return self.num_samples

    def __getitem__(self, idx):
        # Generate a coherent sequence: base + drift * t + noise
        base = torch.randn(self.embed_dim)
        drift = torch.randn(self.embed_dim) * 0.1
        frames = []
        for t in range(self.seq_len):
            noise = torch.randn(self.embed_dim) * 0.05
            frames.append(base + drift * t + noise)
        return {"embeddings": torch.stack(frames)}  # [seq_len, embed_dim]


# ---------------------------------------------------------------------------
# GRPO-style reward computation
# ---------------------------------------------------------------------------

def compute_reward(
    dreamed_embeddings: torch.Tensor,
    actual_embeddings: torch.Tensor,
    dreamed_confidences: torch.Tensor,
    cosine_weight: float = 1.0,
    divergence_penalty: float = 0.5,
) -> torch.Tensor:
    """
    Compute per-step reward for dreamer predictions.

    reward = cosine_weight * cosine_sim - divergence_penalty * divergence

    Args:
        dreamed_embeddings: [B, N_steps, D]
        actual_embeddings: [B, N_steps, D]
        dreamed_confidences: [B, N_steps, 1]
    Returns:
        [B] — reward per sequence
    """
    B, N, D = dreamed_embeddings.shape

    # Per-step cosine similarity
    cos_sim = F.cosine_similarity(
        dreamed_embeddings.reshape(-1, D),
        actual_embeddings.reshape(-1, D),
        dim=-1,
    ).reshape(B, N)  # [B, N]

    # Per-step divergence (MSE)
    divergence = (dreamed_embeddings - actual_embeddings).pow(2).mean(dim=-1)  # [B, N]

    # Per-step reward
    step_reward = cosine_weight * cos_sim - divergence_penalty * divergence  # [B, N]

    # Calibration bonus: confidence should match accuracy
    accuracy = torch.exp(-divergence).detach()  # [B, N]
    conf = dreamed_confidences.squeeze(-1)  # [B, N]
    calibration_bonus = -0.1 * (conf - accuracy).pow(2)  # [B, N]

    # Total reward: mean across steps
    total_reward = (step_reward + calibration_bonus).mean(dim=-1)  # [B]
    return total_reward


# ---------------------------------------------------------------------------
# Training loop
# ---------------------------------------------------------------------------

def parse_args():
    p = argparse.ArgumentParser(description="Stage 7: UniGRPO RL Fine-tuning")
    p.add_argument("--config", type=str, default="configs/scale_1.3b.yaml")
    p.add_argument("--dreamer_config", type=str, default="configs/dreamer.yaml")
    p.add_argument("--hn_config", type=str, default="configs/hypernetwork.yaml")
    p.add_argument("--stage3_ckpt", type=str, default=None)
    p.add_argument("--dreamer_ckpt", type=str, default=None)
    p.add_argument("--hypernet_ckpt", type=str, default=None)
    p.add_argument("--resume", type=str, default=None)
    p.add_argument("--output_dir", type=str, default="checkpoints")
    p.add_argument("--hf_push", action="store_true")
    return p.parse_args()


def main():
    args = parse_args()

    # DDP setup
    dist.init_process_group(backend="nccl")
    rank = dist.get_rank()
    world_size = dist.get_world_size()
    local_rank = int(os.environ.get("LOCAL_RANK", 0))
    device = torch.device(f"cuda:{local_rank}")
    torch.cuda.set_device(device)

    if rank == 0:
        print("=" * 60)
        print("Stage 7: UniGRPO RL Fine-tuning")
        print(f"World size: {world_size}")
        print("=" * 60)

    # Load configs
    with open(args.config) as f:
        config = yaml.safe_load(f)
    if os.path.exists(args.dreamer_config):
        with open(args.dreamer_config) as f:
            config.update(yaml.safe_load(f))
    if os.path.exists(args.hn_config):
        with open(args.hn_config) as f:
            config.update(yaml.safe_load(f))

    rl_cfg = config.get("train_rl", {})
    embed_dim = config.get("predictor", {}).get("embed_dim", 2048)

    # Build dreamer
    dreamer_cfg = config.get("latent_dreamer", {})
    dreamer = LatentDreamer(
        embed_dim=dreamer_cfg.get("embed_dim", embed_dim),
        n_heads=dreamer_cfg.get("n_heads", 16),
        n_layers=dreamer_cfg.get("n_layers", 4),
        max_future_steps=dreamer_cfg.get("max_future_steps", 8),
        max_context_frames=dreamer_cfg.get("max_context_frames", 32),
    ).to(device)

    if args.dreamer_ckpt and os.path.exists(args.dreamer_ckpt):
        ckpt = torch.load(args.dreamer_ckpt, map_location=device)
        dreamer.load_state_dict(ckpt.get("dreamer_state_dict", ckpt), strict=False)
        if rank == 0:
            print(f"Loaded dreamer: {args.dreamer_ckpt}")

    # Wrap in DDP
    dreamer = DDP(dreamer, device_ids=[local_rank])

    # Dataset
    context_frames = 8
    future_steps = 4
    dataset = DummyVideoClipDataset(
        num_samples=5000,
        seq_len=context_frames + future_steps,
        embed_dim=dreamer_cfg.get("embed_dim", embed_dim),
    )
    sampler = DistributedSampler(dataset, num_replicas=world_size, rank=rank)
    dataloader = DataLoader(dataset, batch_size=rl_cfg.get("batch_size", 4), sampler=sampler, num_workers=2)

    # Optimizer
    lr = rl_cfg.get("learning_rate", 5e-6)
    optimizer = torch.optim.AdamW(dreamer.parameters(), lr=lr, weight_decay=0.01)
    max_epochs = rl_cfg.get("max_epochs", 5)
    grad_accum = rl_cfg.get("gradient_accumulation", 16)

    cosine_weight = rl_cfg.get("reward_cosine_weight", 1.0)
    divergence_penalty = rl_cfg.get("reward_divergence_penalty", 0.5)

    if rank == 0:
        print(f"LR: {lr}, Epochs: {max_epochs}, Grad accum: {grad_accum}")
        print(f"Reward: cosine_w={cosine_weight}, divergence_p={divergence_penalty}")

    # Training loop
    dreamer.train()
    global_step = 0

    for epoch in range(max_epochs):
        sampler.set_epoch(epoch)
        epoch_rewards = []

        for batch_idx, batch in enumerate(dataloader):
            embeddings = batch["embeddings"].to(device)  # [B, seq_len, D]
            context = embeddings[:, :context_frames, :]
            actual_future = embeddings[:, context_frames:, :]

            # Dream
            dreamed_embs, dreamed_confs = dreamer.module.dream_sequence(context, n_steps=future_steps)

            # Compute reward
            reward = compute_reward(
                dreamed_embs, actual_future, dreamed_confs,
                cosine_weight=cosine_weight,
                divergence_penalty=divergence_penalty,
            )

            # GRPO-style: maximize reward (minimize negative reward)
            # Use dreamer loss as the differentiable objective
            loss_dict = dreamer.module.compute_dream_loss(
                dreamed_embs, actual_future, dreamed_confs
            )

            # Weight loss by advantage (reward - baseline)
            baseline = reward.mean().detach()
            advantage = (reward - baseline).detach()

            # Scale loss inversely with advantage (better dreams → less loss emphasis)
            weighted_loss = loss_dict["total_loss"] * (1.0 - 0.1 * advantage.mean())
            weighted_loss = weighted_loss / grad_accum

            weighted_loss.backward()

            if (batch_idx + 1) % grad_accum == 0:
                torch.nn.utils.clip_grad_norm_(dreamer.parameters(), 1.0)
                optimizer.step()
                optimizer.zero_grad()
                global_step += 1

            epoch_rewards.append(reward.mean().item())

            if rank == 0 and (batch_idx + 1) % 20 == 0:
                avg_reward = sum(epoch_rewards[-20:]) / min(20, len(epoch_rewards))
                print(f"  Epoch {epoch+1} Step {batch_idx+1}: "
                      f"reward={avg_reward:.4f} "
                      f"loss={loss_dict['total_loss'].item():.4f} "
                      f"cosine={loss_dict['mean_cosine_sim'].item():.4f}")

        if rank == 0:
            avg_epoch_reward = sum(epoch_rewards) / len(epoch_rewards)
            print(f"Epoch {epoch+1}/{max_epochs}: avg_reward={avg_epoch_reward:.4f}")

    # Save checkpoint
    if rank == 0:
        os.makedirs(args.output_dir, exist_ok=True)
        save_path = os.path.join(args.output_dir, "v5_final.pt")
        torch.save({
            "dreamer_state_dict": dreamer.module.state_dict(),
            "optimizer_state_dict": optimizer.state_dict(),
            "epoch": max_epochs,
            "global_step": global_step,
        }, save_path)
        print(f"Saved RL checkpoint: {save_path}")

        if args.hf_push:
            try:
                from huggingface_hub import HfApi
                api = HfApi()
                api.upload_file(
                    path_or_fileobj=save_path,
                    path_in_repo="v5_checkpoints/v5_final.pt",
                    repo_id="hardiksa/arcisvlm",
                    repo_type="model",
                )
                print("✅ Pushed to HuggingFace")
            except Exception as e:
                print(f"⚠️ HF push failed: {e}")

    dist.destroy_process_group()


if __name__ == "__main__":
    main()