#!/usr/bin/env python3
"""
Parse arbitrary physical equations into shape-index entries.
Handles SymPy expressions, plain text, and CALPHAD .tdb format.

Usage:
    python3 equation_shape_parser.py --input /path/to/equations.txt
    python3 equation_shape_parser.py --propnet   # scan Propnet models
"""

import hashlib
import json
import os
import re
import sys
from collections import defaultdict
from datetime import datetime, timezone

Q16_SCALE = 65536

# ─── Shape signature from equation structure ─────────────────────────────────

def classify_equation(eq_str: str, name: str = "", source: str = "") -> dict:
    """Extract FAMM shape signature from an equation string."""
    n_ops = len(re.findall(r'[+\-*/^]', eq_str))
    n_vars = len(set(re.findall(r'\b([a-zα-ω][a-zα-ω0-9]*)\b', eq_str.lower())))
    n_parens = eq_str.count('(') + eq_str.count(')')
    has_trig = len(re.findall(r'sin|cos|tan|exp|log|sqrt', eq_str.lower()))
    has_diff = len(re.findall(r'diff|grad|curl|div|laplacian|\b∂', eq_str))
    complexity = len(eq_str.strip())

    # Sidon graph estimation
    graph_size = max(2, min(n_vars + has_trig + has_diff, 8))
    sidon_addr = [1 << i for i in range(graph_size)]
    n_pairs = graph_size * (graph_size - 1) // 2
    if n_pairs > 0:
        sums = set()
        for i in range(graph_size):
            for j in range(i + 1, graph_size):
                sums.add(sidon_addr[i] + sidon_addr[j])
        sumset_density = len(sums) / n_pairs
        sums_list = sorted(sums)
    else:
        sumset_density = 1.0
        sums_list = []

    # FAMM classification
    structural_noise = min(1.0, n_parens / max(1, complexity * 0.05))
    raw_closure = 1.0 / max(1.0, n_ops * 0.05 + has_trig * 0.1 + has_diff * 0.2)
    closure_frac = min(raw_closure, 1.0)
    scar_pressure = 1.0 - closure_frac

    spectral_raw = int(closure_frac * 4 * Q16_SCALE)
    max_raw = int(4 * Q16_SCALE)
    signal_threshold = int(2 * Q16_SCALE)

    r = Q16_SCALE if spectral_raw >= max_raw else 0
    g = (spectral_raw - signal_threshold) * Q16_SCALE // (max_raw - signal_threshold) if spectral_raw >= signal_threshold else 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 = "HOLD" if scar_pressure > 0.85 else (
        "INSPECT" if scar_pressure > 0.5 else "ACCEPT"
    )

    return {
        "name": name,
        "source": source,
        "equation": eq_str.strip()[:200],
        "constraint_graph": {
            "n_ops": n_ops, "n_vars": n_vars, "n_parens": n_parens,
            "has_trig": bool(has_trig), "has_diff": bool(has_diff),
            "graph_size": graph_size, "complexity_chars": complexity,
        },
        "sidon_signature": {
            "addresses": sidon_addr, "pairwise_sums": sums_list,
            "n_pairs": n_pairs, "unique_sums": len(sums),
            "sumset_density": round(sumset_density, 4),
            "noise_ratio": round(structural_noise, 4),
        },
        "famm": {
            "closure_fraction": round(closure_frac, 4),
            "scar_pressure": round(scar_pressure, 4),
            "famm_state": famm_state, "rrc_shape": rrc_shape,
            "rgb": {"r": r, "g": g, "b": b},
        },
    }


# ─── Propnet parser ──────────────────────────────────────────────────────────

def parse_propnet_models(propnet_path: str) -> list[dict]:
    """Parse SymPy equations from Propnet model files."""
    results = []
    models_dir = os.path.join(propnet_path, "propnet", "models", "python")
    if not os.path.isdir(models_dir):
        return results

    for fname in sorted(os.listdir(models_dir)):
        if not fname.endswith(".py") or fname == "__init__.py":
            continue
        fpath = os.path.join(models_dir, fname)
        with open(fpath) as f:
            content = f.read()
        # Extract equation strings
        eqs = re.findall(r'["\']([^"\']*(?:=|=|≈)[^"\']*)["\']', content)
        eqs += re.findall(r'Eq\([^)]+\)', content)
        for eq in eqs[:5]:
            sig = classify_equation(eq, name=fname.replace(".py", ""), source="propnet")
            results.append(sig)
    return results


# ─── Plain text equation parser ──────────────────────────────────────────────

def parse_equation_file(filepath: str) -> list[dict]:
    """Parse equations from a plain text file (one per line, or name=eq format)."""
    results = []
    with open(filepath) as f:
        for line in f:
            line = line.strip()
            if not line or line.startswith("#") or line.startswith("//"):
                continue
            if "=" in line:
                parts = line.split("=", 1)
                name = parts[0].strip()
                eq = parts[1].strip()
            else:
                name = f"eq_{len(results)}"
                eq = line
            sig = classify_equation(eq, name=name, source=filepath)
            results.append(sig)
    return results


# ─── CALPHAD .tdb parser ────────────────────────────────────────────────────

def parse_tdb_file(filepath: str) -> list[dict]:
    """Parse thermodynamic equations from CALPHAD .tdb format."""
    results = []
    with open(filepath, errors="replace") as f:
        content = f.read()
    # Extract polynomial expressions (G = ...)
    eqs = re.findall(r'G\s*=\s*[^;]+', content)
    for eq in eqs[:50]:
        sig = classify_equation(eq[:200], name="tdb_gibbs", source=filepath)
        results.append(sig)
    return results


# ─── Main ────────────────────────────────────────────────────────────────────

def main():
    import argparse
    ap = argparse.ArgumentParser()
    ap.add_argument("--input", help="Input equation file (plain text)")
    ap.add_argument("--tdb", help="Input CALPHAD .tdb file")
    ap.add_argument("--propnet", action="store_true", help="Scan Propnet models")
    ap.add_argument("--out", default="equation_shape_index.json")
    args = ap.parse_args()

    all_results = []

    if args.propnet:
        for p in ["/home/allaun/propnet_db", "/home/allaun/Research Stack/../propnet_db"]:
            if os.path.isdir(p):
                all_results.extend(parse_propnet_models(p))
                print(f"Parsed {len(all_results)} equations from Propnet")

    if args.input and os.path.isfile(args.input):
        all_results.extend(parse_equation_file(args.input))
        print(f"Parsed {len(all_results)} equations from {args.input}")

    if args.tdb and os.path.isfile(args.tdb):
        all_results.extend(parse_tdb_file(args.tdb))
        print(f"Parsed {len(all_results)} TDB equations")

    # Build shape index
    shape_dist = defaultdict(int)
    for r in all_results:
        shape_dist[r["famm"]["rrc_shape"]] += 1

    index = {
        "meta": {
            "source": "equation_parser",
            "total_equations": len(all_results),
            "shape_distribution": dict(shape_dist),
            "built_at": datetime.now(timezone.utc).isoformat(),
            "schema": "equation_shape_index_v1",
        },
        "equations": all_results,
    }

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

    print(f"\nShape distribution: {dict(shape_dist)}")
    print(f"Output: {args.out}")


if __name__ == "__main__":
    main()