"""
TD Weakness Finder v1 — Diagnose What's Wrong

Takes BOTH benchmark results (what the model gets wrong) AND weight health
(what's broken inside) and outputs a structured weakness report.

This is Step 2 of the self-improvement loop.

Input:
  - Benchmark JSON (from td_benchmark.py)
  - Weight health data (from selfimprove.py health/damage/coherence functions)

Output:
  - Structured weakness report with:
    - Answer weaknesses (categories that scored low)
    - Weight weaknesses (layers that are unhealthy)
    - Search keywords for finding datasets
    - Suggested actions (norm repair, LoRA targets, etc.)
    - Data allocation (how many samples per weakness)
"""

import json
import math
from pathlib import Path
from typing import Dict, List, Optional, Tuple
from dataclasses import dataclass, field


# ============================================================
# DATA STRUCTURES
# ============================================================

@dataclass
class AnswerWeakness:
    """A weakness found from benchmark testing."""
    category: str
    score: float
    target: float
    gap: float  # target - score
    failed_questions: List[str]
    search_keywords: List[str]
    priority: float  # higher = more urgent (based on gap size)


@dataclass
class WeightWeakness:
    """A weakness found from weight health analysis."""
    description: str
    layers: List[int]
    layer_names: List[str]
    severity: str  # "critical", "high", "medium", "low"
    action: str  # what to do about it
    metrics: Dict  # the actual numbers


@dataclass
class WeaknessReport:
    """Complete weakness report combining both sources."""
    answer_weaknesses: List[AnswerWeakness]
    weight_weaknesses: List[WeightWeakness]
    lora_target_layers: List[int]  # which layers to target with LoRA
    norms_to_repair: List[str]  # which norms need surgical repair
    data_allocation: Dict[str, int]  # category -> num_samples
    total_samples_needed: int
    search_plan: List[Dict]  # ordered list of searches to do


# ============================================================
# CATEGORY SEARCH KEYWORDS
# ============================================================

# For each benchmark category, what to search for when it scores low
CATEGORY_SEARCH_MAP = {
    "math": {
        "keywords": [
            "math word problems dataset",
            "grade school math GSM8K",
            "arithmetic reasoning dataset",
            "algebra training data",
            "competition math problems",
            "math chain of thought",
        ],
        "hf_tags": ["math", "arithmetic", "algebra", "gsm8k"],
        "description": "Mathematical reasoning and problem solving",
    },
    "code": {
        "keywords": [
            "code generation dataset",
            "python programming problems",
            "code debugging pairs",
            "competitive programming solutions",
            "code review dataset",
            "function implementation tasks",
        ],
        "hf_tags": ["code", "python", "programming", "code-generation"],
        "description": "Code writing, debugging, and understanding",
    },
    "reasoning": {
        "keywords": [
            "chain of thought reasoning dataset",
            "multi-step logic problems",
            "logical reasoning training",
            "common sense reasoning",
            "deductive reasoning dataset",
            "ARC challenge training",
        ],
        "hf_tags": ["reasoning", "logic", "commonsense", "chain-of-thought"],
        "description": "Logic, multi-step deduction, common sense",
    },
    "creativity": {
        "keywords": [
            "creative writing dataset",
            "story generation training",
            "poetry writing pairs",
            "text rewriting dataset",
            "paraphrase generation",
            "creative instruction following",
        ],
        "hf_tags": ["creative-writing", "story", "paraphrase", "text-generation"],
        "description": "Writing, storytelling, creative text generation",
    },
    "knowledge": {
        "keywords": [
            "MMLU training data",
            "science QA dataset",
            "trivia question answer pairs",
            "general knowledge dataset",
            "encyclopedia QA training",
            "fact verification dataset",
        ],
        "hf_tags": ["qa", "knowledge", "science", "trivia", "mmlu"],
        "description": "Science, history, geography, general facts",
    },
    "instruction_following": {
        "keywords": [
            "instruction following dataset",
            "format constraint training",
            "structured output training",
            "instruction tuning data",
            "precise instruction pairs",
            "constraint satisfaction dataset",
        ],
        "hf_tags": ["instruction-following", "instruction-tuning", "chat"],
        "description": "Following instructions precisely and exactly",
    },
}


# ============================================================
# SCORE TARGETS (what "good" looks like for each category)
# ============================================================

DEFAULT_TARGETS = {
    "math": 0.80,
    "code": 0.70,
    "reasoning": 0.75,
    "creativity": 0.65,
    "knowledge": 0.80,
    "instruction_following": 0.85,
}


