"""8-tension lattice features analyzer.

Compares 6 models on fwd_60 BTC return:
 1. Raw 4 tensions (per-minute)
 2. Raw 8 tensions (4 levels + 4 vels, 4hr window)
 3. Lattice instantaneous (9 channels)
 4. Lattice delta features (21: d60+pr500 per channel + d240 on 3 channels)
 5. Lattice deltas + raw 4
 6. Lattice deltas + raw 8

Splits: chronological 70/30 and 20d/10d block. Standardize on train only.
"""
import pandas as pd, numpy as np, glob
from sklearn.linear_model import LinearRegression
from sklearn.metrics import r2_score

RUN = "/mnt/d/Resonance_Engine/traj/tension_20260608T062437"
T = pd.read_parquet(f"{RUN}/arm_T_tension.parquet")
print(f"loaded arm T: {len(T):,} rows  cols: {len(T.columns)}")

files = sorted(glob.glob("/mnt/d/PaperTrader/research/hl_data/minutes/202604*/*.parquet"))
px_parts = []
for f in files:
    d = pd.read_parquet(f, columns=["minute", "coin", "mid_price"])
    px_parts.append(d[d.coin == "BTC"][["minute", "mid_price"]])
px = pd.concat(px_parts, ignore_index=True).drop_duplicates("minute").sort_values("minute").reset_index(drop=True)
px["fwd_60"] = (np.log(px["mid_price"].shift(-60)) - np.log(px["mid_price"])) * 10000.0
print(f"loaded prices: {len(px):,} minutes")

df = T.merge(px[["minute", "fwd_60"]], on="minute", how="inner").sort_values("minute").reset_index(drop=True)
print(f"joined: {len(df):,} rows")

CHANNELS = ["asymmetry", "coherence", "stress_xx", "stress_yy", "stress_xy",
            "vorticity_mean", "vel_mean", "vel_max", "vel_var"]
TENSIONS_4 = ["funding_bps", "bs_ratio_signed", "activity_excess", "cvd_divergence"]
TENSIONS_8 = (
    [f"{c}_level" for c in TENSIONS_4]
    + [f"{c}_vel" for c in TENSIONS_4]
)

for c in CHANNELS:
    df[f"{c}_inst"] = df[c]
    df[f"{c}_d60"] = df[c] - df[c].shift(60)
    df[f"{c}_pr500"] = df[c].rolling(500, min_periods=100).rank(pct=True) - 0.5

for c in ["asymmetry", "coherence", "vel_max"]:
    df[f"{c}_d240"] = df[c] - df[c].shift(240)

FEAT_RAW4 = TENSIONS_4
FEAT_RAW8 = TENSIONS_8
FEAT_INST = [f"{c}_inst" for c in CHANNELS]
FEAT_DELTA = (
    [f"{c}_d60" for c in CHANNELS]
    + [f"{c}_pr500" for c in CHANNELS]
    + [f"{c}_d240" for c in ["asymmetry", "coherence", "vel_max"]]
)
FEAT_COMBO4 = FEAT_RAW4 + FEAT_DELTA
FEAT_COMBO8 = FEAT_RAW8 + FEAT_DELTA

all_feats = list(set(FEAT_RAW4 + FEAT_RAW8 + FEAT_INST + FEAT_DELTA))
keep = df[all_feats + ["fwd_60", "minute"]].dropna()
print(f"after dropna across all feature sets + y: {len(keep):,} rows")

keep["day"] = pd.to_datetime(keep["minute"] * 60, unit="s", utc=True).dt.date
days = sorted(keep["day"].unique())
print(f"days present: {len(days)}  first={days[0]}  last={days[-1]}")

n = len(keep)
chrono_split = int(n * 0.7)
block_train_days = set(days[:20])
block_test_days = set(days[20:])
block_train_mask = keep["day"].isin(block_train_days)
block_test_mask = keep["day"].isin(block_test_days)
print(f"chrono split row {chrono_split:,}  =>  train {chrono_split:,}  test {n - chrono_split:,}")
print(f"block split  =>  train {block_train_mask.sum():,} ({len(block_train_days)}d)  test {block_test_mask.sum():,} ({len(block_test_days)}d)")


