#!/usr/bin/env python3
"""
chentsov_fusion.py — Parallel proof attack on ChentsovFinite.lean sorries.

DML-style planning + backtracking (DeepClause pattern):
  1. PLAN: For each sorry, generate ranked tactic strategies
  2. EXECUTE: Try first strategy, backtrack on failure
  3. FUSE: Run across qfox (Gemma4, GPU) and neon (Qwen3.6, CPU) in parallel

Usage:
    python3 chentsov_fusion.py                    # all sorries, both endpoints
    python3 chentsov_fusion.py --sorries 1,7,9    # specific sorries
    python3 chentsov_fusion.py --endpoints qfox   # qfox only
    python3 chentsov_fusion.py --max-retries 3    # backtracking depth
"""

import argparse
import json
import os
import re
import sys
import time
from concurrent.futures import ThreadPoolExecutor, as_completed
from datetime import datetime, timezone
from pathlib import Path
from typing import Optional

# Force line-buffered output even when stdout is redirected to a file
sys.stdout.reconfigure(line_buffering=True)

try:
    import requests
    HAS_REQUESTS = True
except ImportError:
    HAS_REQUESTS = False
    import urllib.request
    import urllib.error

# ============================================================
# Configuration
# ============================================================

ENDPOINTS = {
    "qfox": {
        "url": "http://100.88.57.96:8081/v1/chat/completions",
        "model": "gemma4-12b",
        "name": "Gemma4-12B (qfox, GPU)",
        "timeout": 600,
        "max_tokens": 4096,
        "temperature": 0.1,
    },
    "neon-qwen": {
        "url": "http://100.92.88.64:11434/v1/chat/completions",
        "model": "hf.co/llmfan46/Qwen3.6-35B-A3B-uncensored-heretic-GGUF:Q4_K_M",
        "name": "Qwen3.6-35B-A3B (neon, CPU)",
        "timeout": 3600,
        "max_tokens": 4096,
        "temperature": 0.1,
    },
    "neon-step": {
        "url": "http://100.92.88.64:11434/v1/chat/completions",
        "model": "step3.7-flash",
        "name": "Step-3.7-Flash (neon, CPU)",
        "timeout": 3600,
        "max_tokens": 4096,
        "temperature": 0.1,
    },
    # DeepSeek-Prover-V2-7B: purpose-built Lean 4 / math proof model (~4-5 GB Q4_K_M).
    # Best for hard construction sorries; trained on Lean proof search data.
    "neon-deepseek": {
        "url": "http://100.92.88.64:11434/v1/chat/completions",
        "model": "hf.co/irmma/DeepSeek-Prover-V2-7B-Q4_K_M-GGUF:latest",
        "name": "DeepSeek-Prover-V2-7B (neon, CPU)",
        "timeout": 3600,
        "max_tokens": 4096,
        "temperature": 0.05,  # lower temp: prover models are more deterministic
    },
    # FreeLLMAPI: runs on cupfox (100.72.130.76), zero neon RAM impact.
    # Offloads generation to external models; good overflow for hard sorries.
    "freellmapi": {
        "url": "http://46.232.249.226:3001/v1/chat/completions",
        "model": "auto",
        "name": "FreeLLMAPI (cupfox, zero-RAM)",
        "timeout": 120,
        "max_tokens": 4096,
        "temperature": 0.1,
    },
}

# ============================================================
# Proof Strategies (DML-style plan)
# ============================================================
# Each strategy is a tactic family with a focused prompt.
# The planner ranks strategies by difficulty match.
# The executor tries them in order, backtracking on failure.

