#!/usr/bin/env python3
"""Overnight ablation study runner.

Runs 4 ablation training experiments sequentially, then evaluates all models
in a single comprehensive HTML report.

Ablations (all from LaionBox v0.1-wip, 2 epochs, same hyperparameters):
  A: naturalness + quality MLP
  B: naturalness + centroid
  C: naturalness + quality MLP + speaker similarity
  D: naturalness + speaker similarity

After training, generates audio for the best-flow and best-naturalness
checkpoint from each ablation, scores all models (including baselines),
and produces a combined HTML comparison.
"""

import json
import os
import signal
import subprocess
import sys
import time

PYTHON = "/home/deployer/miniconda3/envs/ml-general/bin/python"
VAP_DIR = "/home/deployer/laion/Voice-Acting-Pipeline"
TRAIN_SCRIPT = os.path.join(VAP_DIR, "scripts", "dramabox_finetune_train_multi_aux.py")
EVAL_SCRIPT = os.path.join(VAP_DIR, "scripts", "run_ablation_eval.py")

ABLATIONS = [
    {
        "name": "A_nat_quality",
        "config": "configs/ablation_nat_quality.yaml",
        "output_dir": "finetune_output/ablation_nat_quality",
        "desc_prefix": "Nat+Quality",
        "losses": "naturalness + quality_mlp",
    },
    {
        "name": "B_nat_centroid",
        "config": "configs/ablation_nat_centroid.yaml",
        "output_dir": "finetune_output/ablation_nat_centroid",
        "desc_prefix": "Nat+Centroid",
        "losses": "naturalness + centroid",
    },
    {
        "name": "C_nat_quality_speaker",
        "config": "configs/ablation_nat_quality_speaker.yaml",
        "output_dir": "finetune_output/ablation_nat_quality_speaker",
        "desc_prefix": "Nat+Quality+Speaker",
        "losses": "naturalness + quality_mlp + speaker_sim",
    },
    {
        "name": "D_nat_speaker",
        "config": "configs/ablation_nat_speaker.yaml",
        "output_dir": "finetune_output/ablation_nat_speaker",
        "desc_prefix": "Nat+Speaker",
        "losses": "naturalness + speaker_sim",
    },
]


def log(msg):
    ts = time.strftime("%Y-%m-%d %H:%M:%S")
    line = f"[{ts}] {msg}"
    print(line, flush=True)


def find_best_checkpoints(output_dir):
    """Parse metrics.jsonl to find best flow loss and best naturalness checkpoints."""
    metrics_path = os.path.join(VAP_DIR, output_dir, "metrics.jsonl")
    if not os.path.exists(metrics_path):
        log(f"  WARNING: {metrics_path} not found")
        return None, None

    best_flow_loss = float("inf")
    best_flow_step = None
    best_nat = float("-inf")
    best_nat_step = None

    with open(metrics_path) as f:
        for line in f:
            line = line.strip()
            if not line:
                continue
            try:
                m = json.loads(line)
            except json.JSONDecodeError:
                continue
            step = m.get("step", 0)
            flow = m.get("flow_loss", m.get("loss", float("inf")))
            nat = m.get("naturalness_reward", float("-inf"))

            if flow < best_flow_loss:
                best_flow_loss = flow
                best_flow_step = step
            if nat is not None and nat > best_nat:
                best_nat = nat
                best_nat_step = step

    log(f"  Best flow: step {best_flow_step} (loss={best_flow_loss:.4f})")
    log(f"  Best nat:  step {best_nat_step} (nat={best_nat:.4f})")

    # Find closest saved checkpoint (save_every=10)
    def find_ckpt(step):
        if step is None:
            return None, None
        # Try exact step first, then nearby
        for offset in [0, -1, 1, -2, 2, -3, 3, -4, 4, -5, 5]:
            candidate_step = (step + offset * 10) if offset != 0 else step
            # Try step checkpoint
            fname = f"lora_step_{candidate_step:05d}.safetensors"
            path = os.path.join(VAP_DIR, output_dir, fname)
            if os.path.exists(path):
                return path, candidate_step
            # Also try without zero padding
            fname = f"lora_step_{candidate_step:05d}.safetensors"
            path = os.path.join(VAP_DIR, output_dir, fname)
            if os.path.exists(path):
                return path, candidate_step
        # Try epoch checkpoints
        for epoch in [1, 2]:
            path = os.path.join(VAP_DIR, output_dir, f"lora_epoch{epoch}.safetensors")
            if os.path.exists(path):
                return path, f"epoch{epoch}"
        return None, None

    # Round step to nearest save_every=10
    def snap_step(step):
        if step is None:
            return None
        return round(step / 10) * 10

    flow_path, flow_step = find_ckpt(snap_step(best_flow_step))
    nat_path, nat_step = find_ckpt(snap_step(best_nat_step))

    if flow_path:
        log(f"  Flow checkpoint: {os.path.basename(flow_path)}")
    else:
        log(f"  WARNING: No flow checkpoint found near step {best_flow_step}")

    if nat_path:
        log(f"  Nat checkpoint:  {os.path.basename(nat_path)}")
    else:
        log(f"  WARNING: No nat checkpoint found near step {best_nat_step}")

    return (flow_path, flow_step, best_flow_loss), (nat_path, nat_step, best_nat)


