"""regime_classifier_injector.py — 3-channel lattice regime classifier.

DIRECTIVE FROM CLAUDE DESKTOP — Option B (2026-06-08)
After 6h of 8-tension encoding confirming the predictive-feature path is dead,
abandon "lattice as predictor". Use the lattice as a BINARY REGIME CLASSIFIER
on top of trade_count (the strongest validated raw signal).

INPUTS — three channels, injected per minute (no 4hr aggregation, no levels/vels):
  1. trade_count z-score (rolling 500min one-sided) → CENTER (512, 512)
  2. taker_buy_usd  z-score (rolling 500min one-sided) → (400, 512)
  3. taker_sell_usd z-score (rolling 500min one-sided) → (624, 512)

Strength = z / 3 capped ±1.0. Sigma = 32. No exotic geometry: just three
gaussian blobs along x, with buy at left and sell at right so an imbalance
creates a directional asymmetry in the stress field.

Why these three:
  - trade_count: validated +1.0 Spearman rank stability IS/OOS.
  - taker_buy + taker_sell SEPARATELY (not signed_flow): signed_flow shows
    zero predictive value; the components separately carry volume information
    the lattice can integrate.

OUTPUTS:
  Same 9 telemetry channels as before (asymmetry, coherence, stress_xx, ...).
  Downstream analyzer uses arm A free-running quantiles to set p10/p90
  thresholds for HIGH/LOW/NEUTRAL regime labels.

SUCCESS CRITERION (downstream test, not in this script):
  |fwd_60| volatility ratio for (HIGH TENSION ∩ top-decile trade_count) minutes
  vs (top-decile trade_count alone). Current best is 2.42x. Beat that or the
  lattice is redundant.

OUTPUT FILES:
  /mnt/d/Resonance_Engine/traj/regime_<RUN_ID>/
    meta.json
    arm_A_no_inject.parquet
    arm_T_3ch.parquet
    progress.log
"""
from __future__ import annotations
import argparse, glob, json, queue, threading, time, urllib.request, urllib.error
from pathlib import Path

import numpy as np
import pandas as pd
import zmq

# ───────────────────────── config ─────────────────────────
DATA_ROOT = "/mnt/d/PaperTrader/research/hl_data/minutes"
COIN      = "BTC"

DEFAULT_DAYS_GLOB = "20260415*"

# z-score rolling window (one-sided, past-only via shift(1))
Z_WIN_MIN = 500
Z_MIN_PERIODS = 100

# strength = z / Z_DENOM capped to STR_CAP
Z_DENOM  = 3.0
STR_CAP  = 1.0

# Lattice spatial layout (1024×1024)
CENTER       = (512.0, 512.0)
TAKER_BUY_XY = (400.0, 512.0)   # left of center
TAKER_SELL_XY = (624.0, 512.0)  # right of center
SIGMA        = 32.0

# Timing
PER_MINUTE_MS         = 100     # 10× realtime
WAIT_AFTER_INJECT_MS  = 60      # only 3 injects, less settling time needed
COOLDOWN_BETWEEN_ARMS_S = 1200  # 20 min

# 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"]

# Fractonaut HTTP API (read-only observer, port 28822 — note: 28821 in directive
# was wrong; actual code uses 28822)
FRACTO_URL = "http://127.0.0.1:28822/ask"
FRACTO_TIMEOUT_S = 180        # Ollama can take 30-60s for gemma3:4b responses
FRACTO_QUEUE_MAX = 64         # cap pending queries — drop new ones if full

RUN_ID  = time.strftime("%Y%m%dT%H%M%S")
OUT_DIR = Path(f"/mnt/d/Resonance_Engine/traj/regime_{RUN_ID}")
PROGRESS = OUT_DIR / "progress.log"
FRACTO_LOG = OUT_DIR / "fractonaut_observations.jsonl"


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