PROOF_STRATEGIES = [
    {
        "id": "simp_comm",
        "tactic": "simp",
        "prompt": "Use simp with mul_comm to swap the multiplication order. The goal is symmetric in two variables.",
        "template": "simp [fisherMetric, mul_comm]",
        "applicable_to": ["fisherMetric_sym"],
    },
    {
        "id": "simp_ite",
        "tactic": "simp",
        "prompt": "Use simp with Finset.sum_ite to evaluate the sum of if-then-else expressions. The tangentBasis has two if-then-else layers.",
        "template": "simp [tangentBasis, Finset.sum_ite]",
        "applicable_to": ["tangentBasis_sum"],
    },
    {
        "id": "simp_membership",
        "tactic": "simp",
        "prompt": "Use simp to unfold the set membership and apply the sum lemma.",
        "template": "simp [tangentSpace, tangentBasis_sum]",
        "applicable_to": ["tangentBasis_in_tangentSpace"],
    },
    {
        "id": "constructor_ring",
        "tactic": "constructor; ring",
        "prompt": "For IsLinearMap, use constructor to split into add and smul, then use simp with add_mul/mul_add and Finset.sum_add_distrib, followed by ring.",
        "template": "constructor\n· intro x y; simp [fisherMetric, add_mul, Finset.sum_add_distrib]; ring\n· intro c x; simp [fisherMetric, mul_assoc, Finset.sum_add_distrib]; ring",
        "applicable_to": ["fisherMetric_linear_left", "fisherMetric_linear_right"],
    },
    {
        "id": "div_pos",
        "tactic": "positivity",
        "prompt": "For positive definiteness, each term X i * X i / p.1 i is non-negative (p is in the open simplex so p.1 i > 0). Use div_pos and mul_nonneg, then show at least one term is strictly positive.",
        "template": "simp [fisherMetric]\napply Finset.sum_pos'\n· intro i _; exact div_nonneg (sq_nonneg _) (le_of_lt (p.2.1 i))\n· obtain ⟨i, hi⟩ := not_forall.mp hX\n  exact ⟨i, Finset.mem_unicast.mpr (by simp), div_pos (sq_pos_of_ne_zero hi) (p.2.1 i)⟩",
        "applicable_to": ["fisherMetric_pos_def"],
    },
    {
        "id": "equiv_reindex",
        "tactic": "Finset.sum_equiv",
        "prompt": "For permutation invariance of the sum, use Finset.sum_equiv with the permutation to reindex. The sum ∑ i, p.1 (σ.symm i) = 1 follows from p.2.2 by reindexing.",
        "template": "have := p.2.2\nconv_lhs at this => rw [← Finset.sum_equiv σ.symm (Finset.univ) (by simp)]\nexact this",
        "applicable_to": ["IsPermutationInvariant_sigma_sum"],
    },
    {
        "id": "construction",
        "tactic": "exact/mk",
        "prompt": "This is a definition, not a lemma. Construct the term directly using structure syntax ⟨...⟩.",
        "template": "exact ⟨fun i => if h : i = f.splitIdx then f.hq_pos * p.1 i else if h2 : i = ⟨f.splitIdx, by omega⟩ then (1 - f.hq_pos) * p.1 (f.splitIdx) else p.1 i, ...⟩",
        "applicable_to": ["SplitEmbedding.apply", "SplitEmbedding.pushforward", "uniformDist"],
    },
    {
        "id": "omega_arith",
        "tactic": "omega",
        "prompt": "This is a pure arithmetic goal. Use omega to close it.",
        "template": "omega",
        "applicable_to": [],
    },
    {
        "id": "native_decide",
        "tactic": "native_decide",
        "prompt": "This might be decidable by computation. Try native_decide.",
        "template": "native_decide",
        "applicable_to": [],
    },
    {
        "id": "sorry_fallback",
        "tactic": "sorry",
        "prompt": "Cannot prove this yet. Leave as sorry with a TODO comment.",
        "template": "sorry -- TODO: needs further work",
        "applicable_to": [],
    },
]


# ============================================================
# ChentsovFinite sorries
# ============================================================

