"""inject_continuous_stream.py — 4-arm continuous-streaming experiment.

Question: when real BTC minute data is fed continuously into the lattice,
does field state predict future |price moves| BETTER than the raw data alone?

Arms:
  A — NO-INJECTION CONTROL. Lattice runs free, we sample state every replay-tick.
      Tells us the lattice's natural variability over the experiment window.
  B — STACKED. All 3 vars injected at (512,512), distinct ZMQ messages per var.
      Tests "amplitude only" — does the lattice add anything over data?
  C — SEPARATED. signed_flow @(400,512), vwap_drift @(512,512), wallet_entropy @(624,512).
      Tests "spatial separation creates distinguishable channels".
  D — RAW DATA BASELINE. Computed offline by the analyzer from parquets directly.
      No lattice involved. Reference performance — "how good is the data alone?".

Variables (the 3 confirmed-real channels from varisol overnight):
  signed_flow_usd, vwap_drift, wallet_entropy

Per source minute m:
  1. compute one-sided rolling z-score for each var over past 500 minutes
  2. cap z at +/-3, convert to strength = z/3 * STR_CAP (STR_CAP=0.30, calibrated linear range)
  3. for arms B/C: send 3 inject_density messages per minute
  4. wait WAIT_AFTER_INJECT_MS
  5. snapshot most-recent telemetry as field_state(m)
  6. advance to next minute after PER_MINUTE_MS total

Replay rate: 100ms per source minute (10x realtime). Lattice halflife ~0.5s, so
each minute's injection contributes to a sustained accumulated state. Verified
by _throughput_probe.py — lattice handles 20/sec without packet loss.

Output: /mnt/d/Resonance_Engine/traj/contstream_<RUN_ID>/
    meta.json
    arm_A_no_inject.parquet
    arm_B_stacked.parquet
    arm_C_separated.parquet
    progress.log
"""
from __future__ import annotations
import glob, json, threading, time
from pathlib import Path

import numpy as np
import pandas as pd
import zmq

# ─────── config ───────
DATA_ROOT  = "/mnt/d/PaperTrader/research/hl_data/minutes"
DAYS_GLOB  = "202603*"          # all of March 2026
COIN       = "BTC"

VARIABLES  = ["signed_flow_usd", "vwap_drift", "wallet_entropy"]
ROLL_WIN   = 500                # past-only rolling z-score window
STR_CAP    = 0.30               # stay in calibrated linear range
Z_CLIP     = 3.0

# Spatial coords (used by arm C; B uses XY_CENTER for all)
XY_CENTER  = (512.0, 512.0)
ARM_C_XY   = {
    "signed_flow_usd": (400.0, 512.0),
    "vwap_drift":      (512.0, 512.0),
    "wallet_entropy":  (624.0, 512.0),
}
INJECT_SIG = 32.0

# Timing
PER_MINUTE_MS         = 100     # total per-minute budget (10x realtime)
WAIT_AFTER_INJECT_MS  = 60      # wait after the 3 sends before snapshot
COOLDOWN_BETWEEN_ARMS_S = 1200  # 20 min between arms B/C to let field settle

# ZMQ
TEL_ADDR = "tcp://127.0.0.1:5556"
CMD_ADDR = "tcp://127.0.0.1:5557"

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

RUN_ID  = time.strftime("%Y%m%dT%H%M%S")
OUT_DIR = Path(f"/mnt/d/Resonance_Engine/traj/contstream_{RUN_ID}")
OUT_DIR.mkdir(parents=True, exist_ok=True)
PROGRESS = OUT_DIR / "progress.log"


def log(msg: str) -> None:
    line = f"[{time.strftime('%Y-%m-%dT%H:%M:%S')}] {msg}"
    print(line, flush=True)
    with PROGRESS.open("a") as f:
        f.write(line + "\n")


# ─────── ZMQ telemetry subscriber (keeps most-recent only) ───────
class LatestTel:
    def __init__(self, addr: str):
        self.ctx = zmq.Context.instance()
        self.sock = self.ctx.socket(zmq.SUB)
        self.sock.connect(addr)
        self.sock.setsockopt(zmq.SUBSCRIBE, b"")
        self.sock.setsockopt(zmq.RCVHWM, 2000)
        self.latest: dict | None = None
        self.latest_wall: float = 0.0
        self.n_seen = 0
        self._stop = threading.Event()
        self._t = threading.Thread(target=self._run, daemon=True)
        self._t.start()

    def _run(self):
        while not self._stop.is_set():
            if self.sock.poll(100):
                try:
                    raw = self.sock.recv_string(zmq.NOBLOCK)
                    self.latest = json.loads(raw)
                    self.latest_wall = time.time()
                    self.n_seen += 1
                except Exception:
                    pass

    def snapshot(self) -> tuple[dict | None, float]:
        return self.latest, self.latest_wall

    def stop(self):
        self._stop.set()


