#!/usr/bin/env python3
"""Ablation evaluation — loads models from JSON file instead of hardcoded dict.

Usage:
    python scripts/run_ablation_eval.py --models-json ablation_eval_models.json \
        --num-gpus 8 --output-dir ablation_eval
"""

import argparse
import json
import logging
import os
import sys
import time

# Reuse everything from the comprehensive eval script
sys.path.insert(0, os.path.dirname(__file__))
from run_comprehensive_eval import (
    select_prompts,
    get_ref_audios,
    build_task_list,
    run_generation_phase,
    score_all_audio,
    build_html_report,
    VAP_DIR,
)

logging.basicConfig(level=logging.INFO, format="%(asctime)s %(levelname)s %(message)s")


def main():
    parser = argparse.ArgumentParser(description="Run ablation evaluation with models from JSON")
    parser.add_argument("--models-json", required=True, help="Path to JSON file with models dict")
    parser.add_argument("--output-dir", default=os.path.join(VAP_DIR, "ablation_eval"))
    parser.add_argument("--refs-dir", default="/home/deployer/laion/test-refs")
    parser.add_argument("--num-gpus", type=int, default=8)
    parser.add_argument("--seeds", type=int, nargs="+", default=[42])
    parser.add_argument("--classifiers-dir", default=os.path.join(VAP_DIR, "classifiers"))
    parser.add_argument("--skip-generation", action="store_true")
    parser.add_argument("--skip-scoring", action="store_true")
    args = parser.parse_args()

    # Load models from JSON
    with open(args.models_json) as f:
        models = json.load(f)
    logging.info(f"Loaded {len(models)} models from {args.models_json}")
    for key, minfo in models.items():
        lora_str = os.path.basename(minfo["lora"]) if minfo.get("lora") else "None"
        logging.info(f"  {key}: {minfo['name']} (lora={lora_str})")

    # Verify all LoRA files exist
    for key, minfo in models.items():
        if minfo.get("lora") and not os.path.exists(minfo["lora"]):
            logging.error(f"LoRA file not found for {key}: {minfo['lora']}")
            sys.exit(1)

    args.output_dir = os.path.abspath(args.output_dir)
    os.makedirs(os.path.join(args.output_dir, "wavs"), exist_ok=True)

    # 1. Select prompts
    prompts = select_prompts()
    logging.info(f"Selected {len(prompts)} prompts")

    with open(os.path.join(args.output_dir, "prompts.json"), "w") as f:
        json.dump(prompts, f, indent=2)

    # 2. Get reference audios
    refs = get_ref_audios(args.refs_dir)
    logging.info(f"Found {len(refs)} reference audios")

    # 3. Build task list
    tasks = build_task_list(models, prompts, refs, args.seeds, args.output_dir)
    logging.info(f"Total tasks: {len(tasks)} ({len(models)} models × {len(prompts)} prompts × ({len(refs)} refs + 1 uncond) × {len(args.seeds)} seeds)")

    # 4. Generate (skip tasks where WAV already exists)
    if not args.skip_generation:
        tasks_to_gen = [t for t in tasks if not os.path.exists(t["output"])]
        tasks_existing = [t for t in tasks if os.path.exists(t["output"])]
        logging.info(f"Skipping {len(tasks_existing)} existing WAVs, generating {len(tasks_to_gen)} new samples")
        if tasks_to_gen:
            gen_results = run_generation_phase(tasks_to_gen, args.num_gpus)
        else:
            gen_results = []
        existing_results = [{**t, "success": True, "elapsed": 0, "gpu": -1} for t in tasks_existing]
        results = existing_results + gen_results
        with open(os.path.join(args.output_dir, "gen_results.json"), "w") as f:
            json.dump(results, f, indent=2, default=str)
    else:
        with open(os.path.join(args.output_dir, "gen_results.json")) as f:
            results = json.load(f)

    # 5. Score
    if not args.skip_scoring:
        scored = score_all_audio(results, args.classifiers_dir)
        with open(os.path.join(args.output_dir, "scored_results.json"), "w") as f:
            json.dump(scored, f, indent=2, default=str)
    else:
        with open(os.path.join(args.output_dir, "scored_results.json")) as f:
            scored = json.load(f)

    # 6. Build HTML
    model_means = build_html_report(
        scored, models, prompts, refs,
        os.path.join(args.output_dir, "eval_report.html")
    )

    # Print summary
    print("\n" + "="*70)
    print("ABLATION EVALUATION SUMMARY")
    print("="*70)
    for mk, minfo in models.items():
        means = model_means.get(mk, {})
        print(f"\n{minfo['name']}:")
        for metric in ["naturalness_small", "naturalness_large", "centroid", "quality", "speaker_sim"]:
            val = means.get(metric)
            print(f"  {metric:20s}: {val:.4f}" if val else f"  {metric:20s}: N/A")
    print("="*70)

    # Save summary JSON
    summary = {mk: model_means.get(mk, {}) for mk in models}
    with open(os.path.join(args.output_dir, "summary.json"), "w") as f:
        json.dump(summary, f, indent=2)
    logging.info(f"Summary saved to {os.path.join(args.output_dir, 'summary.json')}")


if __name__ == "__main__":
    main()