#!/usr/bin/env python3
"""
Shape-index mapping for the Lean corpus.
For each theorem, compute its Sidon sumset signature and FAMM shape class,
then index by shape for nearest-neighbor proof template search.

Usage:
    python3 shape_index.py --corpus /path/to/lean_corpus --out index.json

Output:
    index.json  — shape-indexed theorem database
    shapes/     — per-file FAMM classification receipts
"""

import argparse
import hashlib
import json
import os
import re
import sys
import time
from collections import defaultdict, Counter
from pathlib import Path
from datetime import datetime, timezone

Q16_SCALE = 65536


def q16(f):
    return max(-2147483648, min(2147483647, int(f * Q16_SCALE)))


# ─── Sidon signature from theorem structure ───────────────────────────────────

def classify_lean_file(content: str, path: str) -> list[dict]:
    """
    Extract theorems/defs from a .lean file and compute shape signatures.
    """
    results = []
    lines = content.split("\n")

    # Find theorem/lemma/def lines and estimate constraint graph
    for i, line in enumerate(lines):
        stripped = line.strip()
        # Match theorem/lemma/def signatures
        m = re.match(
            r"(theorem|lemma|def)\s+(\w+)\s*(.*?)(?::=|:=)",
            stripped,
        )
        if not m:
            continue
        kind = m.group(1)
        name = m.group(2)
        sig = m.group(3)[:120]

        # Estimate constraint graph from the signature
        # Count variables, types, and arrows — these define the Sidon graph size
        n_vars = len(re.findall(r"\b([a-zα-ω]\w*)\s*:", sig))
        n_types = len(re.findall(r"(ℕ|ℤ|ℚ|ℝ|ℂ|Nat|Int|Rat|Real|Complex|Q16_16|Matrix|Fin\s+\d+)", sig))
        n_arrows = sig.count("→") + sig.count("->")
        n_parens = sig.count("(") + sig.count(")")
        complexity = len(sig.strip())

        # Estimate Sidon graph size from number of variables + types
        # (more interacting parts → larger graph)
        graph_size = max(2, n_vars + n_types // 2)

        # Assign power-of-2 Sidon addresses up to graph_size
        sidon_addrs = [1 << i for i in range(min(graph_size, 8))]
        n_addrs = len(sidon_addrs)

        # Compute pairwise sums (complete graph)
        if n_addrs >= 2:
            sums = [sidon_addrs[i] + sidon_addrs[j]
                    for i in range(n_addrs) for j in range(i + 1, n_addrs)]
            n_pairs = len(sums)
            unique_sums = len(set(sums))
            sumset_density = unique_sums / max(1, n_pairs)
        else:
            n_pairs = 0
            unique_sums = 0
            sumset_density = 1.0

        # Estimate closure fraction from structural features
        # The more variables and arrows, the lower the closure (more constraint collisions)
        raw_closure = 1.0 / max(1.0, n_vars + 0.5 * n_arrows + 0.3 * complexity / 40)
        closure_frac = min(raw_closure, 1.0)
        scar_pressure = 1.0 - closure_frac

        # Compute noise estimate: the fraction of the signature that's
        # unmapped structural overhead (parentheses, type annotations)
        noise_ratio = min(1.0, n_parens / max(1, complexity * 2))

        # Spectral radius (in Q16_16 units, normalized to [0, 4])
        spectral_raw = q16(closure_frac * 4.0)

        # RGB color gate (from PIST/Classify.lean)
        signal_threshold = q16(2.0)
        max_raw = q16(4.0)

        r = Q16_SCALE if spectral_raw >= max_raw else 0
        if spectral_raw >= signal_threshold:
            g = (spectral_raw - signal_threshold) * Q16_SCALE // (max_raw - signal_threshold)
        else:
            g = 0
        b = spectral_raw * Q16_SCALE // signal_threshold if spectral_raw < signal_threshold else 0

        rrc_shape = "CognitiveLoadField" if r > 0 else (
            "SignalShapedRouteCompiler" if g > 0 else "NoiseFloor"
        )

        # FAMM state
        if scar_pressure > 0.85:
            famm_state = "HOLD"
        elif scar_pressure > 0.5:
            famm_state = "INSPECT"
        else:
            famm_state = "ACCEPT"

        results.append({
            "name": name,
            "kind": kind,
            "path": path,
            "line": i + 1,
            "signature": sig,
            "constraint_graph": {
                "n_vars": n_vars,
                "n_types": n_types,
                "n_arrows": n_arrows,
                "graph_size": n_addrs,
                "complexity_chars": complexity,
            },
            "sidon_signature": {
                "addresses": sidon_addrs,
                "pairwise_sums": sorted(list(set(sums))) if n_addrs >= 2 else [],
                "n_pairs": n_pairs,
                "unique_sums": unique_sums,
                "sumset_density": round(sumset_density, 4),
                "noise_ratio": round(noise_ratio, 4),
            },
            "famm": {
                "closure_fraction": round(closure_frac, 4),
                "scar_pressure": round(scar_pressure, 4),
                "famm_state": famm_state,
                "rrc_shape": rrc_shape,
                "spectral_radius_raw": spectral_raw,
                "rgb": {"r": r, "g": g, "b": b},
            },
        })

    return results


# ─── Index building ──────────────────────────────────────────────────────────

def build_index(corpus_dir: str, out_path: str, max_files: int = 0):
    corpus = Path(corpus_dir)
    files = list(corpus.rglob("*.lean"))
    if max_files > 0:
        files = files[:max_files]
    print(f"Scanning {len(files)} Lean files in {corpus_dir}...")

    all_theorems = []
    shape_index = defaultdict(list)  # rrc_shape → [theorem entries]
    stats = Counter()

    start = time.time()
    for fpath in files:
        try:
            content = fpath.read_text(encoding="utf-8", errors="replace")
        except Exception:
            continue
        theorems = classify_lean_file(content, str(fpath.relative_to(corpus)))
        for t in theorems:
            shape = t["famm"]["rrc_shape"]
            shape_index[shape].append(t)
            all_theorems.append(t)
            stats["theorems"] += 1
        stats["files"] += 1

    elapsed = time.time() - start

    # Build output
    index = {
        "meta": {
            "corpus_dir": str(corpus),
            "files_scanned": stats["files"],
            "theorems_indexed": stats["theorems"],
            "elapsed_s": round(elapsed, 1),
            "theorems_per_sec": round(stats["theorems"] / elapsed, 1) if elapsed > 0 else 0,
            "built_at": datetime.now(timezone.utc).isoformat(),
            "schema": "shape_index_v1",
        },
        "shape_summary": {
            shape: len(entries)
            for shape, entries in sorted(shape_index.items())
        },
        "by_shape": {
            shape: entries for shape, entries in shape_index.items()
        },
        "all_theorems": all_theorems,
    }

    with open(out_path, "w") as f:
        json.dump(index, f, indent=2)

    print(f"\nIndexed {stats['theorems']} theorems from {stats['files']} files")
    print(f"Shape distribution: {dict(index['shape_summary'])}")
    print(f"Output: {out_path}")


def build_index_multi(corpus_dirs: list, out_path: str, max_files: int = 0):
    """Build index across multiple corpus directories."""
    all_theorems = []
    shape_index = defaultdict(list)
    stats = Counter()
    corpus_paths = [Path(c) for c in corpus_dirs]

    start = time.time()
    for corpus in corpus_paths:
        if not corpus.exists():
            print(f"Skipping non-existent: {corpus}")
            continue
        files = list(corpus.rglob("*.lean"))
        if max_files > 0:
            files = files[:max_files]
        print(f"  {corpus}: {len(files)} files")

        for fpath in files:
            try:
                content = fpath.read_text(encoding="utf-8", errors="replace")
            except Exception:
                continue
            rel_path = fpath.name  # use just filename to avoid path collisions
            theorems = classify_lean_file(content, rel_path)
            for t in theorems:
                shape = t["famm"]["rrc_shape"]
                shape_index[shape].append(t)
                all_theorems.append(t)
                stats["theorems"] += 1
            stats["files"] += 1

    elapsed = time.time() - start

    index = {
        "meta": {
            "corpus_dirs": [str(c) for c in corpus_paths],
            "files_scanned": stats["files"],
            "theorems_indexed": stats["theorems"],
            "elapsed_s": round(elapsed, 1),
            "theorems_per_sec": round(stats["theorems"] / elapsed, 1) if elapsed > 0 else 0,
            "built_at": datetime.now(timezone.utc).isoformat(),
            "schema": "shape_index_v1",
        },
        "shape_summary": {
            shape: len(entries)
            for shape, entries in sorted(shape_index.items())
        },
        "by_shape": {
            shape: entries for shape, entries in shape_index.items()
        },
        "all_theorems": all_theorems,
    }

    with open(out_path, "w") as f:
        json.dump(index, f, indent=2)

    print(f"\nIndexed {stats['theorems']} theorems from {stats['files']} files")
    print(f"Shape distribution: {dict(index['shape_summary'])}")
    print(f"Output: {out_path}")


# ─── Query: nearest neighbor by shape ────────────────────────────────────────

def nearest_by_shape(unsolved_sig: dict, index_path: str, top_k: int = 3) -> list[dict]:
    """
    Given an unsolved problem's Sidon signature, find the nearest neighbors
    by sumset density distance.
    """
    with open(index_path) as f:
        index = json.load(f)

    target_density = unsolved_sig.get("sumset_density", 0.5)
    target_vars = unsolved_sig.get("n_vars", 1)
    target_state = unsolved_sig.get("target_state", "ACCEPT")

    scored = []
    for t in index["all_theorems"]:
        density = t["sidon_signature"]["sumset_density"]
        n_vars = t["constraint_graph"]["n_vars"]
        famm_state = t["famm"]["famm_state"]

        # Distance: density difference + var difference
        d_density = abs(density - target_density)
        d_vars = abs(n_vars - target_vars) * 0.1
        d_state = 0 if famm_state == target_state else 2.0
        distance = d_density + d_vars + d_state

        scored.append((distance, t))

    scored.sort(key=lambda x: x[0])
    return [t for d, t in scored[:top_k]]


# ─── CLI ─────────────────────────────────────────────────────────────────────

def main():
    ap = argparse.ArgumentParser(description="Shape-index Lean theorems by FAMM class")
    ap.add_argument("--corpus", action="append", default=[],
                    help="Corpus directory (can be repeated)")
    ap.add_argument("--out", default="/home/allaun/lean_corpus/shape_index_full.json")
    ap.add_argument("--nearest", help="Query nearest neighbors by shape signature JSON")
    ap.add_argument("--top-k", type=int, default=3, help="Neighbors to return")
    ap.add_argument("--max-files", type=int, default=0, help="Limit files (0 = all)")
    args = ap.parse_args()

    if args.nearest:
        sig = json.loads(args.nearest)
        results = nearest_by_shape(sig, args.out, args.top_k)
        print(json.dumps(results, indent=2))
        return

    if not args.corpus:
        args.corpus = ["/home/allaun/lean_corpus",
                       "/home/allaun/Research Stack/0-Core-Formalism"]

    build_index_multi(args.corpus, args.out, args.max_files)


if __name__ == "__main__":
    main()