#!/usr/bin/env python3
"""
Test script to diagnose OAT training issues.

This script checks:
1. Forward/backward pass correctness
2. Gradient flow
3. Reconstruction quality
4. Output statistics (detect model collapse)
5. Normalization effects
"""

import argparse
from pathlib import Path
from typing import Dict, Any

import numpy as np
import torch
import torch.nn.functional as F

from pat.models.pat_tokenizer import PATTokenizer, TokenizerType
from pat.models.pat_tokenizer import ActionSpaceConfig, ArmMode, PATConfig


def create_oat_tokenizer(device: str = "cpu") -> PATTokenizer:
    """Create OAT tokenizer with default config."""
    action_space = ActionSpaceConfig(
        action_dim=20,
        horizon=24,
        arm_mode=ArmMode.BILATERAL,
    )

    config = PATConfig(
        tokenizer_type=TokenizerType.OAT,
        action_space=action_space,
        device=device,
        oat_num_registers=16,
        oat_encoder_emb_dim=256,
        oat_encoder_head_dim=64,
        oat_encoder_depth=4,
        oat_encoder_dropout=0.0,
        oat_decoder_emb_dim=256,
        oat_decoder_head_dim=64,
        oat_decoder_depth=4,
        oat_decoder_dropout=0.0,
        oat_token_dropout_mode="drop",
        oat_decoder_use_causal=False,
        oat_fsq_dim=64,
    )

    tokenizer = PATTokenizer(config)
    return tokenizer


def check_gradient_flow(model: torch.nn.Module, inputs: torch.Tensor) -> Dict[str, Any]:
    """Check if gradients are flowing correctly."""
    model.train()
    inputs.requires_grad_(True)

    # Forward pass - match train_pat.py: use encoder/decoder directly
    latent = model.encoder(inputs)
    actions_recon = model.decoder(latent)

    # Compute loss
    loss = F.mse_loss(actions_recon, inputs)

    # Backward pass
    loss.backward()

    # Collect gradient stats
    grad_stats = {
        "loss": loss.item(),
        "input_grad_norm": inputs.grad.norm().item() if inputs.grad is not None else None,
        "param_grads": {},
        "zero_grad_params": 0,
        "nonzero_grad_params": 0,
    }

    for name, param in model.named_parameters():
        if param.grad is not None:
            grad_norm = param.grad.norm().item()
            grad_stats["param_grads"][name] = grad_norm
            if grad_norm > 0:
                grad_stats["nonzero_grad_params"] += 1
            else:
                grad_stats["zero_grad_params"] += 1

    return grad_stats


def check_output_statistics(
    model: torch.nn.Module, inputs: torch.Tensor, num_batches: int = 10
) -> Dict[str, Any]:
    """Check output statistics to detect model collapse."""
    model.eval()

    all_recon_values = []
    all_latent_values = []
    all_recon_stds = []
    all_latent_stds = []

    with torch.no_grad():
        for _ in range(num_batches):
            # Match train_pat.py: use encoder/decoder directly
            latent = model.encoder(inputs)
            actions_recon = model.decoder(latent)

            all_recon_values.append(actions_recon.flatten().cpu().numpy())
            all_latent_values.append(latent.flatten().cpu().numpy())
            all_recon_stds.append(actions_recon.std().item())
            all_latent_stds.append(latent.std().item())

    all_recon = np.concatenate(all_recon_values)
    all_latent = np.concatenate(all_latent_values)

    return {
        "recon_mean": float(all_recon.mean()),
        "recon_std": float(all_recon.std()),
        "recon_min": float(all_recon.min()),
        "recon_max": float(all_recon.max()),
        "latent_mean": float(all_latent.mean()),
        "latent_std": float(all_latent.std()),
        "latent_min": float(all_latent.min()),
        "latent_max": float(all_latent.max()),
        "recon_std_avg": float(np.mean(all_recon_stds)),
        "latent_std_avg": float(np.mean(all_latent_stds)),
    }


