#!/usr/bin/env python3
# -*- coding: utf-8 -*-

"""
DIFT_Shell_Validation_Modes_Trajectory.py

Author: Sławomir Krakowski vel Smoku, Independent Researcher
Creation date: 2026-04-24 UTC

Purpose
-------
Standalone reviewer-oriented script for the next validation step of the
integrated topological-amplitude-gauge DIFT program.

This script performs three tasks on a validated 3D shell-bound candidate:

1. Hypothesis verification
   H1: The candidate remains shell-dominant under controlled real-time evolution.
   H2: No clearly unstable mode is detected in a reduced projected fluctuation sector.
   H3: A parameter-to-observable computational trajectory can be constructed from
       coupling homotopy lambda in [0,1] to the observables
           E*, R*, r_shell, omega*.

2. Reduced linear mode analysis
   It builds a projected Hessian of the frozen-background chi-sector energy
   over a physically interpretable basis of shell/core perturbations and
   low-order modulated shell modes.
   This yields approximate eigenmodes and eigenvalues.

3. Computational trajectory
   It performs a coupling-homotopy relaxation from lambda=0 to lambda=1,
   recording:
       energy proxy E(lambda),
       mean radius R(lambda),
       shell radius r_shell(lambda),
       classification(lambda),
   and then estimates the dominant temporal frequencies omega_n from a
   short real-time evolution of the final state.

Important scope limitation
--------------------------
This script does NOT prove full nonlinear stability of the full backreacted field
(U, chi, A_mu). It validates the chi-sector on a frozen validated carrier background,
including frozen compact gauge background theta when available.

Outputs
-------
Timestamped directory in /Users/mac/Downloads/SKRYPTY with:
    summary.json
    hypothesis_audit.json
    trajectory.csv
    mode_spectrum.csv
    temporal_frequencies.csv
    observables_final.csv
    fig_trajectory.png
    fig_timeseries.png
    fig_mode_spectrum.png
    fig_mode_slices.png
    final_state.npz
    report.md
"""

from __future__ import annotations

import argparse
import csv
import json
from dataclasses import dataclass, asdict
from datetime import datetime, timezone
from pathlib import Path
from typing import Dict, Any, List, Tuple, Optional

import numpy as np
import matplotlib.pyplot as plt


# ============================================================
# Hard defaults
# ============================================================

DEFAULT_OUT_ROOT = Path("/Users/mac/Downloads/SKRYPTY")
DEFAULT_SOURCE_GLOB = "DIFT_3D_SHELL_BINDING_ON_VALIDATED_OBJECT_V7_*/best_candidate.npz"


# ============================================================
# Utilities
# ============================================================

def utc_stamp() -> str:
    return datetime.now(timezone.utc).strftime("%Y%m%dT%H%M%SZ")


def log(msg: str) -> None:
    print(f"[*] [DIFT-SHELL-VALIDATION] {msg}")


def ensure_dir(path: Path) -> None:
    path.mkdir(parents=True, exist_ok=True)


def save_json(path: Path, payload: Dict[str, Any]) -> None:
    with path.open("w", encoding="utf-8") as f:
        json.dump(payload, f, indent=2, ensure_ascii=False)


