#!/usr/bin/env python3
"""
geometric_entropy_explorer.py — Exploration-phase candidate generator
for the Rainbow Raccoon Compiler (RRC).

Places 8 points (braid strands) on a torus and maximizes Shannon entropy
of the pairwise-distance distribution via gradient descent (TF.js-compatible).
Exports candidate BraidReceipt JSON for downstream Lean certification.

Usage:
    python3 geometric_entropy_explorer.py                        # run with defaults
    python3 geometric_entropy_explorer.py --manifold sphere       # try different manifold
    python3 geometric_entropy_explorer.py --export candidates.json

Architecture (per the DP-RRC spec):
    ┌─ Exploration phase ──────────────────────────────────┐
    │  torus → entropy maximization → candidate point cloud │
    │  → export candidate receipt JSON                     │
    └──────────────────────────┬───────────────────────────┘
                               │ candidate (JSON)
                               ▼
    ┌─ Certification phase (Lean) ────────────────────────┐
    │  crossStep verification → eigensolid_convergence     │
    │  → receipt_invertible theorem → AVM stamp             │
    └──────────────────────────────────────────────────────┘

This script is pure I/O + feature extraction (Python-owned per AGENTS.md).
No gating, alignment, or promotion decisions are made here.
"""

import numpy as np
import json
import os
import sys
import argparse
from dataclasses import dataclass, field
from typing import List, Optional, Tuple

# ---------------------------------------------------------------------------
# Q16_16 helpers (fixed-point matching the Lean representation)
# ---------------------------------------------------------------------------

def to_q16(x: float) -> int:
    """Float → Q16_16 signed 32-bit integer."""
    return int(round(x * 65536.0))

def from_q16(q: int) -> float:
    """Q16_16 → float."""
    return q / 65536.0

Q16_ONE = to_q16(1.0)
Q16_PI = to_q16(np.pi)
Q16_PI_4 = to_q16(np.pi / 4)

# ---------------------------------------------------------------------------
# Manifold parameterizations
# ---------------------------------------------------------------------------

class Manifold:
    """Base class for constraint manifolds (torus, sphere, cube, etc.)."""

    def sample(self, n: int, rng: np.random.Generator) -> np.ndarray:
        """Sample n random points on the manifold. Returns (n, 3)."""
        raise NotImplementedError

    def project(self, points: np.ndarray) -> np.ndarray:
        """Project points onto the manifold surface."""
        raise NotImplementedError

    def __str__(self) -> str:
        return self.__class__.__name__


class Torus(Manifold):
    """Torus of revolution: major radius R, minor radius r."""

    def __init__(self, R: float = 2.0, r: float = 1.0):
        self.R = R
        self.r = r

    def _uv_to_xyz(self, u: np.ndarray, v: np.ndarray) -> np.ndarray:
        """(u, v) in [0, 2π)² → (x, y, z) on torus."""
        x = (self.R + self.r * np.cos(v)) * np.cos(u)
        y = (self.R + self.r * np.cos(v)) * np.sin(u)
        z = self.r * np.sin(v)
        return np.stack([x, y, z], axis=-1)

    def sample(self, n: int, rng: np.random.Generator) -> np.ndarray:
        u = rng.uniform(0, 2 * np.pi, size=n)
        v = rng.uniform(0, 2 * np.pi, size=n)
        return self._uv_to_xyz(u, v)

    def project(self, points: np.ndarray) -> np.ndarray:
        """Project (x, y, z) onto nearest point on torus surface."""
        x, y, z = points[..., 0], points[..., 1], points[..., 2]
        phi = np.arctan2(y, x)
        # distance from central circle axis
        d_xy = np.sqrt(x**2 + y**2)
        theta = np.arctan2(z, d_xy - self.R)
        return self._uv_to_xyz(phi, theta)


class Sphere(Manifold):
    """Unit sphere S²."""

    def __init__(self, radius: float = 1.0):
        self.radius = radius

    def sample(self, n: int, rng: np.random.Generator) -> np.ndarray:
        # Normal distribution → normalize to sphere surface
        pts = rng.normal(size=(n, 3))
        norms = np.linalg.norm(pts, axis=-1, keepdims=True)
        return self.radius * pts / norms

    def project(self, points: np.ndarray) -> np.ndarray:
        norms = np.linalg.norm(points, axis=-1, keepdims=True)
        return self.radius * points / norms