SORRIES = [
    {
        "id": 1, "line": 40, "name": "tangentBasis_sum", "difficulty": "easy",
        "strategies": ["simp_ite", "omega_arith"],
        "context": """def tangentBasis {n : ℕ} (i j : Fin n) : Fin n → ℝ :=
  fun k => if k = i then 1 else if k = j then -1 else 0

lemma tangentBasis_sum {n : ℕ} (p : openSimplex n) (i j : Fin n) :
    ∑ k, tangentBasis i j k = 0 := by""",
    },
    {
        "id": 2, "line": 44, "name": "tangentBasis_in_tangentSpace", "difficulty": "easy",
        "strategies": ["simp_membership", "simp_ite"],
        "context": """lemma tangentBasis_in_tangentSpace {n : ℕ} (p : openSimplex n) (i j : Fin n) :
    tangentBasis i j ∈ tangentSpace p := by""",
    },
    {
        "id": 3, "line": 64, "name": "SplitEmbedding.apply", "difficulty": "medium",
        "strategies": ["construction"],
        "context": """structure SplitEmbedding (n : ℕ) where
  splitIdx : Fin n
  q : ℝ
  hq_pos : q > 0
  hq_lt_one : q < 1

def SplitEmbedding.refinedSize {n : ℕ} (_ : SplitEmbedding n) : ℕ := n + 1

def SplitEmbedding.apply {n : ℕ} (f : SplitEmbedding n) (p : openSimplex n) :
    openSimplex (refinedSize f) :=""",
    },
    {
        "id": 4, "line": 68, "name": "SplitEmbedding.pushforward", "difficulty": "medium",
        "strategies": ["construction"],
        "context": """def SplitEmbedding.pushforward {n : ℕ} (f : SplitEmbedding n) (p : openSimplex n)
    (X : Fin n → ℝ) : Fin (refinedSize f) → ℝ :=""",
    },
    {
        "id": 5, "line": 73, "name": "SplitEmbedding.pushforward_sum", "difficulty": "hard",
        "strategies": ["simp_ite", "omega_arith", "sorry_fallback"],
        "context": """lemma SplitEmbedding.pushforward_sum {n : ℕ} (f : SplitEmbedding n) (p : openSimplex n)
    (X : Fin n → ℝ) (hX : ∑ i, X i = 0) :
    ∑ j, f.pushforward p X j = 0 := by""",
    },
    {
        "id": 6, "line": 78, "name": "SplitEmbedding.pushforward_tangent", "difficulty": "hard",
        "strategies": ["equiv_reindex", "sorry_fallback"],
        "context": """lemma SplitEmbedding.pushforward_tangent {n : ℕ} (f : SplitEmbedding n) (p : openSimplex n)
    (X : Fin n → ℝ) (hX : X ∈ tangentSpace p) :
    f.pushforward p X ∈ tangentSpace (f.apply p) := by""",
    },
    {
        "id": 7, "line": 93, "name": "fisherMetric_sym", "difficulty": "easy",
        "strategies": ["simp_comm", "omega_arith"],
        "context": """noncomputable def fisherMetric {n : ℕ} (p : openSimplex n) (X Y : Fin n → ℝ) : ℝ :=
  ∑ i, X i * Y i / p.1 i

lemma fisherMetric_sym {n : ℕ} (p : openSimplex n) (X Y : Fin n → ℝ) :
    fisherMetric p X Y = fisherMetric p Y X := by""",
    },
    {
        "id": 8, "line": 98, "name": "fisherMetric_pos_def", "difficulty": "medium",
        "strategies": ["div_pos", "sorry_fallback"],
        "context": """lemma fisherMetric_pos_def {n : ℕ} (p : openSimplex n) (X : Fin n → ℝ)
    (hX : X ≠ 0) (hXsum : ∑ i, X i = 0) :
    fisherMetric p X X > 0 := by""",
    },
    {
        "id": 9, "line": 102, "name": "fisherMetric_linear_left", "difficulty": "easy",
        "strategies": ["constructor_ring", "omega_arith"],
        "context": """lemma fisherMetric_linear_left {n : ℕ} (p : openSimplex n) (Y : Fin n → ℝ) :
    IsLinearMap ℝ (fun X => fisherMetric p X Y) := by""",
    },
    {
        "id": 10, "line": 106, "name": "fisherMetric_linear_right", "difficulty": "easy",
        "strategies": ["constructor_ring", "omega_arith"],
        "context": """lemma fisherMetric_linear_right {n : ℕ} (p : openSimplex n) (X : Fin n → ℝ) :
    IsLinearMap ℝ (fun Y => fisherMetric p X Y) := by""",
    },
    {
        "id": 11, "line": 136, "name": "IsPermutationInvariant_sigma_sum", "difficulty": "medium",
        "strategies": ["equiv_reindex", "sorry_fallback"],
        "context": """In IsPermutationInvariant definition:
    let σp : openSimplex n :=
      ⟨fun i => p.1 (σ.symm i),
       ⟨fun i => p.2.1 (σ.symm i), by
        have := p.2.2
        sorry⟩⟩
  -- Need: ∑ i, p.1 (σ.symm i) = 1""",
    },
    {
        "id": 12, "line": 156, "name": "fisherMetric_chentsov_invariant", "difficulty": "hard",
        "strategies": ["simp_comm", "sorry_fallback"],
        "context": """lemma fisherMetric_chentsov_invariant {n : ℕ} :
    IsChentsovInvariant (...)  (...) := by
  intro f p X Y hXsum hYsum
  sorry""",
    },
    {
        "id": 13, "line": 169, "name": "uniformDist", "difficulty": "medium",
        "strategies": ["construction", "sorry_fallback"],
        "context": """noncomputable def uniformDist (N : ℕ) (hN : N > 0) : openSimplex N :=""",
    },
    {
        "id": 14, "line": 181, "name": "metric_at_uniform", "difficulty": "hard",
        "strategies": ["sorry_fallback"],
        "context": """lemma metric_at_uniform {N : ℕ} (hN : N ≥ 2) (g : RiemannianMetric N)
    (h_perm : IsPermutationInvariant g) :
    ∃ (lambda_N : ℝ), lambda_N > 0 ∧
    ∀ u v : Fin N → ℝ, ∑ i, u i = 0 → ∑ i, v i = 0 →
      g.toFun (uniformDist N (by linarith)) u v = lambda_N * ∑ i, u i * v i := by""",
    },
    {
        "id": 15, "line": 204, "name": "equal_refinement_const", "difficulty": "hard",
        "strategies": ["sorry_fallback"],
        "context": """lemma equal_refinement_const ... :
    ∃ C : ℝ, C > 0 ∧ ... := by""",
    },
    {
        "id": 16, "line": 228, "name": "fisher_on_rational", "difficulty": "hard",
        "strategies": ["sorry_fallback"],
        "context": """lemma fisher_on_rational ... :
    ∃ C : ℝ, C > 0 ∧ ... := by""",
    },
    {
        "id": 17, "line": 255, "name": "chentsov_theorem", "difficulty": "hard",
        "strategies": ["sorry_fallback"],
        "context": """theorem chentsov_theorem ... :
    ∃ (c : ℝ), c > 0 ∧ ... := by""",
    },
]

