#!/usr/bin/env python3
"""
hachimoji_citation.py — Equation → Hachimoji State → arxiv Citations

Pipeline:
  equation string
      → parse shape (n_vars, n_ops, max_depth, n_quantifiers, n_relations)
      → classify to one of 8 Hachimoji states (mirrors HachimojiCodec.lean)
      → build semantic query from state meaning
      → hybrid_search(query_text, embedding, top_k) via arxiv-pg podman container
      → return top-k arxiv citations for that equation's regime

DB: podman container 'arxiv-pg', database 'arxiv'.
Embedding: sentence-transformers/static-retrieval-mrl-en-v1 (1024-dim).
Search: hybrid_search() — RRF merger of pg_trgm + pgvector HNSW (22ms).

Mirrors formal definitions in:
  formal/CoreFormalism/HachimojiCodec.lean  — classifyEquation, admission
  formal/CoreFormalism/HachimojiLUT.lean    — equationPosition
"""

from __future__ import annotations

import json
import re
import subprocess
from dataclasses import dataclass, field
from typing import Optional

MODEL_NAME = "sentence-transformers/static-retrieval-mrl-en-v1"

# ── State semantic queries ───────────────────────────────────────────────────
# Each Hachimoji state maps to a natural-language query that retrieves
# papers relevant to that regime from the arxiv corpus.
# Source: HachimojiBase.lean §4 semantic cross-reference.

STATE_QUERIES: dict[str, str] = {
    "Φ": (
        "density boundary ordered lattice phase transition Erdős-Rényi "
        "trivial regime fundamental equation constants few variables"
    ),
    "Λ": (
        "Barnes-Wall lattice room regime bounded quantifiers linear algebra "
        "shallow depth ordered exploration"
    ),
    "Ρ": (
        "spectral radius stability boundary tight regime operator norm "
        "high operator count convergence fixed point"
    ),
    "Κ": (
        "complementarity threshold marginal regime softplus many variables "
        "near threshold complementarity duality"
    ),
    "Ω": (
        "Goormaghtigh collision eigensolid terminal fixed point Baker bound "
        "contradiction repunit equation zero lambda"
    ),
    "Σ": (
        "Goormaghtigh symmetric partner solution swap algebraic geometry "
        "symmetric balanced equation equal structure"
    ),
    "Π": (
        "potential violation density probe Baker threshold calculus "
        "high complexity integral limit series path integral"
    ),
    "Ζ": (
        "Riemann zeta near cancellation zero region near zero no integer point "
        "fallback empty undefined bare expression"
    ),
}

# ── Equation shape (mirrors HachimojiCodec.lean EquationShape) ──────────────

@dataclass
class EquationShape:
    n_vars:        int = 0
    n_ops:         int = 0
    max_depth:     int = 0
    n_quantifiers: int = 0
    n_relations:   int = 0


# ── Classifier (mirrors HachimojiCodec.lean classifyEquation) ───────────────
# MUST stay in sync with the Lean definition.  If HachimojiCodec.lean changes,
# update this function and add a note here.

def classify_equation(s: EquationShape) -> str:
    """Return the Hachimoji state label for an EquationShape.

    Exact mirror of HachimojiCodec.lean §8 classifyEquation, order preserved:
      1. Ω  contradiction (n_vars=0, n_ops=0, n_relations≥1)
      2. Λ  bounded quantifiers, shallow depth
      3. Ζ  empty/bare
      4. Φ  fundamental (few vars, few ops, has relation)
      5. Π  high complexity
      6. Σ  symmetric structure
      7. Ρ  high ops, no quantifiers
      8. Κ  many variables, shallow
      9. Ζ  default fallback
    """
    if s.n_vars == 0 and s.n_ops == 0 and s.n_relations >= 1:
        return "Ω"
    if s.n_quantifiers > 0 and s.max_depth <= 2:
        return "Λ"
    if s.n_vars <= 1 and s.n_ops == 0 and s.n_relations == 0:
        return "Ζ"
    if s.n_vars <= 3 and s.n_quantifiers == 0 and s.n_ops <= 5 and s.n_relations >= 1:
        return "Φ"
    if s.n_ops + s.n_vars * s.max_depth + s.n_quantifiers * 2 >= 8 or s.n_ops > 8:
        return "Π"
    is_symmetric = (
        s.n_relations >= 1 and s.n_vars >= 2
        and 1 <= s.n_ops <= 10 and s.n_quantifiers == 0
    )
    if is_symmetric:
        return "Σ"
    if s.n_ops > 5 and s.n_quantifiers == 0:
        return "Ρ"
    if s.n_vars > 5 and s.max_depth <= 1:
        return "Κ"
    return "Ζ"


