"""
Experiment Runner: Adaptive Resource Allocation & Routing Stability.

Evaluates:
  - Marginal Utility vs Static Rank
  - Rank distributions across different token difficulties
  - Routing stability and hysteresis churn reduction
  - Multi-seed variance across 3 seeds
"""

import sys
import os
import json
import argparse
from pathlib import Path

sys.path.insert(0, str(Path(__file__).parent.parent))

import torch
from src.config import ModelConfig
from src.models import QTensorFormer
from src.data import load_synthetic_data


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("--output", type=str, default="outputs/adaptive_results.json")
    parser.add_argument("--seeds", type=int, nargs="+", default=[42, 123, 999])
    args = parser.parse_args()

    os.makedirs(os.path.dirname(args.output) or ".", exist_ok=True)

    print("=" * 65)
    print("EXPERIMENT: Adaptive Allocation & Stability Analysis")
    print("=" * 65)

    cfg = ModelConfig(
        vocab_size=1000,
        d_model=128,
        n_layers=2,
        tt_rank=8,
        use_quantum=True,
    )

    results_by_seed = []
    for seed in args.seeds:
        torch.manual_seed(seed)
        model = QTensorFormer(cfg, preset="QTF_BALANCED")
        model.eval()

        loader = load_synthetic_data(vocab_size=1000, seq_len=32, n_samples=200, batch_size=16)

        ranks_recorded = []
        quantum_ratios = []
        churn_rates = []

        with torch.no_grad():
            for inp, _ in loader:
                _, stats = model(inp, return_stats=True)
                for block_stat in stats:
                    ranks_recorded.append(block_stat["chosen_rank"])
                    quantum_ratios.append(block_stat["quantum_usage_pct"])
                    churn_rates.append(block_stat["routing_churn_rate"])

        mean_rank = sum(ranks_recorded) / len(ranks_recorded)
        mean_quantum = sum(quantum_ratios) / len(quantum_ratios)
        mean_churn = sum(churn_rates) / len(churn_rates)

        rank_counts = {r: ranks_recorded.count(r) / len(ranks_recorded) for r in [1, 2, 4, 8]}

        seed_res = {
            "seed": seed,
            "mean_rank": round(mean_rank, 2),
            "rank_distribution": {k: round(v, 4) for k, v in rank_counts.items()},
            "quantum_usage_pct": round(mean_quantum, 2),
            "routing_churn_rate": round(mean_churn, 4),
            "classification": "MEASURED",
        }
        results_by_seed.append(seed_res)
        print(f"Seed {seed} | Mean Rank: {mean_rank:.2f} | Quantum: {mean_quantum:.1f}% | Churn: {mean_churn:.4f}")

    summary = {
        "seeds": results_by_seed,
        "overall_mean_rank": round(sum(s["mean_rank"] for s in results_by_seed) / len(results_by_seed), 2),
        "overall_quantum_pct": round(sum(s["quantum_usage_pct"] for s in results_by_seed) / len(results_by_seed), 2),
        "overall_churn_rate": round(sum(s["routing_churn_rate"] for s in results_by_seed) / len(results_by_seed), 4),
    }

    with open(args.output, "w") as f:
        json.dump(summary, f, indent=2)

    print(f"\nResults saved to {args.output}")


if __name__ == "__main__":
    main()