"""
Experiment Runner: Adaptive KV Cache Rate-Distortion & Memory Footprint.

Evaluates:
  - Cache size (MB) across sequence lengths (32, 64, 128, 256, 512, 1024)
  - Memory compression under FP16 vs INT8 vs INT4
  - Memory traffic (bytes read / written per token)
  - Eviction fidelity under tight memory constraints
"""

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.kv_cache import AdaptiveKVCache, KVPrecision, QuantizedKVTensor


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

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

    print("=" * 65)
    print("EXPERIMENT: KV Cache Memory & Precision Compression")
    print("=" * 65)

    seq_lengths = [32, 64, 128, 256, 512, 1024]
    precisions = [KVPrecision.FP16, KVPrecision.INT8, KVPrecision.INT4]

    results = []

    B, H, D = 1, 4, 32  # typical head configuration

    for T in seq_lengths:
        row = {"seq_len": T}
        for prec in precisions:
            cache = AdaptiveKVCache(max_capacity=T + 10, default_precision=prec)
            k = torch.randn(B, H, T, D)
            v = torch.randn(B, H, T, D)
            cache.update(k, v)

            footprint_mb = cache.current_mb
            traffic_written = cache.total_bytes_written
            traffic_read = cache.total_bytes_read

            row[f"{prec.value}_mb"] = round(footprint_mb, 4)
            row[f"{prec.value}_bytes_read"] = traffic_read

        # Compute compression vs FP16
        fp16_mb = row["fp16_mb"]
        int8_mb = row["int8_mb"]
        int4_mb = row["int4_mb"]
        row["compression_int8_vs_fp16"] = round(fp16_mb / max(1e-5, int8_mb), 2)
        row["compression_int4_vs_fp16"] = round(fp16_mb / max(1e-5, int4_mb), 2)
        row["classification"] = "MEASURED"

        results.append(row)
        print(f"SeqLen {T:>4} | FP16: {fp16_mb:>6.3f} MB | INT8: {int8_mb:>6.3f} MB ({row['compression_int8_vs_fp16']:.1f}x) | INT4: {int4_mb:>6.3f} MB ({row['compression_int4_vs_fp16']:.1f}x)")

    # Eviction test: simulate context length of 1000 with capacity of 256
    cache_evict = AdaptiveKVCache(max_capacity=256, default_precision=KVPrecision.INT8, window_size=64)
    # Stream in 1000 tokens
    for chunk in range(10):
        k_chunk = torch.randn(B, H, 100, D)
        v_chunk = torch.randn(B, H, 100, D)
        cache_evict.update(k_chunk, v_chunk)

    eviction_summary = {
        "final_seq_len": cache_evict.seq_len,
        "tokens_evicted": cache_evict.evicted_tokens_count,
        "peak_mb": round(cache_evict.current_mb, 3),
        "classification": "MEASURED",
    }
    print(f"\nEviction Test (1000 tokens streamed into capacity 256):")
    print(f"  Final Cache Length: {cache_evict.seq_len} | Evicted Tokens: {cache_evict.evicted_tokens_count} | Footprint: {cache_evict.current_mb:.3f} MB")

    output_data = {
        "scaling_table": results,
        "eviction_test": eviction_summary,
    }

    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()