# ============================================================
# ANSWER WEAKNESS DETECTION
# ============================================================

def find_answer_weaknesses(
    benchmark_path: str,
    targets: Optional[Dict[str, float]] = None,
) -> List[AnswerWeakness]:
    """
    Read benchmark results and find which categories are weak.

    Args:
        benchmark_path: Path to benchmark JSON output
        targets: Score targets per category (default: DEFAULT_TARGETS)

    Returns:
        List of AnswerWeakness sorted by priority (worst first)
    """
    targets = targets or DEFAULT_TARGETS

    # Guard: benchmark file might be missing, empty, or corrupt
    try:
        with open(benchmark_path) as f:
            results = json.load(f)
    except FileNotFoundError:
        print(f"  WARNING: Benchmark file not found: {benchmark_path}")
        return []
    except json.JSONDecodeError as e:
        print(f"  WARNING: Benchmark file corrupt: {e}")
        return []
    except Exception as e:
        print(f"  WARNING: Could not read benchmark file: {e}")
        return []

    if not isinstance(results, dict) or "categories" not in results:
        print(f"  WARNING: Benchmark JSON missing 'categories' key — invalid format")
        return []

    weaknesses = []

    for category, data in results.get("categories", {}).items():
        score = data.get("score", 0.0)
        target = targets.get(category, 0.75)
        gap = target - score

        # Only flag if below target
        if gap <= 0:
            continue

        # BUG FIX: benchmark JSON uses "failed_questions" not "questions"
        # Each failed question has "question", "response", "expected", "type"
        failed = []
        for q in data.get("failed_questions", []):
            failed.append(q.get("question", "unknown")[:80])

        # Look up search keywords for this category
        search_info = CATEGORY_SEARCH_MAP.get(category, {})
        keywords = search_info.get("keywords", [f"{category} training dataset"])

        # Priority is based on how far below target we are
        # Bigger gap = higher priority = more training data
        priority = gap / max(target, 0.01)

        weaknesses.append(AnswerWeakness(
            category=category,
            score=score,
            target=target,
            gap=gap,
            failed_questions=failed,
            search_keywords=keywords,
            priority=priority,
        ))

    # Sort by priority (worst gaps first)
    weaknesses.sort(key=lambda w: w.priority, reverse=True)
    return weaknesses


# ============================================================
# WEIGHT WEAKNESS DETECTION
# ============================================================

