"""Compare iter100 vs iter150 checkpoints: action_mean diff and refine_mask equality."""
import numpy as np
import torch
from torch_geometric.data import Batch

from src.network import create_model
from src.utils import load_checkpoint, setup_helmholtz_config


def load_config():
    from src.utils import load_config as _lc
    from pathlib import Path
    cfg_path = Path(__file__).resolve().parent / "src" / "config.yaml"
    return _lc(str(cfg_path))


def compare_checkpoints(ckpt_a, ckpt_b, label_a="iter100", label_b="iter150"):
    config = load_config()
    setup_helmholtz_config(config)
    algo = config.get("algorithm", {})

    from environment.mesh_refinement import MeshRefinement

    env = MeshRefinement(
        environment_config=config.get("environment", {}).get("mesh_refinement", {}),
        seed=99,
    )

    # ── Load both models ──
    model_a = create_model(env, config.get("network", {}), algo.get("ppo", {}))
    load_checkpoint(model_a, ckpt_a)
    model_a.eval()

    model_b = create_model(env, config.get("network", {}), algo.get("ppo", {}))
    load_checkpoint(model_b, ckpt_b)
    model_b.eval()

    # ── Get same initial observation ──
    env.reset()
    obs = env.reset()  # second reset ensures same state

    with torch.no_grad():
        batch = Batch.from_data_list([obs])

        # Model A
        shared_a, batch_a = model_a._encode(batch)
        latent_pi_a = model_a.policy_mlp(shared_a)
        action_mean_a = model_a.action_out(latent_pi_a).cpu().numpy().flatten()
        dist_a = model_a._make_distribution(latent_pi_a)
        actions_a = dist_a.get_actions(deterministic=True).cpu().numpy().flatten()

        # Model B
        shared_b, batch_b = model_b._encode(batch)
        latent_pi_b = model_b.policy_mlp(shared_b)
        action_mean_b = model_b.action_out(latent_pi_b).cpu().numpy().flatten()
        dist_b = model_b._make_distribution(latent_pi_b)
        actions_b = dist_b.get_actions(deterministic=True).cpu().numpy().flatten()

    # ── Compare action_mean ──
    diff = action_mean_a - action_mean_b
    print(f"\n{'='*60}")
    print(f"  1. action_mean comparison")
    print(f"{'='*60}")
    print(f"  {label_a} action_mean:  min={action_mean_a.min():.6f}  max={action_mean_a.max():.6f}  mean={action_mean_a.mean():.6f}  std={action_mean_a.std():.6f}")
    print(f"  {label_b} action_mean:  min={action_mean_b.min():.6f}  max={action_mean_b.max():.6f}  mean={action_mean_b.mean():.6f}  std={action_mean_b.std():.6f}")
    print(f"  ---")
    print(f"  |diff|:              min={np.abs(diff).min():.8f}  max={np.abs(diff).max():.8f}  mean={np.abs(diff).mean():.8f}")
    print(f"  diff = 0 exactly:    {int(np.sum(diff == 0))} / {len(diff)}  ({100 * np.sum(diff == 0) / len(diff):.2f}%)")
    print(f"  |diff| < 1e-6:       {int(np.sum(np.abs(diff) < 1e-6))} / {len(diff)}")
    print(f"  |diff| < 1e-4:       {int(np.sum(np.abs(diff) < 1e-4))} / {len(diff)}")
    print(f"  cos similarity:      {np.dot(action_mean_a, action_mean_b) / (np.linalg.norm(action_mean_a) * np.linalg.norm(action_mean_b) + 1e-12):.8f}")

    # ── Compare refine_mask (action > 0) ──
    mask_a = actions_a > 0.0
    mask_b = actions_b > 0.0
    mask_equal = np.array_equal(mask_a, mask_b)

    print(f"\n{'='*60}")
    print(f"  2. refine_mask comparison")
    print(f"{'='*60}")
    print(f"  {label_a} refine_mask: sum={mask_a.sum()} / {len(mask_a)}  ({100 * mask_a.sum() / len(mask_a):.1f}%)")
    print(f"  {label_b} refine_mask: sum={mask_b.sum()} / {len(mask_b)}  ({100 * mask_b.sum() / len(mask_b):.1f}%)")
    print(f"  refine_mask exactly equal: {mask_equal}")
    print(f"  mask XOR sum:              {(mask_a ^ mask_b).sum()} / {len(mask_a)}")

    if not mask_equal:
        diff_idx = np.where(mask_a != mask_b)[0]
        print(f"  First 20 differing indices: {diff_idx[:20].tolist()}")
        print(f"  At those indices, {label_a} action_mean: {action_mean_a[diff_idx[:10]]}")
        print(f"  At those indices, {label_b} action_mean: {action_mean_b[diff_idx[:10]]}")

    # ── 3. Parameter-level diff ──
    print(f"\n{'='*60}")
    print(f"  3. Model parameter weight diff (L2 norm)")
    print(f"{'='*60}")
    sd_a = torch.load(ckpt_a, map_location="cpu")["model_state_dict"]
    sd_b = torch.load(ckpt_b, map_location="cpu")["model_state_dict"]
    for k in sorted(sd_a.keys()):
        w_a = sd_a[k].float()
        w_b = sd_b[k].float()
        l2 = torch.norm(w_a - w_b).item()
        rel = l2 / (torch.norm(w_a).item() + 1e-12)
        print(f"  {k:55s}  |Δ|₂={l2:.6e}  rel={rel:.6e}")


if __name__ == "__main__":
    import sys
    d1 = sys.argv[1] if len(sys.argv) > 1 else "checkpoints/model_iter0100.pt"
    d2 = sys.argv[2] if len(sys.argv) > 2 else "checkpoints/model_iter0150.pt"
    l1 = sys.argv[3] if len(sys.argv) > 3 else "iter100"
    l2 = sys.argv[4] if len(sys.argv) > 4 else "iter150"
    compare_checkpoints(d1, d2, l1, l2)