"""
TD Self-Improvement Loop v1 — The Full Autonomous Pipeline

Ties everything together into an endless self-improvement loop:

  Step 1: Test (benchmark + weight health)
  Step 2: Diagnose (weakness finder)
  Step 3: Find data (dataset finder)
  Step 4: Train + Repair (selfimprove.py)
  Step 5: Verify (re-benchmark)
  → Loop

This is the script you start and walk away from.

Usage:
  python td_loop.py                          # Start from latest checkpoint
  python td_loop.py --model path/to/model    # Start from specific model
  python td_loop.py --max-cycles 10          # Stop after 10 cycles
  python td_loop.py --dry-run                # Just test+diagnose, don't train
"""

import json
import time
import shutil
import gc
import sys
import argparse
from pathlib import Path
from typing import Dict, List, Optional, Tuple
from dataclasses import dataclass


# ============================================================
# LOOP CONFIGURATION
# ============================================================

@dataclass
class LoopConfig:
    """Configuration for the self-improvement loop."""

    # Model
    model_path: str = ""  # Auto-detected if empty
    reference_model: str = "Qwen/Qwen3-VL-8B-Instruct"

    # Directories
    output_dir: str = "td_fuse_outputs/self_improve"
    benchmark_dir: str = "td_fuse_outputs/benchmarks"
    weakness_dir: str = "td_fuse_outputs/weakness_reports"
    data_dir: str = "td_fuse_outputs/training_data"

    # Loop settings
    max_cycles: int = 0  # 0 = infinite
    dry_run: bool = False  # Test + diagnose only
    training_budget: int = 8000  # Samples per cycle

    # Score targets (when to stop)
    targets: Dict = None

    # HuggingFace
    hf_token: str = ""

    # Regression protection
    max_regression_pct: float = 0.05  # 5% drop = rollback
    max_consecutive_regressions: int = 2  # 2 in a row = stop

    # Quick test mode (compressed run to verify code paths)
    quick_test: bool = False

    # Safety limits
    max_training_hours: float = 4.0  # Warn if training exceeds this (hours per cycle)

    def __post_init__(self):
        if self.targets is None:
            self.targets = {
                "math": 0.80,
                "code": 0.70,
                "reasoning": 0.75,
                "creativity": 0.65,
                "knowledge": 0.80,
                "instruction_following": 0.85,
            }


# ============================================================
# CYCLE STATE TRACKING
# ============================================================

@dataclass
class CycleResult:
    """Results from one improvement cycle."""
    cycle_num: int
    benchmark_scores: Dict[str, float]
    weight_health_summary: Dict
    weaknesses_found: int
    data_collected: int
    training_loss: float
    improved: bool  # Overall, did we improve?
    regressions: List[str]  # Categories that got worse
    duration_seconds: float
    model_path: str


def load_cycle_history(output_dir: str) -> List[CycleResult]:
    """Load previous cycle results. Returns empty list if file is corrupt."""
    history_file = Path(output_dir) / "loop_history.json"
    if not history_file.exists():
        return []

    try:
        with open(history_file) as f:
            data = json.load(f)
        return [CycleResult(**c) for c in data]
    except (json.JSONDecodeError, TypeError, KeyError) as e:
        print(f"  WARNING: History file corrupt ({e}) — starting fresh")
        # Back up the corrupt file
        backup = history_file.with_suffix(".json.bak")
        try:
            shutil.copy2(str(history_file), str(backup))
            print(f"  Backed up corrupt history to {backup}")
        except Exception:
            pass
        return []


def save_cycle_history(history: List[CycleResult], output_dir: str):
    """Save cycle history to disk. Atomic write to prevent corruption."""
    history_file = Path(output_dir) / "loop_history.json"
    tmp_file = Path(output_dir) / "loop_history.json.tmp"
    Path(output_dir).mkdir(parents=True, exist_ok=True)

    data = []
    for c in history:
        d = {
            "cycle_num": c.cycle_num,
            "benchmark_scores": c.benchmark_scores,
            "weight_health_summary": c.weight_health_summary,
            "weaknesses_found": c.weaknesses_found,
            "data_collected": c.data_collected,
            "training_loss": c.training_loss,
            "improved": c.improved,
            "regressions": c.regressions,
            "duration_seconds": c.duration_seconds,
            "model_path": c.model_path,
        }
        data.append(d)

    try:
        # Write to temp file first, then atomic rename
        with open(tmp_file, "w") as f:
            json.dump(data, f, indent=2)
        # Atomic rename — if this fails, the old file is still intact
        tmp_file.replace(history_file)
    except OSError as e:
        print(f"  WARNING: Failed to save history ({e}). Disk may be full!")
        # Try to clean up temp file
        try:
            tmp_file.unlink(missing_ok=True)
        except Exception:
            pass


# ============================================================
# THE LOOP
# ============================================================