# ============================================================
# System prompt
# ============================================================

SYSTEM_PROMPT = """You are a Lean 4 formal verification expert. You will be given:
1. A theorem statement with `sorry` to prove
2. A specific tactic strategy to try

Output ONLY the Lean tactic proof that replaces `sorry`. No explanation, no markdown fences.

Key definitions:
- openSimplex n = { p | (∀ i, p i > 0) ∧ (∑ i, p i = 1) }
- tangentSpace p = { X | ∑ i, X i = 0 }
- fisherMetric p X Y = ∑ i, X i * Y i / p.1 i
- tangentBasis i j = fun k => if k = i then 1 else if k = j then -1 else 0

Available tactics: simp, ring, omega, exact, apply, constructor, intro, ext,
funext, rcases, calc, nlinarith, norm_num, decide, native_decide,
Finset.sum_ite, Finset.sum_add_distrib, add_mul, mul_add, mul_comm,
Finset.sum_equiv, div_pos, mul_pos, sq_nonneg, sq_pos_of_ne_zero."""


# ============================================================
# API calls
# ============================================================


def call_model(endpoint: dict, prompt: str, sorry_id: int) -> dict:
    """Call a model endpoint and return the result."""
    start = time.time()
    try:
        payload = {
            "model": endpoint["model"],
            "messages": [
                {"role": "system", "content": SYSTEM_PROMPT},
                {"role": "user", "content": prompt},
            ],
            "max_tokens": endpoint["max_tokens"],
            "temperature": endpoint["temperature"],
            "stream": False,
        }

        if HAS_REQUESTS:
            resp = requests.post(endpoint["url"], json=payload, timeout=endpoint["timeout"])
            resp.raise_for_status()
            data = resp.json()
        else:
            req = urllib.request.Request(
                endpoint["url"],
                data=json.dumps(payload).encode("utf-8"),
                headers={"Content-Type": "application/json"},
                method="POST",
            )
            with urllib.request.urlopen(req, timeout=endpoint["timeout"]) as resp:
                data = json.loads(resp.read().decode("utf-8"))

        elapsed = time.time() - start
        content = data["choices"][0]["message"]["content"]
        # Some models (Gemma4) put reasoning in reasoning_content
        reasoning = data["choices"][0]["message"].get("reasoning_content", "")
        # Strip markdown code fences
        if content:
            content = re.sub(r"^```(?:lean)?\s*\n?", "", content.strip())
            content = re.sub(r"\n?```\s*$", "", content.strip())
        return {
            "sorry_id": sorry_id,
            "endpoint": endpoint["name"],
            "response": content,
            "reasoning": reasoning[:500] if reasoning else None,
            "elapsed_s": round(elapsed, 1),
            "success": True,
            "error": None,
        }
    except Exception as e:
        elapsed = time.time() - start
        return {
            "sorry_id": sorry_id,
            "endpoint": endpoint["name"],
            "response": None,
            "elapsed_s": round(elapsed, 1),
            "success": False,
            "error": str(e),
        }