class CubeShell(Manifold):
    """Cube surface (6 faces)."""

    def __init__(self, side: float = 2.0):
        self.side = side
        self.half = side / 2

    def sample(self, n: int, rng: np.random.Generator) -> np.ndarray:
        # Pick a random face, then random 2D coordinates on it
        pts = np.zeros((n, 3))
        faces = rng.integers(0, 6, size=n)
        for i, f in enumerate(faces):
            coord = rng.uniform(-self.half, self.half, size=2)
            if f == 0: pts[i] = [self.half, coord[0], coord[1]]
            if f == 1: pts[i] = [-self.half, coord[0], coord[1]]
            if f == 2: pts[i] = [coord[0], self.half, coord[1]]
            if f == 3: pts[i] = [coord[0], -self.half, coord[1]]
            if f == 4: pts[i] = [coord[0], coord[1], self.half]
            if f == 5: pts[i] = [coord[0], coord[1], -self.half]
        return pts

    def project(self, points: np.ndarray) -> np.ndarray:
        # Clamp to cube surface (closest face)
        return np.clip(points, -self.half, self.half)


# ---------------------------------------------------------------------------
# Entropy computation (matching the geometric-entropy-lab approach)
# ---------------------------------------------------------------------------

def pairwise_distances(points: np.ndarray) -> np.ndarray:
    """(n, 3) → (n, n) Euclidean distance matrix."""
    diff = points[:, np.newaxis, :] - points[np.newaxis, :, :]
    return np.sqrt(np.sum(diff**2, axis=-1))


def gaussian_kde_entropy(
    distances: np.ndarray,
    bandwidth: float = 0.3,
    temperature: float = 1.0,
) -> float:
    """
    Shannon entropy of the pairwise-distance distribution using Gaussian KDE,
    matching the geometric-entropy-lab approach:

        G = D² (Gram matrix of squared distances)
        ρ = softmax(G / τ)  (density via softmax)
        H = -Σ p·log(p)     (Shannon entropy)

    The lab uses dot products; we use squared distances, which is equivalent
    for centered point clouds.
    """
    n = distances.shape[0]
    # Gaussian kernel over squared distances
    D2 = distances**2
    K = np.exp(-D2 / (2 * bandwidth**2))
    # Density via softmax over kernel matrix
    K_scaled = K / temperature
    K_max = np.max(K_scaled, axis=-1, keepdims=True)
    K_stable = K_scaled - K_max
    exp_K = np.exp(K_stable)
    rho = exp_K / np.sum(exp_K, axis=-1, keepdims=True)
    # Shannon entropy: H = -Σ p·log(p)
    p = np.mean(rho, axis=0)
    p = p / np.sum(p)
    H = -np.sum(p * np.log(p + 1e-30))
    return float(H)


def entropy_gradient(
    points: np.ndarray,
    bandwidth: float = 0.3,
    temperature: float = 1.0,
    eps: float = 1e-6,
) -> np.ndarray:
    """
    Numerical gradient of entropy w.r.t. point positions via central differences.
    Returns (n, 3) gradient: dH/dx_i.
    """
    grad = np.zeros_like(points)
    D = pairwise_distances(points)
    H0 = gaussian_kde_entropy(D, bandwidth, temperature)
    for i in range(points.shape[0]):
        for j in range(3):
            points[i, j] += eps
            Dp = pairwise_distances(points)
            Hp = gaussian_kde_entropy(Dp, bandwidth, temperature)
            points[i, j] -= 2 * eps
            Dm = pairwise_distances(points)
            Hm = gaussian_kde_entropy(Dm, bandwidth, temperature)
            points[i, j] += eps
            grad[i, j] = (Hp - Hm) / (2 * eps)
    return grad


# ---------------------------------------------------------------------------
# BraidStrand mapping: point on manifold → BraidStrand parameters
# ---------------------------------------------------------------------------