def _setup_logging(output_dir: str):
    """
    Tee all stdout/stderr to a log file so diagnostics survive terminal disconnects.
    Returns the log file path.
    """
    import sys as _sys

    Path(output_dir).mkdir(parents=True, exist_ok=True)
    log_path = Path(output_dir) / f"loop_{time.strftime('%Y%m%d_%H%M%S')}.log"

    class Tee:
        def __init__(self, original, log_file):
            self.original = original
            self.log_file = log_file
        def write(self, data):
            self.original.write(data)
            try:
                self.log_file.write(data)
                self.log_file.flush()
            except Exception:
                pass  # Don't crash if log write fails
        def flush(self):
            self.original.flush()
            try:
                self.log_file.flush()
            except Exception:
                pass

    log_fh = open(log_path, "a")
    _sys.stdout = Tee(_sys.stdout, log_fh)
    _sys.stderr = Tee(_sys.stderr, log_fh)
    return str(log_path)


def _acquire_lockfile(output_dir: str) -> Optional[str]:
    """
    Prevent two td_loop instances from running simultaneously.
    Returns lockfile path if acquired, None if another instance is running.
    """
    import os
    lock_path = Path(output_dir) / "td_loop.lock"
    Path(output_dir).mkdir(parents=True, exist_ok=True)

    if lock_path.exists():
        # Check if the PID in the lockfile is still alive
        try:
            with open(lock_path) as f:
                old_pid = int(f.read().strip())
            os.kill(old_pid, 0)  # Signal 0 = check if alive
            # Process is alive — another instance is running
            return None
        except (ValueError, ProcessLookupError, PermissionError, OSError):
            # PID invalid or dead — stale lockfile, safe to remove
            print(f"  Removing stale lockfile (PID {lock_path.read_text().strip()} is dead)")

    with open(lock_path, "w") as f:
        f.write(str(os.getpid()))
    return str(lock_path)


def _release_lockfile(output_dir: str):
    """Release the lockfile."""
    lock_path = Path(output_dir) / "td_loop.lock"
    try:
        lock_path.unlink(missing_ok=True)
    except Exception:
        pass


def run_loop(cfg: LoopConfig):
    """
    The main self-improvement loop.

    Start it and walk away. It will:
    1. Test the model (benchmark + weight health)
    2. Find weaknesses
    3. Search for training data
    4. Train + repair
    5. Verify improvement
    6. Repeat
    """
    # ── LOCKFILE: Prevent duplicate instances ──
    lock = _acquire_lockfile(cfg.output_dir)
    if lock is None:
        print("ERROR: Another td_loop instance is already running!")
        print(f"  Lock file: {Path(cfg.output_dir) / 'td_loop.lock'}")
        print("  If this is wrong (crashed process), delete the lock file and retry.")
        return

    # ── LOG TEE: Save all output to file ──
    log_path = _setup_logging(cfg.output_dir)

    # ── SIGINT HANDLER: Graceful shutdown ──
    import signal
    shutdown_requested = [False]  # Mutable container for closure

    def _signal_handler(signum, frame):
        if shutdown_requested[0]:
            print("\n  FORCED SHUTDOWN — exiting immediately!")
            _release_lockfile(cfg.output_dir)
            raise SystemExit(1)
        shutdown_requested[0] = True
        print(f"\n  SHUTDOWN REQUESTED (Ctrl+C) — will stop after current step completes.")
        print(f"  Press Ctrl+C again to force quit.")

    signal.signal(signal.SIGINT, _signal_handler)
    signal.signal(signal.SIGTERM, _signal_handler)

    print("\n" + "=" * 60)
    print("TD SELF-IMPROVEMENT LOOP v1")
    print("=" * 60)
    print(f"  Max cycles: {'infinite' if cfg.max_cycles == 0 else cfg.max_cycles}")
    print(f"  Training budget: {cfg.training_budget} samples/cycle")
    print(f"  Dry run: {cfg.dry_run}")
    print(f"  Started: {time.strftime('%Y-%m-%d %H:%M:%S')}")
    print(f"  Log file: {log_path}")
    print("=" * 60)

    try:
        _run_loop_inner(cfg, shutdown_requested)
    finally:
        _release_lockfile(cfg.output_dir)
        print(f"\n  Lock released. Log saved to: {log_path}")


