#!/usr/bin/env python3
"""
Equation Ingestion Pipeline — Markdown → Parse → Classify → Compute → Emit

Accepts a .md file containing mathematical equations.
Runs each through the spectral pipeline. If the file is too dense,
pauses and writes an intermediate decision document.

Usage:
  python3 scripts/ingest.py equations.md
  python3 scripts/ingest.py equations.md --max-eqns 20
"""
import json, re, sys, argparse
from pathlib import Path
from dataclasses import dataclass, field
from typing import List, Dict, Optional, Tuple

# ── Stage 0: Equation Parser ────────────────────────────────────────

@dataclass
class ParsedEquation:
    text: str
    line: int
    is_block: bool
    classification: str = "unknown"  # braid, spectral, cartan, unknown
    metadata: Dict = field(default_factory=dict)

def parse_markdown(text: str) -> List[ParsedEquation]:
    """Extract all LaTeX equations from markdown.
    Returns list with line numbers and block/inline type."""
    equations = []
    lines = text.split('\n')
    
    # Block equations: $$...$$
    for i, line in enumerate(lines):
        blocks = re.findall(r'\$\$(.+?)\$\$', line)
        for eq in blocks:
            equations.append(ParsedEquation(text=eq.strip(), line=i+1, is_block=True))
    
    # Inline equations: $...$
    for i, line in enumerate(lines):
        inlines = re.findall(r'\$(.+?)\$', line)
        for eq in inlines:
            # Skip if it was already captured as block
            if f"$${eq}$$" in line:
                continue
            equations.append(ParsedEquation(text=eq.strip(), line=i+1, is_block=False))
    
    return equations

# ── Stage 1: Equation Classifier ────────────────────────────────────

SPECTRAL_KEYWORDS = [
    r'\sigma', r'\tau', r'\Delta', r'\lambda', r'\gamma', r'\chi',
    r'spectral', r'eigenvalue', r'gap', r'radius', r'threshold',
    r'\d\{1792\}', r'\d\{256\}', r'\d\{273\}', r'\d\{17\}', r'\d\{28\}',
    r'Cartan', r'Sidon', r'chiral', r'rossby', r'kelvin', r'braid',
]

BRAID_KEYWORDS = [
    r'braid', r'strand', r'cross', r'Sidon', r'eigensolid', r'bracket',
    r'Yang-Baxter', r'crossStep', r'residue', r'phase',
]

CARTAN_KEYWORDS = [
    r'Cartan', r'weight', r'matrix', r'diagonal', r'adjacent',
    r'block', r'2\\times2', r'8\\times8', r'Gram',
]

def classify_equation(eq: ParsedEquation) -> ParsedEquation:
    """Classify equation as braid/spectral/cartan based on keyword detection."""
    text = eq.text.lower()
    
    scores = {"spectral": 0, "braid": 0, "cartan": 0}
    
    for kw in SPECTRAL_KEYWORDS:
        if re.search(kw, text, re.IGNORECASE):
            scores["spectral"] += 1
    
    for kw in BRAID_KEYWORDS:
        if re.search(kw, text, re.IGNORECASE):
            scores["braid"] += 1
    
    for kw in CARTAN_KEYWORDS:
        if re.search(kw, text, re.IGNORECASE):
            scores["cartan"] += 1
    
    best = max(scores, key=scores.get)
    if scores[best] > 0:
        eq.classification = best
    else:
        eq.classification = "unknown"
    
    eq.metadata["keyword_scores"] = scores
    return eq

# ── Stage 2: Density Check ──────────────────────────────────────────

def check_density(equations: List[ParsedEquation], max_eqns: int = 20, 
                  max_unknown: int = 5, max_lines: int = 500) -> Dict:
    """Check if the input is too dense/complex to auto-process.
    Returns decision dict with whether to proceed or pause."""
    total = len(equations)
    unknown = sum(1 for e in equations if e.classification == "unknown")
    
    should_pause = (total > max_eqns) or (unknown > max_unknown)
    
    return {
        "total_equations": total,
        "classified": total - unknown,
        "unknown": unknown,
        "max_allowed": max_eqns,
        "should_pause": should_pause,
        "reason": f"too many equations ({total} > {max_eqns})" if total > max_eqns else
                  f"too many unknowns ({unknown} > {max_unknown})" if unknown > max_unknown else
                  "proceed"
    }