def find_weight_weaknesses(
    health_scores: Dict[str, Dict],
    damage_scores: Optional[Dict[str, float]] = None,
    coherence_scores: Optional[Dict[int, float]] = None,
    clogged_norms: Optional[List[str]] = None,
) -> List[WeightWeakness]:
    """
    Analyze weight health data and find structural problems.

    Args:
        health_scores: From selfimprove.compute_health_scores()
        damage_scores: From selfimprove.compute_damage_scores()
        coherence_scores: From selfimprove.compute_coherence_scores()
        clogged_norms: From selfimprove.detect_clogged_norms()

    Returns:
        List of WeightWeakness sorted by severity
    """
    weaknesses = []

    # Guard: health_scores might be empty or have unexpected structure
    if not health_scores:
        print("  WARNING: No health_scores provided — skipping weight weakness detection")
        return []

    # --- 1. Clogged norms (highest priority) ---
    if clogged_norms:
        # Group by layer number
        norm_layers = _extract_layer_numbers(clogged_norms)

        # Find the kurtosis values (with guards for missing keys)
        kurtosis_vals = {}
        for name in clogged_norms:
            if name in health_scores:
                stats = health_scores[name].get("stats", {})
                if "kurtosis" in stats:
                    kurtosis_vals[name] = abs(stats["kurtosis"])

        avg_kurt = sum(kurtosis_vals.values()) / max(len(kurtosis_vals), 1)
        max_kurt = max(kurtosis_vals.values()) if kurtosis_vals else 0

        severity = "critical" if max_kurt > 100 else "high" if max_kurt > 50 else "medium"

        weaknesses.append(WeightWeakness(
            description=f"{len(clogged_norms)} clogged norms (avg kurtosis {avg_kurt:.0f}, max {max_kurt:.0f})",
            layers=sorted(set(norm_layers)),
            layer_names=clogged_norms,
            severity=severity,
            action="surgical_norm_repair",
            metrics={"avg_kurtosis": avg_kurt, "max_kurtosis": max_kurt, "count": len(clogged_norms)},
        ))

    # --- 2. High condition number layers (numerically unstable) ---
    high_condition = []
    for name, h in health_scores.items():
        stats = h.get("stats", {}) if isinstance(h, dict) else {}
        sv_ratio = stats.get("sv_ratio", 0)
        if sv_ratio > 10000:
            high_condition.append((name, sv_ratio))

    if high_condition:
        layers = _extract_layer_numbers([n for n, _ in high_condition])
        max_cond = max(v for _, v in high_condition)

        weaknesses.append(WeightWeakness(
            description=f"{len(high_condition)} layers with extreme condition numbers (max {max_cond:.0f})",
            layers=sorted(set(layers)),
            layer_names=[n for n, _ in high_condition],
            severity="critical" if max_cond > 50000 else "high",
            action="lora_target_high_priority",
            metrics={"max_condition": max_cond, "count": len(high_condition)},
        ))

    # --- 3. Low rank utilization (layers not using their capacity) ---
    low_rank = []
    for name, h in health_scores.items():
        stats = h.get("stats", {}) if isinstance(h, dict) else {}
        rank_util = stats.get("sv_rank_utilization", -1)
        if 0 < rank_util < 0.1:
            low_rank.append((name, rank_util))

    if low_rank:
        layers = _extract_layer_numbers([n for n, _ in low_rank])
        avg_util = sum(v for _, v in low_rank) / len(low_rank)

        weaknesses.append(WeightWeakness(
            description=f"{len(low_rank)} layers with very low rank utilization (avg {avg_util:.1%})",
            layers=sorted(set(layers)),
            layer_names=[n for n, _ in low_rank],
            severity="high",
            action="lora_target_complex_tasks",
            metrics={"avg_rank_util": avg_util, "count": len(low_rank)},
        ))

    # --- 4. Dead neurons ---
    dead_neuron_layers = []
    for name, h in health_scores.items():
        stats = h.get("stats", {}) if isinstance(h, dict) else {}
        dead = stats.get("near_zero_pct", 0)
        if dead > 0.3:
            dead_neuron_layers.append((name, dead))

    if dead_neuron_layers:
        layers = _extract_layer_numbers([n for n, _ in dead_neuron_layers])
        avg_dead = sum(v for _, v in dead_neuron_layers) / len(dead_neuron_layers)

        weaknesses.append(WeightWeakness(
            description=f"{len(dead_neuron_layers)} layers with >30% dead neurons (avg {avg_dead:.1%})",
            layers=sorted(set(layers)),
            layer_names=[n for n, _ in dead_neuron_layers],
            severity="high",
            action="lora_target_activation",
            metrics={"avg_dead_pct": avg_dead, "count": len(dead_neuron_layers)},
        ))

    # --- 5. Coherence issues (layer-to-layer jumps) ---
    if coherence_scores:
        bad_coherence = [(l, s) for l, s in coherence_scores.items() if s > 0.3]
        if bad_coherence:
            weaknesses.append(WeightWeakness(
                description=f"{len(bad_coherence)} layers with coherence issues",
                layers=sorted([l for l, _ in bad_coherence]),
                layer_names=[],
                severity="medium",
                action="lora_target_smooth_transition",
                metrics={"worst_score": max(s for _, s in bad_coherence)},
            ))

    # --- 6. Damage vs reference (merge artifacts) ---
    if damage_scores:
        high_damage = [(n, s) for n, s in damage_scores.items() if s > 0.5]
        if high_damage:
            layers = _extract_layer_numbers([n for n, _ in high_damage])
            avg_damage = sum(s for _, s in high_damage) / len(high_damage)

            weaknesses.append(WeightWeakness(
                description=f"{len(high_damage)} layers with high merge damage (avg {avg_damage:.2f})",
                layers=sorted(set(layers)),
                layer_names=[n for n, _ in high_damage],
                severity="high" if avg_damage > 0.8 else "medium",
                action="lora_target_repair",
                metrics={"avg_damage": avg_damage, "count": len(high_damage)},
            ))

    # Sort by severity
    severity_order = {"critical": 0, "high": 1, "medium": 2, "low": 3}
    weaknesses.sort(key=lambda w: severity_order.get(w.severity, 99))

    return weaknesses


def _extract_layer_numbers(names: List[str]) -> List[int]:
    """Pull layer numbers from parameter names like 'model.layers.5.self_attn.q_proj'."""
    layers = []
    for name in names:
        parts = name.split(".")
        for i, part in enumerate(parts):
            if part == "layers" and i + 1 < len(parts) and parts[i + 1].isdigit():
                layers.append(int(parts[i + 1]))
                break
    return layers