def _run_loop_inner(cfg: LoopConfig, shutdown_requested: list):
    """Inner loop logic, separated so run_loop can handle cleanup."""

    # ── GPU HEALTH CHECK ──
    try:
        import torch
        if torch.cuda.is_available():
            gpu_name = torch.cuda.get_device_name(0)
            gpu_mem = torch.cuda.get_device_properties(0).total_memory / 1e9
            print(f"\n  GPU: {gpu_name} ({gpu_mem:.1f} GB)")
            if gpu_mem < 20:
                print(f"  WARNING: GPU has only {gpu_mem:.1f} GB — Qwen3-VL-8B needs ~24 GB for training!")
            # Quick allocation test — try to allocate 1GB tensor
            test = torch.empty(256, 1024, 1024, dtype=torch.bfloat16, device="cuda")
            del test
            torch.cuda.empty_cache()
            print(f"  GPU health: OK (allocation test passed)")
        else:
            print("\n  WARNING: No CUDA GPU detected! Training will be extremely slow or fail.")
    except Exception as e:
        print(f"\n  WARNING: GPU check failed: {e}")

    # ── DISK HEALTH CHECK ──
    try:
        disk = shutil.disk_usage(cfg.output_dir if Path(cfg.output_dir).exists() else ".")
        free_gb = disk.free / 1e9
        print(f"  Disk: {free_gb:.1f} GB free")
        if free_gb < 50:
            print(f"  WARNING: Low disk space! Need ~25 GB per cycle for checkpoints.")
    except Exception:
        pass

    # Load history
    history = load_cycle_history(cfg.output_dir)
    cycle_num = len(history) + 1
    cycles_run_this_session = 0
    consecutive_regressions = 0

    # Auto-detect model if not specified
    if not cfg.model_path:
        cfg.model_path = _auto_detect_model()

    print(f"\n  Starting model: {cfg.model_path}")
    print(f"  Starting at cycle: {cycle_num}")

    # The loop
    while True:
        # ── GRACEFUL SHUTDOWN CHECK ──
        if shutdown_requested[0]:
            print(f"\n  Shutdown requested — stopping before cycle {cycle_num}")
            _print_final_summary(history)
            return
        if cfg.max_cycles > 0 and cycles_run_this_session >= cfg.max_cycles:
            print(f"\n  Reached max cycles ({cfg.max_cycles}). Stopping.")
            break

        cycle_start = time.time()
        print(f"\n{'#' * 60}")
        print(f"# CYCLE {cycle_num}")
        print(f"# {time.strftime('%Y-%m-%d %H:%M:%S')}")
        print(f"{'#' * 60}")

        # ============================================
        # STEP 1: TEST (benchmark + weight health)
        # ============================================
        print(f"\n{'='*50}")
        print(f"STEP 1/{5 if not cfg.dry_run else 2}: TEST MODEL")
        print(f"{'='*50}")

        # 1A: Run benchmark (with crash recovery)
        try:
            benchmark_scores = _run_benchmark(cfg.model_path, cfg.benchmark_dir, cycle_num)
        except Exception as e:
            print(f"\n  CRITICAL: Benchmark crashed: {e}")
            import traceback
            traceback.print_exc()
            # Can't proceed without benchmark scores — try to clean GPU and retry once
            gc.collect()
            try:
                import torch
                if torch.cuda.is_available():
                    torch.cuda.empty_cache()
            except ImportError:
                pass
            try:
                print("  Retrying benchmark after GPU cleanup...")
                benchmark_scores = _run_benchmark(cfg.model_path, cfg.benchmark_dir, cycle_num)
            except Exception as e2:
                print(f"  Benchmark failed twice: {e2}")
                print(f"  Skipping cycle {cycle_num}")
                consecutive_regressions += 1
                if consecutive_regressions >= cfg.max_consecutive_regressions:
                    print(f"\n  {consecutive_regressions} consecutive failures — STOPPING LOOP")
                    break
                cycle_num += 1
                continue

        # 1B: Analyze weight health
        health_data = _analyze_weights(cfg.model_path, cfg.reference_model, cfg.output_dir)

        # ============================================
        # STEP 2: DIAGNOSE (find weaknesses)
        # ============================================
        print(f"\n{'='*50}")
        print(f"STEP 2/{5 if not cfg.dry_run else 2}: DIAGNOSE WEAKNESSES")
        print(f"{'='*50}")

        try:
            weakness_report = _find_weaknesses(
                benchmark_path=str(Path(cfg.benchmark_dir) / f"cycle{cycle_num}" / "results.json"),
                health_data=health_data,
                targets=cfg.targets,
                budget=cfg.training_budget,
                output_dir=cfg.weakness_dir,
                cycle_num=cycle_num,
            )
        except Exception as e:
            print(f"\n  WARNING: Weakness analysis crashed: {e}")
            print(f"  Using fallback: train all categories equally")
            weakness_report = {
                "search_plan": [{"category": c, "budget": cfg.training_budget // 6,
                                 "keywords": [f"{c} training dataset"],
                                 "hf_tags": [], "priority": 0.5, "score": 0.5,
                                 "target": 0.75, "gap": 0.25}
                                for c in ["math", "code", "reasoning", "creativity",
                                          "knowledge", "instruction_following"]],
                "data_allocation": {c: cfg.training_budget // 6
                                    for c in ["math", "code", "reasoning", "creativity",
                                              "knowledge", "instruction_following"]},
                "lora_target_layers": list(range(36)),
                "norms_to_repair": [],
                "total_weaknesses": 6,
            }

        if cfg.dry_run:
            print("\n  DRY RUN — stopping after diagnosis.")
            result = CycleResult(
                cycle_num=cycle_num,
                benchmark_scores=benchmark_scores,
                weight_health_summary=health_data.get("summary", {}),
                weaknesses_found=weakness_report.get("total_weaknesses", 0),
                data_collected=0,
                training_loss=0.0,
                improved=False,
                regressions=[],
                duration_seconds=time.time() - cycle_start,
                model_path=cfg.model_path,
            )
            history.append(result)
            save_cycle_history(history, cfg.output_dir)
            break

        # ============================================
        # STEP 3: FIND TRAINING DATA
        # ============================================
        print(f"\n{'='*50}")
        print(f"STEP 3/5: FIND TRAINING DATA")
        print(f"{'='*50}")

        # Guard: If no weaknesses found, skip data search + training
        if not weakness_report.get("search_plan"):
            print("\n  No weaknesses found — all targets met! Skipping training.")
            print("  (This shouldn't happen unless the model is already great)")
            training_data = {"total_samples": 0, "categories_covered": [], "categories_missing": []}
        else:
            try:
                training_data = _find_training_data(
                    weakness_report=weakness_report,
                    output_dir=cfg.data_dir,
                    hf_token=cfg.hf_token,
                    cycle_num=cycle_num,
                )
            except Exception as e:
                print(f"\n  WARNING: Data finder crashed: {e}")
                print(f"  Continuing with generic datasets (selfimprove.py fallback)")
                training_data = {"total_samples": 0, "categories_covered": [], "categories_missing": []}

        # ============================================
        # STEP 4: TRAIN + REPAIR
        # ============================================
        print(f"\n{'='*50}")
        print(f"STEP 4/5: TRAIN + REPAIR")
        print(f"{'='*50}")

        # ── DISK SPACE PRE-CHECK ──
        try:
            disk = shutil.disk_usage(cfg.output_dir if Path(cfg.output_dir).exists() else ".")
            free_gb = disk.free / 1e9
            if free_gb < 10:
                print(f"\n  CRITICAL: Only {free_gb:.1f} GB disk free — skipping training!")
                print(f"  Need at least 10 GB to safely save a checkpoint.")
                consecutive_regressions += 1
                cycle_num += 1
                continue
        except Exception:
            pass  # Can't check disk — proceed anyway

        try:
            new_model_path, training_loss = _train_model(
                model_path=cfg.model_path,
                training_data_path=str(Path(cfg.data_dir) / f"cycle{cycle_num}" / "training_data.jsonl"),
                lora_targets=weakness_report.get("lora_target_layers", list(range(36))),
                norms_to_repair=weakness_report.get("norms_to_repair", []),
                reference_model=cfg.reference_model,
                output_dir=cfg.output_dir,
                cycle_num=cycle_num,
                quick_test=cfg.quick_test,
            )
        except Exception as e:
            print(f"\n  CRITICAL: Training crashed: {e}")
            print(f"  Skipping this cycle and retrying with current model...")
            import traceback
            traceback.print_exc()
            # Save partial results and continue to next cycle
            cycle_duration = time.time() - cycle_start
            result = CycleResult(
                cycle_num=cycle_num,
                benchmark_scores=benchmark_scores,
                weight_health_summary=health_data.get("summary", {}),
                weaknesses_found=weakness_report.get("total_weaknesses", 0),
                data_collected=training_data.get("total_samples", 0),
                training_loss=0.0,
                improved=False,
                regressions=["TRAINING_CRASH"],
                duration_seconds=cycle_duration,
                model_path=cfg.model_path,
            )
            history.append(result)
            save_cycle_history(history, cfg.output_dir)
            consecutive_regressions += 1
            if consecutive_regressions >= cfg.max_consecutive_regressions:
                print(f"\n  {consecutive_regressions} consecutive failures — STOPPING LOOP")
                break
            cycle_num += 1
            # Free GPU memory before retry
            gc.collect()
            try:
                import torch
                if torch.cuda.is_available():
                    torch.cuda.empty_cache()
            except ImportError:
                pass
            continue

        # ── TRAINING DURATION WATCHDOG ──
        train_elapsed_h = (time.time() - cycle_start) / 3600
        if train_elapsed_h > cfg.max_training_hours:
            print(f"\n  WARNING: Training took {train_elapsed_h:.1f}h (limit: {cfg.max_training_hours}h)")
            print(f"  This may indicate a stuck process or oversized dataset.")

        # ============================================
        # STEP 5: VERIFY
        # ============================================
        print(f"\n{'='*50}")
        print(f"STEP 5/5: VERIFY IMPROVEMENT")
        print(f"{'='*50}")

        try:
            new_scores = _run_benchmark(new_model_path, cfg.benchmark_dir, cycle_num, suffix="_post")
        except Exception as e:
            print(f"\n  CRITICAL: Post-training benchmark crashed: {e}")
            print(f"  Cannot verify improvement — treating as regression for safety.")
            import traceback
            traceback.print_exc()
            new_scores = {}  # Empty = all categories missing = treated as regression

        # Compare scores
        improved, regressions = _compare_scores(benchmark_scores, new_scores, cfg.max_regression_pct)

        if improved:
            print(f"\n  ✓ IMPROVEMENT CONFIRMED — cycle {cycle_num} successful!")
            consecutive_regressions = 0
            cfg.model_path = new_model_path  # Use improved model next cycle
        else:
            print(f"\n  ✗ REGRESSION DETECTED in: {', '.join(regressions)}")
            consecutive_regressions += 1

            if consecutive_regressions >= cfg.max_consecutive_regressions:
                print(f"\n  {consecutive_regressions} consecutive regressions — STOPPING LOOP")
                print(f"  The model may have hit its improvement ceiling.")
                break

            # Still use the new model — small regressions can recover
            cfg.model_path = new_model_path

        # Save cycle results
        cycle_duration = time.time() - cycle_start
        result = CycleResult(
            cycle_num=cycle_num,
            benchmark_scores=new_scores,
            weight_health_summary=health_data.get("summary", {}),
            weaknesses_found=weakness_report.get("total_weaknesses", 0),
            data_collected=training_data.get("total_samples", 0),
            training_loss=training_loss,
            improved=improved,
            regressions=regressions,
            duration_seconds=cycle_duration,
            model_path=new_model_path,
        )
        history.append(result)
        save_cycle_history(history, cfg.output_dir)

        # Print cycle summary
        _print_cycle_summary(result, history)

        # Check if all targets met
        if _all_targets_met(new_scores, cfg.targets):
            print(f"\n{'='*60}")
            print(f"TARGET PERFORMANCE REACHED AFTER {cycle_num} CYCLES!")
            print(f"{'='*60}")
            break

        # ── PERIODIC FULL GC between cycles ──
        # After a full cycle, Python has accumulated lots of garbage
        # (deleted models, tokenizers, datasets). Force collection.
        gc.collect()
        try:
            import torch
            if torch.cuda.is_available():
                torch.cuda.empty_cache()
                # Print GPU memory status for debugging
                allocated = torch.cuda.memory_allocated() / 1e9
                reserved = torch.cuda.memory_reserved() / 1e9
                print(f"\n  GPU memory: {allocated:.1f} GB allocated, {reserved:.1f} GB reserved")
        except ImportError:
            pass

        # ── SHUTDOWN CHECK before starting next cycle ──
        cycles_run_this_session += 1
        if shutdown_requested[0]:
            print(f"\n  Shutdown requested — stopping after cycle {cycle_num}")
            break

        cycle_num += 1

    # Final summary
    _print_final_summary(history)