# ============================================================
# DML-style planning + backtracking
# ============================================================


def plan_proof(sorry: dict) -> list[dict]:
    """Generate ranked proof strategies for a sorry (DML planner phase)."""
    strategy_ids = sorry.get("strategies", [])
    ranked = []
    for sid in strategy_ids:
        strategy = next((s for s in PROOF_STRATEGIES if s["id"] == sid), None)
        if strategy:
            ranked.append(strategy)
    # Add fallback if not already present
    if not any(s["id"] == "sorry_fallback" for s in ranked):
        ranked.append(next(s for s in PROOF_STRATEGIES if s["id"] == "sorry_fallback"))
    return ranked


def build_strategy_prompt(sorry: dict, strategy: dict) -> str:
    """Build a focused prompt for one sorry + one strategy."""
    return f"""Prove this Lean 4 theorem using the specified tactic strategy.

Theorem:
{sorry['context']}

Strategy: {strategy['prompt']}
Suggested tactic: {strategy['template']}

Provide ONLY the tactic proof (the `by` block). No markdown fences."""


def execute_with_backtracking(
    sorry: dict,
    endpoint: dict,
    max_retries: int = 3,
) -> dict:
    """Try proof strategies in order, backtracking on failure (DML executor phase)."""
    strategies = plan_proof(sorry)
    attempts = []

    for i, strategy in enumerate(strategies[:max_retries]):
        prompt = build_strategy_prompt(sorry, strategy)
        result = call_model(endpoint, prompt, sorry["id"])
        result["strategy_id"] = strategy["id"]
        result["strategy_tactic"] = strategy["tactic"]
        result["attempt"] = i + 1
        attempts.append(result)

        # Check if the response looks like a valid proof (not sorry)
        if result["success"] and result["response"]:
            response = result["response"].strip()
            if "sorry" not in response.lower() or response.startswith("sorry -- TODO"):
                # Looks like a real proof attempt
                result["looks_valid"] = True
                break
            else:
                result["looks_valid"] = False
                # Backtrack: try next strategy
                continue
        else:
            result["looks_valid"] = False
            continue

    # Return the best attempt
    best = attempts[-1] if attempts else {
        "sorry_id": sorry["id"],
        "endpoint": endpoint["name"],
        "response": None,
        "elapsed_s": 0,
        "success": False,
        "error": "No attempts made",
        "strategy_id": None,
        "attempt": 0,
        "looks_valid": False,
    }
    best["total_attempts"] = len(attempts)
    best["all_attempts"] = [
        {
            "strategy": a.get("strategy_id"),
            "success": a["success"],
            "elapsed_s": a["elapsed_s"],
            "looks_valid": a.get("looks_valid", False),
            "preview": (a.get("response") or "")[:100],
        }
        for a in attempts
    ]
    return best