def check_reconstruction_quality(
    tokenizer: PATTokenizer, device: str = "cpu"
) -> Dict[str, Any]:
    """Check reconstruction quality with known inputs."""
    tokenizer.eval()

    # Create test patterns
    test_patterns = {
        "zeros": torch.zeros((2, 24, 20), dtype=torch.float32),
        "ones": torch.ones((2, 24, 20), dtype=torch.float32),
        "random": torch.randn(2, 24, 20, dtype=torch.float32) * 0.5,
        "large": torch.randn(2, 24, 20, dtype=torch.float32) * 5.0,
    }

    results = {}

    for name, actions in test_patterns.items():
        # Check if normalizer is fitted
        if not tokenizer.normalizer.action.is_fitted():
            # Fit normalizer with dummy data
            dummy_actions = torch.randn(100, 24, 20)
            tokenizer.normalizer.action.fit(dummy_actions)

        tokens = tokenizer.encode(actions.numpy())
        actions_recon, mask = tokenizer.decode(tokens)

        mse = F.mse_loss(
            torch.from_numpy(actions_recon), actions
        ).item()
        mae = F.l1_loss(
            torch.from_numpy(actions_recon), actions
        ).item()

        results[name] = {
            "input_mean": float(actions.mean()),
            "input_std": float(actions.std()),
            "recon_mean": float(actions_recon.mean()),
            "recon_std": float(actions_recon.std()),
            "mse": mse,
            "mae": mae,
        }

    return results


def check_normalizer(tokenizer: PATTokenizer) -> Dict[str, Any]:
    """Check normalizer statistics."""
    normalizer = tokenizer.normalizer.action

    if not normalizer.is_fitted():
        return {"fitted": False}

    return {
        "fitted": True,
        "scale_mean": float(normalizer.scale.mean().item()),
        "scale_std": float(normalizer.scale.std().item()),
        "scale_min": float(normalizer.scale.min().item()),
        "scale_max": float(normalizer.scale.max().item()),
        "offset_mean": float(normalizer.offset.mean().item()),
        "offset_std": float(normalizer.offset.std().item()),
    }


def run_training_simulation(
    tokenizer: PATTokenizer,
    device: str = "cpu",
    num_steps: int = 100,
    batch_size: int = 32,
) -> Dict[str, Any]:
    """Simulate training to check loss curve."""
    tokenizer.train()
    model = tokenizer.model

    # Create optimizer
    optimizer = torch.optim.AdamW(
        model.parameters(), lr=1e-4, weight_decay=0.01, betas=(0.9, 0.95)
    )

    # Fit normalizer
    if not tokenizer.normalizer.action.is_fitted():
        dummy_actions = torch.randn(100, 24, 20)
        tokenizer.normalizer.action.fit(dummy_actions)

    losses = []

    for step in range(num_steps):
        # Generate random batch
        actions = torch.randn(batch_size, 24, 20, device=device)

        optimizer.zero_grad()

        # Forward pass - match train_pat.py: use encoder/decoder directly
        latent = model.encoder(actions)
        actions_recon = model.decoder(latent)

        # Compute loss
        loss = F.mse_loss(actions_recon, actions)

        # Backward pass
        loss.backward()

        # Gradient clipping
        torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)

        # Optimizer step
        optimizer.step()

        losses.append(loss.item())

        if (step + 1) % 10 == 0:
            print(f"Step {step + 1}/{num_steps} | Loss: {loss.item():.6f}")

    return {
        "losses": losses,
        "initial_loss": losses[0],
        "final_loss": losses[-1],
        "loss_ratio": losses[0] / max(losses[-1], 1e-8),
    }


