#!/usr/bin/env python3
"""
rrc_self_classify.py — Self-classifying RRC pipeline.

Takes a new equation (name + LaTeX text + route_hint), runs it through
all 6 kernel stages, assigns manifold location + regime, and emits a receipt.

Usage:
  # Classify a single equation
  python3 4-Infrastructure/shim/rrc_self_classify.py \\
    --name "my_sidon_test" \\
    --equation "|A| ≤ √(2N) + 1" \\
    --route "number_theory"

  # Batch classify from JSONL
  python3 4-Infrastructure/shim/rrc_self_classify.py --batch new_eqs.jsonl

  # Self-test: classify all kernel titles against themselves
  python3 4-Infrastructure/shim/rrc_self_classify.py --self-test
"""
from __future__ import annotations

import json
import re
import subprocess
import sys
import time
from collections import defaultdict
from pathlib import Path
from typing import Any

NEON_HOST = "neon-64gb"
CONTAINER = "arxiv-pg"
DB = "arxiv"

# Import the kernel detection logic
sys.path.insert(0, str(Path(__file__).resolve().parent))
from rrc_arxiv_kernel_refine import (
    DIOPHANTINE_KERNEL, COMBINATORICS_KERNEL,
    detect_diophantine_type, detect_combinatorics_type,
    detect_obscure_type, detect_dataset_type, detect_sidon_type,
    detect_geometry_type, detect_reconstruction_type,
    extract_keywords, search_papers,
    RECEIPT_PATH,
)

RECEIPT_PATH = Path("archive/experimental-shim-probes/rrc_equation_classifier_receipt.json")


def classify_equation(name: str, eq_text: str, route_hint: str = "") -> dict:
    """Run an equation through all 6 kernel stages and return the best match."""
    stages = [
        # v0: graph-reconstruction kernel — fires FIRST (matches the batch
        # kernel_refine ordering), so reconstruction-conjecture equations are
        # tagged before the generic combinatorics/sidon stages claim them.
        ("kernel_refine_v0", lambda: detect_reconstruction_type(name, eq_text)),
        ("kernel_refine_v6", lambda: detect_sidon_type(name, eq_text)),
        ("kernel_refine_v3", lambda: detect_combinatorics_type(name, eq_text)),
        ("kernel_refine_v4", lambda: detect_dataset_type(name, eq_text)),
        ("kernel_refine_v2", lambda: detect_diophantine_type(name, eq_text)),
        ("kernel_refine_v5", lambda: detect_obscure_type(name, eq_text)),
        # v7: geometry/topology kernel — last kernel stage before the keyword
        # fallback, so existing number-theory matches are untouched and only
        # otherwise-unmatched eqs (e.g. the geodesic equation) reach it.
        ("kernel_refine_v7", lambda: detect_geometry_type(name, eq_text)),
    ]

    best_paper_id = None
    best_match = None
    best_stage = None

    for stage_name, detector in stages:
        try:
            matches = detector()
            if matches:
                best = matches[0]
                best_paper_id = best.get("paper_id")
                best_match = best
                best_stage = stage_name
                break
        except Exception:
            continue

    # Fallback: generic keyword search
    if not best_paper_id:
        kw = extract_keywords(name + " " + eq_text + " " + route_hint)
        results = search_papers(kw)
        if results and results[0]["score"] >= 3:
            best_paper_id = results[0]["paper_id"]
            best_match = results[0]
            best_stage = "kernel_refine_v1"

    # Compute manifold assignment
    slack, regime = _compute_regime(best_match)
    sidon_label, strand = _assign_label(name)
    manifold_route = _guess_route(name, eq_text, route_hint)
    # Coherence: a geometry-kernel match implies the geometry/topology route,
    # overriding the keyword route-guess (which lacks geometry vocabulary like
    # "perelman"/"chern"). Only applied when the caller gave no explicit hint.
    if (best_match and str(best_match.get("match_type", "")).startswith("geometry")
            and (not route_hint or route_hint == "?")):
        manifold_route = "geometry_topology"

    result = {
        "name": name,
        "equation_snippet": eq_text[:100],
        "classified_at": time.strftime("%Y-%m-%dT%H:%M:%SZ"),
        "match": {
            "paper_id": best_paper_id,
            "title": best_match.get("title", "") if best_match else "",
            "score": best_match.get("score", 0) if best_match else 0,
            "stage": best_stage,
            "match_type": best_match.get("match_type", "") if best_match else "",
            "signals": best_match.get("signals", []) if best_match else [],
        },
        "manifold": {
            "route": manifold_route,
            "regime": regime,
            "slack": slack,
            "sidon_label": sidon_label,
            "strand": strand,
        },
        "classification": "classified" if best_paper_id else "unmatched",
    }

    return result


def _compute_regime(match: dict | None) -> tuple[int, str]:
    if match is None:
        return 0, "unclassified"
    score = match.get("score", 0)
    if isinstance(score, str):
        try:
            score = int(score)
        except ValueError:
            score = 0
    # Geometry/topology eqs sit off the Sidon diophantine axis — label them
    # by curvature regime instead of (anti_)diophantine slack.
    if str(match.get("match_type", "")).startswith("geometry"):
        return (64 if score >= 6 else 16), "riemannian"
    if score >= 100:
        return 128, "anti_diophantine"
    elif score >= 20:
        return 64, "transition"
    elif score >= 5:
        return 16, "transition_tight"
    return 4, "diophantine"


