#!/usr/bin/env python3
"""
RAIL / RomanAILabs - "Run the script in 4D" (practical meaning)

This script runs your FPMF computation on a 4D state-space trajectory:
- A 4D vector (x,y,z,w) evolves along an outer parameter σ (your contour Γ).
- Inner evolution parameter τ is used to compute frenetic activity:
      A(σ) = ∫_R γ(ψdot, φ) dτ
- Phase:
      Θ(σ) = λ α Γ_term(σ) + β κ(σ) Π_term(σ)
- Your bridge (axiom):
      exp(iΘ) = exp(-A) * K   =>   K_bridge = exp(A) * exp(iΘ)
- Also prints a physically useful damped propagator:
      P_prop = exp(-A) * exp(iΘ)

Outputs:
- F (complex), |F|
- E_mis (real mismatch energy)
- CI (coherence index)
- K_bridge stats and P_prop stats

IMPORTANT:
If A gets large, K_bridge magnitude MUST explode because |K_bridge| = exp(A).
That is not a bug; it's your axiom. If you want K to stay bounded, use P_prop
as the transport/propagator object, or keep A small via gamma normalization.

Copyright Daniel Harding - RomanAILabs
Credits: Drafted with assistance from Nova (GPT-5.2 Thinking)
"""

from __future__ import annotations

import argparse
import json
import os
from dataclasses import dataclass
from typing import Dict, Tuple

import numpy as np


# -----------------------------
# 4D math primitives
# -----------------------------

@dataclass(frozen=True)
class Vec4:
    x: float
    y: float
    z: float
    w: float

    def as_np(self) -> np.ndarray:
        return np.array([self.x, self.y, self.z, self.w], dtype=float)

    @staticmethod
    def from_np(v: np.ndarray) -> "Vec4":
        v = np.asarray(v, dtype=float)
        return Vec4(float(v[0]), float(v[1]), float(v[2]), float(v[3]))


def rot4_plane(i: int, j: int, theta: float) -> np.ndarray:
    """
    4D rotation matrix in the (i, j) plane.
    Indices 0..3 correspond to (x,y,z,w).
    """
    R = np.eye(4, dtype=float)
    c = float(np.cos(theta))
    s = float(np.sin(theta))
    R[i, i] = c
    R[j, j] = c
    R[i, j] = -s
    R[j, i] = s
    return R


# -----------------------------
# FPMF core computations (scalar phase version)
# -----------------------------

def compute_F(A: np.ndarray, theta: np.ndarray, delta_sigma: np.ndarray) -> complex:
    amp_damp = np.exp(-A)
    amp_phase = np.exp(1j * theta)
    return complex(np.sum((amp_damp - amp_phase) * delta_sigma))


def mismatch_energy(A: np.ndarray, theta: np.ndarray, delta_sigma: np.ndarray) -> float:
    amp_damp = np.exp(-A)
    amp_phase = np.exp(1j * theta)
    mis = np.abs(amp_damp - amp_phase) ** 2
    return float(np.sum(mis * delta_sigma))


def coherence_index(A: np.ndarray, theta: np.ndarray, delta_sigma: np.ndarray, eps: float = 1e-12) -> float:
    w = np.exp(-A)
    num = np.sum(w * np.cos(theta) * delta_sigma)
    den = np.sum(w * delta_sigma) + eps
    return float(num / den)


def compute_K_bridge(A: np.ndarray, theta: np.ndarray) -> np.ndarray:
    """
    Your axiom:
      exp(iΘ) = exp(-A) * K   =>  K = exp(A) * exp(iΘ)
    """
    return np.exp(A) * np.exp(1j * theta)


def compute_P_prop(A: np.ndarray, theta: np.ndarray) -> np.ndarray:
    """
    A physically useful damped propagator:
      P = exp(-A) * exp(iΘ)
    """
    return np.exp(-A) * np.exp(1j * theta)


# -----------------------------
# 4D trajectory + inner fields
# -----------------------------

def simulate_4d_trajectory(N: int, seed: int = 0, drift: float = 0.01) -> Tuple[np.ndarray, np.ndarray]:
    """
    Returns:
      sigma: (N,)
      X:     (N,4)  4D state along σ
    """
    rng = np.random.default_rng(seed)
    sigma = np.linspace(0.0, 1.0, N)
    X = np.zeros((N, 4), dtype=float)

    v = np.array([1.0, 0.0, 0.0, 0.0], dtype=float)

    for k in range(N):
        s = sigma[k]

        th_xy = 2.0 * np.pi * 1.0 * s
        th_zw = 2.0 * np.pi * 0.7 * s
        th_xw = 2.0 * np.pi * 0.3 * s

        R = rot4_plane(0, 1, th_xy) @ rot4_plane(2, 3, th_zw) @ rot4_plane(0, 3, th_xw)

        v = (R @ v) + float(drift) * rng.normal(size=4)

        # Keep bounded
        v = v / (np.linalg.norm(v) + 1e-12)
        X[k, :] = v

    return sigma, X


