Warrior_EA/research/wyckoff_gold.py

"""
Gold-set generator for Wyckoff structure labelling (2026-10-05).

Each sample = the last W H1 bars ending at a decision bar t.  CAUSAL: nothing after t is drawn, so a
labeller (a vision model or a human) sees exactly what a trader sees at t.  ANONYMOUS: no dates, no
symbol, price rescaled to ATR units from the window's first close, volume to the window's median, so
the labeller cannot recall the history.  The key (symbol, timestamp) goes to a separate CSV.
Sealed: nothing dated >= 2025-01-01 is ever sampled.
  python wyckoff_gold.py SP500,NAS100 40 pilot      # symbols, n per symbol, set name
"""
import os, sys, numpy as np, pandas as pd
import matplotlib; matplotlib.use("Agg")
import matplotlib.pyplot as plt
sys.path.insert(0, os.path.dirname(__file__))
import discover as D

W = 120
OUT = os.path.join(os.path.dirname(__file__), "wyckoff_gold")


def render(b: pd.DataFrame, path: str):
    c0 = b.c.iloc[0]
    atr = (b.h - b.l).rolling(14, min_periods=1).mean().iloc[-1]
    sc = lambda x: (x - c0) / atr
    v = b.tv / b.tv.median()
    fig, (a, v_ax) = plt.subplots(2, 1, figsize=(11, 6.5), sharex=True, gridspec_kw={"height_ratios": [3, 1]})
    x = np.arange(len(b))
    up = b.c >= b.o
    a.vlines(x, sc(b.l), sc(b.h), color="k", lw=0.8)
    a.bar(x, (sc(b.c) - sc(b.o)).abs().clip(lower=0.02), bottom=np.minimum(sc(b.o), sc(b.c)), width=0.7,
          color=np.where(up, "#2a9d4a", "#d62f2f"))
    a.set_ylabel("price (ATR units)"); a.grid(alpha=.25)
    v_ax.bar(x, v, color=np.where(up, "#2a9d4a", "#d62f2f"), width=0.8); v_ax.set_ylabel("rel. volume")
    v_ax.set_xlabel("bar (last bar = now; nothing to the right exists)"); v_ax.grid(alpha=.25)
    fig.tight_layout(); fig.savefig(path, dpi=80); plt.close(fig)


def main():
    syms = sys.argv[1].split(","); n = int(sys.argv[2]); name = sys.argv[3]
    d = os.path.join(OUT, name); os.makedirs(d, exist_ok=True)
    rng = np.random.default_rng(int(os.environ.get("SEED","7"))); key = []
    for s in syms:
        h = D.load_h1(s); h = h[h.index < D.SEAL]
        # event-driven sampling: decision bars where price moved >=2 ATR from the 20-bar mean or at random
        idx = rng.choice(np.arange(W, len(h)), n, replace=False)
        for i in sorted(idx):
            sid = f"{len(key):04d}"
            render(h.iloc[i - W + 1:i + 1], os.path.join(d, sid + ".png"))
            key.append((sid, s, h.index[i], h.c.iloc[i]))
    pd.DataFrame(key, columns=["id", "sym", "t", "close"]).to_csv(os.path.join(d, "_key.csv"), index=False)
    print(len(key), "charts ->", d)


if __name__ == "__main__":
    main()