# ───────────────────────── Fractonaut querier ─────────────────────────
class FractoQuerier:
    """Fire-and-forget background querier. Submit() returns immediately;
    a worker thread serializes calls to /ask (Ollama only handles one
    request at a time anyway) and appends responses to fractonaut_observations.jsonl.
    """

    def __init__(self, url: str = FRACTO_URL):
        self.url = url
        self.q: queue.Queue = queue.Queue(maxsize=FRACTO_QUEUE_MAX)
        self.dropped = 0
        self.submitted = 0
        self.responded = 0
        self.errors = 0
        self._stop = threading.Event()
        self._t = threading.Thread(target=self._worker, daemon=True)
        self._t.start()

    def submit(self, question: str, ctx: dict) -> None:
        """Non-blocking. Drops the query if the queue is full."""
        item = {
            "submit_wall": time.time(),
            "question": question,
            "context": ctx,
        }
        try:
            self.q.put_nowait(item)
            self.submitted += 1
        except queue.Full:
            self.dropped += 1

    def _worker(self) -> None:
        while not self._stop.is_set():
            try:
                item = self.q.get(timeout=1.0)
            except queue.Empty:
                continue
            self._call(item)
            self.q.task_done()

    def _call(self, item: dict) -> None:
        body = json.dumps({"question": item["question"]}).encode()
        req = urllib.request.Request(
            self.url, data=body,
            headers={"Content-Type": "application/json"}, method="POST",
        )
        t0 = time.time()
        resp_text = None
        err = None
        try:
            with urllib.request.urlopen(req, timeout=FRACTO_TIMEOUT_S) as r:
                resp_json = json.loads(r.read())
            resp_text = resp_json.get("response", "")
            self.responded += 1
        except Exception as e:
            err = str(e)
            self.errors += 1
        elapsed = time.time() - t0

        entry = {
            "submit_wall": item["submit_wall"],
            "response_wall": time.time(),
            "elapsed_s": elapsed,
            "context": item["context"],
            "question": item["question"],
            "response": resp_text,
            "error": err,
        }
        OUT_DIR.mkdir(parents=True, exist_ok=True)
        with FRACTO_LOG.open("a") as f:
            f.write(json.dumps(entry) + "\n")

    def drain(self, max_wait_s: float = 600.0) -> None:
        """Block until queue is empty or timeout."""
        deadline = time.time() + max_wait_s
        while time.time() < deadline:
            if self.q.empty():
                # let in-flight call finish
                time.sleep(1)
                if self.q.empty():
                    return
            time.sleep(2)

    def stop(self) -> None:
        self._stop.set()
        self._t.join(timeout=5)

    def stats(self) -> str:
        return (f"fracto submitted={self.submitted} responded={self.responded} "
                f"dropped={self.dropped} errors={self.errors} pending={self.q.qsize()}")


# ───────────────────────── ZMQ telemetry ─────────────────────────
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()
        self._t.join(timeout=2)
        try:
            self.sock.close(0)
        except Exception:
            pass


# ───────────────────────── data loading ─────────────────────────
def load_btc(days_glob: str) -> 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:
        for f in sorted(glob.glob(f"{d}/*.parquet")):
            dfs.append(pd.read_parquet(f))
    df = pd.concat(dfs, ignore_index=True)
    df = df[df.coin == COIN].sort_values("minute").drop_duplicates("minute").reset_index(drop=True)
    log(f"loaded {len(df)} unique minutes of {COIN}")
    return df


def compute_zscores(df: pd.DataFrame) -> pd.DataFrame:
    """Rolling 500-min one-sided z-scores for trade_count, taker_buy_usd,
    taker_sell_usd. Past-only (shift(1) before rolling)."""
    out = df.copy()
    for col in ["trade_count", "taker_buy_usd", "taker_sell_usd"]:
        s = out[col].astype(float).shift(1)
        mu = s.rolling(Z_WIN_MIN, min_periods=Z_MIN_PERIODS).mean()
        sd = s.rolling(Z_WIN_MIN, min_periods=Z_MIN_PERIODS).std()
        out[f"{col}_z"] = (out[col].astype(float) - mu) / sd.replace(0, np.nan)
    return out


# ───────────────────────── injection ─────────────────────────
def _send_inject(pub: zmq.Socket, x: float, y: float,
                 strength: float, sigma: float = SIGMA) -> None:
    payload = {"cmd": "inject_density",
               "x": float(x), "y": float(y),
               "sigma": float(sigma),
               "strength": float(strength)}
    pub.send_string(json.dumps(payload))


def _z_to_strength(z: float) -> float | None:
    if not np.isfinite(z):
        return None
    s = float(np.clip(z / Z_DENOM, -STR_CAP, +STR_CAP))
    if abs(s) < 1e-6:
        return None
    return s


def encode_trade_count(pub: zmq.Socket, z: float) -> None:
    s = _z_to_strength(z)
    if s is None: return
    cx, cy = CENTER
    _send_inject(pub, cx, cy, s)


