"""data_loader.py -- the ONE way to read an M1 quote file in this package: true clock, UTC, split-enforcing, session cubes.

Phase 0 finding (out/PHASE0_DATA_AUDIT.md): every liquid file is stamped at the END of the minute under the zone in lab_utc.TZ_MAP, in summer and winter.
`load_true_clock` converts the zone to UTC and subtracts 1 minute, so the index is the START of the minute (the clock the live MT5 bot sees).
Reading a window after the Discovery end requires an unlock record written by the stage-transition step (see `unlock`); until then those rows are not
returned, whatever the caller asks for. `build_cubes` turns the minute table into [days, minutes] arrays for one pinned session window.
"""
from __future__ import annotations

import hashlib
import json
import os
import sys
from dataclasses import dataclass, field

import numpy as np
import pandas as pd
import pyarrow.parquet as pq

HERE = os.path.dirname(os.path.abspath(__file__))
sys.path.insert(0, HERE)
sys.path.insert(0, os.path.dirname(HERE))
import config as C  # noqa: E402
from lab_utc import TZ_MAP  # noqa: E402

_COLS = ["date", "bid_open", "ask_open", "bid_high", "ask_high", "bid_low", "ask_low", "bid_close", "ask_close"]
STAMP_CORRECTION_MIN = 1          # file stamp = end of minute (Phase 0, 433 releases, both seasons)
UNLOCK_FILE = os.path.join(C.OUT, "unlock.json")
_STAGES = ("discovery", "validation", "holdout")
# [withheld: refers to the lab's private live strategy]


def _unlocked() -> dict:
    try:
        return json.load(open(UNLOCK_FILE))
    except Exception:
        return {}


def unlock(stage: str, cells_sha256: str, note: str = "") -> None:
    """Called ONLY by the stage-transition step, together with a ledger row and the hashed frozen list of cells that go forward."""
    assert stage in ("validation", "holdout")
    u = _unlocked()
    u[stage] = dict(cells_sha256=cells_sha256, note=note, utc=pd.Timestamp.now("UTC").isoformat())
    json.dump(u, open(UNLOCK_FILE, "w"), indent=1)


def stage_bounds(stage: str) -> tuple[pd.Timestamp, pd.Timestamp]:
    far = pd.Timestamp("2100-01-01", tz="UTC")
    if stage == "discovery":
        return pd.Timestamp("1990-01-01", tz="UTC"), pd.Timestamp(C.DISCOVERY_END, tz="UTC") + pd.Timedelta(days=1)
    if stage == "validation":
        return pd.Timestamp(C.VALIDATION_START, tz="UTC"), pd.Timestamp(C.VALIDATION_END, tz="UTC") + pd.Timedelta(days=1)
    if stage == "holdout":
        return pd.Timestamp(C.HOLDOUT_START, tz="UTC"), far
    raise ValueError(stage)