def build_inner_fields_from_4d(Xk: np.ndarray, tau: np.ndarray) -> Tuple[np.ndarray, np.ndarray]:
    """
    Given a 4D state Xk at contour sample k, create inner functions ψ(τ), φ(τ).

    Returns:
      dotpsi(τ): (M,)
      phi(τ):    (M,)
    """
    x, y, z, w = float(Xk[0]), float(Xk[1]), float(Xk[2]), float(Xk[3])

    psi = (0.6 + 0.2 * x) * np.sin(2 * np.pi * tau) + (0.4 + 0.2 * y) * np.sin(4 * np.pi * tau)
    dotpsi = (0.6 + 0.2 * x) * (2 * np.pi) * np.cos(2 * np.pi * tau) + (0.4 + 0.2 * y) * (4 * np.pi) * np.cos(4 * np.pi * tau)

    phi = (0.7 + 0.2 * z) * np.cos(2 * np.pi * tau) + (0.3 + 0.2 * w) * np.cos(6 * np.pi * tau)

    return dotpsi, phi


def gamma_frenetic_normalized(
    dotpsi: np.ndarray,
    phi: np.ndarray,
    gamma0: float,
    a: float,
    b: float,
    dot_scale: float,
    phi_scale: float,
) -> np.ndarray:
    """
    Normalized γ so A doesn't explode.

    A = ∫ γ dτ should be dimensionless.
    If dotpsi gets large (because of 2π factors), A will blow up unless we scale.

    Default behavior targets A ~ O(0.1 .. 0.5) rather than O(5 .. 10).

    γ = gamma0 + a*(dotpsi/dot_scale)^2 + b*(phi/phi_scale)^2
    """
    dp = dotpsi / float(dot_scale)
    ph = phi / float(phi_scale)
    return float(gamma0) + float(a) * (dp ** 2) + float(b) * (ph ** 2)


def compute_A_over_contour(
    X: np.ndarray,
    tau: np.ndarray,
    gamma0: float,
    a: float,
    b: float,
    dot_scale: float,
    phi_scale: float,
) -> np.ndarray:
    """
    For each contour sample k, compute A_k = ∫_R γ(ψdot, φ) dτ.
    """
    M = tau.size
    dt = float(tau[1] - tau[0]) if M > 1 else 1.0
    A = np.zeros((X.shape[0],), dtype=float)

    for k in range(X.shape[0]):
        dotpsi, phi = build_inner_fields_from_4d(X[k, :], tau)
        g = gamma_frenetic_normalized(dotpsi, phi, gamma0, a, b, dot_scale, phi_scale)
        A[k] = float(np.sum(g) * dt)

    return A


def compute_Gamma_kappa_Pi_from_4d(sigma: np.ndarray, X: np.ndarray) -> Tuple[np.ndarray, np.ndarray, np.ndarray]:
    """
    Scalar proxies Γ_term(σ), κ(σ), Π_term(σ) from a 4D trajectory.
    Swap these for your real definitions later.
    """
    Gamma_term = np.arctan2(X[:, 1], X[:, 0])  # angle in xy

    dX = np.gradient(X, axis=0)
    speed = np.linalg.norm(dX, axis=1)
    kappa = 0.8 + 0.6 * (speed / (np.max(speed) + 1e-12))

    Pi_term = (X[:, 2] * np.cos(8 * np.pi * sigma)) + (X[:, 3] * np.sin(6 * np.pi * sigma))
    return Gamma_term, kappa, Pi_term


def compute_theta(lam: float, alpha: float, Gamma_term: np.ndarray, beta: float, kappa: np.ndarray, Pi_term: np.ndarray) -> np.ndarray:
    return float(lam) * float(alpha) * Gamma_term + float(beta) * kappa * Pi_term


def stats(arr: np.ndarray) -> Dict[str, float]:
    arr = np.asarray(arr)
    return {
        "min": float(np.min(arr)),
        "max": float(np.max(arr)),
        "mean": float(np.mean(arr)),
        "std": float(np.std(arr)),
    }