# ============================================================
# COMBINE INTO REPORT
# ============================================================

def build_weakness_report(
    answer_weaknesses: List[AnswerWeakness],
    weight_weaknesses: List[WeightWeakness],
    total_budget: int = 8000,
) -> WeaknessReport:
    """
    Combine answer + weight weaknesses into a single action plan.

    Args:
        answer_weaknesses: From find_answer_weaknesses()
        weight_weaknesses: From find_weight_weaknesses()
        total_budget: Total training samples to allocate

    Returns:
        WeaknessReport with everything the training loop needs
    """
    # --- Figure out LoRA target layers ---
    lora_layers = set()
    for ww in weight_weaknesses:
        if "lora" in ww.action:
            lora_layers.update(ww.layers)

    # If no weight issues pointed to specific layers, target all layers
    if not lora_layers:
        lora_layers = set(range(36))  # Qwen3-VL has 36 layers

    # --- Figure out norms to repair ---
    norms_to_repair = []
    for ww in weight_weaknesses:
        if ww.action == "surgical_norm_repair":
            norms_to_repair.extend(ww.layer_names)

    # --- Allocate training data ---
    data_allocation = _allocate_data(answer_weaknesses, total_budget)

    # --- Build search plan ---
    search_plan = _build_search_plan(answer_weaknesses)

    return WeaknessReport(
        answer_weaknesses=answer_weaknesses,
        weight_weaknesses=weight_weaknesses,
        lora_target_layers=sorted(lora_layers),
        norms_to_repair=norms_to_repair,
        data_allocation=data_allocation,
        total_samples_needed=sum(data_allocation.values()),
        search_plan=search_plan,
    )


