"""Analyze trajectory experiment: are the 6 variables distinguishable in lattice trajectory-space?

For each variable: load its trace.parquet, restrict to t in [PRE, PRE+SEQ], i.e. while pulses were active.
Compute trajectory in (asymmetry, coherence, stress_xy, stress_xx, stress_yy, vorticity_mean) space.
Compare pairwise distances between variable trajectories vs. self-consistency.

Output:
  - per-variable summary table (mean/std/range during pulse window)
  - pairwise trajectory distance matrix
  - decision: are trajectories distinguishable above field noise?
"""
import json
from pathlib import Path
import numpy as np
import pandas as pd

OUT_DIR = Path("/mnt/d/Resonance_Engine/traj/20260606T193626")
META = json.loads((OUT_DIR / "meta.json").read_text())
PRE = META["pre_record_s"]
N = META["n_pulses"]
SPACING = META["pulse_spacing_s"]
SEQ_END = PRE + N * SPACING  # end of pulse window in rel_t
VARS = META["vars"]

CHANNELS = ["asymmetry", "coherence", "stress_xx", "stress_yy", "stress_xy",
            "vorticity_mean", "vel_mean", "vel_var"]

print(f"=== run {META['run_id']}  pulse window [{PRE:.0f}s .. {SEQ_END:.0f}s] ===\n")

# load all traces
traces = {}
for v in VARS:
    df = pd.read_parquet(OUT_DIR / f"{v}.parquet")
    traces[v] = df

# helper: select pulse-active window
def pulse_window(df):
    return df[(df.rel_t >= PRE) & (df.rel_t <= SEQ_END)]

def pre_window(df):
    return df[df.rel_t < PRE]

# ---- per-variable stats during pulse window ----
print("=== per-variable telemetry during pulse window ===")
print(f"{'var':18s} {'n':>4} {'asym_mean':>10} {'asym_std':>9} {'asym_max':>9}  "
      f"{'sxy_mean':>10} {'sxy_std':>10} {'sxy_amp':>10}  "
      f"{'coh_mean':>9} {'coh_min':>9}  {'vort_mean':>10}")
print("-" * 130)
summary = {}
for v in VARS:
    pw = pulse_window(traces[v])
    s = {
        "n": len(pw),
        "asym_mean": pw.asymmetry.mean(),
        "asym_std":  pw.asymmetry.std(),
        "asym_max":  pw.asymmetry.max(),
        "asym_min":  pw.asymmetry.min(),
        "sxy_mean":  pw.stress_xy.mean(),
        "sxy_std":   pw.stress_xy.std(),
        "sxy_amp":   pw.stress_xy.max() - pw.stress_xy.min(),
        "sxx_mean":  pw.stress_xx.mean(),
        "syy_mean":  pw.stress_yy.mean(),
        "coh_mean":  pw.coherence.mean(),
        "coh_min":   pw.coherence.min(),
        "vort_mean": pw.vorticity_mean.mean(),
        "vort_std":  pw.vorticity_mean.std(),
        "vel_var_mean": pw.vel_var.mean(),
    }
    summary[v] = s
    print(f"{v:18s} {s['n']:>4d} {s['asym_mean']:>10.4f} {s['asym_std']:>9.4f} "
          f"{s['asym_max']:>9.4f}  "
          f"{s['sxy_mean']:>+10.6f} {s['sxy_std']:>10.6f} {s['sxy_amp']:>10.6f}  "
          f"{s['coh_mean']:>9.5f} {s['coh_min']:>9.5f}  {s['vort_mean']:>10.5f}")

# Compute spread across variables (between-variable variance)
print("\n=== spread ACROSS variables (between-variable std / within-variable std) ===")
between = {}
within = {}
for ch in ["asymmetry", "stress_xy", "stress_xx", "stress_yy", "coherence",
           "vorticity_mean", "vel_var"]:
    var_means = [pulse_window(traces[v])[ch].mean() for v in VARS]
    var_stds = [pulse_window(traces[v])[ch].std() for v in VARS]
    between[ch] = np.std(var_means)
    within[ch] = np.mean(var_stds)
    ratio = between[ch] / within[ch] if within[ch] > 0 else 0
    print(f"  {ch:18s}  between_std={between[ch]:.6f}  within_std={within[ch]:.6f}  "
          f"ratio={ratio:.3f}")