def run_training(ablation):
    """Run a single ablation training."""
    name = ablation["name"]
    config = ablation["config"]
    log_file = f"/tmp/ablation_{name}.log"

    log(f"{'='*70}")
    log(f"STARTING ABLATION: {name} ({ablation['losses']})")
    log(f"Config: {config}")
    log(f"Log: {log_file}")
    log(f"{'='*70}")

    cmd = [
        "env", "-u", "LD_LIBRARY_PATH",
        "accelerate", "launch", "--num_processes=8",
        TRAIN_SCRIPT,
        "--config", config,
    ]

    env = os.environ.copy()
    env.pop("LD_LIBRARY_PATH", None)
    # Ensure accelerate from ml-general is used
    env["PATH"] = "/home/deployer/miniconda3/envs/ml-general/bin:" + env.get("PATH", "")

    with open(log_file, "w") as lf:
        proc = subprocess.Popen(
            cmd, stdout=lf, stderr=subprocess.STDOUT,
            cwd=VAP_DIR, env=env,
        )

    log(f"Training PID: {proc.pid}")

    # Wait for completion
    t0 = time.time()
    while True:
        ret = proc.poll()
        if ret is not None:
            break
        elapsed = time.time() - t0
        # Print progress every 5 minutes
        if int(elapsed) % 300 == 0 and elapsed > 10:
            # Read last line of status.json if it exists
            status_path = os.path.join(VAP_DIR, ablation["output_dir"], "status.json")
            if os.path.exists(status_path):
                try:
                    with open(status_path) as sf:
                        status = json.load(sf)
                    step = status.get("step", "?")
                    total = status.get("total_steps", "?")
                    loss = status.get("flow_loss", status.get("loss", "?"))
                    eta = status.get("eta_sec", 0)
                    eta_m = int(eta // 60) if isinstance(eta, (int, float)) else "?"
                    log(f"  Progress: step {step}/{total}, loss={loss}, ETA={eta_m}m")
                except Exception:
                    pass
        time.sleep(10)

    elapsed_min = (time.time() - t0) / 60
    if ret == 0:
        log(f"Training {name} COMPLETED in {elapsed_min:.1f} min (exit code 0)")
    else:
        log(f"Training {name} FAILED with exit code {ret} after {elapsed_min:.1f} min")
        log(f"  Check log: {log_file}")

    return ret == 0


def write_eval_models_json(all_models, path):
    """Write the models dict as JSON for the eval script."""
    with open(path, "w") as f:
        json.dump(all_models, f, indent=2)
    log(f"Wrote {len(all_models)} models to {path}")


def upload_to_hf(abl_name, output_dir, flow_info, nat_info):
    """Upload best checkpoints + metrics to HF."""
    try:
        from huggingface_hub import HfApi
        api = HfApi(token="HF_TOKEN_REDACTED")
        repo = "TTS-AGI/laionbox-ablation-checkpoints"

        uploads = []
        metrics_path = os.path.join(VAP_DIR, output_dir, "metrics.jsonl")
        if os.path.exists(metrics_path):
            uploads.append((metrics_path, f"{abl_name}/metrics.jsonl"))

        if flow_info and flow_info[0]:
            path, step, _ = flow_info
            uploads.append((path, f"{abl_name}/best_flow_step{step}.safetensors"))
        if nat_info and nat_info[0]:
            path, step, _ = nat_info
            uploads.append((path, f"{abl_name}/best_nat_step{step}.safetensors"))

        for local, remote in uploads:
            if os.path.exists(local):
                sz = os.path.getsize(local) / 1024 / 1024
                log(f"  HF upload: {remote} ({sz:.0f} MB)")
                api.upload_file(path_or_fileobj=local, path_in_repo=remote, repo_id=repo)
        log(f"  HF upload complete for {abl_name}")
    except Exception as e:
        log(f"  HF upload FAILED for {abl_name}: {e}")


def run_eval(models_json_path, output_dir):
    """Run the comprehensive evaluation."""
    log(f"{'='*70}")
    log("STARTING COMPREHENSIVE EVALUATION")
    log(f"{'='*70}")

    log_file = "/tmp/ablation_eval.log"

    cmd = [
        "env", "-u", "LD_LIBRARY_PATH",
        PYTHON, EVAL_SCRIPT,
        "--models-json", models_json_path,
        "--num-gpus", "8",
        "--output-dir", output_dir,
    ]

    env = os.environ.copy()
    env.pop("LD_LIBRARY_PATH", None)

    with open(log_file, "w") as lf:
        proc = subprocess.Popen(
            cmd, stdout=lf, stderr=subprocess.STDOUT,
            cwd=VAP_DIR, env=env,
        )

    log(f"Eval PID: {proc.pid}, log: {log_file}")

    proc.wait()
    elapsed_min = proc.returncode
    if proc.returncode == 0:
        log("Evaluation COMPLETED successfully")
    else:
        log(f"Evaluation FAILED with exit code {proc.returncode}")
        log(f"  Check log: {log_file}")

    return proc.returncode == 0


def main():
    log("="*70)
    log("ABLATION STUDY: Auxiliary Loss Comparison")
    log("4 training runs + comprehensive evaluation")
    log("="*70)
    log("")
    log("Ablations:")
    for i, abl in enumerate(ABLATIONS):
        log(f"  {abl['name']}: {abl['losses']}")
    log("")

    # Phase 1: Training
    results = {}
    for abl in ABLATIONS:
        # Check if this ablation already completed (metrics.jsonl has >= 20 entries)
        metrics_path = os.path.join(VAP_DIR, abl["output_dir"], "metrics.jsonl")
        if os.path.exists(metrics_path):
            with open(metrics_path) as f:
                n_lines = sum(1 for _ in f)
            if n_lines >= 20:
                log(f"SKIPPING {abl['name']}: already has {n_lines} metric entries")
                log(f"Finding best checkpoints for {abl['name']}...")
                flow_info, nat_info = find_best_checkpoints(abl["output_dir"])
                results[abl["name"]] = {
                    "ablation": abl,
                    "flow": flow_info,
                    "nat": nat_info,
                }
                log("")
                continue

        success = run_training(abl)
        if success:
            log(f"\nFinding best checkpoints for {abl['name']}...")
            flow_info, nat_info = find_best_checkpoints(abl["output_dir"])
            results[abl["name"]] = {
                "ablation": abl,
                "flow": flow_info,  # (path, step, loss)
                "nat": nat_info,    # (path, step, nat_score)
            }
            # Upload to HF
            log(f"Uploading {abl['name']} checkpoints to HF...")
            upload_to_hf(abl["name"], abl["output_dir"], flow_info, nat_info)
        else:
            log(f"\nWARNING: {abl['name']} failed, skipping checkpoint extraction")
            results[abl["name"]] = None
        log("")

    # Phase 2: Build models dict for evaluation
    log("\n" + "="*70)
    log("TRAINING PHASE COMPLETE — BUILDING EVALUATION")
    log("="*70)

    models = {
        "vanilla": {
            "name": "Vanilla DramaBox",
            "lora": None,
            "desc": "Base DramaBox model without any fine-tuning",
        },
        "laionbox_v01": {
            "name": "LaionBox v0.1-wip",
            "lora": "/home/deployer/.cache/huggingface/hub/models--laion--laionbox-v0.1-wip/snapshots/66176d2a653a013a7b71c1ccb7a7a4d4cf514b0d/lora_epoch5.safetensors",
            "desc": "Previous best LoRA (5-epoch diff reward, DramaBox+Emolia data)",
        },
        "nat_only_best_flow": {
            "name": "Nat-Only Best-Flow (s160)",
            "lora": os.path.join(VAP_DIR, "finetune_output/nat_only_2ep/lora_step_00160.safetensors"),
            "desc": "Naturalness-only (CLAP-7B), best flow=0.528 @s160",
        },
        "nat_only_best_nat": {
            "name": "Nat-Only Best-Nat (s190)",
            "lora": os.path.join(VAP_DIR, "finetune_output/nat_only_2ep/lora_step_00190.safetensors"),
            "desc": "Naturalness-only (CLAP-7B), best nat=0.111 @s190",
        },
    }

    for abl_name, res in results.items():
        if res is None:
            continue
        abl = res["ablation"]
        prefix = abl["desc_prefix"]
        losses_short = abl["losses"]

        if res["flow"] and res["flow"][0]:
            path, step, loss = res["flow"]
            key = f"{abl_name}_best_flow"
            models[key] = {
                "name": f"{prefix} Best-Flow (s{step})",
                "lora": path,
                "desc": f"{losses_short}, best flow={loss:.4f} @s{step}",
            }

        if res["nat"] and res["nat"][0]:
            path, step, nat_score = res["nat"]
            key = f"{abl_name}_best_nat"
            models[key] = {
                "name": f"{prefix} Best-Nat (s{step})",
                "lora": path,
                "desc": f"{losses_short}, best nat={nat_score:.4f} @s{step}",
            }

    log(f"\nTotal models for evaluation: {len(models)}")
    for key, minfo in models.items():
        log(f"  {key}: {minfo['name']}")

    # Write models JSON
    models_json = os.path.join(VAP_DIR, "ablation_eval_models.json")
    write_eval_models_json(models, models_json)

    # Phase 3: Evaluation
    eval_output = os.path.join(VAP_DIR, "ablation_eval")
    success = run_eval(models_json, eval_output)

    # Phase 4: Summary
    log("\n" + "="*70)
    log("ABLATION STUDY COMPLETE")
    log("="*70)

    if success:
        report_path = os.path.join(eval_output, "eval_report.html")
        log(f"HTML report: {report_path}")
        log("Serve with: python -m http.server 8780 --directory " + eval_output)
    else:
        log("Evaluation failed — check /tmp/ablation_eval.log")

    # Print training summary
    log("\nTraining Summary:")
    for abl_name, res in results.items():
        if res is None:
            log(f"  {abl_name}: FAILED")
            continue
        flow_str = f"flow={res['flow'][2]:.4f}@s{res['flow'][1]}" if res["flow"] and res["flow"][0] else "N/A"
        nat_str = f"nat={res['nat'][2]:.4f}@s{res['nat'][1]}" if res["nat"] and res["nat"][0] else "N/A"
        log(f"  {abl_name}: {flow_str}, {nat_str}")


if __name__ == "__main__":
    main()