def encode_taker_buy(pub: zmq.Socket, z: float) -> None:
    s = _z_to_strength(z)
    if s is None: return
    x, y = TAKER_BUY_XY
    _send_inject(pub, x, y, s)


def encode_taker_sell(pub: zmq.Socket, z: float) -> None:
    s = _z_to_strength(z)
    if s is None: return
    x, y = TAKER_SELL_XY
    _send_inject(pub, x, y, s)


# ───────────────────────── per-arm runner ─────────────────────────
def classify_regime(asym: float, coh: float,
                    asym_p10: float | None, asym_p90: float | None,
                    coh_p10: float | None, coh_p90: float | None) -> str:
    """Return one of HIGH / LOW / NEUTRAL. Returns NEUTRAL if thresholds
    not yet known or telemetry invalid."""
    if any(v is None or not np.isfinite(v) for v in
           (asym, coh, asym_p10, asym_p90, coh_p10, coh_p90)):
        return "NEUTRAL"
    if asym > asym_p90 or coh < coh_p10:
        return "HIGH"
    if asym < asym_p10 and coh > coh_p90:
        return "LOW"
    return "NEUTRAL"

def snapshot_row(tel: LatestTel, minute: int, arm: str,
                 row: pd.Series) -> dict:
    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,
        # raw values for downstream analysis
        "trade_count":     float(row.get("trade_count", np.nan)),
        "taker_buy_usd":   float(row.get("taker_buy_usd", np.nan)),
        "taker_sell_usd":  float(row.get("taker_sell_usd", np.nan)),
        "mid_price":       float(row.get("mid_price", np.nan)),
        # z-scores actually injected
        "trade_count_z":     row.get("trade_count_z"),
        "taker_buy_usd_z":   row.get("taker_buy_usd_z"),
        "taker_sell_usd_z":  row.get("taker_sell_usd_z"),
    }
    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,
            fracto: FractoQuerier | None = None,
            thresholds: dict | None = None) -> pd.DataFrame:
    """Run one arm.

    If `fracto` and `thresholds` are supplied, regime labels (HIGH/LOW/NEUTRAL)
    are computed each minute, transitions are detected, and a Fractonaut /ask
    query is fired (non-blocking) on every transition.
    """
    log(f"\n=== arm {arm}: starting ({len(df)} minutes) ===")
    if arm == "A":
        log("  arm A: NO INJECTIONS — sampling state at cadence only")
    elif arm == "T":
        log("  arm T: 3 Z-SCORE INJECTIONS per minute (trade_count center, buy left, sell right)")
    if fracto is not None and thresholds is not None:
        log(f"  fractonaut transitions ENABLED  thresholds={thresholds}")

    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

    prev_regime = "NEUTRAL"
    n_transitions = 0

    for i in range(n_minutes):
        t_tick = time.time()
        row = df.iloc[i]

        if arm == "T" and pub is not None:
            encode_trade_count(pub, row.get("trade_count_z"))
            encode_taker_buy  (pub, row.get("taker_buy_usd_z"))
            encode_taker_sell (pub, row.get("taker_sell_usd_z"))
            time.sleep(wait_inject_s)
        else:
            time.sleep(wait_inject_s)

        rec = snapshot_row(tel, int(row.minute), arm, row)

        # Regime classification & transition detection (arm T only — arm A is the
        # threshold source, we don't ask Fractonaut about its own baseline)
        if arm == "T" and fracto is not None and thresholds is not None:
            asym = rec.get("asymmetry")
            coh  = rec.get("coherence")
            label = classify_regime(
                asym, coh,
                thresholds.get("asym_p10"), thresholds.get("asym_p90"),
                thresholds.get("coh_p10"),  thresholds.get("coh_p90"),
            )
            rec["regime"] = label
            if label != prev_regime and label in ("HIGH", "LOW"):
                n_transitions += 1
                tc_z   = row.get("trade_count_z", float("nan"))
                buy_z  = row.get("taker_buy_usd_z", float("nan"))
                sell_z = row.get("taker_sell_usd_z", float("nan"))
                question = (
                    f"Field just entered {label} TENSION state "
                    f"(from {prev_regime}). "
                    f"asymmetry={asym:.1f} coherence={coh:.4f} "
                    f"trade_count_z={tc_z:.2f} buy_z={buy_z:.2f} sell_z={sell_z:.2f}. "
                    f"Have you seen this configuration before? "
                    f"What pattern does this resemble and what followed last time?"
                )
                fracto.submit(question, {
                    "event": "regime_transition",
                    "minute": int(row.minute),
                    "from": prev_regime, "to": label,
                    "asymmetry": float(asym) if asym is not None else None,
                    "coherence": float(coh) if coh is not None else None,
                    "trade_count_z": float(tc_z) if np.isfinite(tc_z) else None,
                    "taker_buy_usd_z": float(buy_z) if np.isfinite(buy_z) else None,
                    "taker_sell_usd_z": float(sell_z) if np.isfinite(sell_z) else None,
                    "stress_xx": rec.get("stress_xx"),
                    "stress_yy": rec.get("stress_yy"),
                })
            elif label != prev_regime and label == "NEUTRAL" and prev_regime in ("HIGH", "LOW"):
                n_transitions += 1
                question = (
                    f"Field just exited {prev_regime} TENSION back to NEUTRAL. "
                    f"asymmetry={asym:.1f} coherence={coh:.4f}. "
                    f"How did the configuration evolve while in {prev_regime}? "
                    f"Did anything notable change before the return to baseline?"
                )
                fracto.submit(question, {
                    "event": "regime_transition",
                    "minute": int(row.minute),
                    "from": prev_regime, "to": label,
                    "asymmetry": float(asym) if asym is not None else None,
                    "coherence": float(coh) if coh is not None else None,
                })
            prev_regime = label

        records.append(rec)

        spent = time.time() - t_tick
        remaining = step_s - spent
        if remaining > 0:
            time.sleep(remaining)

        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 {}
            asym = cur_field.get('asymmetry', float('nan'))
            extras = ""
            if fracto is not None and arm == "T":
                extras = f"  trans={n_transitions}  " + fracto.stats()
            log(f"  arm {arm} {i+1}/{n_minutes} ({pct:.1f}%)  "
                f"rate={rate_per_s:.1f}min/s  ETA={eta_s/60:.1f}min  "
                f"asym={asym:.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}{extras}")
            last_status = time.time()

    log(f"  arm {arm} COMPLETE  records={len(records)}  total={(time.time()-t_arm_start)/60:.1f}min")
    if fracto is not None and arm == "T":
        log(f"  arm T transitions={n_transitions}  {fracto.stats()}")
    return pd.DataFrame(records)