# ============================================================
# STEP IMPLEMENTATIONS
# ============================================================

def _auto_detect_model() -> str:
    """Find the best starting model."""
    base = Path("td_fuse_outputs/self_improve")

    # Check for highest cycle output
    if base.exists():
        for n in range(50, 0, -1):
            cycle_dir = base / f"improved_cycle{n}"
            if cycle_dir.exists() and list(cycle_dir.glob("*.safetensors")):
                print(f"  Found cycle {n} output at {cycle_dir}")
                return str(cycle_dir)

    # Check for old "improved" folder
    old_improved = base / "improved"
    if old_improved.exists() and list(old_improved.glob("*.safetensors")):
        return str(old_improved)

    # Fall back to reasoning_healed
    healed = Path("td_fuse_outputs/reasoning_healed")
    if healed.exists() and list(healed.glob("*.safetensors")):
        return str(healed)

    raise FileNotFoundError("No model found! Need a starting model checkpoint.")


def _run_benchmark(
    model_path: str,
    benchmark_dir: str,
    cycle_num: int,
    suffix: str = "",
) -> Dict[str, float]:
    """Run td_benchmark.py and return scores per category."""
    from td_benchmark import BenchmarkConfig, run_benchmark

    output_dir = str(Path(benchmark_dir) / f"cycle{cycle_num}{suffix}")
    results_path = str(Path(output_dir) / "results.json")
    Path(output_dir).mkdir(parents=True, exist_ok=True)

    # Validate checkpoint before loading
    mp = Path(model_path)
    if mp.is_dir() and not (mp / "config.json").exists():
        raise FileNotFoundError(f"No config.json in {model_path} — stale or corrupt checkpoint")

    config = BenchmarkConfig(
        model_path=model_path,
        output_dir=output_dir,
        output_path=results_path,
        disable_thinking=True,  # SPEED: Skip thinking for benchmark
    )

    print(f"  Running benchmark on {model_path}...")
    results = run_benchmark(config)

    # Guard: results might be None or not a dict if benchmark crashed internally
    if not isinstance(results, dict):
        print(f"  WARNING: Benchmark returned {type(results)} instead of dict — treating as empty")
        results = {}

    # Extract scores per category (already 0.0-1.0 after our fix)
    scores = {}
    for category, data in results.get("categories", {}).items():
        if isinstance(data, dict):
            scores[category] = data.get("score", 0.0)
        else:
            scores[category] = 0.0

    print(f"  Benchmark scores:")
    for cat, score in sorted(scores.items()):
        print(f"    {cat}: {score:.0%}")

    # Clean up model from GPU memory
    gc.collect()
    try:
        import torch
        if torch.cuda.is_available():
            torch.cuda.empty_cache()
    except ImportError:
        pass

    return scores