# ─────── data loading ───────
def load_btc_month() -> pd.DataFrame:
    day_dirs = sorted(glob.glob(f"{DATA_ROOT}/{DAYS_GLOB}"))
    if not day_dirs:
        raise RuntimeError(f"no day dirs matching {DAYS_GLOB} in {DATA_ROOT}")
    log(f"loading {len(day_dirs)} day dirs ({day_dirs[0].split('/')[-1]} .. {day_dirs[-1].split('/')[-1]})")
    dfs = []
    for d in day_dirs:
        files = sorted(glob.glob(f"{d}/*.parquet"))
        for f in files:
            dfs.append(pd.read_parquet(f))
    df = pd.concat(dfs, ignore_index=True)
    df = df[df.coin == COIN].sort_values("minute").reset_index(drop=True)
    df = df.drop_duplicates(subset=["minute"]).reset_index(drop=True)
    log(f"loaded {len(df)} unique minutes of {COIN} data, "
        f"minute range {df.minute.min()}..{df.minute.max()}")
    return df


def compute_one_sided_zscores(df: pd.DataFrame, vars_: list[str], win: int) -> dict[str, np.ndarray]:
    """For each var: rolling z-score using only PAST `win` minutes (no look-ahead)."""
    out = {}
    for v in vars_:
        s = df[v].astype(float)
        # shift(1) so the window for minute m uses minutes m-win..m-1 (strictly past)
        roll = s.shift(1).rolling(window=win, min_periods=50)
        mu = roll.mean()
        sd = roll.std()
        z = (s - mu) / sd.replace(0, np.nan)
        out[v] = z.values
    return out


# ─────── inject + snapshot ───────
def inject_one(pub: zmq.Socket, x: float, y: float, strength: float) -> None:
    payload = {"cmd": "inject_density",
               "x": float(x), "y": float(y),
               "sigma": INJECT_SIG, "strength": float(strength)}
    pub.send_string(json.dumps(payload))


def snapshot_row(tel: LatestTel, minute: int, strengths: dict[str, float],
                 arm: str) -> dict:
    """Pull latest telemetry, build per-minute record."""
    snap, snap_wall = tel.snapshot()
    rec = {
        "minute":  int(minute),
        "arm":     arm,
        "snap_wall": snap_wall,
        "snap_age_ms": (time.time() - snap_wall) * 1000 if snap_wall else None,
    }
    for v in VARIABLES:
        rec[f"str_{v}"] = strengths.get(v)
    if snap is None:
        for ch in CHANNELS:
            rec[ch] = None
    else:
        for ch in CHANNELS:
            rec[ch] = snap.get(ch)
    return rec


def run_arm(tel: LatestTel, pub: zmq.Socket | None, arm: str,
            df: pd.DataFrame, zscores: dict[str, np.ndarray]) -> pd.DataFrame:
    """Run one arm end to end. Returns a DataFrame of per-minute records."""
    log(f"\n=== arm {arm}: starting ({len(df)} minutes) ===")
    if arm == "A":
        log("  arm A: NO INJECTIONS — sampling state at cadence only")
    elif arm == "B":
        log(f"  arm B: all 3 vars at {XY_CENTER}")
    elif arm == "C":
        log(f"  arm C: spatially separated — {ARM_C_XY}")

    records: list[dict] = []
    t_arm_start = time.time()
    last_status = t_arm_start
    n_minutes = len(df)
    step_s = PER_MINUTE_MS / 1000.0
    wait_inject_s = WAIT_AFTER_INJECT_MS / 1000.0

    for i in range(n_minutes):
        t_tick = time.time()
        minute = int(df.minute.iat[i])
        strengths = {}
        for v in VARIABLES:
            z = zscores[v][i]
            if not np.isfinite(z):
                strengths[v] = None
            else:
                z_c = float(np.clip(z, -Z_CLIP, Z_CLIP))
                strengths[v] = (z_c / Z_CLIP) * STR_CAP

        # inject (skip if arm A, or if strength is None/insufficient warmup)
        if arm in ("B", "C") and pub is not None:
            for v in VARIABLES:
                s = strengths[v]
                if s is None:
                    continue
                if arm == "B":
                    x, y = XY_CENTER
                else:
                    x, y = ARM_C_XY[v]
                inject_one(pub, x, y, s)
            # let field react
            time.sleep(wait_inject_s)
        else:
            # arm A — just wait the inject-equivalent gap so cadence matches
            time.sleep(wait_inject_s)

        rec = snapshot_row(tel, minute, strengths, arm)
        records.append(rec)

        # pace to step_s per minute total
        spent = time.time() - t_tick
        remaining = step_s - spent
        if remaining > 0:
            time.sleep(remaining)

        # status every 60s wall-time
        if time.time() - last_status > 60:
            elapsed = time.time() - t_arm_start
            pct = (i + 1) / n_minutes * 100
            rate_per_s = (i + 1) / elapsed if elapsed > 0 else 0
            eta_s = (n_minutes - i - 1) / rate_per_s if rate_per_s > 0 else 0
            cur_field = tel.latest or {}
            log(f"  arm {arm} {i+1}/{n_minutes} ({pct:.1f}%)  rate={rate_per_s:.1f}min/s  "
                f"ETA={eta_s/60:.1f}min  asym={cur_field.get('asymmetry','?'):.3f}  "
                f"latest_tel_age={(time.time()-tel.latest_wall)*1000 if tel.latest_wall else -1:.0f}ms  "
                f"tel_seen={tel.n_seen}")
            last_status = time.time()

    log(f"  arm {arm} COMPLETE  records={len(records)}  total={(time.time()-t_arm_start)/60:.1f}min")
    out = pd.DataFrame(records)
    return out