def main() -> None:
    ap = argparse.ArgumentParser(description="RAIL FPMF 4D runner (4D state-space test).")
    ap.add_argument("--N", type=int, default=800, help="Contour samples along σ.")
    ap.add_argument("--M", type=int, default=200, help="Inner samples along τ.")
    ap.add_argument("--seed", type=int, default=0, help="Random seed for 4D drift.")
    ap.add_argument("--drift", type=float, default=0.01, help="4D drift/noise per step.")

    ap.add_argument("--lam", type=float, default=1.0)
    ap.add_argument("--alpha", type=float, default=0.75)
    ap.add_argument("--beta", type=float, default=0.90)

    # Gamma normalization knobs (these are the key fix)
    ap.add_argument("--gamma0", type=float, default=0.10, help="Base gamma offset.")
    ap.add_argument("--a", type=float, default=0.25, help="Weight for (dotpsi/dot_scale)^2.")
    ap.add_argument("--b", type=float, default=0.15, help="Weight for (phi/phi_scale)^2.")
    ap.add_argument("--dot_scale", type=float, default=10.0, help="Scale for dotpsi to prevent A blow-up.")
    ap.add_argument("--phi_scale", type=float, default=2.0, help="Scale for phi to prevent A blow-up.")

    ap.add_argument("--out", default=None, help="Optional JSON output path.")
    args = ap.parse_args()

    N = int(args.N)
    M = int(args.M)

    sigma, X = simulate_4d_trajectory(N=N, seed=int(args.seed), drift=float(args.drift))
    delta_sigma = np.full(N, float(sigma[1] - sigma[0]) if N > 1 else 1.0, dtype=float)

    tau = np.linspace(0.0, 1.0, M)

    A = compute_A_over_contour(
        X=X,
        tau=tau,
        gamma0=float(args.gamma0),
        a=float(args.a),
        b=float(args.b),
        dot_scale=float(args.dot_scale),
        phi_scale=float(args.phi_scale),
    )

    Gamma_term, kappa, Pi_term = compute_Gamma_kappa_Pi_from_4d(sigma, X)
    theta = compute_theta(float(args.lam), float(args.alpha), Gamma_term, float(args.beta), kappa, Pi_term)

    F = compute_F(A, theta, delta_sigma)
    E = mismatch_energy(A, theta, delta_sigma)
    CI = coherence_index(A, theta, delta_sigma)

    K_bridge = compute_K_bridge(A, theta)
    P_prop = compute_P_prop(A, theta)

    print("RAIL FPMF 4D Runner Report")
    print("==========================")
    print(f"params: lam={args.lam}, alpha={args.alpha}, beta={args.beta}")
    print(f"gamma:  gamma0={args.gamma0}, a={args.a}, b={args.b}, dot_scale={args.dot_scale}, phi_scale={args.phi_scale}")
    print()
    print(f"A stats         : {stats(A)}")
    print(f"theta stats     : {stats(theta)}")
    print(f"exp(-A) stats   : {stats(np.exp(-A))}")
    print()
    print(f"K_bridge |abs|  : {stats(np.abs(K_bridge))}   (your axiom; |K|=exp(A))")
    print(f"P_prop   |abs|  : {stats(np.abs(P_prop))}     (damped propagator; |P|=exp(-A))")
    print()
    print(f"F               : {F.real:.6f} + {F.imag:.6f} i")
    print(f"|F|             : {abs(F):.6f}")
    print(f"E_mis           : {E:.6f}")
    print(f"CI              : {CI:.6f}")

    summary = {
        "params": {"lam": float(args.lam), "alpha": float(args.alpha), "beta": float(args.beta)},
        "gamma_params": {
            "gamma0": float(args.gamma0),
            "a": float(args.a),
            "b": float(args.b),
            "dot_scale": float(args.dot_scale),
            "phi_scale": float(args.phi_scale),
        },
        "stats": {
            "A": stats(A),
            "theta": stats(theta),
            "exp_minus_A": stats(np.exp(-A)),
            "K_bridge_abs": stats(np.abs(K_bridge)),
            "P_prop_abs": stats(np.abs(P_prop)),
        },
        "results": {
            "F_real": float(F.real),
            "F_imag": float(F.imag),
            "F_abs": float(abs(F)),
            "E_mis": float(E),
            "CI": float(CI),
        },
    }

    if args.out:
        out_path = os.path.expanduser(args.out)
        os.makedirs(os.path.dirname(out_path), exist_ok=True)
        with open(out_path, "w", encoding="utf-8") as f:
            json.dump(summary, f, indent=2, sort_keys=True)
        print()
        print(f"Saved JSON to: {out_path}")


if __name__ == "__main__":
    main()