def points_to_braid_state(
    points: np.ndarray,
    slots: Optional[List[int]] = None,
) -> dict:
    """
    Map (n, 3) point cloud on torus to BraidState-compatible dict.

    Encoding (per DP-RRC spec):
        - point spherical angles → PhaseVec (x, y)
        - Sidon labels from toroidal coordinate quanta
        - kappa from pairwise distance entropy gradient
        - bracket from PhaseVec via fromPhaseVec equivalent
    """
    n = points.shape[0]
    assert n == 8, f"BraidStorm requires exactly 8 strands, got {n}"

    if slots is None:
        slots = [1, 2, 4, 8, 16, 32, 64, 128]

    # Normalize to unit sphere for phase angles
    norms = np.linalg.norm(points, axis=-1, keepdims=True)
    unit = points / (norms + 1e-30)

    # Theta (polar) and phi (azimuthal) as PhaseVec (x, y) in Q16_16
    theta = np.arccos(np.clip(unit[:, 2], -1.0, 1.0))
    phi = np.arctan2(unit[:, 1], unit[:, 0])

    # Pairwise distances for kappa computation (octagonal norm analog)
    D = pairwise_distances(points)
    # kappa = normalized mean distance to nearest neighbor (like octagonal norm)
    diag_mask = np.eye(n, dtype=bool)
    D_masked = D.copy()
    D_masked[diag_mask] = np.inf
    min_d = np.min(D_masked, axis=1)
    kappa_vals = min_d / np.max(min_d + 1e-30)

    # Compute bracket kappa as octagonal norm equivalent
    bracket_kappa = float(np.mean(D[~diag_mask]))

    strands = []
    for i in range(n):
        phase_vec = {
            "x": to_q16(float(np.sin(theta[i]) * np.cos(phi[i]))),
            "y": to_q16(float(np.sin(theta[i]) * np.sin(phi[i]))),
        }
        mu = slots[i]
        kappa_q = to_q16(float(kappa_vals[i]))
        # BraidBracket: lower = κ - μ, upper = κ + μ, gap = 2μ
        lower_q = to_q16(float(kappa_vals[i] - 0.1 * mu / 128.0))
        upper_q = to_q16(float(kappa_vals[i] + 0.1 * mu / 128.0))
        gap_q = to_q16(float(0.2 * mu / 128.0))
        admissible = lower_q <= upper_q

        strands.append({
            "phaseAcc": phase_vec,
            "parity": bool(i % 2),
            "slot": slots[i],
            "residue": 0,
            "jitter": 0,
            "bracket": {
                "lower": lower_q,
                "upper": upper_q,
                "gap": gap_q,
                "kappa": kappa_q,
                "phi": Q16_PI_4 if kappa_q != 0 else 0,
                "admissible": admissible,
            }
        })

    # Sidon slack: budget - max label used
    sidon_slack = 128 - max(slots)

    return {
        "strands": strands,
        "bracket_kappa": to_q16(float(bracket_kappa)),
        "sidon_slack": sidon_slack,
    }


# ---------------------------------------------------------------------------
# Gradient descent optimizer (entropy maximization)
# ---------------------------------------------------------------------------

def optimize_entropy(
    manifold: Manifold,
    n_points: int = 8,
    n_steps: int = 200,
    lr: float = 0.1,
    bandwidth: float = 0.3,
    temperature: float = 1.0,
    seed: Optional[int] = None,
    verbose: bool = True,
    cluster_init: bool = True,
) -> Tuple[np.ndarray, List[float]]:
    """
    Run gradient descent to maximize Shannon entropy of pairwise distances
    on the given manifold.

    Strategy: start with a clustered initialization (low entropy), then
    maximize entropy to spread points out. This gives a clear gradient signal.

    Returns:
        points: (n, 3) optimized point cloud
        history: [H_0, H_1, ..., H_n_steps] entropy trace
    """
    rng = np.random.default_rng(seed)
    if cluster_init:
        # Start all points in a tight cluster → low entropy → strong gradient
        center = manifold.sample(1, rng)[0]
        points = center + rng.normal(0, 0.05, size=(n_points, 3))
        points = manifold.project(points)
    else:
        points = manifold.sample(n_points, rng)
    history = []

    for step in range(n_steps):
        D = pairwise_distances(points)
        H = gaussian_kde_entropy(D, bandwidth, temperature)
        history.append(H)

        if step % 20 == 0 and verbose:
            print(f"  step {step:4d}: H = {H:.6f}  (spread: {float(np.mean(D[~np.eye(n_points, dtype=bool)])):.4f})")

        if step == n_steps - 1:
            break

        grad = entropy_gradient(points, bandwidth, temperature)
        # Gradient ascent (maximize entropy)
        points = points + lr * grad
        # Project back onto manifold
        points = manifold.project(points)
        # Repulsion regularizer: prevent collapse
        D_self = pairwise_distances(points)
        np.fill_diagonal(D_self, np.inf)
        min_sep = np.min(D_self)
        if min_sep < 0.05:
            for i in range(n_points):
                for j in range(n_points):
                    if i != j:
                        diff = points[i] - points[j]
                        dist = np.linalg.norm(diff)
                        if 0 < dist < 0.2:
                            repel = 0.02 * diff / (dist + 1e-30)
                            points[i] += repel
                            points[j] -= repel
            points = manifold.project(points)

    return points, history