def main():
    parser = argparse.ArgumentParser(description="Test OAT training")
    parser.add_argument(
        "--device",
        type=str,
        default="cuda" if torch.cuda.is_available() else "cpu",
        help="Device for testing",
    )
    parser.add_argument(
        "--test",
        type=str,
        default="all",
        choices=["all", "gradient", "output", "recon", "normalizer", "train"],
        help="Which test to run",
    )
    args = parser.parse_args()

    print("=" * 80)
    print("OAT Training Diagnostic Test")
    print("=" * 80)
    print(f"Device: {args.device}")
    print(f"Test: {args.test}")
    print("=" * 80)

    # Create tokenizer
    tokenizer = create_oat_tokenizer(device=args.device)
    model = tokenizer.model

    # Test 1: Gradient flow
    if args.test in ["all", "gradient"]:
        print("\n" + "=" * 80)
        print("TEST 1: Gradient Flow")
        print("=" * 80)

        inputs = torch.randn(4, 24, 20, device=args.device, requires_grad=False)
        grad_stats = check_gradient_flow(model, inputs)

        print(f"Loss: {grad_stats['loss']:.6f}")
        print(f"Input grad norm: {grad_stats['input_grad_norm']:.6f}")
        print(f"Non-zero grad params: {grad_stats['nonzero_grad_params']}")
        print(f"Zero grad params: {grad_stats['zero_grad_params']}")

        if grad_stats["zero_grad_params"] > 0:
            print("WARNING: Some parameters have zero gradients!")
            for name, grad_norm in grad_stats["param_grads"].items():
                if grad_norm == 0:
                    print(f"  - {name}: zero gradient")

    # Test 2: Output statistics
    if args.test in ["all", "output"]:
        print("\n" + "=" * 80)
        print("TEST 2: Output Statistics (Model Collapse Detection)")
        print("=" * 80)

        inputs = torch.randn(8, 24, 20, device=args.device)
        output_stats = check_output_statistics(model, inputs)

        print("Reconstruction statistics:")
        print(f"  Mean: {output_stats['recon_mean']:.6f}")
        print(f"  Std:  {output_stats['recon_std']:.6f}")
        print(f"  Min:  {output_stats['recon_min']:.6f}")
        print(f"  Max:  {output_stats['recon_max']:.6f}")
        print(f"  Per-batch std avg: {output_stats['recon_std_avg']:.6f}")

        print("\nLatent statistics:")
        print(f"  Mean: {output_stats['latent_mean']:.6f}")
        print(f"  Std:  {output_stats['latent_std']:.6f}")
        print(f"  Min:  {output_stats['latent_min']:.6f}")
        print(f"  Max:  {output_stats['latent_max']:.6f}")

        # Check for collapse
        if output_stats["recon_std"] < 0.01:
            print("WARNING: Very low reconstruction std - possible model collapse!")
        if output_stats["recon_std_avg"] < 0.01:
            print("WARNING: Very low per-batch std - model outputting constant values!")

    # Test 3: Reconstruction quality
    if args.test in ["all", "recon"]:
        print("\n" + "=" * 80)
        print("TEST 3: Reconstruction Quality")
        print("=" * 80)

        recon_results = check_reconstruction_quality(tokenizer, device=args.device)

        for name, stats in recon_results.items():
            print(f"\n{name}:")
            print(f"  Input:  mean={stats['input_mean']:.4f}, std={stats['input_std']:.4f}")
            print(f"  Recon:  mean={stats['recon_mean']:.4f}, std={stats['recon_std']:.4f}")
            print(f"  MSE: {stats['mse']:.6f}, MAE: {stats['mae']:.6f}")

    # Test 4: Normalizer
    if args.test in ["all", "normalizer"]:
        print("\n" + "=" * 80)
        print("TEST 4: Normalizer Statistics")
        print("=" * 80)

        norm_stats = check_normalizer(tokenizer)

        if not norm_stats["fitted"]:
            print("Normalizer is NOT fitted - this could cause issues!")
        else:
            print(f"Scale:  mean={norm_stats['scale_mean']:.4f}, std={norm_stats['scale_std']:.4f}")
            print(f"        min={norm_stats['scale_min']:.4f}, max={norm_stats['scale_max']:.4f}")
            print(f"Offset: mean={norm_stats['offset_mean']:.4f}, std={norm_stats['offset_std']:.4f}")

    # Test 5: Training simulation
    if args.test in ["all", "train"]:
        print("\n" + "=" * 80)
        print("TEST 5: Training Simulation")
        print("=" * 80)

        train_results = run_training_simulation(
            tokenizer, device=args.device, num_steps=100, batch_size=32
        )

        print(f"\nInitial loss: {train_results['initial_loss']:.6f}")
        print(f"Final loss:   {train_results['final_loss']:.6f}")
        print(f"Loss ratio:   {train_results['loss_ratio']:.2f}x")

        # Check for suspicious loss drop
        if train_results["loss_ratio"] > 100:
            print("WARNING: Loss dropped too quickly (>100x) - possible data leakage or collapse!")
        elif train_results["loss_ratio"] > 10:
            print("CAUTION: Loss dropped significantly (>10x) - verify this is expected")

        # Plot loss curve (if matplotlib available)
        try:
            import matplotlib.pyplot as plt

            plt.figure(figsize=(10, 4))
            plt.plot(train_results["losses"])
            plt.xlabel("Step")
            plt.ylabel("Loss")
            plt.title("Training Loss Curve")
            plt.grid(True, alpha=0.3)
            plt.savefig("oat_test_loss_curve.png")
            print("\nLoss curve saved to: oat_test_loss_curve.png")
        except ImportError:
            pass

    print("\n" + "=" * 80)
    print("Test Complete")
    print("=" * 80)


if __name__ == "__main__":
    main()