def admission(state: str) -> str:
    """Mirror of HachimojiCodec.lean §4 admission.

    ADMIT:      Φ, Λ, Ρ, Κ, Σ
    QUARANTINE: Ω, Π, Ζ
    """
    return "ADMIT" if state in {"Φ", "Λ", "Ρ", "Κ", "Σ"} else "QUARANTINE"


# ── Simple equation shape parser ────────────────────────────────────────────

def parse_shape(equation: str) -> EquationShape:
    """Heuristic shape extraction from a raw equation string.

    Good enough for classification; not a full parse.  Replace with a proper
    AST parser when available.
    """
    ops      = len(re.findall(r'[+\-*/^√∫∑∏∂∇]', equation))
    vars_    = len(set(re.findall(r'\b[a-zA-Z]\b', equation)))
    rels     = len(re.findall(r'[=<>≤≥≠≈]', equation))
    quants   = len(re.findall(r'[∀∃]', equation))
    # depth: count nested parens / brackets
    depth    = 0
    max_d    = 0
    for ch in equation:
        if ch in '([{':
            depth += 1
            max_d = max(max_d, depth)
        elif ch in ')]}':
            depth = max(0, depth - 1)
    return EquationShape(
        n_vars=vars_,
        n_ops=ops,
        max_depth=max_d,
        n_quantifiers=quants,
        n_relations=rels,
    )


# ── DB helpers ───────────────────────────────────────────────────────────────

def psql(sql: str, timeout: int = 30) -> tuple[str, str]:
    """Run SQL in the arxiv-pg podman container."""
    r = subprocess.run(
        ['podman', 'exec', 'arxiv-pg', 'psql',
         '-U', 'postgres', '-d', 'arxiv', '-t', '-A', '-c', sql],
        capture_output=True, text=True, timeout=timeout,
    )
    return r.stdout.strip(), r.stderr.strip()


def embed(text: str) -> list[float]:
    """Embed a query string with the canonical model (1024-dim, normalised)."""
    from sentence_transformers import SentenceTransformer
    model = SentenceTransformer(MODEL_NAME)
    vec = model.encode(text, normalize_embeddings=True)
    return vec.tolist()


# ── Citation dataclass ───────────────────────────────────────────────────────

@dataclass
class Citation:
    paper_id:      str
    title:         str
    rrf_score:     float
    trigram_rank:  int
    vector_rank:   int


# ── Main entry point ─────────────────────────────────────────────────────────

@dataclass
class EquationCitation:
    equation:    str
    shape:       EquationShape
    state:       str
    admission:   str
    phase_deg:   int
    query:       str
    citations:   list[Citation] = field(default_factory=list)


STATE_PHASE: dict[str, int] = {
    "Φ": 0, "Λ": 45, "Ρ": 90, "Κ": 135,
    "Ω": 180, "Σ": 225, "Π": 270, "Ζ": 315,
}


def cite_equation(
    equation: str,
    shape: Optional[EquationShape] = None,
    top_k: int = 5,
    use_embedding: bool = True,
) -> EquationCitation:
    """Classify an equation and retrieve its top-k arxiv citations.

    Args:
        equation:      Raw equation string (LaTeX or ASCII math).
        shape:         Pre-computed shape; parsed from equation if None.
        top_k:         Number of citations to retrieve.
        use_embedding: If True, use semantic embedding (requires model);
                       if False, trigram-only (faster, no model needed).

    Returns:
        EquationCitation with state, phase, admission, and top-k papers.
    """
    if shape is None:
        shape = parse_shape(equation)

    state = classify_equation(shape)
    adm   = admission(state)
    phase = STATE_PHASE[state]
    query = STATE_QUERIES[state]

    # Augment query with the equation itself (first 200 chars)
    full_query = f"{query} {equation[:200]}".strip()
    full_query_esc = full_query.replace("'", "''")

    if use_embedding:
        emb   = embed(full_query)
        emb_s = '[' + ','.join(f'{x:.6f}' for x in emb) + ']'
        sql   = (
            f"SELECT paper_id, title, trigram_rank, vector_rank, rrf_score "
            f"FROM hybrid_search('{full_query_esc}', '{emb_s}'::vector(1024), {top_k})"
        )
    else:
        # trigram-only fallback via ts_rank
        sql = (
            f"SELECT paper_id, title, "
            f"row_number() OVER () AS trigram_rank, "
            f"0 AS vector_rank, "
            f"ts_rank(to_tsvector('english', title || ' ' || COALESCE(abstract, '')), "
            f"        plainto_tsquery('english', '{full_query_esc}')) AS rrf_score "
            f"FROM arxiv_papers "
            f"WHERE to_tsvector('english', title || ' ' || COALESCE(abstract, '')) "
            f"      @@ plainto_tsquery('english', '{full_query_esc}') "
            f"ORDER BY rrf_score DESC LIMIT {top_k}"
        )

    out, err = psql(sql)
    citations: list[Citation] = []
    if err and "ERROR" in err:
        print(f"DB error: {err[:120]}")
    else:
        for line in out.split("\n"):
            if not line.strip():
                continue
            parts = line.split("|")
            if len(parts) >= 5:
                try:
                    trig = int(parts[2]) if parts[2].strip() else 0
                    vec = int(parts[3]) if parts[3].strip() else 0
                    citations.append(Citation(
                        paper_id     = parts[0].strip(),
                        title        = parts[1].strip()[:200],
                        trigram_rank = trig,
                        vector_rank  = vec,
                        rrf_score    = float(parts[4].strip()),
                    ))
                except (ValueError, IndexError):
                    continue

    return EquationCitation(
        equation  = equation,
        shape     = shape,
        state     = state,
        admission = adm,
        phase_deg = phase,
        query     = full_query[:120],
        citations = citations,
    )