def _analyze_weights(
    model_path: str,
    reference_model: str,
    output_dir: str,
) -> Dict:
    """Run weight health analysis from selfimprove.py."""
    try:
        from selfimprove import (
            analyze_model_weights,
            compute_health_scores,
            compute_damage_scores,
            compute_coherence_scores,
            detect_clogged_norms,
        )

        print(f"  Analyzing model weights...")
        model_analysis = analyze_model_weights(model_path, "TD model")
        health_scores = compute_health_scores(model_analysis)
        coherence_scores = compute_coherence_scores(model_analysis)
        clogged_norms = detect_clogged_norms(health_scores)

        # Reference comparison — skip if no reference cached (saves ~5 min)
        # selfimprove.py downloads it during training anyway
        damage_scores = {}
        ref_path = Path(output_dir) / "reference_weights"
        if ref_path.exists() and list(ref_path.glob("*.safetensors")):
            print(f"  Found cached reference at {ref_path}, comparing...")
            ref_analysis = analyze_model_weights(str(ref_path), "Reference")
            damage_scores = compute_damage_scores(model_analysis, ref_analysis)
        else:
            print(f"  No cached reference — skipping damage comparison (saves time)")

        # Build summary
        unhealthy = sum(1 for h in health_scores.values() if h["health_score"] > 0.2)
        summary = {
            "total_tensors": len(health_scores),
            "unhealthy_tensors": unhealthy,
            "clogged_norms": len(clogged_norms),
            "coherence_issues": sum(1 for s in coherence_scores.values() if s > 0.3),
        }

        print(f"  Weight health: {unhealthy}/{len(health_scores)} tensors with issues")
        print(f"  Clogged norms: {len(clogged_norms)}")

        # Clean up
        gc.collect()

        return {
            "health_scores": health_scores,
            "damage_scores": damage_scores,
            "coherence_scores": coherence_scores,
            "clogged_norms": clogged_norms,
            "summary": summary,
        }

    except Exception as e:
        print(f"  Warning: Weight analysis failed: {e}")
        return {"summary": {}, "health_scores": {}, "clogged_norms": []}


