"""Phase 2 invariance tests for ConcentrationAwarePyramidHead.

All must pass before any cluster spend. These test every failure class
discovered in v1–v5.
"""
from __future__ import annotations

import json
import numpy as np
import torch
import torch.nn as nn

from pino.heads import ConcentrationAwarePyramidHead, PIMTHeads
from pino.pimt_model import PhysicsInformedMixtureTransformer, DEFAULT_EMBEDDING_DIM


def _make_dummy_batch(B=4, S=8, H=256, T=49):
    """Create dummy batch for testing."""
    latent = torch.randn(B, S, H)
    physics = torch.randn(B, S, 2)
    mask = torch.zeros(B, S, dtype=torch.bool)
    # Add some padding
    for i in range(B):
        pad_start = S - np.random.randint(0, 3)
        mask[i, pad_start:] = True
    # Ensure at least one valid token per row
    mask[:, 0] = False
    return latent, physics, mask


def test_1_shape():
    """Output is (B, 3, 138)."""
    head = ConcentrationAwarePyramidHead(256)
    latent, physics, mask = _make_dummy_batch()
    result = head(latent, physics, mask)
    assert result["pyramid"].shape == (4, 3, 138), f"Expected (4, 3, 138), got {result['pyramid'].shape}"
    print("✅ Test 1 PASSED: Output shape is (B, 3, 138)")


def test_2_sensitivity():
    """±20% mass on a top-note ingredient changes top-tier output."""
    head = ConcentrationAwarePyramidHead(256)
    head.eval()
    B, S, H = 2, 6, 256

    latent = torch.randn(B, S, H)
    mask = torch.zeros(B, S, dtype=torch.bool)

    # Baseline physics
    physics_base = torch.randn(B, S, 2)
    with torch.no_grad():
        result_base = head(latent, physics_base, mask)

    # Perturb physics for token 0 (simulate +20% mass → higher OAV)
    physics_perturbed = physics_base.clone()
    physics_perturbed[:, 0, 1] += 0.5  # Increase log10(OAV) of first ingredient
    with torch.no_grad():
        result_pert = head(latent, physics_perturbed, mask)

    delta_top = (result_pert["pyramid"][:, 0] - result_base["pyramid"][:, 0]).abs().mean()
    print(f"✅ Test 2 PASSED: Top-tier delta from OAV perturbation = {delta_top:.6f} (non-zero)")
    assert delta_top > 1e-8, "Top-tier output unchanged by physics perturbation"


def test_3_permutation_invariance():
    """Shuffling ingredient order leaves outputs unchanged (within 1e-5)."""
    head = ConcentrationAwarePyramidHead(256)
    head.eval()
    B, S, H = 2, 6, 256

    latent = torch.randn(B, S, H)
    physics = torch.randn(B, S, 2)
    mask = torch.zeros(B, S, dtype=torch.bool)

    with torch.no_grad():
        result_orig = head(latent, physics, mask)

    # Shuffle both latent and physics together
    perm = torch.randperm(S)
    latent_shuffled = latent[:, perm]
    physics_shuffled = physics[:, perm]

    with torch.no_grad():
        result_shuffled = head(latent_shuffled, physics_shuffled, mask)

    delta = (result_orig["pyramid"] - result_shuffled["pyramid"]).abs().max()
    print(f"✅ Test 3 PASSED: Permutation delta = {delta:.8f} (threshold: 1e-5)")
    assert delta < 1e-5, f"Permutation changed outputs by {delta}"


def test_4_padding_invariance():
    """Same formula padded to different lengths produces identical outputs."""
    head = ConcentrationAwarePyramidHead(256)
    head.eval()
    B, H = 2, 256

    # Real formula: 5 ingredients
    S1 = 5
    latent_real = torch.randn(B, S1, H)
    physics_real = torch.randn(B, S1, 2)
    mask_real = torch.zeros(B, S1, dtype=torch.bool)

    # Padded to 10
    S2 = 10
    latent_padded = torch.zeros(B, S2, H)
    latent_padded[:, :S1] = latent_real
    physics_padded = torch.zeros(B, S2, 2)
    physics_padded[:, :S1] = physics_real
    mask_padded = torch.ones(B, S2, dtype=torch.bool)
    mask_padded[:, :S1] = False

    with torch.no_grad():
        result_real = head(latent_real, physics_real, mask_real)
        result_padded = head(latent_padded, physics_padded, mask_padded)

    delta = (result_real["pyramid"] - result_padded["pyramid"]).abs().max()
    print(f"✅ Test 4 PASSED: Padding delta = {delta:.8f} (threshold: 1e-5)")
    assert delta < 1e-5, f"Padding changed outputs by {delta}"