def load_true_clock(symbol: str, stage: str = "discovery", control_token: str | None = None) -> pd.DataFrame:
    """UTC ns index (start of minute), columns bid_*/ask_* OHLC. Rows after the allowed stage window are never returned.
    # [withheld: refers to the lab's private live strategy]
    if stage not in _STAGES + ("all", "control", "through_validation"):
        raise ValueError(stage)
    if stage == "control":
        if control_token != _CONTROL_TOKEN:
            # [withheld: refers to the lab's private live strategy]
        need = []
    else:
        need = {"discovery": [], "validation": ["validation"], "holdout": ["holdout"], "all": ["validation", "holdout"],
                "through_validation": ["validation"]}[stage]
    u = _unlocked()
    missing = [s for s in need if s not in u]
    if missing:
        raise PermissionError(f"stage(s) {missing} are locked: no unlock record in {UNLOCK_FILE}")
    _, tz = TZ_MAP[symbol]
    tbl = pq.read_table(C.parquet_path(symbol), columns=_COLS).to_pandas()
    if "date" in tbl.columns:
        tbl = tbl.set_index("date")
    idx = pd.DatetimeIndex(tbl.index)
    loc = idx if idx.tz is not None else idx.tz_localize(tz, ambiguous="NaT", nonexistent="NaT")
    ok = ~pd.isna(loc)
    tbl = tbl[np.asarray(ok)]
    tbl.index = loc[ok].tz_convert("UTC").as_unit("ns") - pd.Timedelta(minutes=STAMP_CORRECTION_MIN)
    tbl = tbl[~tbl.index.duplicated(keep="last")].sort_index()
    tbl = tbl[(tbl["ask_close"] >= tbl["bid_close"]) | tbl["ask_close"].isna()]
    if stage in _STAGES:
        lo, hi = stage_bounds(stage)
        tbl = tbl[(tbl.index >= lo) & (tbl.index < hi)]
    elif stage == "through_validation":                  # warm-up history + validation days; nothing after the validation end
        _, hi = stage_bounds("validation")
        tbl = tbl[tbl.index < hi]
    return tbl


def mid(df: pd.DataFrame) -> pd.DataFrame:
    return pd.DataFrame({k: 0.5 * (df[f"bid_{k}"] + df[f"ask_{k}"]) for k in ("open", "high", "low", "close")}, index=df.index)


def embargo_mask(index: pd.DatetimeIndex, stage: str, n_days: int = C.EMBARGO_TRADING_DAYS) -> np.ndarray:
    """True for rows that may be used: drops `n_days` weekdays at the start and at the end of the validation window."""
    if stage != "validation":
        return np.ones(len(index), dtype=bool)
    lo, hi = stage_bounds("validation")
    days = pd.bdate_range(lo.tz_localize(None), (hi - pd.Timedelta(days=1)).tz_localize(None))
    a, b = days[n_days], days[-n_days - 1] + pd.Timedelta(days=1)
    t = index.tz_convert("UTC").tz_localize(None)
    return np.asarray((t >= a) & (t < b))


def file_sha256(symbol: str) -> str:
    h = hashlib.sha256()
    with open(C.parquet_path(symbol), "rb") as f:
        for chunk in iter(lambda: f.read(1 << 20), b""):
            h.update(chunk)
    return h.hexdigest()


# ----------------------------------------------------------------------------------------------------------------- session cubes
@dataclass
class Quotes:
    """Minute table as flat numpy arrays (shared by every window of one asset)."""
    symbol: str
    idx_ns: np.ndarray                      # UTC ns, start of minute, sorted unique
    mo: np.ndarray; mh: np.ndarray; ml: np.ndarray; mc: np.ndarray      # mid
    sp: np.ndarray                          # ask_open - bid_open
    bo: np.ndarray; bh: np.ndarray; bl: np.ndarray; bc: np.ndarray      # bid
    ao: np.ndarray; ah: np.ndarray; al: np.ndarray; ac: np.ndarray      # ask


def to_quotes(symbol: str, df: pd.DataFrame) -> Quotes:
    a = lambda c: df[c].to_numpy(dtype=np.float64)
    bo, bh, bl, bc = a("bid_open"), a("bid_high"), a("bid_low"), a("bid_close")
    ao, ah, al, ac = a("ask_open"), a("ask_high"), a("ask_low"), a("ask_close")
    return Quotes(symbol, df.index.asi8, 0.5 * (bo + ao), 0.5 * (bh + ah), 0.5 * (bl + al), 0.5 * (bc + ac), ao - bo, bo, bh, bl, bc, ao, ah, al, ac)


@dataclass
class Cubes:
    """[days, minutes] arrays for one pinned session window. Minute index 0 = window start, index T_excl = the flat minute."""
    symbol: str
    window: str
    wcfg: dict
    T_excl: int
    dates: np.ndarray                       # datetime64[D], weekdays
    year: np.ndarray
    utc0: np.ndarray                        # int64 ns of the window start per day
    o: np.ndarray; h: np.ndarray; l: np.ndarray; c: np.ndarray; sp: np.ndarray       # mid + spread, NaN where the minute is missing
    strict: dict = field(default_factory=dict)                                       # bid/ask arrays when requested

    @property
    def D(self) -> int:
        return len(self.dates)

    @property
    def T(self) -> int:
        return self.T_excl + 1


def hm_to_min(s: str) -> int:
    return int(s[:2]) * 60 + int(s[3:])


def build_cubes(q: Quotes, window: str, wcfg: dict, strict: bool = False, dates: pd.DatetimeIndex | None = None) -> Cubes:
    tz = wcfg["tz"]
    T_excl = hm_to_min(wcfg["end"]) - hm_to_min(wcfg["start"])
    T = T_excl + 1
    if dates is None:
        d0 = pd.Timestamp(int(q.idx_ns[0]), tz="UTC").tz_localize(None).normalize()
        d1 = pd.Timestamp(int(q.idx_ns[-1]), tz="UTC").tz_localize(None).normalize()
        dates = pd.bdate_range(d0, d1)
    utc0 = np.array([pd.Timestamp(f"{d.date()} {wcfg['start']}", tz=tz).tz_convert("UTC").value for d in dates], dtype=np.int64)
    grid = utc0[:, None] + (np.arange(T, dtype=np.int64) * 60_000_000_000)[None, :]
    pos = np.searchsorted(q.idx_ns, grid.ravel()).reshape(grid.shape)
    posc = np.minimum(pos, len(q.idx_ns) - 1)
    ok = q.idx_ns[posc] == grid
    take = lambda a: np.where(ok, a[posc], np.nan)
    cub = Cubes(q.symbol, window, wcfg, T_excl, dates.values.astype("datetime64[D]"), dates.year.values.astype(np.int16), utc0,
                take(q.mo), take(q.mh), take(q.ml), take(q.mc), take(q.sp))
    if strict:
        cub.strict = dict(bo=take(q.bo), bh=take(q.bh), bl=take(q.bl), bc=take(q.bc), ao=take(q.ao), ah=take(q.ah), al=take(q.al), ac=take(q.ac))
    return cub


def split_masks(dates: np.ndarray) -> dict:
    """Boolean day masks for the three stages (validation keeps the 5-trading-day embargo at both ends)."""
    d = pd.DatetimeIndex(dates)
    disc = d <= pd.Timestamp(C.DISCOVERY_END)
    val_all = (d >= pd.Timestamp(C.VALIDATION_START)) & (d <= pd.Timestamp(C.VALIDATION_END))
    val = val_all.copy()
    vi = np.flatnonzero(val_all)
    if len(vi) > 2 * C.EMBARGO_TRADING_DAYS:
        val[vi[:C.EMBARGO_TRADING_DAYS]] = False
        val[vi[-C.EMBARGO_TRADING_DAYS:]] = False
    hold = d >= pd.Timestamp(C.HOLDOUT_START)
    hi = np.flatnonzero(hold)
    # embargo on the holdout side of the validation boundary is the last 5 validation days (above); holdout itself starts clean
    return dict(discovery=np.asarray(disc), validation=np.asarray(val), holdout=np.asarray(hold))


if __name__ == "__main__":
    d = load_true_clock("nas100usd")
    assert d.index.max() < pd.Timestamp(C.DISCOVERY_END, tz="UTC") + pd.Timedelta(days=1), "discovery leak"
    try:
        load_true_clock("nas100usd", "holdout")
        raise SystemExit("FAIL: holdout readable without unlock")
    except PermissionError:
        pass
    q = to_quotes("nas100usd", d)
    cb = build_cubes(q, "US_RTH", C.CONTROL_WINDOW)
    print("cube", cb.o.shape, "valid share", float(np.isfinite(cb.o).mean()), "first day", cb.dates[0], "last", cb.dates[-1])
    ny = pd.Timestamp(int(cb.utc0[np.flatnonzero(cb.dates == np.datetime64("2019-06-03"))[0]]), tz="UTC").tz_convert("America/New_York")
    print("2019-06-03 window start local:", ny)
    print("loader self-test OK")