def mid_slice(arr: np.ndarray) -> np.ndarray:
    return arr[:, :, arr.shape[2] // 2]


def parse_json_like(x: Any) -> Dict[str, Any]:
    try:
        return json.loads(str(x))
    except Exception:
        return {}


def discover_default_source(out_root: Path) -> Path:
    candidates = sorted(out_root.glob(DEFAULT_SOURCE_GLOB))
    if not candidates:
        raise FileNotFoundError(
            f"Could not auto-discover best_candidate.npz under {out_root}. "
            f"Expected pattern: {DEFAULT_SOURCE_GLOB}"
        )
    return candidates[-1].resolve()


def flatten_complex(z: np.ndarray) -> np.ndarray:
    return np.concatenate([np.real(z).ravel(), np.imag(z).ravel()])


def unflatten_complex(v: np.ndarray, shape: Tuple[int, int, int]) -> np.ndarray:
    n = np.prod(shape)
    re = v[:n].reshape(shape)
    im = v[n:].reshape(shape)
    return re + 1j * im


def mass_of(chi: np.ndarray) -> float:
    return float(np.sum(np.abs(chi) ** 2))


def renormalize_mass(chi: np.ndarray, target_mass: float, tiny: float) -> np.ndarray:
    m = mass_of(chi)
    if m <= tiny:
        return chi
    return chi * np.sqrt(target_mass / m)


# ============================================================
# Parameters
# ============================================================

@dataclass
class Params:
    source: Optional[str] = None
    out_root: str = str(DEFAULT_OUT_ROOT)

    dx: Optional[float] = None

    # classification
    edge_margin_cells: int = 3
    shell_top_quantile: float = 0.85
    core_top_quantile: float = 0.80
    shell_bound_threshold: float = 0.45
    core_bound_threshold: float = 0.45
    edge_bound_threshold: float = 0.35
    shell_score_threshold: float = 0.05

    # real-time hypothesis test
    rt_steps: int = 1200
    rt_dt: float = 0.01
    rt_sample_every: int = 20
    rt_gamma_damp: float = 5.0e-4
    perturb_strength: float = 0.01
    add_phase_noise: bool = True
    seed: int = 12345

    # trajectory / relaxation
    lambda_grid_n: int = 9
    relax_steps_per_lambda: int = 80
    relax_tau: float = 0.03

    # temporal frequency estimation
    omega_steps: int = 800
    omega_dt: float = 0.01
    omega_sample_every: int = 2

    # reduced mode analysis
    hessian_eps: float = 2.0e-3
    projected_basis_max: int = 12

    # chi model
    D_chi: float = 1.0
    lambda_chi: float = 0.5

    tiny: float = 1.0e-14


# ============================================================
# Loading
# ============================================================

def load_best_candidate(path: Path) -> Dict[str, Any]:
    data = np.load(path, allow_pickle=True)

    required = ["phi", "chi", "S_n", "G_n", "W_n"]
    missing = [k for k in required if k not in data]
    if missing:
        raise RuntimeError(f"Missing required arrays in NPZ: {missing}")

    phi = data["phi"]
    chi = data["chi"]
    S_n = data["S_n"]
    G_n = data["G_n"]
    W_n = data["W_n"]
    theta = data["theta"] if "theta" in data else None

    best_row = {}
    if "best_row" in data:
        best_row = parse_json_like(data["best_row"])

    cfg = {}
    if "cfg" in data:
        cfg = parse_json_like(data["cfg"])

    meta = {}
    if "meta" in data:
        meta = parse_json_like(data["meta"])

    if phi.ndim != 4 or phi.shape[0] != 4:
        raise RuntimeError(f"Expected phi with shape (4,Nx,Ny,Nz), got {phi.shape}")
    if chi.ndim != 3:
        raise RuntimeError(f"Expected chi with shape (Nx,Ny,Nz), got {chi.shape}")
    if S_n.shape != chi.shape or G_n.shape != chi.shape or W_n.shape != chi.shape:
        raise RuntimeError("S_n, G_n, W_n must match chi shape")

    if theta is not None:
        if theta.ndim != 4 or theta.shape[0] != 3 or theta.shape[1:] != chi.shape:
            raise RuntimeError(f"theta has invalid shape: {theta.shape}")

    return {
        "phi": phi,
        "chi": chi,
        "theta": theta,
        "S_n": S_n,
        "G_n": G_n,
        "W_n": W_n,
        "best_row": best_row,
        "cfg": cfg,
        "meta": meta,
        "npz_keys": list(data.keys()),
        "raw_npz": data,
    }


def resolve_dx(loaded: Dict[str, Any], cli_dx: Optional[float]) -> Tuple[float, str]:
    if cli_dx is not None:
        return float(cli_dx), "cli"

    raw_npz = loaded["raw_npz"]
    if "dx" in loaded["npz_keys"]:
        try:
            return float(raw_npz["dx"]), "npz:dx"
        except Exception:
            pass

    cfg = loaded.get("cfg", {})
    phi = loaded["phi"]
    n = int(phi.shape[1])

    if isinstance(cfg, dict) and "dx" in cfg:
        try:
            return float(cfg["dx"]), "cfg.dx"
        except Exception:
            pass

    if isinstance(cfg, dict) and "length" in cfg:
        try:
            length = float(cfg["length"])
            return length / n, "cfg.length/N"
        except Exception:
            pass

    meta = loaded.get("meta", {})
    if isinstance(meta, dict) and "dx" in meta:
        try:
            return float(meta["dx"]), "meta.dx"
        except Exception:
            pass

    return 1.0, "unit-normalized-fallback"


# ============================================================
# Diagnostics / masks
# ============================================================

def border_mask(shape: Tuple[int, int, int], margin: int) -> np.ndarray:
    nx, ny, nz = shape
    m = np.zeros(shape, dtype=bool)
    m[:margin, :, :] = True
    m[-margin:, :, :] = True
    m[:, :margin, :] = True
    m[:, -margin:, :] = True
    m[:, :, :margin] = True
    m[:, :, -margin:] = True
    return m


def quantile_mask(field: np.ndarray, q: float) -> np.ndarray:
    thr = float(np.quantile(field, q))
    return field >= thr


def classify_shell_state(
    rho: np.ndarray,
    S_n: np.ndarray,
    W_n: np.ndarray,
    edge_mask: np.ndarray,
    p: Params,
) -> Dict[str, float]:
    total = max(float(np.sum(rho)), p.tiny)

    core_mask = quantile_mask(S_n, p.core_top_quantile)
    shell_mask = quantile_mask(W_n, p.shell_top_quantile)

    eta_core = float(np.sum(rho[core_mask]) / total)
    eta_shell = float(np.sum(rho[shell_mask]) / total)
    eta_edge = float(np.sum(rho[edge_mask]) / total)
    shell_score = eta_shell - max(eta_core, eta_edge)

    if eta_shell >= p.shell_bound_threshold and shell_score > p.shell_score_threshold:
        cls = "shell_bound"
    elif eta_core >= p.core_bound_threshold:
        cls = "core_bound"
    elif eta_edge >= p.edge_bound_threshold:
        cls = "edge_bound"
    else:
        cls = "diffuse_or_unclassified"

    return {
        "classification": cls,
        "eta_core": eta_core,
        "eta_shell": eta_shell,
        "eta_edge": eta_edge,
        "shell_score": shell_score,
    }


def carrier_center(phi: np.ndarray) -> np.ndarray:
    core_signal = np.maximum(0.0, 1.0 - phi[0])
    total = float(np.sum(core_signal))
    idx = np.indices(core_signal.shape)
    center = np.array([np.sum(idx[d] * core_signal) / max(total, 1e-14) for d in range(3)], dtype=float)
    return center


def radial_geometry(shape: Tuple[int, int, int], center: np.ndarray, dx: float) -> np.ndarray:
    idx = np.indices(shape).astype(float)
    for d in range(3):
        idx[d] = (idx[d] - center[d]) * dx
    r = np.sqrt(idx[0] ** 2 + idx[1] ** 2 + idx[2] ** 2)
    return r


def radial_shell_radius(r: np.ndarray, rho: np.ndarray, nbins: int = 80) -> float:
    rmax = float(np.max(r))
    bins = np.linspace(0.0, rmax, nbins + 1)
    shell_density = np.zeros(nbins, dtype=float)
    for i in range(nbins):
        m = (r >= bins[i]) & (r < bins[i + 1])
        if np.any(m):
            shell_density[i] = float(np.sum(rho[m]))
    k = int(np.argmax(shell_density))
    return 0.5 * (bins[k] + bins[k + 1])


def mean_radius(r: np.ndarray, rho: np.ndarray, tiny: float) -> float:
    total = max(float(np.sum(rho)), tiny)
    return float(np.sum(r * rho) / total)


# ============================================================
# Gauge-aware operators
# ============================================================

def covariant_laplacian(chi: np.ndarray, theta: Optional[np.ndarray], dx: float) -> np.ndarray:
    if theta is None:
        out = np.zeros_like(chi, dtype=np.complex128)
        for axis in range(3):
            out += np.roll(chi, -1, axis=axis) + np.roll(chi, 1, axis=axis) - 2.0 * chi
        return out / (dx * dx)

    out = np.zeros_like(chi, dtype=np.complex128)
    for axis in range(3):
        U = np.exp(1j * theta[axis])
        chi_fwd = np.roll(chi, -1, axis=axis)
        chi_bwd = np.roll(chi, 1, axis=axis)
        U_bwd = np.roll(np.conj(U), 1, axis=axis)
        out += U * chi_fwd + U_bwd * chi_bwd - 2.0 * chi
    return out / (dx * dx)


def effective_potential(
    rho: np.ndarray,
    S_n: np.ndarray,
    G_n: np.ndarray,
    W_n: np.ndarray,
    lambda_chi: float,
    m2: float,
    g_core: float,
    g_shell: float,
    g_nonlocal: float,
) -> np.ndarray:
    return (
        m2
        - g_core * S_n
        - g_shell * G_n
        - g_nonlocal * W_n
        + 2.0 * lambda_chi * rho
    )


def energy_proxy(
    chi: np.ndarray,
    theta: Optional[np.ndarray],
    dx: float,
    S_n: np.ndarray,
    G_n: np.ndarray,
    W_n: np.ndarray,
    D_chi: float,
    lambda_chi: float,
    m2: float,
    g_core: float,
    g_shell: float,
    g_nonlocal: float,
) -> float:
    rho = np.abs(chi) ** 2
    lap = covariant_laplacian(chi, theta, dx)
    kinetic = -np.real(np.conj(chi) * (D_chi * lap))
    potential = (m2 - g_core * S_n - g_shell * G_n - g_nonlocal * W_n) * rho
    nonlinear = lambda_chi * rho * rho
    return float(np.sum(kinetic + potential + nonlinear) * (dx ** 3))


def kinetic_step_midpoint(chi: np.ndarray, theta: Optional[np.ndarray], dt: float, dx: float, D_chi: float) -> np.ndarray:
    lap1 = covariant_laplacian(chi, theta, dx)
    chi_mid = chi + 0.5 * dt * (1j * D_chi * lap1)
    lap2 = covariant_laplacian(chi_mid, theta, dx)
    return chi + dt * (1j * D_chi * lap2)


def real_time_step(
    chi: np.ndarray,
    theta: Optional[np.ndarray],
    S_n: np.ndarray,
    G_n: np.ndarray,
    W_n: np.ndarray,
    dx: float,
    p: Params,
    g_core: float,
    g_shell: float,
    g_nonlocal: float,
    m2: float,
    target_mass: float,
) -> np.ndarray:
    rho = np.abs(chi) ** 2
    V = effective_potential(rho, S_n, G_n, W_n, p.lambda_chi, m2, g_core, g_shell, g_nonlocal)
    chi = chi * np.exp((-1j * V - p.rt_gamma_damp) * (0.5 * p.rt_dt))
    chi = kinetic_step_midpoint(chi, theta, p.rt_dt, dx, p.D_chi)
    rho = np.abs(chi) ** 2
    V = effective_potential(rho, S_n, G_n, W_n, p.lambda_chi, m2, g_core, g_shell, g_nonlocal)
    chi = chi * np.exp((-1j * V - p.rt_gamma_damp) * (0.5 * p.rt_dt))
    chi = renormalize_mass(chi, target_mass, p.tiny)
    return chi


def relax_step(
    chi: np.ndarray,
    theta: Optional[np.ndarray],
    S_n: np.ndarray,
    G_n: np.ndarray,
    W_n: np.ndarray,
    dx: float,
    p: Params,
    g_core: float,
    g_shell: float,
    g_nonlocal: float,
    m2: float,
    target_mass: float,
) -> np.ndarray:
    rho = np.abs(chi) ** 2
    V = effective_potential(rho, S_n, G_n, W_n, p.lambda_chi, m2, g_core, g_shell, g_nonlocal)
    gradE = -p.D_chi * covariant_laplacian(chi, theta, dx) + V * chi
    chi = chi - p.relax_tau * gradE
    chi = renormalize_mass(chi, target_mass, p.tiny)
    return chi


# ============================================================
# Observables
# ============================================================

def compute_observables(
    chi: np.ndarray,
    phi: np.ndarray,
    theta: Optional[np.ndarray],
    S_n: np.ndarray,
    G_n: np.ndarray,
    W_n: np.ndarray,
    dx: float,
    p: Params,
    g_core: float,
    g_shell: float,
    g_nonlocal: float,
    m2: float,
) -> Dict[str, Any]:
    rho = np.abs(chi) ** 2
    edge_m = border_mask(rho.shape, p.edge_margin_cells)
    diag = classify_shell_state(rho, S_n, W_n, edge_m, p)

    c = carrier_center(phi)
    r = radial_geometry(rho.shape, c, dx)
    r_shell = radial_shell_radius(r, rho)
    R_mean = mean_radius(r, rho, p.tiny)
    E = energy_proxy(chi, theta, dx, S_n, G_n, W_n, p.D_chi, p.lambda_chi, m2, g_core, g_shell, g_nonlocal)

    out = {
        "energy_proxy": E,
        "mean_radius": R_mean,
        "shell_radius": r_shell,
        "mass": mass_of(chi),
        **diag,
    }
    return out


# ============================================================
# Hypothesis H1: real-time persistence
# ============================================================

def run_hypothesis_H1(
    chi0: np.ndarray,
    phi: np.ndarray,
    theta: Optional[np.ndarray],
    S_n: np.ndarray,
    G_n: np.ndarray,
    W_n: np.ndarray,
    dx: float,
    p: Params,
    best_row: Dict[str, Any],
) -> Dict[str, Any]:
    g_core = float(best_row.get("g_core", 0.0))
    g_shell = float(best_row.get("g_shell", 0.0))
    g_nonlocal = float(best_row.get("g_nonlocal", 0.0))
    m2 = float(best_row.get("m2", 1.0))

    rng = np.random.default_rng(p.seed)
    amp_noise = 1.0 + p.perturb_strength * rng.standard_normal(size=chi0.shape)
    chi = chi0 * amp_noise
    if p.add_phase_noise:
        chi = chi * np.exp(1j * p.perturb_strength * rng.standard_normal(size=chi0.shape))

    target_mass = mass_of(chi)
    edge_m = border_mask(chi.shape, p.edge_margin_cells)

    ts: Dict[str, List[Any]] = {
        "step": [], "time": [], "mass_rel": [], "energy_rel": [],
        "eta_core": [], "eta_shell": [], "eta_edge": [], "shell_score": [],
        "classification": []
    }

    E0 = energy_proxy(chi, theta, dx, S_n, G_n, W_n, p.D_chi, p.lambda_chi, m2, g_core, g_shell, g_nonlocal)

    for step in range(p.rt_steps + 1):
        rho = np.abs(chi) ** 2
        diag = classify_shell_state(rho, S_n, W_n, edge_m, p)

        if step % p.rt_sample_every == 0 or step == p.rt_steps:
            E = energy_proxy(chi, theta, dx, S_n, G_n, W_n, p.D_chi, p.lambda_chi, m2, g_core, g_shell, g_nonlocal)
            ts["step"].append(step)
            ts["time"].append(step * p.rt_dt)
            ts["mass_rel"].append(mass_of(chi) / max(target_mass, p.tiny))
            ts["energy_rel"].append(E / max(abs(E0), p.tiny))
            ts["eta_core"].append(diag["eta_core"])
            ts["eta_shell"].append(diag["eta_shell"])
            ts["eta_edge"].append(diag["eta_edge"])
            ts["shell_score"].append(diag["shell_score"])
            ts["classification"].append(diag["classification"])

        if step == p.rt_steps:
            break

        chi = real_time_step(
            chi, theta, S_n, G_n, W_n, dx, p,
            g_core, g_shell, g_nonlocal, m2, target_mass
        )

    counts = {
        "shell_bound": sum(1 for c in ts["classification"] if c == "shell_bound"),
        "core_bound": sum(1 for c in ts["classification"] if c == "core_bound"),
        "edge_bound": sum(1 for c in ts["classification"] if c == "edge_bound"),
        "diffuse_or_unclassified": sum(1 for c in ts["classification"] if c == "diffuse_or_unclassified"),
    }
    frac_shell = counts["shell_bound"] / max(len(ts["classification"]), 1)

    final_obs = compute_observables(
        chi, phi, theta, S_n, G_n, W_n, dx, p,
        g_core, g_shell, g_nonlocal, m2
    )

    passed = bool(
        frac_shell >= 0.70
        and final_obs["classification"] == "shell_bound"
        and final_obs["shell_score"] >= p.shell_score_threshold
        and final_obs["eta_edge"] < p.edge_bound_threshold
    )

    return {
        "timeseries": ts,
        "counts": counts,
        "shell_bound_fraction": frac_shell,
        "passed": passed,
        "final_observables": final_obs,
        "final_state": chi,
    }


# ============================================================
# Hypothesis H3: computational trajectory lambda -> observables
# ============================================================

def run_computational_trajectory(
    chi0: np.ndarray,
    phi: np.ndarray,
    theta: Optional[np.ndarray],
    S_n: np.ndarray,
    G_n: np.ndarray,
    W_n: np.ndarray,
    dx: float,
    p: Params,
    best_row: Dict[str, Any],
) -> Dict[str, Any]:
    base_g_core = float(best_row.get("g_core", 0.0))
    base_g_shell = float(best_row.get("g_shell", 0.0))
    base_g_nonlocal = float(best_row.get("g_nonlocal", 0.0))
    m2 = float(best_row.get("m2", 1.0))

    lam_grid = np.linspace(0.0, 1.0, p.lambda_grid_n)
    target_mass = mass_of(chi0)
    chi = chi0.copy()

    rows: List[Dict[str, Any]] = []

    for lam in lam_grid:
        g_core = lam * base_g_core
        g_shell = lam * base_g_shell
        g_nonlocal = lam * base_g_nonlocal

        for _ in range(p.relax_steps_per_lambda):
            chi = relax_step(
                chi, theta, S_n, G_n, W_n, dx, p,
                g_core, g_shell, g_nonlocal, m2, target_mass
            )

        obs = compute_observables(
            chi, phi, theta, S_n, G_n, W_n, dx, p,
            g_core, g_shell, g_nonlocal, m2
        )
        row = {
            "lambda": float(lam),
            "g_core": g_core,
            "g_shell": g_shell,
            "g_nonlocal": g_nonlocal,
            "m2": m2,
            "energy_proxy": obs["energy_proxy"],
            "mean_radius": obs["mean_radius"],
            "shell_radius": obs["shell_radius"],
            "mass": obs["mass"],
            "eta_core": obs["eta_core"],
            "eta_shell": obs["eta_shell"],
            "eta_edge": obs["eta_edge"],
            "shell_score": obs["shell_score"],
            "classification": obs["classification"],
        }
        rows.append(row)
        log(
            f"trajectory lambda={lam:.2f} "
            f"E={row['energy_proxy']:.6f} R={row['mean_radius']:.6f} "
            f"r_shell={row['shell_radius']:.6f} class={row['classification']}"
        )

    final_row = rows[-1]
    return {
        "rows": rows,
        "final_state": chi,
        "E_star": final_row["energy_proxy"],
        "R_star": final_row["mean_radius"],
        "r_shell_star": final_row["shell_radius"],
    }


# ============================================================
# Temporal frequencies omega*
# ============================================================

def estimate_temporal_frequencies(
    chi_ref: np.ndarray,
    chi_start: np.ndarray,
    theta: Optional[np.ndarray],
    S_n: np.ndarray,
    G_n: np.ndarray,
    W_n: np.ndarray,
    dx: float,
    p: Params,
    best_row: Dict[str, Any],
) -> Dict[str, Any]:
    g_core = float(best_row.get("g_core", 0.0))
    g_shell = float(best_row.get("g_shell", 0.0))
    g_nonlocal = float(best_row.get("g_nonlocal", 0.0))
    m2 = float(best_row.get("m2", 1.0))

    shell_mask = quantile_mask(W_n, p.shell_top_quantile)
    chi = chi_start.copy()
    target_mass = mass_of(chi)

    sig_t: List[float] = []
    times: List[float] = []

    for step in range(p.omega_steps + 1):
        if step % p.omega_sample_every == 0:
            overlap = np.sum(np.conj(chi_ref[shell_mask]) * chi[shell_mask])
            sig_t.append(float(np.real(overlap)))
            times.append(step * p.omega_dt)

        if step == p.omega_steps:
            break

        chi = real_time_step(
            chi, theta, S_n, G_n, W_n, dx, p,
            g_core, g_shell, g_nonlocal, m2, target_mass
        )

    y = np.asarray(sig_t, dtype=float)
    t = np.asarray(times, dtype=float)
    y = y - np.mean(y)
    dt_eff = t[1] - t[0] if len(t) > 1 else p.omega_dt

    freqs = np.fft.rfftfreq(len(y), d=dt_eff)
    spec = np.abs(np.fft.rfft(y)) ** 2

    if len(spec) > 0:
        spec[0] = 0.0

    idx = np.argsort(spec)[::-1][:5]
    peaks = []
    for k in idx:
        if k < len(freqs) and spec[k] > 0:
            peaks.append({
                "frequency_cycles": float(freqs[k]),
                "omega": float(2.0 * np.pi * freqs[k]),
                "power": float(spec[k]),
            })

    omega_star = peaks[0]["omega"] if peaks else 0.0

    return {
        "times": t,
        "signal": np.asarray(sig_t, dtype=float),
        "freqs": freqs,
        "spectrum": spec,
        "peaks": peaks,
        "omega_star": omega_star,
    }


# ============================================================
# Hypothesis H2: reduced projected mode analysis
# ============================================================

def gram_schmidt_complex(vecs: List[np.ndarray], tiny: float) -> List[np.ndarray]:
    ortho: List[np.ndarray] = []
    for v in vecs:
        w = v.astype(np.complex128).copy()
        for u in ortho:
            proj = np.sum(np.conj(u) * w)
            w = w - proj * u
        nrm = np.sqrt(np.sum(np.abs(w) ** 2))
        if nrm > tiny:
            ortho.append(w / nrm)
    return ortho


def build_reduced_basis(
    chi: np.ndarray,
    phi: np.ndarray,
    S_n: np.ndarray,
    W_n: np.ndarray,
    dx: float,
    p: Params,
) -> List[np.ndarray]:
    shape = chi.shape
    nx, ny, nz = shape
    center = carrier_center(phi)
    r = radial_geometry(shape, center, dx)
    r_shell = radial_shell_radius(r, np.abs(chi) ** 2)

    shell_w = W_n / max(float(np.max(W_n)), p.tiny)
    core_signal = np.maximum(0.0, 1.0 - phi[0])
    core_w = core_signal / max(float(np.max(core_signal)), p.tiny)

    x = (np.arange(nx) - center[0]) * dx
    y = (np.arange(ny) - center[1]) * dx
    z = (np.arange(nz) - center[2]) * dx
    X, Y, Z = np.meshgrid(x, y, z, indexing="ij")

    Lx = max(float(np.max(x) - np.min(x)), dx)
    Ly = max(float(np.max(y) - np.min(y)), dx)
    Lz = max(float(np.max(z) - np.min(z)), dx)

    templates: List[np.ndarray] = []

    # shell-amplitude / breathing
    templates.append(shell_w.astype(np.complex128))
    templates.append(shell_w * ((r - r_shell) / max(r_shell, dx)))
    templates.append(core_w.astype(np.complex128))

    # low-order shell modulations
    templates.append(shell_w * np.cos(2.0 * np.pi * X / Lx))
    templates.append(shell_w * np.cos(2.0 * np.pi * Y / Ly))
    templates.append(shell_w * np.cos(2.0 * np.pi * Z / Lz))
    templates.append(shell_w * np.sin(2.0 * np.pi * X / Lx))
    templates.append(shell_w * np.sin(2.0 * np.pi * Y / Ly))
    templates.append(shell_w * np.sin(2.0 * np.pi * Z / Lz))

    # phase-following shell mode
    templates.append(1j * shell_w)
    templates.append(1j * shell_w * np.cos(2.0 * np.pi * X / Lx))
    templates.append(1j * shell_w * np.cos(2.0 * np.pi * Y / Ly))

    basis = gram_schmidt_complex(templates[:p.projected_basis_max], p.tiny)
    return basis


def constrained_energy_along(
    chi0: np.ndarray,
    delta: np.ndarray,
    eps: float,
    theta: Optional[np.ndarray],
    dx: float,
    S_n: np.ndarray,
    G_n: np.ndarray,
    W_n: np.ndarray,
    p: Params,
    best_row: Dict[str, Any],
) -> float:
    g_core = float(best_row.get("g_core", 0.0))
    g_shell = float(best_row.get("g_shell", 0.0))
    g_nonlocal = float(best_row.get("g_nonlocal", 0.0))
    m2 = float(best_row.get("m2", 1.0))

    target_mass = mass_of(chi0)
    chi = chi0 + eps * delta
    chi = renormalize_mass(chi, target_mass, p.tiny)
    return energy_proxy(chi, theta, dx, S_n, G_n, W_n, p.D_chi, p.lambda_chi, m2, g_core, g_shell, g_nonlocal)


def reduced_hessian(
    chi0: np.ndarray,
    basis: List[np.ndarray],
    theta: Optional[np.ndarray],
    dx: float,
    S_n: np.ndarray,
    G_n: np.ndarray,
    W_n: np.ndarray,
    p: Params,
    best_row: Dict[str, Any],
) -> np.ndarray:
    n = len(basis)
    H = np.zeros((n, n), dtype=float)
    e0 = constrained_energy_along(chi0, np.zeros_like(chi0), 0.0, theta, dx, S_n, G_n, W_n, p, best_row)

    for i in range(n):
        for j in range(i, n):
            ei = basis[i]
            ej = basis[j]
            epm = constrained_energy_along(chi0, ei + ej, +p.hessian_eps, theta, dx, S_n, G_n, W_n, p, best_row)
            emm = constrained_energy_along(chi0, ei + ej, -p.hessian_eps, theta, dx, S_n, G_n, W_n, p, best_row)

            if i == j:
                H[i, i] = (epm - 2.0 * e0 + emm) / (p.hessian_eps ** 2)
            else:
                epp = constrained_energy_along(chi0, ei + ej, +p.hessian_eps, theta, dx, S_n, G_n, W_n, p, best_row)
                epn = constrained_energy_along(chi0, ei - ej, +p.hessian_eps, theta, dx, S_n, G_n, W_n, p, best_row)
                enp = constrained_energy_along(chi0, -ei + ej, +p.hessian_eps, theta, dx, S_n, G_n, W_n, p, best_row)
                enn = constrained_energy_along(chi0, -ei - ej, +p.hessian_eps, theta, dx, S_n, G_n, W_n, p, best_row)
                H[i, j] = (epp - epn - enp + enn) / (4.0 * p.hessian_eps ** 2)
                H[j, i] = H[i, j]

    return H


def run_hypothesis_H2_modes(
    chi0: np.ndarray,
    phi: np.ndarray,
    theta: Optional[np.ndarray],
    S_n: np.ndarray,
    G_n: np.ndarray,
    W_n: np.ndarray,
    dx: float,
    p: Params,
    best_row: Dict[str, Any],
) -> Dict[str, Any]:
    basis = build_reduced_basis(chi0, phi, S_n, W_n, dx, p)
    H = reduced_hessian(chi0, basis, theta, dx, S_n, G_n, W_n, p, best_row)

    evals, evecs = np.linalg.eigh(H)
    mode_fields: List[np.ndarray] = []
    for k in range(len(evals)):
        field = np.zeros_like(chi0, dtype=np.complex128)
        for j, b in enumerate(basis):
            field = field + evecs[j, k] * b
        nrm = np.sqrt(np.sum(np.abs(field) ** 2))
        if nrm > p.tiny:
            field = field / nrm
        mode_fields.append(field)

    num_negative = int(np.sum(evals < -1.0e-8))
    passed = bool(num_negative == 0)

    return {
        "basis_size": len(basis),
        "hessian": H,
        "eigenvalues": evals,
        "eigenvectors": evecs,
        "mode_fields": mode_fields,
        "num_negative_modes": num_negative,
        "passed": passed,
    }


# ============================================================
# Writers
# ============================================================

def write_csv(path: Path, rows: List[Dict[str, Any]], fieldnames: List[str]) -> None:
    with path.open("w", encoding="utf-8", newline="") as fh:
        writer = csv.DictWriter(fh, fieldnames=fieldnames)
        writer.writeheader()
        for row in rows:
            writer.writerow(row)


def write_observables_final(path: Path, payload: Dict[str, Any]) -> None:
    rows = [{"metric": k, "value": v} for k, v in payload.items()]
    with path.open("w", encoding="utf-8", newline="") as fh:
        writer = csv.DictWriter(fh, fieldnames=["metric", "value"])
        writer.writeheader()
        writer.writerows(rows)


def write_timeseries_csv(path: Path, ts: Dict[str, List[Any]]) -> None:
    n = len(ts["step"])
    with path.open("w", encoding="utf-8", newline="") as fh:
        writer = csv.writer(fh)
        writer.writerow(["step", "time", "mass_rel", "energy_rel", "eta_core", "eta_shell", "eta_edge", "shell_score", "classification"])
        for i in range(n):
            writer.writerow([
                ts["step"][i], ts["time"][i], ts["mass_rel"][i], ts["energy_rel"][i],
                ts["eta_core"][i], ts["eta_shell"][i], ts["eta_edge"][i],
                ts["shell_score"][i], ts["classification"][i]
            ])


def write_mode_spectrum_csv(path: Path, evals: np.ndarray) -> None:
    rows = [{"mode_index": int(i), "eigenvalue": float(v)} for i, v in enumerate(evals)]
    write_csv(path, rows, ["mode_index", "eigenvalue"])


def write_temporal_freqs_csv(path: Path, peaks: List[Dict[str, Any]]) -> None:
    write_csv(path, peaks, ["frequency_cycles", "omega", "power"])


def make_fig_trajectory(out_dir: Path, rows: List[Dict[str, Any]]) -> None:
    lam = np.array([r["lambda"] for r in rows], dtype=float)
    E = np.array([r["energy_proxy"] for r in rows], dtype=float)
    R = np.array([r["mean_radius"] for r in rows], dtype=float)
    RS = np.array([r["shell_radius"] for r in rows], dtype=float)
    SH = np.array([r["eta_shell"] for r in rows], dtype=float)

    fig, axes = plt.subplots(2, 2, figsize=(12, 8))
    axes[0, 0].plot(lam, E, marker="o")
    axes[0, 0].set_title("Energy proxy vs lambda")
    axes[0, 0].grid(True, alpha=0.3)

    axes[0, 1].plot(lam, R, marker="o", label="R")
    axes[0, 1].plot(lam, RS, marker="s", label="r_shell")
    axes[0, 1].set_title("R and r_shell vs lambda")
    axes[0, 1].legend()
    axes[0, 1].grid(True, alpha=0.3)

    axes[1, 0].plot(lam, SH, marker="o")
    axes[1, 0].set_title("eta_shell vs lambda")
    axes[1, 0].grid(True, alpha=0.3)

    classes = [r["classification"] for r in rows]
    class_num = []
    mapping = {"diffuse_or_unclassified": 0, "edge_bound": 1, "core_bound": 2, "shell_bound": 3}
    for c in classes:
        class_num.append(mapping.get(c, 0))
    axes[1, 1].plot(lam, class_num, marker="o")
    axes[1, 1].set_yticks([0, 1, 2, 3])
    axes[1, 1].set_yticklabels(["diffuse", "edge", "core", "shell"])
    axes[1, 1].set_title("Classification vs lambda")
    axes[1, 1].grid(True, alpha=0.3)

    for ax in axes.flat:
        ax.set_xlabel("lambda")

    fig.tight_layout()
    fig.savefig(out_dir / "fig_trajectory.png", dpi=180)
    plt.close(fig)


def make_fig_timeseries(out_dir: Path, ts: Dict[str, List[Any]]) -> None:
    t = np.array(ts["time"], dtype=float)

    fig, axes = plt.subplots(2, 2, figsize=(12, 8))
    axes[0, 0].plot(t, ts["mass_rel"])
    axes[0, 0].set_title("Relative mass")
    axes[0, 0].grid(True, alpha=0.3)

    axes[0, 1].plot(t, ts["energy_rel"])
    axes[0, 1].set_title("Relative energy")
    axes[0, 1].grid(True, alpha=0.3)

    axes[1, 0].plot(t, ts["eta_shell"], label="eta_shell")
    axes[1, 0].plot(t, ts["eta_core"], label="eta_core")
    axes[1, 0].plot(t, ts["eta_edge"], label="eta_edge")
    axes[1, 0].set_title("Localization fractions")
    axes[1, 0].legend()
    axes[1, 0].grid(True, alpha=0.3)

    axes[1, 1].plot(t, ts["shell_score"])
    axes[1, 1].axhline(0.05, linestyle="--")
    axes[1, 1].set_title("Shell score")
    axes[1, 1].grid(True, alpha=0.3)

    for ax in axes.flat:
        ax.set_xlabel("time")

    fig.tight_layout()
    fig.savefig(out_dir / "fig_timeseries.png", dpi=180)
    plt.close(fig)


def make_fig_mode_spectrum(out_dir: Path, evals: np.ndarray, freqs: np.ndarray, spec: np.ndarray) -> None:
    fig, axes = plt.subplots(1, 2, figsize=(12, 4.8))

    axes[0].plot(np.arange(len(evals)), evals, marker="o")
    axes[0].axhline(0.0, linestyle="--")
    axes[0].set_title("Projected Hessian eigenvalues")
    axes[0].set_xlabel("mode index")
    axes[0].set_ylabel("eigenvalue")
    axes[0].grid(True, alpha=0.3)

    axes[1].plot(freqs, spec)
    axes[1].set_title("Temporal spectrum")
    axes[1].set_xlabel("frequency (cycles / unit time)")
    axes[1].set_ylabel("power")
    axes[1].grid(True, alpha=0.3)

    fig.tight_layout()
    fig.savefig(out_dir / "fig_mode_spectrum.png", dpi=180)
    plt.close(fig)


def make_fig_mode_slices(out_dir: Path, phi: np.ndarray, mode_fields: List[np.ndarray], evals: np.ndarray) -> None:
    # show first 4 modes by ascending eigenvalue
    nm = min(4, len(mode_fields))
    fig, axes = plt.subplots(2, nm, figsize=(4 * nm, 7))
    if nm == 1:
        axes = np.array([[axes[0]], [axes[1]]])

    for j in range(nm):
        fld = np.abs(mode_fields[j]) ** 2
        axes[0, j].imshow(mid_slice(fld), origin="lower")
        axes[0, j].set_title(f"|mode {j}|^2\nlambda={evals[j]:.4e}")

        overlay = np.maximum(0.0, 1.0 - phi[0])
        axes[1, j].imshow(mid_slice(overlay), origin="lower")
        axes[1, j].contour(mid_slice(fld), levels=5, colors="white", linewidths=0.7)
        axes[1, j].set_title(f"mode {j} on carrier")

    fig.tight_layout()
    fig.savefig(out_dir / "fig_mode_slices.png", dpi=180)
    plt.close(fig)


def write_report_md(path: Path, summary: Dict[str, Any]) -> None:
    lines = [
        "# DIFT shell validation / modes / trajectory report",
        "",
        "## Core outcome",
        "",
        f"- **H1 persistence passed**: {summary['H1_passed']}",
        f"- **H2 reduced-mode stability passed**: {summary['H2_passed']}",
        f"- **H3 trajectory constructed**: {summary['H3_constructed']}",
        "",
        "## Final observables",
        "",
        f"- **E_star**: {summary['E_star']}",
        f"- **R_star**: {summary['R_star']}",
        f"- **r_shell_star**: {summary['r_shell_star']}",
        f"- **omega_star**: {summary['omega_star']}",
        "",
        "## Reduced mode analysis",
        "",
        f"- **num_negative_modes**: {summary['num_negative_modes']}",
        f"- **basis_size**: {summary['basis_size']}",
        "",
        "## Interpretation",
        "",
        "H1 refers to shell-dominant persistence under real-time evolution on the frozen validated carrier.",
        "H2 refers to the absence of negative eigenvalues in the reduced projected fluctuation sector only.",
        "H3 refers to an explicit homotopy trajectory from lambda to the observables E*, R*, r_shell, and omega*.",
        "This is a serious validation step, but not yet a full proof of global nonlinear stability of the complete integrated field.",
        "",
    ]
    path.write_text("\n".join(lines), encoding="utf-8")


# ============================================================
# Main
# ============================================================

def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(description="DIFT shell validation, reduced mode analysis, and computational trajectory")
    parser.add_argument("--source", type=str, default=None)
    parser.add_argument("--out-root", type=str, default=str(DEFAULT_OUT_ROOT))
    parser.add_argument("--dx", type=float, default=None)
    return parser.parse_args()


def main() -> None:
    args = parse_args()

    p = Params(
        source=args.source,
        out_root=args.out_root,
        dx=args.dx,
    )

    out_root = Path(p.out_root).expanduser().resolve()
    ensure_dir(out_root)

    if p.source is None:
        source = discover_default_source(out_root)
        log(f"Auto-discovered source: {source}")
    else:
        source = Path(p.source).expanduser().resolve()

    out_dir = out_root / f"DIFT_SHELL_VALIDATION_MODES_TRAJECTORY_{utc_stamp()}"
    ensure_dir(out_dir)

    log(f"source={source}")
    log(f"out_dir={out_dir}")

    loaded = load_best_candidate(source)
    phi = loaded["phi"]
    chi0 = loaded["chi"]
    theta = loaded["theta"]
    S_n = loaded["S_n"]
    G_n = loaded["G_n"]
    W_n = loaded["W_n"]
    best_row = loaded["best_row"]

    dx, dx_source = resolve_dx(loaded, p.dx)
    theta_used = theta is not None

    log(f"theta_used={theta_used}")
    log(f"dx={dx} ({dx_source})")

    # H1 persistence
    h1 = run_hypothesis_H1(chi0, phi, theta, S_n, G_n, W_n, dx, p, best_row)
    log(f"H1 persistence passed = {h1['passed']}")

    # H3 computational trajectory
    h3 = run_computational_trajectory(chi0, phi, theta, S_n, G_n, W_n, dx, p, best_row)
    log("H3 trajectory constructed")

    # omega*
    omega_data = estimate_temporal_frequencies(
        chi_ref=h3["final_state"],
        chi_start=h3["final_state"],
        theta=theta,
        S_n=S_n,
        G_n=G_n,
        W_n=W_n,
        dx=dx,
        p=p,
        best_row=best_row,
    )
    log(f"omega_star = {omega_data['omega_star']:.6f}")

    # H2 reduced mode analysis around final trajectory state
    h2 = run_hypothesis_H2_modes(
        chi0=h3["final_state"],
        phi=phi,
        theta=theta,
        S_n=S_n,
        G_n=G_n,
        W_n=W_n,
        dx=dx,
        p=p,
        best_row=best_row,
    )
    log(f"H2 reduced-mode stability passed = {h2['passed']}")

    summary = {
        "source": str(source),
        "script_name": "DIFT_Shell_Validation_Modes_Trajectory.py",
        "theta_used": theta_used,
        "dx": dx,
        "dx_source": dx_source,
        "params": asdict(p),
        "best_row": best_row,
        "H1_passed": bool(h1["passed"]),
        "H2_passed": bool(h2["passed"]),
        "H3_constructed": True,
        "shell_bound_fraction": float(h1["shell_bound_fraction"]),
        "final_classification": str(h1["final_observables"]["classification"]),
        "final_eta_shell": float(h1["final_observables"]["eta_shell"]),
        "final_eta_core": float(h1["final_observables"]["eta_core"]),
        "final_eta_edge": float(h1["final_observables"]["eta_edge"]),
        "final_shell_score": float(h1["final_observables"]["shell_score"]),
        "E_star": float(h3["E_star"]),
        "R_star": float(h3["R_star"]),
        "r_shell_star": float(h3["r_shell_star"]),
        "omega_star": float(omega_data["omega_star"]),
        "num_negative_modes": int(h2["num_negative_modes"]),
        "basis_size": int(h2["basis_size"]),
        "leading_eigenvalues": [float(x) for x in h2["eigenvalues"][:8]],
        "dominant_temporal_peaks": omega_data["peaks"],
        "interpretation_note": (
            "H1 validates shell-dominant persistence on the frozen validated carrier. "
            "H2 reports the projected reduced-sector fluctuation spectrum around the final state. "
            "H3 provides an explicit lambda-to-observable computational trajectory yielding "
            "E*, R*, r_shell, and omega*. This is not yet a full 3D backreacted proof of global "
            "nonlinear stability of the complete integrated field."
        ),
    }

    hypothesis_audit = {
        "H1_shell_persistence": {
            "passed": bool(h1["passed"]),
            "shell_bound_fraction": h1["shell_bound_fraction"],
            "counts": h1["counts"],
            "final_observables": h1["final_observables"],
        },
        "H2_reduced_mode_stability": {
            "passed": bool(h2["passed"]),
            "num_negative_modes": h2["num_negative_modes"],
            "basis_size": h2["basis_size"],
            "leading_eigenvalues": [float(x) for x in h2["eigenvalues"][:12]],
        },
        "H3_computational_trajectory": {
            "constructed": True,
            "E_star": h3["E_star"],
            "R_star": h3["R_star"],
            "r_shell_star": h3["r_shell_star"],
            "omega_star": omega_data["omega_star"],
        },
    }

    save_json(out_dir / "summary.json", summary)
    save_json(out_dir / "hypothesis_audit.json", hypothesis_audit)
    write_timeseries_csv(out_dir / "timeseries_h1.csv", h1["timeseries"])
    write_csv(out_dir / "trajectory.csv", h3["rows"], list(h3["rows"][0].keys()))
    write_mode_spectrum_csv(out_dir / "mode_spectrum.csv", h2["eigenvalues"])
    write_temporal_freqs_csv(out_dir / "temporal_frequencies.csv", omega_data["peaks"])

    final_obs_row = {
        "E_star": h3["E_star"],
        "R_star": h3["R_star"],
        "r_shell_star": h3["r_shell_star"],
        "omega_star": omega_data["omega_star"],
        "final_classification": h1["final_observables"]["classification"],
        "final_eta_shell": h1["final_observables"]["eta_shell"],
        "final_eta_core": h1["final_observables"]["eta_core"],
        "final_eta_edge": h1["final_observables"]["eta_edge"],
        "final_shell_score": h1["final_observables"]["shell_score"],
    }
    write_observables_final(out_dir / "observables_final.csv", final_obs_row)

    make_fig_trajectory(out_dir, h3["rows"])
    make_fig_timeseries(out_dir, h1["timeseries"])
    make_fig_mode_spectrum(out_dir, h2["eigenvalues"], omega_data["freqs"], omega_data["spectrum"])
    make_fig_mode_slices(out_dir, phi, h2["mode_fields"], h2["eigenvalues"])
    write_report_md(out_dir / "report.md", summary)

    np.savez_compressed(
        out_dir / "final_state.npz",
        phi=phi,
        chi_initial=chi0,
        chi_h1_final=h1["final_state"],
        chi_h3_final=h3["final_state"],
        theta=theta if theta is not None else np.zeros((3, *chi0.shape), dtype=float),
        S_n=S_n,
        G_n=G_n,
        W_n=W_n,
        mode_eigenvalues=h2["eigenvalues"],
        omega_freqs=omega_data["freqs"],
        omega_spectrum=omega_data["spectrum"],
        best_row=json.dumps(best_row),
        summary=json.dumps(summary),
    )

    log("==============================================")
    log(f"H1 persistence passed: {summary['H1_passed']}")
    log(f"H2 reduced-mode stability passed: {summary['H2_passed']}")
    log(f"H3 trajectory constructed: {summary['H3_constructed']}")
    log(f"E_star = {summary['E_star']:.6f}")
    log(f"R_star = {summary['R_star']:.6f}")
    log(f"r_shell_star = {summary['r_shell_star']:.6f}")
    log(f"omega_star = {summary['omega_star']:.6f}")
    log(f"num_negative_modes = {summary['num_negative_modes']}")
    log(f"summary.json -> {out_dir / 'summary.json'}")
    log("==============================================")


if __name__ == "__main__":
    main()