#!/usr/bin/env python3
"""
pist_trace_classify_offline.py
==============================

Offline, token-free proof trace classifier.
Computes transition matrix spectra, finds nearest neighbors from the 57-theorem
flexure library (offline vectors), determines the RRC shape proxy using
color-space thresholds matching Lean's `Semantics.PIST.Classify`, and invokes
the local Lean `rrc-watchdog` binary in the podman container to verify alignment.

Usage:
    python3 pist_trace_classify_offline.py trace.json --rrc-shape signalShapedRouteCompiler
"""

import argparse
import json
import math
import os
import subprocess
import sys
from pathlib import Path
from collections import Counter

REPO_ROOT = Path(__file__).resolve().parents[2]
VECTORS_PATH = REPO_ROOT / "shared-data" / "pist_trace_scaled_vectors.jsonl"
WATCHDOG_PATH = "/home/researcher/stack/0-Core-Formalism/lean/Semantics/.lake/build/bin/rrc-watchdog"

FEATURE_KEYS = [
    "matrix_size",
    "rank",
    "spectral_gap",
    "laplacian_zero_count",
    "density",
    "adjacency_eigenvalue_max",
]


def power_iteration(matrix, max_iter=100):
    n = len(matrix)
    if n == 0:
        return 0.0
    v = [1.0 / math.sqrt(n)] * n
    for _ in range(max_iter):
        vn = [sum(matrix[i][j] * v[j] for j in range(n)) for i in range(n)]
        nm = math.sqrt(sum(x * x for x in vn))
        if nm < 1e-12:
            return 0.0
        v = [x / nm for x in vn]
    num = sum(v[i] * sum(matrix[i][j] * v[j] for j in range(n)) for i in range(n))
    den = sum(v[i] * v[i] for i in range(n))
    return num / den if den > 0 else 0.0


def compute_spectral(matrix):
    n = len(matrix)
    if n == 0:
        return {}
    sym = [[(matrix[i][j] + matrix[j][i]) / 2.0 for j in range(n)] for i in range(n)]
    lap = [[sum(sym[i]) if i == j else -sym[i][j] for j in range(n)] for i in range(n)]
    ev_max = power_iteration(sym)
    shifted = [[sym[i][j] - 0.9 * ev_max * (1 if i == j else 0) for j in range(n)] for i in range(n)]
    ev_shift = power_iteration(shifted)
    ev_second = max(0, ev_max - ev_shift) if ev_shift < ev_max else ev_max
    gap = ev_max - ev_second
    lap_max = power_iteration(lap)
    neg_lap = [[-lap[i][j] for j in range(n)] for i in range(n)]
    lap_min = -power_iteration(neg_lap)
    ata = [[sum(matrix[k][i] * matrix[k][j] for k in range(n)) for j in range(n)] for i in range(n)]
    sv_max = math.sqrt(max(0, power_iteration(ata)))
    rank = sum(1 for row in matrix if sum(row) > 0)
    total = sum(sum(row) for row in matrix)
    frob = math.sqrt(sum(cell * cell for row in matrix for cell in row))
    lap_zero = sum(1 for i in range(n) if abs(sum(matrix[i]) - matrix[i][i]) < 1e-9)
    return {
        "matrix_size": n,
        "rank": rank,
        "spectral_gap": round(gap, 6),
        "density": round(total / max(n * n, 1), 6),
        "trace": sum(matrix[i][i] for i in range(n)),
        "frobenius_norm": round(frob, 6),
        "laplacian_zero_count": lap_zero,
        "adjacency_eigenvalue_max": round(ev_max, 6),
        "laplacian_eigenvalue_max": round(lap_max, 6),
        "singular_value_max": round(sv_max, 6),
    }


def classify_tactic_family(name: str) -> str:
    n = name.lower()
    if "rw" in n:
        return "rewrite"
    if "simp" in n:
        return "normalization"
    if "omega" in n:
        return "arithmetic"
    if "induct" in n:
        return "induction"
    if "ring" in n or "calc" in n:
        return "algebraic"
    if "cases" in n or "constructor" in n:
        return "case_analysis"
    if any(k in n for k in ["apply", "intro", "have", "logic"]):
        return "discharge"
    if "rfl" in n:
        return "reflexivity"
    return "unknown"