def write_decision_doc(equations: List[ParsedEquation], density: Dict, input_path: Path):
    """Write an intermediate decision document for excessive density."""
    out = input_path.with_suffix(".decision.md")
    
    lines = [
        f"# Equation Pipeline — Decision Required",
        f"",
        f"**File:** `{input_path}`",
        f"**Date:** auto-generated",
        f"",
        f"## Density Report",
        f"",
        f"| Metric | Value |",
        f"|--------|-------|",
        f"| Total equations | {density['total_equations']} |",
        f"| Classified | {density['classified']} |",
        f"| Unknown | {density['unknown']} |",
        f"| Max auto-process | {density['max_allowed']} |",
        f"",
        f"**Reason for pause:** {density['reason']}",
        f"",
        f"## Detected Equations",
        f"",
        f"| Line | Type | Classification | Equation |",
        f"|------|------|---------------|----------|",
    ]
    
    for eq in equations:
        cls = eq.classification
        lines.append(f"| {eq.line} | {'block' if eq.is_block else 'inline'} | {cls} | `{eq.text[:60]}{'...' if len(eq.text)>60 else ''}` |")
    
    lines += [
        "",
        "## Decision Required",
        "",
        f"To proceed, run with `--force` or reduce equation count:",
        f"```",
        f"python3 scripts/ingest.py {input_path} --force",
        f"```",
        f"Or process only classified equations:",
        f"```",
        f"python3 scripts/ingest.py {input_path} --classified-only",
        f"```",
    ]
    
    out.write_text("\n".join(lines))
    return out

# ── Stage 3: Compute (delegates to fundamental force pipeline) ───────

def run_pipeline(equations: List[ParsedEquation]) -> List[Dict]:
    """Run each classified equation through the spectral pipeline."""
    from fundamental_force import pipeline as fp_pipeline
    
    results = []
    for eq in equations:
        if eq.classification == "unknown":
            results.append({"equation": eq.text, "line": eq.line, 
                          "result": "skipped", "reason": "unknown classification"})
            continue
        
        # Build input from equation classification
        inp = {
            "strand_count": 8,
            "crossing_pairs": [(0,1),(2,3),(4,5),(6,7)],
            "chiral_labels": ["A"]*8,
            "strand_phases": [0]*8,
            "source_equation": eq.text,
            "source_line": eq.line,
        }
        
        # Spectral equations: vary chiral labels based on detected keywords
        if "rossby" in eq.text.lower() or "chiral" in eq.text.lower():
            inp["chiral_labels"] = ["L"]*8
        
        if "scarred" in eq.text.lower():
            inp["chiral_labels"] = ["S"]*8
        
        result = fp_pipeline(inp)
        results.append({
            "equation": eq.text, "line": eq.line,
            "classification": eq.classification,
            "fingerprint": result["derived"]
        })
    
    return results

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

def main():
    parser = argparse.ArgumentParser(description="Equation ingestion pipeline")
    parser.add_argument("input", type=str, help="Markdown file with equations")
    parser.add_argument("--max-eqns", type=int, default=20, help="Max equations before pausing")
    parser.add_argument("--max-unknown", type=int, default=5, help="Max unknown equations before pausing")
    parser.add_argument("--force", action="store_true", help="Process even if too dense")
    parser.add_argument("--classified-only", action="store_true", help="Skip unknown equations")
    
    args = parser.parse_args()
    input_path = Path(args.input)
    
    if not input_path.exists():
        print(f"Error: {input_path} not found")
        sys.exit(1)
    
    text = input_path.read_text()
    
    # Parse
    equations = parse_markdown(text)
    print(f"Parsed {len(equations)} equations from {input_path}")
    
    # Classify
    equations = [classify_equation(e) for e in equations]
    spectral = sum(1 for e in equations if e.classification == "spectral")
    braid = sum(1 for e in equations if e.classification == "braid")
    cartan = sum(1 for e in equations if e.classification == "cartan")
    unknown = sum(1 for e in equations if e.classification == "unknown")
    print(f"  Spectral: {spectral}, Braid: {braid}, Cartan: {cartan}, Unknown: {unknown}")
    
    # Density check
    density = check_density(equations, args.max_eqns, args.max_unknown)
    
    if density["should_pause"] and not args.force:
        doc = write_decision_doc(equations, density, input_path)
        print(f"\n⚠️ Too dense — wrote decision document:")
        print(f"   {doc}")
        print(f"   Review and re-run with --force to proceed.")
        return
    
    # Filter
    if args.classified_only:
        equations = [e for e in equations if e.classification != "unknown"]
        print(f"Processing {len(equations)} classified equations")
    
    # Compute
    results = run_pipeline(equations)
    
    # Emit
    receipt = {
        "schema": "equation_ingestion_v1",
        "source_file": str(input_path),
        "total_parsed": len(equations),
        "total_processed": len(results),
        "density": density,
        "results": results
    }
    
    out_path = input_path.with_suffix(".receipt.json")
    out_path.write_text(json.dumps(receipt, indent=2))
    
    print(f"\n✅ Pipeline complete")
    print(f"   Processed: {len(results)} equations")
    print(f"   Receipt: {out_path}")
    
    # Summary
    computed = [r for r in results if "fingerprint" in r]
    if computed:
        gaps = [r["fingerprint"]["delta"][0] for r in computed]
        regimes = {r["fingerprint"]["delta"][0]: "CANONICAL" if r["fingerprint"]["delta"][0] > 0 else "ROSSBY" 
                   for r in computed}
        print(f"   Gaps: {[f'{g:.6f}' for g in gaps]}")
        print(f"   Regimes: {list(regimes.values())}")

if __name__ == "__main__":
    main()