# ---------------------------------------------------------------------------
# Candidate export (bridge to Lean certification pipeline)
# ---------------------------------------------------------------------------

def export_candidate(
    points: np.ndarray,
    manifold: Manifold,
    entropy_history: List[float],
    equation_id: str = "rrc_eq_entropy_explorer",
    output_path: Optional[str] = None,
    bandwidth: float = 0.3,
    temperature: float = 1.0,
) -> dict:
    """
    Export a candidate receipt JSON that the Lean pipeline can consume.

    Format matches the BraidReceipt structure from BraidEigensolid.lean
    plus provenance metadata for the exploration phase.
    """
    braid_state = points_to_braid_state(points)
    final_entropy = entropy_history[-1] if entropy_history else 0.0

    # Serialize manifold params safely
    if isinstance(manifold, Torus):
        mparams = {"R": manifold.R, "r": manifold.r}
    elif isinstance(manifold, Sphere):
        mparams = {"radius": manifold.radius}
    elif isinstance(manifold, CubeShell):
        mparams = {"side": manifold.side}
    else:
        mparams = {}

    candidate = {
        "schema": "rrc_candidate_entropy_v1",
        "claim_boundary": "exploration-phase-only;not-certified",
        "genesis": {
            "method": "entropy_maximization",
            "manifold": str(manifold),
            "manifold_params": mparams,
            "entropy_final": round(final_entropy, 6),
            "entropy_history": [round(h, 6) for h in entropy_history],
            "bandwidth": bandwidth,
            "temperature": temperature,
        },
        "braid_state": braid_state,
        "equation_id": equation_id,
        "sidon_slack": braid_state["sidon_slack"],
    }

    if output_path:
        os.makedirs(os.path.dirname(output_path) or ".", exist_ok=True)
        with open(output_path, "w") as f:
            json.dump(candidate, f, indent=2)
        print(f"Exported candidate to {output_path}")

    return candidate


# ---------------------------------------------------------------------------
# Main
# ---------------------------------------------------------------------------