def _find_weaknesses(
    benchmark_path: str,
    health_data: Dict,
    targets: Dict,
    budget: int,
    output_dir: str,
    cycle_num: int,
) -> Dict:
    """Run td_weakness.py analysis."""
    from td_weakness import run_weakness_analysis

    output_path = str(Path(output_dir) / f"cycle{cycle_num}" / "weakness_report.json")
    Path(output_path).parent.mkdir(parents=True, exist_ok=True)

    report = run_weakness_analysis(
        benchmark_path=benchmark_path,
        health_scores=health_data.get("health_scores") or {},
        damage_scores=health_data.get("damage_scores") or {},
        coherence_scores=health_data.get("coherence_scores") or {},
        clogged_norms=health_data.get("clogged_norms") or [],
        targets=targets,
        total_budget=budget,
        output_path=output_path,
    )

    return {
        "search_plan": report.search_plan,
        "data_allocation": report.data_allocation,
        "lora_target_layers": report.lora_target_layers,
        "norms_to_repair": report.norms_to_repair,
        "total_weaknesses": len(report.answer_weaknesses) + len(report.weight_weaknesses),
    }


def _find_training_data(
    weakness_report: Dict,
    output_dir: str,
    hf_token: str,
    cycle_num: int,
) -> Dict:
    """Run td_data_finder.py to search for and download data."""
    from td_data_finder import find_training_data

    cycle_data_dir = str(Path(output_dir) / f"cycle{cycle_num}")

    result = find_training_data(
        search_plan=weakness_report["search_plan"],
        data_allocation=weakness_report["data_allocation"],
        output_dir=cycle_data_dir,
        hf_token=hf_token if hf_token else None,
    )

    return {
        "total_samples": result.total_samples,
        "categories_covered": result.categories_covered,
        "categories_missing": result.categories_missing,
    }