def cite_batch(
    equations: list[str],
    top_k: int = 5,
    use_embedding: bool = True,
) -> list[EquationCitation]:
    """Classify and cite a batch of equations.

    Loads the embedding model once for efficiency.
    """
    if use_embedding:
        from sentence_transformers import SentenceTransformer
        model = SentenceTransformer(MODEL_NAME)
    else:
        model = None

    results = []
    for eq in equations:
        shape = parse_shape(eq)
        state = classify_equation(shape)
        query = f"{STATE_QUERIES[state]} {eq[:200]}".strip()
        query_esc = query.replace("'", "''")

        if model is not None:
            emb   = model.encode(query, normalize_embeddings=True).tolist()
            emb_s = '[' + ','.join(f'{x:.6f}' for x in emb) + ']'
            sql   = (
                f"SELECT paper_id, title, trigram_rank, vector_rank, rrf_score "
                f"FROM hybrid_search('{query_esc}', '{emb_s}'::vector(1024), {top_k})"
            )
        else:
            sql = (
                f"SELECT paper_id, title, 0, 0, "
                f"ts_rank(to_tsvector('english', title), "
                f"plainto_tsquery('english', '{query_esc}')) "
                f"FROM arxiv_papers "
                f"WHERE to_tsvector('english', title) @@ plainto_tsquery('english', '{query_esc}') "
                f"ORDER BY 5 DESC LIMIT {top_k}"
            )

        out, err = psql(sql)
        citations: list[Citation] = []
        for line in out.split("\n"):
            parts = line.split("|")
            if len(parts) >= 5:
                try:
                    trig = int(parts[2]) if parts[2].strip() else 0
                    vec = int(parts[3]) if parts[3].strip() else 0
                    citations.append(Citation(
                        paper_id=parts[0].strip(), title=parts[1].strip()[:200],
                        trigram_rank=trig,
                        vector_rank=vec,
                        rrf_score=float(parts[4].strip()),
                    ))
                except (ValueError, IndexError):
                    continue

        results.append(EquationCitation(
            equation=eq, shape=shape, state=state,
            admission=admission(state),
            phase_deg=STATE_PHASE[state],
            query=query[:120], citations=citations,
        ))
    return results


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

if __name__ == "__main__":
    import argparse, sys

    parser = argparse.ArgumentParser(
        description="Classify an equation and retrieve arxiv citations.")
    parser.add_argument("equation", nargs="?",
                        help="Equation string (quote it)")
    parser.add_argument("--top-k", type=int, default=5)
    parser.add_argument("--no-embed", action="store_true",
                        help="Trigram-only; skip embedding model")
    parser.add_argument("--json", action="store_true",
                        help="Output JSON")
    args = parser.parse_args()

    if not args.equation:
        # Demo batch
        demo = [
            "E = mc^2",
            "a^2 + b^2 = c^2",
            "∀x. P(x) → Q(x)",
            "0 = 1",
            "∫ f(x) dx = F(x) + C",
        ]
        results = cite_batch(demo, top_k=args.top_k,
                             use_embedding=not args.no_embed)
    else:
        results = [cite_equation(args.equation, top_k=args.top_k,
                                 use_embedding=not args.no_embed)]

    if args.json:
        def _serial(obj):
            if hasattr(obj, '__dict__'):
                return obj.__dict__
            return str(obj)
        print(json.dumps([r.__dict__ for r in results], default=_serial, indent=2))
    else:
        for r in results:
            print(f"\n{'='*60}")
            print(f"Equation : {r.equation}")
            print(f"State    : {r.state}  phase={r.phase_deg}°  admission={r.admission}")
            print(f"Query    : {r.query[:80]}…")
            print(f"Citations:")
            if not r.citations:
                print("  (none)")
            for i, c in enumerate(r.citations, 1):
                print(f"  {i}. [{c.rrf_score:.4f}] {c.title} ({c.paper_id})")