def evaluate(name, feats):
    X = keep[feats].values
    y = keep["fwd_60"].values

    def fit_eval(Xtr, ytr, Xte, yte):
        mu, sd = Xtr.mean(axis=0), Xtr.std(axis=0) + 1e-12
        Xtr_z = (Xtr - mu) / sd
        Xte_z = (Xte - mu) / sd
        reg = LinearRegression().fit(Xtr_z, ytr)
        pred_te = reg.predict(Xte_z)
        pred_tr = reg.predict(Xtr_z)
        r2_tr = r2_score(ytr, pred_tr)
        r2_te = r2_score(yte, pred_te)
        dacc = (np.sign(pred_te) == np.sign(yte)).mean() * 100
        return r2_tr, r2_te, dacc, reg

    r2_tr_c, r2_te_c, dacc_c, _ = fit_eval(X[:chrono_split], y[:chrono_split],
                                            X[chrono_split:], y[chrono_split:])
    Xtr = X[block_train_mask.values]
    ytr = y[block_train_mask.values]
    Xte = X[block_test_mask.values]
    yte = y[block_test_mask.values]
    r2_tr_b, r2_te_b, dacc_b, _ = fit_eval(Xtr, ytr, Xte, yte)
    return dict(name=name, n_feat=len(feats),
                chrono_r2_tr=r2_tr_c, chrono_r2_te=r2_te_c, chrono_dacc=dacc_c,
                block_r2_tr=r2_tr_b, block_r2_te=r2_te_b, block_dacc=dacc_b)


results = []
results.append(evaluate("Raw 4 tensions (per-minute)", FEAT_RAW4))
results.append(evaluate("Raw 8 tensions (4lvl+4vel)", FEAT_RAW8))
results.append(evaluate("Lattice instantaneous (9)",  FEAT_INST))
results.append(evaluate("Lattice delta features (21)", FEAT_DELTA))
results.append(evaluate("Lattice deltas + raw 4 (25)", FEAT_COMBO4))
results.append(evaluate("Lattice deltas + raw 8 (29)", FEAT_COMBO8))

print()
print("=" * 120)
print("FORWARD-60-MINUTE BTC RETURN PREDICTION  (R^2 and directional accuracy, OOS)")
print("=" * 120)
hdr = (f"{'Model':<32} {'feat':>5}   "
       f"{'chrono R^2 tr':>14}  {'chrono R^2 te':>14}  {'chrono dacc%':>13}   "
       f"{'block R^2 tr':>13}  {'block R^2 te':>13}  {'block dacc%':>12}")
print(hdr)
print("-" * len(hdr))
for r in results:
    print(f"{r['name']:<32} {r['n_feat']:>5}   "
          f"{r['chrono_r2_tr']:>+14.5f}  {r['chrono_r2_te']:>+14.5f}  {r['chrono_dacc']:>12.2f}%   "
          f"{r['block_r2_tr']:>+13.5f}  {r['block_r2_te']:>+13.5f}  {r['block_dacc']:>11.2f}%")

print()
print("=== top 10 |beta| in 'Lattice deltas + raw 8' (chronological train, standardized) ===")
X = keep[FEAT_COMBO8].values
y = keep["fwd_60"].values
Xtr = X[:chrono_split]
ytr = y[:chrono_split]
mu, sd = Xtr.mean(axis=0), Xtr.std(axis=0) + 1e-12
reg = LinearRegression().fit((Xtr - mu) / sd, ytr)
betas = pd.Series(reg.coef_, index=FEAT_COMBO8).sort_values(key=lambda s: s.abs(), ascending=False)
for k, v in betas.head(10).items():
    print(f"  beta[{k:<28}] = {v:+.4f}")

print()
print("=== top 8 |beta| in 'Raw 8 tensions' alone (chronological train, standardized) ===")
X = keep[FEAT_RAW8].values
y = keep["fwd_60"].values
Xtr = X[:chrono_split]
ytr = y[:chrono_split]
mu, sd = Xtr.mean(axis=0), Xtr.std(axis=0) + 1e-12
reg = LinearRegression().fit((Xtr - mu) / sd, ytr)
betas = pd.Series(reg.coef_, index=FEAT_RAW8).sort_values(key=lambda s: s.abs(), ascending=False)
for k, v in betas.head(8).items():
    print(f"  beta[{k:<28}] = {v:+.4f}")

print()
print(f"y autocorr lag1   = {pd.Series(y).autocorr(1):+.4f}   (overlapping fwd_60 windows)")
print(f"y autocorr lag60  = {pd.Series(y).autocorr(60):+.4f}")
print(f"y std            = {y.std():.2f} bps   y mean = {y.mean():+.2f} bps")