def _train_model(
    model_path: str,
    training_data_path: str,
    lora_targets: List[int],
    norms_to_repair: List[str],
    reference_model: str,
    output_dir: str,
    cycle_num: int,
    quick_test: bool = False,
) -> Tuple[str, float]:
    """
    Run training using selfimprove.py's training pipeline.

    Returns: (new_model_path, final_training_loss)
    """
    from selfimprove import SelfImproveConfig, run_cycle

    cfg = SelfImproveConfig()
    # CRITICAL: Set model_path BEFORE run_cycle so auto_detect_model_path
    # doesn't pick a different checkpoint
    cfg.model_path = model_path
    cfg.reference_model = reference_model
    cfg.output_dir = output_dir

    if quick_test:
        # ── COMPRESSED TEST MODE ──
        # Exercises every code path but finishes in ~30-45 min
        print("  *** QUICK TEST: Compressed training settings ***")
        cfg.mega_mode = True  # Still mega mode so norm unclogging runs
        cfg.top_n_layers = 8  # Only 8 layers (not 36) — 4x faster LoRA setup
        cfg.min_improvement_score = 0.0
        cfg.lora_r = 32  # Smaller LoRA — 4x less trainable params
        cfg.lora_alpha = 64
        cfg.train_batch = 4
        cfg.train_grad_accum = 2  # Effective batch 8 (not 16)
        cfg.num_train_samples = 500  # 16x fewer samples
        cfg.train_epochs = 1  # 1 epoch (not 3) — 3x faster
        cfg.learning_rate = 1e-5
        cfg.norm_learning_rate = 1e-3
        cfg.save_steps = 9999  # No mid-training saves — saves ~2 min
    else:
        # ── FULL MEGA MODE ──
        cfg.mega_mode = True
        cfg.top_n_layers = 36
        cfg.min_improvement_score = 0.0
        cfg.lora_r = 128
        cfg.lora_alpha = 256
        cfg.train_batch = 4
        cfg.train_grad_accum = 4
        cfg.num_train_samples = 8000
        cfg.learning_rate = 1e-5
        cfg.norm_learning_rate = 1e-3
        cfg.save_steps = 200

    # Pass targeted training data from td_data_finder
    # selfimprove.py will use this INSTEAD of its hardcoded datasets
    training_data_file = Path(training_data_path)
    if training_data_file.exists() and training_data_file.stat().st_size > 0:
        cfg.custom_data_path = str(training_data_file)
        print(f"  Using targeted training data: {training_data_file}")
    else:
        print(f"  No targeted data found — selfimprove.py will use generic datasets")

    # Figure out what cycle_num selfimprove.py will ACTUALLY use
    # so we know where the output will be saved
    import re as _re
    si_match = _re.search(r'improved_cycle(\d+)', model_path)
    if si_match:
        si_cycle_num = int(si_match.group(1)) + 1
    elif "improve" in model_path:
        si_cycle_num = 2
    else:
        si_cycle_num = 1

    print(f"  Starting training (loop cycle {cycle_num}, selfimprove cycle {si_cycle_num})...")
    print(f"  Model: {model_path}")
    print(f"  Output will be: {output_dir}/improved_cycle{si_cycle_num}")
    run_cycle(cfg)

    # Find the output model — use selfimprove's cycle numbering
    new_model_path = str(Path(output_dir) / f"improved_cycle{si_cycle_num}")
    if not Path(new_model_path).exists():
        # Fallback: try loop's cycle number
        alt_path = str(Path(output_dir) / f"improved_cycle{cycle_num}")
        if Path(alt_path).exists():
            new_model_path = alt_path
        else:
            # Last resort: scan for any improved_cycle* dir
            base = Path(output_dir)
            found = False
            for n in range(50, 0, -1):
                check = base / f"improved_cycle{n}"
                if check.exists() and list(check.glob("*.safetensors")):
                    new_model_path = str(check)
                    print(f"  Output path mismatch — found model at {check}")
                    found = True
                    break
            if not found:
                old_improved = base / "improved"
                if old_improved.exists():
                    new_model_path = str(old_improved)
                else:
                    # Training produced nothing — return original model
                    print(f"  ERROR: No output model found! Training may have failed silently.")
                    print(f"  Returning original model: {model_path}")
                    return model_path, 0.0

    # Validate the output model has actual weights
    out_p = Path(new_model_path)
    if out_p.is_dir():
        st_files = list(out_p.glob("*.safetensors"))
        if not st_files:
            print(f"  ERROR: Output model has no safetensors files! Returning original.")
            return model_path, 0.0
        total_sz = sum(f.stat().st_size for f in st_files) / 1e9
        if total_sz < 1.0:
            print(f"  ERROR: Output model only {total_sz:.2f} GB — likely corrupt. Returning original.")
            return model_path, 0.0

    # Try to get final loss from training log
    training_loss = 0.0
    log_file = Path(new_model_path) / "training_log.json"
    if log_file.exists():
        try:
            with open(log_file) as f:
                log = json.load(f)
            training_loss = log.get("final_loss", 0.0)
        except Exception:
            pass

    return new_model_path, training_loss