# Pairwise distance in normalized multi-channel space
print("\n=== pairwise trajectory distance ===")
# Build a feature vector per variable: time series of each channel during pulse window,
# resampled to common length, then concatenated and normalized per channel by the
# WITHIN-variable std (so we measure cross-variable difference in units of own noise)
NRESAMPLE = 100
feature_vecs = {}
for v in VARS:
    pw = pulse_window(traces[v]).sort_values("rel_t")
    # resample to NRESAMPLE points uniformly across [PRE, SEQ_END]
    times = np.linspace(PRE, SEQ_END, NRESAMPLE)
    vec = []
    for ch in ["asymmetry", "stress_xy", "stress_xx", "stress_yy", "coherence",
               "vorticity_mean", "vel_var"]:
        resampled = np.interp(times, pw.rel_t.values, pw[ch].values)
        # normalize by within-channel std across all variables
        if within[ch] > 0:
            resampled = (resampled - resampled.mean()) / within[ch]
        vec.append(resampled)
    feature_vecs[v] = np.concatenate(vec)

print(f"{'':18s}  " + "  ".join(f"{v[:12]:>12s}" for v in VARS))
for v1 in VARS:
    row = []
    for v2 in VARS:
        d = np.linalg.norm(feature_vecs[v1] - feature_vecs[v2])
        row.append(d)
    print(f"  {v1:16s}  " + "  ".join(f"{d:>12.2f}" for d in row))

# Are the distances large enough to call distinguishable?
off_diag = []
for i, v1 in enumerate(VARS):
    for j, v2 in enumerate(VARS):
        if i < j:
            off_diag.append(np.linalg.norm(feature_vecs[v1] - feature_vecs[v2]))

print(f"\n  pairwise distances:  min={min(off_diag):.2f}  max={max(off_diag):.2f}  "
      f"mean={np.mean(off_diag):.2f}")

# Sanity baseline: distance between two halves of the SAME variable's trace
print("\n=== sanity: distance within same variable (split halves of pulse window) ===")
for v in VARS:
    pw = pulse_window(traces[v]).sort_values("rel_t")
    mid = (PRE + SEQ_END) / 2
    a = pw[pw.rel_t < mid]
    b = pw[pw.rel_t >= mid]
    if len(a) < 5 or len(b) < 5: continue
    times_a = np.linspace(PRE, mid, NRESAMPLE)
    times_b = np.linspace(mid, SEQ_END, NRESAMPLE)
    vec_a, vec_b = [], []
    for ch in ["asymmetry", "stress_xy", "stress_xx", "stress_yy", "coherence",
               "vorticity_mean", "vel_var"]:
        ra = np.interp(times_a, a.rel_t.values, a[ch].values)
        rb = np.interp(times_b, b.rel_t.values, b[ch].values)
        if within[ch] > 0:
            ra = (ra - ra.mean()) / within[ch]
            rb = (rb - rb.mean()) / within[ch]
        vec_a.append(ra)
        vec_b.append(rb)
    d_self = np.linalg.norm(np.concatenate(vec_a) - np.concatenate(vec_b))
    print(f"  {v:18s}  self_half_distance={d_self:.2f}")

print()
print("=== INTERPRETATION ===")
mean_cross = np.mean(off_diag)

# clean within-variable half-distance
half_dists = []
for v in VARS:
    pw = pulse_window(traces[v]).sort_values("rel_t")
    mid = (PRE + SEQ_END) / 2
    a = pw[pw.rel_t < mid]
    b = pw[pw.rel_t >= mid]
    if len(a) < 5 or len(b) < 5:
        continue
    times_a = np.linspace(PRE, mid, NRESAMPLE)
    times_b = np.linspace(mid, SEQ_END, NRESAMPLE)
    vec_a, vec_b = [], []
    for ch in ["asymmetry", "stress_xy", "stress_xx", "stress_yy", "coherence",
               "vorticity_mean", "vel_var"]:
        ra = np.interp(times_a, a.rel_t.values, a[ch].values)
        rb = np.interp(times_b, b.rel_t.values, b[ch].values)
        if within[ch] > 0:
            ra = (ra - ra.mean()) / within[ch]
            rb = (rb - rb.mean()) / within[ch]
        vec_a.append(ra)
        vec_b.append(rb)
    half_dists.append(np.linalg.norm(np.concatenate(vec_a) - np.concatenate(vec_b)))
mean_self = np.mean(half_dists)

print(f"  mean cross-variable distance: {mean_cross:.2f}")
print(f"  mean within-variable half-distance: {mean_self:.2f}")
ratio = mean_cross / mean_self if mean_self > 0 else float('inf')
print(f"  RATIO: {ratio:.2f}x  (>1 = variables more different from each other")
print(f"          than each is from itself between halves)")