# ─────── main ───────
def main():
    log(f"=== inject_continuous_stream run_id={RUN_ID} ===")
    log(f"out_dir={OUT_DIR}")
    log(f"config: per_minute_ms={PER_MINUTE_MS} wait_after_inject_ms={WAIT_AFTER_INJECT_MS} "
        f"str_cap={STR_CAP} roll_win={ROLL_WIN}")

    # load data
    df = load_btc_month()
    # compute one-sided z-scores
    zscores = compute_one_sided_zscores(df, VARIABLES, ROLL_WIN)
    n_finite = {v: int(np.isfinite(zscores[v]).sum()) for v in VARIABLES}
    log(f"zscore finite counts: {n_finite}")

    # write meta
    meta = {
        "run_id": RUN_ID,
        "coin":   COIN,
        "data_root": DATA_ROOT,
        "days_glob": DAYS_GLOB,
        "n_minutes": int(len(df)),
        "minute_range": [int(df.minute.min()), int(df.minute.max())],
        "variables": VARIABLES,
        "roll_win": ROLL_WIN,
        "z_clip": Z_CLIP,
        "str_cap": STR_CAP,
        "xy_center": list(XY_CENTER),
        "arm_C_xy": {k: list(v) for k, v in ARM_C_XY.items()},
        "inject_sigma": INJECT_SIG,
        "per_minute_ms": PER_MINUTE_MS,
        "wait_after_inject_ms": WAIT_AFTER_INJECT_MS,
        "cooldown_between_arms_s": COOLDOWN_BETWEEN_ARMS_S,
        "wall_iso_start": time.strftime("%Y-%m-%dT%H:%M:%S"),
    }
    (OUT_DIR / "meta.json").write_text(json.dumps(meta, indent=2))

    # ZMQ
    tel = LatestTel(TEL_ADDR)
    ctx = zmq.Context.instance()
    pub = ctx.socket(zmq.PUB)
    pub.connect(CMD_ADDR)
    time.sleep(0.8)

    # wait for telemetry up
    t0 = time.time()
    while tel.latest is None and time.time() - t0 < 15:
        time.sleep(0.2)
    if tel.latest is None:
        log("FATAL: no telemetry on 5556 within 15s")
        return
    log(f"telemetry up — initial asym={tel.latest['asymmetry']:.3f}")

    # ── ARM A: no-injection control FIRST (preserves natural state) ──
    df_A = run_arm(tel, None, "A", df, zscores)
    df_A.to_parquet(OUT_DIR / "arm_A_no_inject.parquet")
    log(f"  saved arm_A_no_inject.parquet  ({len(df_A)} rows)")

    # ── cooldown ──
    log(f"\ncooldown {COOLDOWN_BETWEEN_ARMS_S}s before arm B")
    time.sleep(COOLDOWN_BETWEEN_ARMS_S)

    # ── ARM B: stacked at center ──
    df_B = run_arm(tel, pub, "B", df, zscores)
    df_B.to_parquet(OUT_DIR / "arm_B_stacked.parquet")
    log(f"  saved arm_B_stacked.parquet  ({len(df_B)} rows)")

    # ── cooldown ──
    log(f"\ncooldown {COOLDOWN_BETWEEN_ARMS_S}s before arm C")
    time.sleep(COOLDOWN_BETWEEN_ARMS_S)

    # ── ARM C: spatially separated ──
    df_C = run_arm(tel, pub, "C", df, zscores)
    df_C.to_parquet(OUT_DIR / "arm_C_separated.parquet")
    log(f"  saved arm_C_separated.parquet  ({len(df_C)} rows)")

    log(f"\n=== ALL ARMS COMPLETE ===  total wall time {(time.time()-t0)/3600:.2f}h")
    log(f"output: {OUT_DIR}")
    tel.stop()


if __name__ == "__main__":
    main()