def _compare_scores(
    old_scores: Dict[str, float],
    new_scores: Dict[str, float],
    max_regression: float,
) -> Tuple[bool, List[str]]:
    """
    Compare benchmark scores before and after training.
    Returns: (overall_improved, list_of_regressed_categories)
    """
    regressions = []
    improvements = 0

    for category in old_scores:
        if category not in new_scores:
            print(f"    ? {category}: MISSING from new scores — treating as regression!")
            regressions.append(category)
            continue

        old = old_scores[category]
        new = new_scores[category]
        diff = new - old

        if diff > 0.01:  # Meaningful improvement (>1%)
            improvements += 1
            print(f"    ↑ {category}: {old:.1%} → {new:.1%} (+{diff:.1%})")
        elif diff < -max_regression:  # Regression beyond threshold
            regressions.append(category)
            print(f"    ↓ {category}: {old:.1%} → {new:.1%} ({diff:.1%}) REGRESSION!")
        else:
            print(f"    = {category}: {old:.1%} → {new:.1%} (stable)")

    # Improved if: at least one category improved AND improvements outnumber regressions.
    # Old logic (improvements > 0 AND regressions == 0) was too strict —
    # with 6 categories, some random fluctuation is expected. If math goes
    # up 5% but creativity dips 2%, that's still a net win.
    overall_improved = improvements > len(regressions)
    return overall_improved, regressions


def _all_targets_met(scores: Dict[str, float], targets: Dict[str, float]) -> bool:
    """Check if all scores meet or exceed targets."""
    for category, target in targets.items():
        if category not in scores:
            return False
        if scores[category] < target:
            return False
    return True


# ============================================================
# SUMMARIES
# ============================================================

def _print_cycle_summary(result: CycleResult, history: List[CycleResult]):
    """Print a summary of the completed cycle."""
    try:
        hours = result.duration_seconds / 3600
        print(f"\n{'='*50}")
        print(f"CYCLE {result.cycle_num} SUMMARY")
        print(f"{'='*50}")
        print(f"  Duration: {hours:.1f} hours")
        print(f"  Weaknesses found: {result.weaknesses_found}")
        print(f"  Data collected: {result.data_collected}")
        print(f"  Training loss: {result.training_loss:.4f}")
        print(f"  Improved: {'Yes' if result.improved else 'No'}")
        if result.regressions:
            print(f"  Regressions: {', '.join(result.regressions)}")
        print(f"  Model saved: {result.model_path}")

        # Score history
        if len(history) > 1:
            print(f"\n  Score progression:")
            for h in history[-5:]:  # Last 5 cycles
                scores = h.benchmark_scores or {}
                avg = sum(scores.values()) / max(len(scores), 1) if scores else 0.0
                print(f"    Cycle {h.cycle_num}: avg {avg:.1%}")
    except Exception as e:
        print(f"  (Could not print cycle summary: {e})")


def _print_final_summary(history: List[CycleResult]):
    """Print final summary after loop ends."""
    if not history:
        return

    try:
        total_hours = sum(r.duration_seconds for r in history) / 3600
        improved_cycles = sum(1 for r in history if r.improved)

        print(f"\n{'#' * 60}")
        print(f"# LOOP COMPLETE")
        print(f"{'#' * 60}")
        print(f"  Total cycles: {len(history)}")
        print(f"  Improved cycles: {improved_cycles}/{len(history)}")
        print(f"  Total time: {total_hours:.1f} hours")

        if len(history) >= 2:
            first = history[0]
            last = history[-1]
            first_scores = first.benchmark_scores or {}
            last_scores = last.benchmark_scores or {}
            if first_scores and last_scores:
                print(f"\n  Score changes (first → last):")
                for cat in sorted(first_scores.keys()):
                    if cat in last_scores:
                        old = first_scores[cat]
                        new = last_scores[cat]
                        diff = new - old
                        arrow = "↑" if diff > 0 else "↓" if diff < 0 else "="
                        print(f"    {arrow} {cat}: {old:.1%} → {new:.1%} ({diff:+.1%})")

        print(f"\n  Final model: {history[-1].model_path}")
    except Exception as e:
        print(f"  (Could not print final summary: {e})")


# ============================================================
# CLI
# ============================================================

def main():
    parser = argparse.ArgumentParser(description="TD Self-Improvement Loop")
    parser.add_argument("--model", default="", help="Starting model path")
    parser.add_argument("--max-cycles", type=int, default=0, help="Max cycles (0=infinite)")
    parser.add_argument("--budget", type=int, default=8000, help="Training samples per cycle")
    parser.add_argument("--dry-run", action="store_true", help="Test + diagnose only")
    parser.add_argument("--quick-test", action="store_true",
                        help="Compressed test run (~30-45 min): 500 samples, 1 epoch, LoRA r=32, 8 layers. "
                             "Tests every code path but finishes fast.")
    parser.add_argument("--hf-token", default="", help="HuggingFace token")
    args = parser.parse_args()

    cfg = LoopConfig(
        model_path=args.model,
        max_cycles=args.max_cycles,
        training_budget=args.budget,
        dry_run=args.dry_run,
        hf_token=args.hf_token,
    )

    # Quick test mode: override to compressed settings
    if args.quick_test:
        cfg.quick_test = True
        cfg.max_cycles = 1
        cfg.training_budget = 500
        print("\n  *** QUICK TEST MODE ***")
        print("  500 samples, 1 epoch, LoRA r=32, 8 layers")
        print("  Expected time: ~30-45 minutes")
        print("  Purpose: verify every code path works before real run\n")

    run_loop(cfg)


if __name__ == "__main__":
    main()