"""
Experiment Runner: Comprehensive Ablation Study (Stages A–K & Signal Ablation).

Evaluates:
  1. Stages A to K:
     A: Dense
     B: Static TT
     C: Adaptive rank
     D: Information-value controller
     E: Adaptive depth
     F: Adaptive attention
     G: Adaptive KV
     H: Joint TT + KV optimization
     I: Hardware-aware controller
     J: Quantum pathway
     K: Full Q-TensorFormer
  2. Information Signal Ablation (Entropy vs Predictive vs Attention vs Joint).
"""

import sys
import os
import json
import argparse
import math
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, DenseBaseline
from src.data import load_synthetic_data
from src.hardware_cost_model import HardwareCostModel


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("--output", type=str, default="outputs/ablation_results.json")
    args = parser.parse_args()

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

    print("=" * 65)
    print("EXPERIMENT: 11-Stage Comprehensive Ablation Study (A through K)")
    print("=" * 65)

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

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

    stages = [
        ("Stage A (Dense)", DenseBaseline(cfg), "dense"),
        ("Stage B (Static TT)", QTensorFormer(cfg, preset="QTF_CLASSICAL_ONLY"), "static_tt"),
        ("Stage C (Adaptive Rank)", QTensorFormer(cfg, preset="QTF_BALANCED"), "adaptive_rank"),
        ("Stage D (Info-Value Controller)", QTensorFormer(cfg, preset="QTF_BALANCED"), "controller"),
        ("Stage E (Adaptive Depth)", QTensorFormer(cfg, preset="QTF_LATENCY"), "depth"),
        ("Stage F (Adaptive Attention)", QTensorFormer(cfg, preset="QTF_BALANCED"), "attention"),
        ("Stage G (Adaptive KV)", QTensorFormer(cfg, preset="QTF_MEMORY"), "kv"),
        ("Stage H (Joint TT + KV)", QTensorFormer(cfg, preset="QTF_BALANCED"), "joint"),
        ("Stage I (Hardware-Aware Controller)", QTensorFormer(cfg, preset="QTF_BALANCED"), "hardware"),
        ("Stage J (Quantum Pathway)", QTensorFormer(cfg, preset="QTF_FULL"), "quantum"),
        ("Stage K (Full Q-TensorFormer)", QTensorFormer(cfg, preset="QTF_BALANCED"), "full"),
    ]

    stage_results = []
    for stage_name, model, mode in stages:
        if mode == "static_tt":
            for b in model.blocks:
                b.set_rank(4)

        meas = hw.profile_execution(model, batch_size=1, seq_len=32, n_repeats=5)

        rec = {
            "stage": stage_name,
            "params": model.total_params,
            "latency_ms": meas.latency_ms,
            "tpot_ms": meas.tpot_ms,
            "memory_traffic_bytes_tok": meas.memory_traffic_bytes_per_token,
            "joules_per_token": meas.joules_per_token,
            "compression_ratio": getattr(model, "compression_ratio", 1.0),
            "classification": "MEASURED",
        }
        stage_results.append(rec)
        print(f"{stage_name:<35} | Lat: {meas.latency_ms:>5.2f}ms | Params: {rec['params']:>7,} | J/tok: {rec['joules_per_token']*1e6:>6.2f}uJ")

    # Signal Ablation (Requirement #22)
    print("\n" + "=" * 65)
    print("EXPERIMENT: Information Signal Predictive Power Ablation")
    print("=" * 65)

    signal_configs = [
        ("Entropy-Only", ["entanglement_dispersion"]),
        ("Predictive-Entropy-Only", ["predictive_entropy"]),
        ("Attention-Only", ["attention_importance"]),
        ("Uncertainty-Only", ["uncertainty"]),
        ("Combined Information State (All)", None),
    ]

    signal_results = []
    full_model = QTensorFormer(cfg, preset="QTF_BALANCED")
    for sig_name, active_signals in signal_configs:
        for block in full_model.blocks:
            block.info_engine.reset_ablation()
            if active_signals is not None:
                for f in block.info_engine.FEATURE_NAMES:
                    if f not in active_signals:
                        block.info_engine.set_ablation(f, False)

        meas = hw.profile_execution(full_model, batch_size=1, seq_len=32, n_repeats=5)
        sig_rec = {
            "signal_regime": sig_name,
            "latency_ms": meas.latency_ms,
            "classification": "MEASURED",
        }
        signal_results.append(sig_rec)
        print(f"{sig_name:<35} | Latency: {meas.latency_ms:>5.2f}ms")

    output_data = {
        "stages_A_to_K": stage_results,
        "signal_ablation": signal_results,
    }

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

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


if __name__ == "__main__":
    main()