def _allocate_data(
    weaknesses: List[AnswerWeakness],
    total_budget: int,
) -> Dict[str, int]:
    """
    Allocate training samples proportional to weakness severity.
    Worse categories get more data.

    IMPORTANT: Also allocates a FLOOR (10% of budget per category) to ALL
    categories, even those above target. This prevents catastrophic forgetting
    where training on weak categories makes strong categories worse.

    Example:
      math: 42% (gap 38%) → gets ~2x more data (from weakness pool)
      code: 65% (gap 5%)  → gets less data (from weakness pool)
      knowledge: 85% (above target) → still gets floor allocation
    """
    ALL_CATEGORIES = ["math", "code", "reasoning", "creativity",
                      "knowledge", "instruction_following"]

    # FLOOR: 20% of budget split equally across ALL categories
    # This prevents catastrophic forgetting on categories above target
    # Increased from 10% to 20% — model at 95%+ needs more retention data
    floor_pct = 0.20
    floor_total = int(total_budget * floor_pct)
    floor_per_cat = max(floor_total // len(ALL_CATEGORIES), 100)

    # Start with floor for every category
    allocation = {cat: floor_per_cat for cat in ALL_CATEGORIES}

    if not weaknesses:
        return allocation

    # Remaining budget goes to weak categories proportionally
    weakness_budget = total_budget - (floor_per_cat * len(ALL_CATEGORIES))
    if weakness_budget <= 0:
        return allocation

    # Weight by gap size (bigger gap = more data)
    total_gap = sum(w.gap for w in weaknesses)
    if total_gap <= 0:
        per_cat = weakness_budget // len(weaknesses)
        for w in weaknesses:
            allocation[w.category] = allocation.get(w.category, 0) + per_cat
        return allocation

    remaining = weakness_budget

    for i, w in enumerate(weaknesses):
        if i == len(weaknesses) - 1:
            # Last category gets whatever is left — but never negative
            allocation[w.category] = allocation.get(w.category, 0) + max(remaining, 0)
        else:
            share = int(weakness_budget * (w.gap / total_gap))
            share = min(max(share, 500), remaining)  # Don't exceed remaining budget
            allocation[w.category] = allocation.get(w.category, 0) + share
            remaining -= share

    # Make sure we don't exceed budget
    total_alloc = sum(allocation.values())
    if total_alloc > total_budget:
        scale = total_budget / total_alloc
        allocation = {k: max(int(v * scale), 50) for k, v in allocation.items()}

    return allocation


def _build_search_plan(weaknesses: List[AnswerWeakness]) -> List[Dict]:
    """
    Build an ordered list of dataset searches to perform.
    Each entry has: category, keywords, hf_tags, priority, num_samples_needed.

    Includes ALL categories (even above target) because the data allocation
    now has a floor for every category to prevent catastrophic forgetting.
    Weak categories come first (higher priority).
    """
    ALL_CATEGORIES = ["math", "code", "reasoning", "creativity",
                      "knowledge", "instruction_following"]

    plan = []
    weak_cats = set()

    # First: weak categories (high priority, targeted search)
    for w in weaknesses:
        search_info = CATEGORY_SEARCH_MAP.get(w.category, {})
        plan.append({
            "category": w.category,
            "keywords": w.search_keywords,
            "hf_tags": search_info.get("hf_tags", []),
            "priority": w.priority,
            "score": w.score,
            "target": w.target,
            "gap": w.gap,
        })
        weak_cats.add(w.category)

    # Then: strong categories (low priority, just use fallback datasets)
    for cat in ALL_CATEGORIES:
        if cat not in weak_cats:
            search_info = CATEGORY_SEARCH_MAP.get(cat, {})
            plan.append({
                "category": cat,
                "keywords": search_info.get("keywords", [f"{cat} training dataset"]),
                "hf_tags": search_info.get("hf_tags", []),
                "priority": 0.1,  # Low priority — just floor allocation
                "score": 1.0,  # Above target
                "target": DEFAULT_TARGETS.get(cat, 0.75),
                "gap": 0.0,
            })

    return plan


# ============================================================
# PRETTY PRINT
# ============================================================

def print_report(report: WeaknessReport):
    """Print a human-readable weakness report."""
    print("\n" + "=" * 60)
    print("TD WEAKNESS REPORT")
    print("=" * 60)

    # Answer weaknesses
    if report.answer_weaknesses:
        print("\n  ANSWER WEAKNESSES (what it gets wrong):")
        print("  " + "-" * 50)
        for w in report.answer_weaknesses:
            status = "CRITICAL" if w.gap > 0.3 else "HIGH" if w.gap > 0.15 else "MEDIUM"
            print(f"    [{status}] {w.category}: {w.score:.0%} (target: {w.target:.0%}, gap: {w.gap:.0%})")
            print(f"      Failed: {len(w.failed_questions)} questions")
            print(f"      Search: {', '.join(w.search_keywords[:3])}")
    else:
        print("\n  ANSWER WEAKNESSES: None! All categories at or above target.")

    # Weight weaknesses
    if report.weight_weaknesses:
        print(f"\n  WEIGHT WEAKNESSES (what's broken inside):")
        print("  " + "-" * 50)
        for w in report.weight_weaknesses:
            print(f"    [{w.severity.upper()}] {w.description}")
            if w.layers:
                print(f"      Layers: {w.layers}")
            print(f"      Action: {w.action}")
    else:
        print("\n  WEIGHT WEAKNESSES: None! All layers healthy.")

    # Data allocation
    if report.data_allocation:
        print(f"\n  DATA ALLOCATION ({report.total_samples_needed} total samples):")
        print("  " + "-" * 50)
        for cat, num in sorted(report.data_allocation.items(), key=lambda x: -x[1]):
            pct = num / max(report.total_samples_needed, 1)
            bar = "█" * int(pct * 30)
            print(f"    {cat:25s} {num:5d} samples ({pct:5.1%}) {bar}")

    # LoRA targets
    if report.lora_target_layers:
        print(f"\n  LORA TARGETS: {len(report.lora_target_layers)} layers")
        print(f"    Layers: {report.lora_target_layers}")

    # Norm repair
    if report.norms_to_repair:
        print(f"\n  NORMS TO REPAIR: {len(report.norms_to_repair)} parameters")
        for n in report.norms_to_repair[:10]:  # Show first 10
            print(f"    - {n}")
        if len(report.norms_to_repair) > 10:
            print(f"    ... and {len(report.norms_to_repair) - 10} more")

    # Search plan
    if report.search_plan:
        print(f"\n  SEARCH PLAN ({len(report.search_plan)} categories to search):")
        print("  " + "-" * 50)
        for i, s in enumerate(report.search_plan, 1):
            print(f"    {i}. {s['category']} (score: {s['score']:.0%}, need: {s['gap']:.0%} improvement)")
            print(f"       Keywords: {', '.join(s['keywords'][:2])}")

    print("\n" + "=" * 60)


# ============================================================
# SAVE / LOAD
# ============================================================

def save_report(report: WeaknessReport, path: str):
    """Save weakness report to JSON."""
    data = {
        "answer_weaknesses": [
            {
                "category": w.category,
                "score": w.score,
                "target": w.target,
                "gap": w.gap,
                "failed_questions": w.failed_questions,
                "search_keywords": w.search_keywords,
                "priority": w.priority,
            }
            for w in report.answer_weaknesses
        ],
        "weight_weaknesses": [
            {
                "description": w.description,
                "layers": w.layers,
                "layer_names": w.layer_names,
                "severity": w.severity,
                "action": w.action,
                "metrics": w.metrics,
            }
            for w in report.weight_weaknesses
        ],
        "lora_target_layers": report.lora_target_layers,
        "norms_to_repair": report.norms_to_repair,
        "data_allocation": report.data_allocation,
        "total_samples_needed": report.total_samples_needed,
        "search_plan": report.search_plan,
    }

    Path(path).parent.mkdir(parents=True, exist_ok=True)
    with open(path, "w") as f:
        json.dump(data, f, indent=2)
    print(f"\nWeakness report saved to {path}")


def load_report(path: str) -> WeaknessReport:
    """Load weakness report from JSON."""
    with open(path) as f:
        data = json.load(f)

    answer_weaknesses = [
        AnswerWeakness(**aw) for aw in data["answer_weaknesses"]
    ]
    weight_weaknesses = [
        WeightWeakness(**ww) for ww in data["weight_weaknesses"]
    ]

    return WeaknessReport(
        answer_weaknesses=answer_weaknesses,
        weight_weaknesses=weight_weaknesses,
        lora_target_layers=data["lora_target_layers"],
        norms_to_repair=data["norms_to_repair"],
        data_allocation=data["data_allocation"],
        total_samples_needed=data["total_samples_needed"],
        search_plan=data["search_plan"],
    )


# ============================================================
# MAIN — run from command line
# ============================================================

def run_weakness_analysis(
    benchmark_path: str,
    health_scores: Optional[Dict] = None,
    damage_scores: Optional[Dict] = None,
    coherence_scores: Optional[Dict] = None,
    clogged_norms: Optional[List] = None,
    targets: Optional[Dict[str, float]] = None,
    total_budget: int = 8000,
    output_path: Optional[str] = None,
) -> WeaknessReport:
    """
    Run the full weakness analysis pipeline.

    This is the main entry point that td_loop.py calls.

    Args:
        benchmark_path: Path to benchmark results JSON
        health_scores: From selfimprove.compute_health_scores()
        damage_scores: From selfimprove.compute_damage_scores()
        coherence_scores: From selfimprove.compute_coherence_scores()
        clogged_norms: From selfimprove.detect_clogged_norms()
        targets: Score targets per category
        total_budget: Total training samples to allocate
        output_path: Where to save the report JSON

    Returns:
        Complete WeaknessReport
    """
    print("\n" + "=" * 60)
    print("TD WEAKNESS FINDER v1")
    print("=" * 60)

    # Step 1: Find answer weaknesses
    print("\nAnalyzing benchmark results...")
    answer_weaknesses = find_answer_weaknesses(benchmark_path, targets)
    print(f"  Found {len(answer_weaknesses)} answer weaknesses")

    # Step 2: Find weight weaknesses
    print("\nAnalyzing weight health...")
    weight_weaknesses = []
    if health_scores:
        weight_weaknesses = find_weight_weaknesses(
            health_scores=health_scores,
            damage_scores=damage_scores,
            coherence_scores=coherence_scores,
            clogged_norms=clogged_norms,
        )
        print(f"  Found {len(weight_weaknesses)} weight weaknesses")
    else:
        print("  No weight health data provided — skipping weight analysis")

    # Step 3: Build combined report
    print("\nBuilding combined report...")
    report = build_weakness_report(answer_weaknesses, weight_weaknesses, total_budget)

    # Step 4: Print it
    print_report(report)

    # Step 5: Save if requested
    if output_path:
        save_report(report, output_path)

    return report


if __name__ == "__main__":
    import sys

    if len(sys.argv) < 2:
        print("Usage: python td_weakness.py <benchmark_results.json> [output_path.json]")
        print("\nThis reads benchmark results and outputs a weakness report.")
        print("Weight health data is passed programmatically from td_loop.py.")
        sys.exit(1)

    benchmark_path = sys.argv[1]
    output_path = sys.argv[2] if len(sys.argv) > 2 else None

    # When run standalone, only uses benchmark data (no weight health)
    report = run_weakness_analysis(
        benchmark_path=benchmark_path,
        output_path=output_path,
    )