def get_rrc_shape(ev_max: float) -> str:
    # Q16.16 conversion
    lam = int(ev_max * 65536)
    if lam >= 262144:  # 4.0
        return "cognitiveLoadField"
    elif lam >= 131072:  # 2.0
        return "signalShapedRouteCompiler"
    else:
        return "holdForUnlawfulOrUnderspecifiedShape"


def load_library():
    library = []
    if not VECTORS_PATH.exists():
        print(f"Warning: local vectors file not found at {VECTORS_PATH}", file=sys.stderr)
        return library
    with open(VECTORS_PATH) as f:
        for line in f:
            if line.strip():
                r = json.loads(line)
                library.append(r)
    return library


def find_nearest_neighbors(features, library, top_k=3):
    if not library:
        return []
    scored = []
    for r in library:
        # Distance over the FEATURE_KEYS
        dist = 0.0
        for k in FEATURE_KEYS:
            val_a = features.get(k, 0.0)
            val_b = r.get(k, 0.0)
            dist += (val_a - val_b) ** 2
        dist = math.sqrt(dist)
        scored.append({"dist": dist, "record": r})
    scored.sort(key=lambda x: x["dist"])
    return scored[:top_k]


def run_lean_watchdog(predicted_shape: str, expected_shape: str) -> dict:
    cmd = [
        "podman",
        "exec",
        "research-stack",
        WATCHDOG_PATH,
        "--pist-label",
        predicted_shape,
        "--exact-label",
        predicted_shape,
        "--rrc-shape",
        expected_shape,
    ]
    try:
        res = subprocess.run(cmd, capture_output=True, text=True, timeout=15)
        if res.returncode not in (0, 1):
            return {"error": f"watchdog returned exit code {res.returncode}", "stderr": res.stderr}
        return json.loads(res.stdout)
    except Exception as e:
        return {"error": f"failed to run rrc-watchdog: {e}"}


def main():
    parser = argparse.ArgumentParser(description="Offline RRC Trace Classifier")
    parser.add_argument("trace_path", help="Path to ProofTraceReceipt v2 JSON file")
    parser.add_argument(
        "--rrc-shape",
        default="signalShapedRouteCompiler",
        help="Expected shape for Lean alignment gate check",
    )
    args = parser.parse_args()

    # Load trace
    with open(args.trace_path) as f:
        trace = json.load(f)

    name = trace.get("name", "unnamed")
    matrix = trace.get("transition_matrix", [])
    if not matrix:
        print(json.dumps({"error": "Empty transition matrix"}))
        sys.exit(1)

    # 1. Compute spectral profile
    spectral = compute_spectral(matrix)
    ev_max = spectral["adjacency_eigenvalue_max"]

    # 2. Map color domain (Lean-anchored logic)
    rrc_shape = get_rrc_shape(ev_max)

    # 3. K-NN tactic family matching (free offline)
    library = load_library()
    neighbors = find_nearest_neighbors(spectral, library)

    tactic_family = classify_tactic_family(name)
    predicted_tactic_family = "unknown"
    predicted_status = "failed"
    knn_support = 0

    if neighbors:
        families = [classify_tactic_family(n["record"].get("name", "")) for n in neighbors]
        statuses = [n["record"].get("status", "failed") for n in neighbors]
        predicted_tactic_family = Counter(families).most_common(1)[0][0]
        predicted_status = Counter(statuses).most_common(1)[0][0]
        knn_support = len(neighbors)

    # 4. Lean RRC alignment verification
    watchdog_res = run_lean_watchdog(rrc_shape, args.rrc_shape)

    output = {
        "theorem_name": name,
        "spectral_radius": ev_max,
        "spectral_radius_q16": int(ev_max * 65536),
        "predicted_rrc_shape": rrc_shape,
        "tactic_family_heuristic": tactic_family,
        "knn_predictions": {
            "tactic_family": predicted_tactic_family,
            "status": predicted_status,
            "support": knn_support,
        },
        "lean_alignment": watchdog_res,
    }

    print(json.dumps(output, indent=2))


if __name__ == "__main__":
    main()