def _assign_label(name: str) -> tuple[int, int]:
    labels = [1, 2, 4, 8, 16, 32, 64, 128]
    idx = hash(name) % len(labels)
    return labels[idx], idx


def _guess_route(name: str, eq_text: str, route_hint: str) -> str:
    if route_hint and route_hint != "?":
        return route_hint
    combined = (name + " " + eq_text).lower()
    route_patterns = [
        ("thermodynamic_energy", ["energy", "entropy", "heat", "temperature", "thermo"]),
        ("geometry_topology", ["geometry", "metric", "manifold", "curvature", "geodesic"]),
        ("cognitive_load", ["cognitive", "load", "emotional", "signal", "gate"]),
        ("compression_route", ["compress", "encoding", "codec", "hutter", "entropy"]),
        ("magnetic_signal", ["magnetic", "field", "plasma", "wave"]),
        ("control_signal", ["control", "overflow", "gain", "threshold", "tuning"]),
        ("number_theory", ["prime", "modulo", "sidon", "sumset", "additive", "bound"]),
        ("chaotic_couch", ["chaotic", "couch", "soliton", "turbulence"]),
    ]
    best_route, best_score = "unclassified", 0
    for route, kws in route_patterns:
        score = sum(3 for kw in kws if kw in combined)
        if score > best_score:
            best_score = score
            best_route = route
    return best_route


def self_test():
    """Self-test: classify a set of known equations to verify pipeline."""
    test_cases = [
        {"name": "sidon_maximum_bound", "equation": "|A| ≤ √(2N) + 1", "route": "number_theory"},
        {"name": "sumset_growth", "equation": "|A+A| ≥ |A|(|A|−1)/2", "route": "combinatorics"},
        {"name": "baker_lower_bound", "equation": "log|Λ| > −C·log(H₁)·log(H₂)", "route": "number_theory"},
        {"name": "entropy_rate", "equation": "H(X|Y) = H(X) − I(X;Y)", "route": "thermodynamic_energy"},
        {"name": "geodesic_equation", "equation": "d²x^i/ds² + Γ^i_jk dx^j/ds dx^k/ds = 0", "route": "geometry_topology"},
        {"name": "sidon_set_collision", "equation": "a + b = c + d ⇒ {a,b} = {c,d}", "route": "number_theory"},
        {"name": "singer_construction", "equation": "|D| = q+1, D ⊂ ℤ_{q²+q+1}", "route": "number_theory"},
        {"name": "cap_set_bound", "equation": "|A| ≤ 3·(2.756)^n", "route": "number_theory"},
        {"name": "CAUCHY_DAVENPORT", "equation": "|A+B| ≥ min(p, |A|+|B|−1)", "route": "number_theory"},
        {"name": "emotional_gate", "equation": "G_em = max(0, L_em − T_em)", "route": "cognitive_load"},
    ]

    print("=" * 60)
    print("RRC Self-Classification Test")
    print("=" * 60)

    results = []
    for tc in test_cases:
        result = classify_equation(tc["name"], tc["equation"], tc["route"])
        results.append(result)

        stage = result["match"]["stage"] or "NONE"
        paper = result["match"]["paper_id"] or "—"
        route = result["manifold"]["route"]
        regime = result["manifold"]["regime"]
        status = "✓" if result["classification"] == "classified" else "✗"

        print(f"\n  {status} {tc['name']:35s} {stage:20s} {route:25s} {regime:15s}")
        print(f"      → paper={paper}")
        print(f"      → {tc['equation'][:60]}")
        signals = result["match"].get("signals", [])
        if signals:
            print(f"      → signals: {', '.join(signals)}")

    # Summary
    classified = sum(1 for r in results if r["classification"] == "classified")
    print(f"\n{'='*60}")
    print(f"  Classified: {classified}/{len(results)}")
    for stage in set(r["match"]["stage"] for r in results if r["match"]["stage"]):
        cnt = sum(1 for r in results if r["match"]["stage"] == stage)
        print(f"    {stage:25s} {cnt}")


def main():
    import argparse
    ap = argparse.ArgumentParser(description="RRC Self-Classifying Pipeline")
    ap.add_argument("--name", type=str, help="Equation name")
    ap.add_argument("--equation", type=str, help="Equation LaTeX")
    ap.add_argument("--route", type=str, default="", help="Route hint")
    ap.add_argument("--self-test", action="store_true", help="Run self-test")
    args = ap.parse_args()

    if args.self_test:
        self_test()
        return

    if not args.name or not args.equation:
        print("ERROR: --name and --equation required (or --self-test)", file=sys.stderr)
        sys.exit(1)

    result = classify_equation(args.name, args.equation, args.route)
    print(json.dumps(result, indent=2, ensure_ascii=False))


if __name__ == "__main__":
    main()