# ============================================================
# Main fusion logic
# ============================================================


def run_fusion(
    sorries: list[dict],
    endpoints: dict[str, dict],
    output_dir: Path,
    max_retries: int = 3,
    max_workers: int = 6,
    difficulty_split: bool = False,
) -> dict:
    """Run parallel DML-style fusion across all endpoints and sorries."""
    output_dir.mkdir(parents=True, exist_ok=True)
    timestamp = datetime.now(timezone.utc).strftime("%Y%m%d_%H%M%S")

    results = []
    tasks = []

    # Difficulty-aware routing: when enabled and all 3 standard endpoints are
    # present, route each sorry to the endpoint best suited to its difficulty.
    # easy → qfox (GPU, fast iteration)
    # medium → neon-step (StepFun, strong math reasoning)
    # hard → neon-qwen (Qwen3.6-35B, large context)
    # Fallback: each endpoint still tries the sorry if its primary isn't present.
    difficulty_map = {
        # easy: GPU first (fast), fallback to step
        "easy":   ["qfox", "neon-step", "neon-deepseek"],
        # medium: StepFun strong at math reasoning; GPU backup
        "medium": ["neon-step", "qfox", "neon-deepseek"],
        # hard: purpose-built Lean prover first, then FreeLLMAPI overflow, then GPU
        "hard":   ["neon-deepseek", "freellmapi", "qfox"],
    }

    use_split = difficulty_split and len(endpoints) >= 2

    for sorry in sorries:
        if use_split:
            preferred = difficulty_map.get(sorry.get("difficulty", "hard"), list(endpoints.keys()))
            # Assign to the first preferred endpoint that's available
            ep_key = next((k for k in preferred if k in endpoints), list(endpoints.keys())[0])
            tasks.append((ep_key, endpoints[ep_key], sorry))
        else:
            for ep_key, ep_config in endpoints.items():
                tasks.append((ep_key, ep_config, sorry))

    print(f"=== ChentsovFinite DML Fusion Run ===")
    print(f"Timestamp: {timestamp}")
    print(f"Sorries: {len(sorries)}")
    print(f"Endpoints: {len(endpoints)}")
    print(f"Max retries per sorry: {max_retries}")
    print(f"Total tasks: {len(tasks)}")
    print(f"Max workers: {max_workers}")
    print("---")

    with ThreadPoolExecutor(max_workers=max_workers) as executor:
        futures = {}
        for ep_key, ep_config, sorry in tasks:
            future = executor.submit(
                execute_with_backtracking, sorry, ep_config, max_retries
            )
            futures[future] = (ep_key, sorry)

        for future in as_completed(futures):
            ep_key, sorry = futures[future]
            result = future.result()
            result["sorry_name"] = sorry["name"]
            result["difficulty"] = sorry["difficulty"]
            result["endpoint_key"] = ep_key
            results.append(result)

            valid = "VALID" if result.get("looks_valid") else "SORRY"
            attempts = result.get("total_attempts", 0)
            print(
                f"  [{ep_key}] SORRY {result['sorry_id']:2d} ({sorry['name']:30s}) "
                f"{result['elapsed_s']:6.1f}s  {attempts} attempts  {valid}"
            )

    # Save results
    output_file = output_dir / f"chentsov_fusion_{timestamp}.json"
    with open(output_file, "w") as f:
        json.dump(
            {
                "timestamp": timestamp,
                "sorries": len(sorries),
                "endpoints": list(endpoints.keys()),
                "max_retries": max_retries,
                "results": results,
            },
            f,
            indent=2,
        )

    # Generate summary
    summary_file = output_dir / f"chentsov_fusion_{timestamp}_summary.md"
    with open(summary_file, "w") as f:
        f.write(f"# ChentsovFinite DML Fusion Results — {timestamp}\n\n")
        f.write(f"**Sorries attacked:** {len(sorries)}\n")
        f.write(f"**Endpoints:** {', '.join(endpoints.keys())}\n")
        f.write(f"**Max retries:** {max_retries}\n\n")

        # Stats
        valid_count = sum(1 for r in results if r.get("looks_valid"))
        total_attempts = sum(r.get("total_attempts", 0) for r in results)
        f.write(f"**Valid proofs found:** {valid_count}/{len(results)}\n")
        f.write(f"**Total attempts:** {total_attempts}\n\n")

        by_sorry = {}
        for r in results:
            sid = r["sorry_id"]
            if sid not in by_sorry:
                by_sorry[sid] = []
            by_sorry[sid].append(r)

        for sid in sorted(by_sorry.keys()):
            sorry_info = next(s for s in sorries if s["id"] == sid)
            f.write(f"## SORRY {sid}: {sorry_info['name']} ({sorry_info['difficulty']})\n\n")

            for r in sorted(by_sorry[sid], key=lambda x: x["elapsed_s"]):
                valid = "VALID" if r.get("looks_valid") else "SORRY"
                f.write(f"### {r['endpoint']} — {r['elapsed_s']}s — {valid}\n\n")
                f.write(f"Strategy: `{r.get('strategy_id')}` ({r.get('total_attempts', 0)} attempts)\n\n")

                if r.get("all_attempts"):
                    f.write("| # | Strategy | Valid | Time | Preview |\n")
                    f.write("|---|----------|-------|------|----------|\n")
                    for a in r["all_attempts"]:
                        v = "Y" if a["looks_valid"] else "N"
                        f.write(f"| {a.get('attempt', '?')} | {a['strategy']} | {v} | {a['elapsed_s']}s | `{a['preview'][:60]}...` |\n")
                    f.write("\n")

                if r["response"]:
                    f.write(f"```lean\n{r['response']}\n```\n\n")
                else:
                    f.write(f"Error: {r['error']}\n\n")

    print(f"---")
    print(f"Results: {output_file}")
    print(f"Summary: {summary_file}")
    print(f"Valid proofs: {valid_count}/{len(results)}")

    return {"output_file": str(output_file), "summary_file": str(summary_file)}