# ───────────────────────── main ─────────────────────────
def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--days-glob", default=DEFAULT_DAYS_GLOB,
                    help="day-dir glob under hl_data/minutes/")
    ap.add_argument("--max-minutes", type=int, default=None,
                    help="cap minutes (shakedown)")
    ap.add_argument("--arms", default="A,T",
                    help="comma-separated arms to run (default A,T)")
    ap.add_argument("--cooldown-s", type=int, default=COOLDOWN_BETWEEN_ARMS_S,
                    help="seconds between arms (only after arm T)")
    args = ap.parse_args()

    log(f"=== regime_classifier_injector run_id={RUN_ID} ===")
    log(f"out_dir={OUT_DIR}")
    log(f"args: days_glob={args.days_glob} max_minutes={args.max_minutes} "
        f"arms={args.arms}")

    df = load_btc(args.days_glob)
    if args.max_minutes:
        df = df.head(args.max_minutes).reset_index(drop=True)
        log(f"capped to first {args.max_minutes} minutes")

    df = compute_zscores(df)
    log("z-scores computed (past-only 500min rolling):")
    for col in ["trade_count_z", "taker_buy_usd_z", "taker_sell_usd_z"]:
        s = df[col]
        finite = int(np.isfinite(s.astype(float)).sum())
        log(f"  {col:<22}  n_finite={finite:>5}/{len(s)}  "
            f"mean={s.mean():+.4f}  std={s.std():+.4f}  "
            f"min={s.min():+.4f}  max={s.max():+.4f}")

    meta = {
        "run_id": RUN_ID,
        "coin": COIN,
        "data_root": DATA_ROOT,
        "days_glob": args.days_glob,
        "n_minutes": int(len(df)),
        "minute_range": [int(df.minute.min()), int(df.minute.max())],
        "z_win_min": Z_WIN_MIN,
        "z_denom": Z_DENOM,
        "str_cap": STR_CAP,
        "spatial": {
            "center": list(CENTER),
            "taker_buy_xy": list(TAKER_BUY_XY),
            "taker_sell_xy": list(TAKER_SELL_XY),
            "sigma": SIGMA,
        },
        "timing": {
            "per_minute_ms": PER_MINUTE_MS,
            "wait_after_inject_ms": WAIT_AFTER_INJECT_MS,
            "cooldown_between_arms_s": COOLDOWN_BETWEEN_ARMS_S,
        },
        "arms": args.arms.split(","),
    }
    (OUT_DIR / "meta.json").write_text(json.dumps(meta, indent=2))

    tel = LatestTel(TEL_ADDR)
    ctx = zmq.Context.instance()
    pub = ctx.socket(zmq.PUB)
    pub.connect(CMD_ADDR)
    log(f"ZMQ connected: SUB {TEL_ADDR}, PUB {CMD_ADDR}")
    log("warming up SUB socket (3s) ...")
    time.sleep(3)
    log(f"  initial tel_seen={tel.n_seen}  latest_age_ms="
        f"{(time.time()-tel.latest_wall)*1000 if tel.latest_wall else -1:.0f}")

    # Fractonaut querier (fire-and-forget). It is OK if the HTTP service is down
    # — the worker will just log errors; the run keeps going.
    fracto = FractoQuerier()
    log(f"FractoQuerier ready; HTTP target = {FRACTO_URL}")

    arm_A_telemetry: pd.DataFrame | None = None
    thresholds: dict | None = None

    try:
        arms_to_run = [a.strip().upper() for a in args.arms.split(",")]
        for ai, arm in enumerate(arms_to_run):
            if arm not in ("A", "T"):
                log(f"skipping unknown arm: {arm}")
                continue
            if ai > 0 and arms_to_run[ai - 1] == "T":
                log(f"--- cooldown {args.cooldown_s}s ---")
                time.sleep(args.cooldown_s)

            if arm == "T" and arm_A_telemetry is not None:
                # Derive thresholds from arm A telemetry distribution
                arm_A_clean = arm_A_telemetry.dropna(subset=["asymmetry", "coherence"])
                thresholds = {
                    "asym_p10": float(arm_A_clean["asymmetry"].quantile(0.10)),
                    "asym_p90": float(arm_A_clean["asymmetry"].quantile(0.90)),
                    "coh_p10":  float(arm_A_clean["coherence"].quantile(0.10)),
                    "coh_p90":  float(arm_A_clean["coherence"].quantile(0.90)),
                }
                log(f"derived thresholds from arm A: {thresholds}")
                # START-OF-T summary query
                fracto.submit(
                    ("Arm A free-running baseline complete. "
                     "Arm T market injection beginning now. "
                     "Three channels: trade_count at centre (512,512), "
                     "taker_buy at x=400, taker_sell at x=624, each as "
                     "rolling-500min z-score capped ±1. "
                     "Summarise what you observed during arm A. "
                     "What is the baseline field character?"),
                    {"event": "arm_T_start", "thresholds": thresholds,
                     "arm_A_rows": int(len(arm_A_clean))},
                )

            arm_df = run_arm(
                tel,
                pub if arm == "T" else None,
                arm, df,
                fracto=fracto if arm == "T" else None,
                thresholds=thresholds if arm == "T" else None,
            )
            outf = OUT_DIR / f"arm_{arm}_{'3ch' if arm=='T' else 'no_inject'}.parquet"
            arm_df.to_parquet(outf)
            log(f"  saved {outf.name}  ({len(arm_df)} rows)")
            if arm == "A":
                arm_A_telemetry = arm_df
            elif arm == "T":
                # END-OF-T summary query
                fracto.submit(
                    ("Arm T complete. Injection has been running for the full "
                     "April dataset. What patterns did you observe during the "
                     "injection period that were absent or different from the "
                     "arm A baseline? What field configurations appeared "
                     "repeatedly at HIGH TENSION moments?"),
                    {"event": "arm_T_end", "arm_T_rows": int(len(arm_df))},
                )
        log(f"\n=== ALL ARMS COMPLETE ===")
        log(f"output: {OUT_DIR}")
        log(f"draining FractoQuerier (timeout 600s) ...  {fracto.stats()}")
        fracto.drain(max_wait_s=600)
        log(f"fracto drain complete  {fracto.stats()}")
    finally:
        fracto.stop()
        tel.stop()
        pub.close(0)
        ctx.term()


if __name__ == "__main__":
    main()