def main():
    parser = argparse.ArgumentParser(
        description="Geometric Entropy Explorer — RRC candidate generator"
    )
    parser.add_argument(
        "--manifold", choices=["torus", "sphere", "cube"],
        default="torus", help="Constraint manifold"
    )
    parser.add_argument("--n-strands", type=int, default=8)
    parser.add_argument("--steps", type=int, default=200)
    parser.add_argument("--lr", type=float, default=0.1)
    parser.add_argument("--seed", type=int, default=None)
    parser.add_argument("--export", type=str, default=None,
                        help="Export candidate JSON to path (single)")
    parser.add_argument("--batch", type=int, default=None,
                        help="Run N random seeds, export best candidates")
    parser.add_argument("--equation-id", type=str,
                        default="rrc_eq_entropy_explorer",
                        help="Equation ID for the candidate")
    parser.add_argument("--output-dir", type=str,
                        default=None,
                        help="Output directory for candidates")
    args = parser.parse_args()

    # Default output dir
    if args.output_dir is None:
        args.output_dir = (
            "/home/allaun/Research Stack/shared-data/data/stack_solidification/candidates"
        )

    manifolds = {
        "torus": Torus(R=2.0, r=1.0),
        "sphere": Sphere(radius=2.0),
        "cube": CubeShell(side=3.0),
    }
    manifold = manifolds[args.manifold]

    if args.batch:
        # Batch mode: run multiple seeds, pick best by final entropy
        print(f"Batch exploration: {args.batch} runs on {manifold}")
        all_candidates = []
        for trial in range(args.batch):
            trial_seed = (args.seed or 0) + trial
            pts, hist = optimize_entropy(
                manifold=manifold,
                n_points=args.n_strands,
                n_steps=args.steps,
                lr=args.lr,
                seed=trial_seed,
                verbose=False,
            )
            final_H = hist[-1]
            cand = export_candidate(
                points=pts,
                manifold=manifold,
                entropy_history=hist,
                equation_id=f"{args.equation_id}_seed{trial_seed}",
                bandwidth=0.3,
                temperature=1.0,
            )
            all_candidates.append((final_H, cand, pts))
            print(f"  trial {trial:3d} (seed {trial_seed:3d}): H = {final_H:.6f}")

        # Sort by entropy descending
        all_candidates.sort(key=lambda x: -x[0])

        # Export all to batch dir
        batch_dir = os.path.join(args.output_dir, f"batch_{args.manifold}")
        os.makedirs(batch_dir, exist_ok=True)

        best = []
        for rank, (h, cand, pts) in enumerate(all_candidates):
            fname = f"candidate_{args.manifold}_rank{rank:03d}_seed{args.seed + rank if args.seed else rank}.json"
            path = os.path.join(batch_dir, fname)
            with open(path, "w") as f:
                json.dump(cand, f, indent=2)
            best.append({
                "rank": rank,
                "entropy": round(h, 6),
                "file": fname,
                "sidon_slack": cand["sidon_slack"],
            })

        # Write manifest
        manifest = {
            "schema": "rrc_candidate_batch_manifest_v1",
            "claim_boundary": "exploration-phase-only;not-certified",
            "manifold": str(manifold),
            "n_trials": args.batch,
            "candidates": best,
        }
        manifest_path = os.path.join(batch_dir, "manifest.json")
        with open(manifest_path, "w") as f:
            json.dump(manifest, f, indent=2)

        print(f"\nBatch complete. {args.batch} candidates -> {batch_dir}/")
        print(f"Best entropy: H = {best[0]['entropy']}")
        print(f"Worst entropy: H = {best[-1]['entropy']}")
        print(f"Manifest: {manifest_path}")

    else:
        # Single run
        print(f"Running entropy exploration on {manifold} with {args.n_strands} strands")

        points, history = optimize_entropy(
            manifold=manifold,
            n_points=args.n_strands,
            n_steps=args.steps,
            lr=args.lr,
            seed=args.seed,
        )
        print(f"\nFinal entropy: H = {history[-1]:.6f}")

        if args.export:
            candidate = export_candidate(
                points=points,
                manifold=manifold,
                entropy_history=history,
                equation_id=args.equation_id,
                output_path=args.export,
            )
        else:
            os.makedirs(args.output_dir, exist_ok=True)
            candidate = export_candidate(
                points=points,
                manifold=manifold,
                entropy_history=history,
                equation_id=args.equation_id,
                output_path=os.path.join(
                    args.output_dir,
                    f"candidate_{manifold}_{args.seed or 0}.json",
                )
            )

        print(f"\nCandidate braid state:")
        print(f"  Sidon slack: σ = {candidate['braid_state']['sidon_slack']}")
        slots = [s["slot"] for s in candidate["braid_state"]["strands"]]
        print(f"  Slot labels: {slots}")
        admissibility = [
            "✓" if s["bracket"]["admissible"] else "✗"
            for s in candidate["braid_state"]["strands"]
        ]
        print(f"  Admissible:  {''.join(admissibility)}")

    print(f"\nTo certify candidates:")
    print(f"  lake build Semantics.AVMIsa.Emit")
    print(f"  python3 4-Infrastructure/shim/emit250_extract.py")


if __name__ == "__main__":
    main()