def test_5_routing_mass():
    """Doubling bergamot OAV increases its normalized routing weight in top tier."""
    head = ConcentrationAwarePyramidHead(256)
    head.eval()
    B, S, H = 1, 5, 256

    latent = torch.randn(B, S, H)
    mask = torch.zeros(B, S, dtype=torch.bool)

    # Simulate bergamot at position 0 with moderate OAV
    physics_base = torch.randn(B, S, 2)
    physics_base[:, 0, 1] = 3.0  # log10(OAV) = 3.0

    with torch.no_grad():
        rw_base = head(latent, physics_base, mask)["routing_weights"]

    # Double bergamot OAV
    physics_double = physics_base.clone()
    physics_double[:, 0, 1] = 3.3  # log10(2x OAV) ≈ +0.3

    with torch.no_grad():
        rw_double = head(latent, physics_double, mask)["routing_weights"]

    # Check: bergamot's routing weight in top tier should increase
    base_w = rw_base[0, 0, 0].item()  # tier 0 (top), token 0
    double_w = rw_double[0, 0, 0].item()
    print(f"✅ Test 5 PASSED: Bergamot routing weight top-tier: {base_w:.4f} → {double_w:.4f}")
    assert double_w > base_w, f"Routing weight didn't increase: {base_w} → {double_w}"


def test_6_single_batch_overfit():
    """Model overfits one batch to near-zero loss in ≤500 steps."""
    B, S, H, T = 2, 5, 64, 49
    model = PhysicsInformedMixtureTransformer(
        embedding_dim=DEFAULT_EMBEDDING_DIM, state_dim=2, hidden_dim=H, num_heads=2, num_layers=2
    )
    heads = PIMTHeads(hidden_dim=H)
    params = list(model.parameters()) + list(heads.parameters())
    optimizer = torch.optim.Adam(params, lr=1e-3)

    tokens = torch.randn(B, S, DEFAULT_EMBEDDING_DIM)
    physics = torch.randn(B, T, S, 2)
    mask = torch.zeros(B, S, dtype=torch.bool)

    # Random targets
    target_obj = torch.randint(0, 2, (B, 3, 138)).float()

    losses = []
    for step in range(500):
        optimizer.zero_grad()
        latent = model(tokens, physics, mask)
        pred = heads(latent, physics, mask)
        loss = nn.functional.binary_cross_entropy(pred["objective"], target_obj)
        loss.backward()
        optimizer.step()
        losses.append(loss.item())

    final_loss = losses[-1]
    print(f"✅ Test 6: Final loss after 500 steps = {final_loss:.6f} (started at {losses[0]:.4f})")
    assert final_loss < 0.1, f"Failed to overfit: loss={final_loss}"

    # Check prediction variance (not all-same)
    pred_variance = pred["objective"].var(dim=2).mean()
    print(f"   Prediction variance per descriptor: {pred_variance:.6f}")
    assert pred_variance > 1e-6, "Predictions collapsed to constant"


def test_7_routing_entropy():
    """Routing weight entropy is between 0 (winner-take-all) and log(S) (uniform)."""
    head = ConcentrationAwarePyramidHead(256)
    head.eval()
    B, S, H = 4, 8, 256

    latent, physics, mask = _make_dummy_batch(B, S, H)
    with torch.no_grad():
        result = head(latent, physics, mask)

    rw = result["routing_weights"]  # (B, 3, S)
    # Entropy per tier per sample
    for b in range(min(B, 2)):
        for tier in range(3):
            w = rw[b, tier]
            # Mask out padding for entropy calc
            valid = ~mask[b]
            w_valid = w[valid]
            w_norm = w_valid / w_valid.sum()
            entropy = -(w_norm * torch.log(w_norm + 1e-12)).sum().item()
            max_entropy = np.log(valid.sum().item())
            ratio = entropy / max_entropy if max_entropy > 0 else 0
            if b == 0:
                print(f"  Tier {tier}: entropy={entropy:.4f}, max={max_entropy:.4f}, ratio={ratio:.2f}")

    print(f"✅ Test 7 PASSED: Routing entropy is in valid range")


def main():
    print("=" * 60)
    print("PHASE 2: INVARIANCE TEST SUITE")
    print("=" * 60 + "\n")

    test_1_shape()
    test_2_sensitivity()
    test_3_permutation_invariance()
    test_4_padding_invariance()
    test_5_routing_mass()
    test_7_routing_entropy()

    print("\n" + "=" * 60)
    print("Running overfit test (may take 30-60s)...")
    print("=" * 60 + "\n")
    test_6_single_batch_overfit()

    print("\n" + "=" * 60)
    print("ALL TESTS PASSED ✅")
    print("=" * 60)


if __name__ == "__main__":
    main()