"""Lattice features analyzer. Comparison table requested by Claude Desktop.

Constraints:
- Chronological splits only (random would leak via lag-1 autocorr 0.999)
- Use delta and rolling-rank features, not raw instantaneous values
- Two splits: 70/30 chronological + 20-day/10-day block
- Comparison: raw tensions vs lattice instantaneous vs lattice deltas vs combined
"""
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_20260607T130445"
T = pd.read_parquet(f"{RUN}/arm_T_tension.parquet")
print(f"loaded arm T: {len(T):,} rows")

# load BTC prices
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")

# join
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 = ["funding_bps","bs_ratio_signed","activity_excess","cvd_divergence"]

# === FEATURE CONSTRUCTION ===
# Set B: raw instantaneous channels
for c in CHANNELS:
    df[f"{c}_inst"] = df[c]

# Set C: delta_60 + pct_rank_500 per channel
for c in CHANNELS:
    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

# also a couple of multi-scale deltas to catch slow drift
for c in ["asymmetry","coherence","vel_max"]:
    df[f"{c}_d240"] = df[c] - df[c].shift(240)

FEAT_RAW     = TENSIONS
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_COMBO   = FEAT_RAW + FEAT_DELTA

# drop rows where any feature or y is NaN (rolling warm-up + price tail)
all_feats = list(set(FEAT_RAW + FEAT_INST + FEAT_DELTA))
keep = df[all_feats + ["fwd_60","minute"]].dropna()
print(f"after dropna across all feature sets + y: {len(keep):,} rows")

# minute -> calendar day for block split
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]}")

# === split definitions ===
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
    # standardize on train only to avoid leakage
    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
    # chrono
    r2_tr_c, r2_te_c, dacc_c, reg_c = fit_eval(X[:chrono_split], y[:chrono_split],
                                               X[chrono_split:], y[chrono_split:])
    # block
    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, reg_b = fit_eval(Xtr, ytr, Xte, yte)
    return {
        "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 tensions (4)",                FEAT_RAW))
results.append(evaluate("Lattice instantaneous (9)",       FEAT_INST))
results.append(evaluate("Lattice delta features (21)",     FEAT_DELTA))
results.append(evaluate("Lattice deltas + raw (25)",       FEAT_COMBO))

print()
print("="*112)
print("FORWARD-60-MINUTE BTC RETURN PREDICTION  (R^2 and directional accuracy, OOS)")
print("="*112)
hdr = f"{'Model':<32} {'feat':>5}   {'chrono R^2 tr':>14}  {'chrono R^2 te':>14}  {'chrono dacc%':>13}   {'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}%")

# sanity: report which delta features have the largest absolute betas in the combined fit
print()
print("=== top 8 |beta| in 'Lattice deltas + raw' (chronological train, standardized) ===")
X = keep[FEAT_COMBO].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_COMBO).sort_values(key=lambda s: s.abs(), ascending=False)
for k,v in betas.head(8).items():
    print(f"  beta[{k:<28}] = {v:+.4f}")

# triviality reminder
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")