# ============================================================
# Main
# ============================================================


def main():
    parser = argparse.ArgumentParser(description="ChentsovFinite DML fusion")
    parser.add_argument("--sorries", type=str, default="all")
    parser.add_argument("--output", type=str, default="/tmp/chentsov_fusion")
    parser.add_argument("--endpoints", type=str, default="qfox,neon-qwen")
    parser.add_argument("--max-retries", type=int, default=3)
    parser.add_argument("--workers", type=int, default=6)
    parser.add_argument("--difficulty-split", action="store_true",
                        help="Route easy→qfox, medium→neon-step, hard→neon-qwen (17 tasks) "
                             "instead of all endpoints trying all sorries (51 tasks)")
    args = parser.parse_args()

    if args.sorries == "all":
        selected = SORRIES
    else:
        ids = [int(x.strip()) for x in args.sorries.split(",")]
        selected = [s for s in SORRIES if s["id"] in ids]

    ep_keys = [x.strip() for x in args.endpoints.split(",")]
    selected_eps = {k: ENDPOINTS[k] for k in ep_keys if k in ENDPOINTS}

    if not selected_eps:
        print(f"Error: no valid endpoints in {args.endpoints}")
        print(f"Available: {', '.join(ENDPOINTS.keys())}")
        sys.exit(1)

    run_fusion(selected, selected_eps, Path(args.output), args.max_retries, args.workers,
               difficulty_split=args.difficulty_split)


if __name__ == "__main__":
    main()