#!/usr/bin/env python3
"""Data layer for the umar-audit "Initial Balance" backtest (tools/ib_backtest.py).

Source: Dukascopy Bank's free historical datafeed (datafeed.dukascopy.com),
instrument USATECHIDXUSD (Dukascopy's "US Tech 100" CFD on the NASDAQ-100),
BID and ASK 1-minute candles, one LZMA .bi5 file per UTC day and side.
Raw files are cached byte-for-byte under projects/umar-audit/data_raw/dukascopy/
(month in the path is 0-based, exactly as in the feed). Prices are never edited.

Only NYSE trading sessions are fetched (exchange_calendars XNYS: holidays,
early closes and special closures handled). A NY cash session (09:30-16:00 ET)
always falls inside one UTC day, so one daily file per side covers it.

Usage:
  .venv/bin/python tools/ib_data.py fetch [--start 2015-01-01 --end 2026-09-30 --threads 24]
  .venv/bin/python tools/ib_data.py build      # -> data_raw/usatech_session_1m.pkl (+ gaps report)
  .venv/bin/python tools/ib_data.py status     # how many files are cached / missing
  .venv/bin/python tools/ib_data.py histdata   # 2nd source for cross-checks: HistData.com NSXUSD M1 (BID only)

Second source: HistData.com free "NSXUSD" (NASDAQ 100) 1-minute ASCII files, one zip per
year (monthly for the current year). HistData says "EST without DST", but for this symbol
the stamps behave as New York local time (see cmd_histdata). Cached as downloaded under
data_raw/histdata/.
"""
from __future__ import annotations

import argparse
import concurrent.futures as cf
import json
import lzma
import os
import random
import sys
import threading
import time
import urllib.parse
import urllib.request
from datetime import date, datetime, timezone
from pathlib import Path

import numpy as np
import pandas as pd

ROOT = Path(__file__).resolve().parents[1]
RAW = ROOT / "projects/umar-audit/data_raw"
SYMBOL = "USATECHIDXUSD"
DUKA_RAW = RAW / "dukascopy" / SYMBOL
DUKA_URL = f"https://datafeed.dukascopy.com/datafeed/{SYMBOL}"
PRICE_SCALE = 1e3  # Dukascopy index CFDs are stored in 1/1000 points (verified against known NDX levels)
NY = "America/New_York"
SESSION_PKL = RAW / "usatech_session_1m.pkl"
GAPS_JSON = RAW / "usatech_gaps.json"
DEFAULT_START, DEFAULT_END = "2015-01-01", "2026-09-30"


def nyse_schedule(start: str, end: str) -> pd.DataFrame:
    import exchange_calendars as xc
    cal = xc.get_calendar("XNYS", start="2015-01-02", end="2026-12-31")
    sch = cal.schedule.loc[max(pd.Timestamp(start), pd.Timestamp("2015-01-02")):pd.Timestamp(end)]
    out = pd.DataFrame({"open": sch["open"].dt.tz_convert(NY), "close": sch["close"].dt.tz_convert(NY)})
    out.index = pd.DatetimeIndex(out.index.date, name="day")
    return out


def rel_path(d: date, side: str) -> str:
    return f"{d.year}/{d.month - 1:02d}/{d.day:02d}/{side}_candles_min_1.bi5"


# ------------------------------------------------------------------ fetch ---
_lock = threading.Lock()
_stats = {"ok": 0, "cached": 0, "fail": 0, "retries": 0}


def fetch_one(rel: str, max_attempts: int = 30) -> bool:
    path = DUKA_RAW / rel
    if path.exists():
        with _lock:
            _stats["cached"] += 1
        return True
    url = f"{DUKA_URL}/{rel}"
    for attempt in range(max_attempts):
        try:
            req = urllib.request.Request(url, headers={"User-Agent": "Mozilla/5.0"})
            with urllib.request.urlopen(req, timeout=120) as r:
                body = r.read()
            if body:
                lzma.decompress(body)  # validate before caching
            path.parent.mkdir(parents=True, exist_ok=True)
            tmp = path.with_suffix(f".part{os.getpid()}_{threading.get_ident()}")
            tmp.write_bytes(body)
            tmp.rename(path)
            with _lock:
                _stats["ok"] += 1
            return True
        except Exception:  # 503 "no server available" / timeouts are common on the free feed
            with _lock:
                _stats["retries"] += 1
            time.sleep(min(60, 2 + attempt * 2) * random.uniform(0.5, 1.5))
    with _lock:
        _stats["fail"] += 1
    print(f"FAILED {url}", flush=True)
    return False


def cmd_fetch(a) -> None:
    sch = nyse_schedule(a.start, a.end)
    # all BID files first (enough to run the backtest with a modelled spread), then ASK
    jobs = [rel_path(d.date(), s) for s in ("BID", "ASK") for d in sch.index]
    todo = [j for j in jobs if not (DUKA_RAW / j).exists()]
    if a.reverse:  # newest first within each side, so two workers meet in the middle
        todo = [j for j in reversed(todo) if "BID_" in j] + [j for j in reversed(todo) if "ASK_" in j]
    print(f"{len(jobs)} files for {len(sch)} NYSE sessions, {len(todo)} to download", flush=True)
    t0 = time.time()
    with cf.ThreadPoolExecutor(a.threads) as ex:
        futs = [ex.submit(fetch_one, j) for j in todo]
        for i, _ in enumerate(cf.as_completed(futs), 1):
            if i % 50 == 0 or i == len(todo):
                el = time.time() - t0
                print(f"{i}/{len(todo)} done in {el/60:.1f} min ({i/el:.2f}/s) stats={_stats}", flush=True)
    print("fetch finished", _stats, flush=True)


def cmd_status(a) -> None:
    sch = nyse_schedule(a.start, a.end)
    n = miss = 0
    for d in sch.index:
        for s in ("BID", "ASK"):
            n += 1
            miss += not (DUKA_RAW / rel_path(d.date(), s)).exists()
    print(f"{n - miss}/{n} cached, {miss} missing")


# --------------------------------------------------------------- histdata ---
HD_RAW = RAW / "histdata"
HD_PAGE = "https://www.histdata.com/download-free-forex-historical-data/?/ascii/1-minute-bar-quotes/nsxusd/"


def histdata_fetch(period: str) -> Path:
    """period '2019' (a year) or '2026/9' (a month of the current year). Retries on network errors."""
    for attempt in range(6):
        try:
            return _histdata_fetch(period)
        except (OSError, AttributeError) as e:
            print(f"  histdata {period}: {e!r}, retrying", flush=True)
            time.sleep(5 * (attempt + 1))
    raise SystemExit(f"HistData {period} failed")


def _histdata_fetch(period: str) -> Path:
    import re
    import http.cookiejar
    path = HD_RAW / f"NSXUSD_M1_{period.replace('/', '_')}.zip"
    if path.exists():
        return path
    cj = http.cookiejar.CookieJar()
    op = urllib.request.build_opener(urllib.request.HTTPCookieProcessor(cj))
    op.addheaders = [("User-Agent", "Mozilla/5.0")]
    page = op.open(HD_PAGE + period, timeout=60).read().decode("utf-8", "replace")
    form = {k: re.search(rf'name="{k}" id="{k}" value="([^"]*)"', page).group(1)
            for k in ("tk", "date", "datemonth", "platform", "timeframe", "fxpair")}
    data = urllib.parse.urlencode(form).encode()
    req = urllib.request.Request("https://www.histdata.com/get.php", data=data,
                                 headers={"Referer": HD_PAGE + period})
    body = op.open(req, timeout=120).read()
    if body[:2] != b"PK":
        raise SystemExit(f"HistData {period}: not a zip ({body[:80]!r})")
    HD_RAW.mkdir(parents=True, exist_ok=True)
    path.write_bytes(body)
    time.sleep(2)
    return path


def cmd_histdata(a) -> None:
    import zipfile
    periods = [str(y) for y in range(2015, 2026)] + [f"2026/{m}" for m in range(1, 10)]
    frames = []
    for per in periods:
        z = zipfile.ZipFile(histdata_fetch(per))
        name = next(n for n in z.namelist() if n.endswith(".csv"))
        df = pd.read_csv(z.open(name), sep=";", header=None, names=["ts", "o", "h", "l", "c", "v"])
        frames.append(df)
        print(per, len(df), flush=True)
    df = pd.concat(frames)
    # HistData documents "EST without DST", but for NSXUSD the minute returns line up with
    # Dukascopy only if the stamps are read as New York LOCAL time (with DST): in summer the
    # "+5h" reading is 60 min off (lag-60 return correlation 0.97-1.00, lag-0 ~0). Checked in
    # BACKTEST.md; times in the DST switch hour (Sunday night, outside the session) are dropped.
    t = pd.DatetimeIndex(pd.to_datetime(df["ts"], format="%Y%m%d %H%M%S"))
    df.index = t.tz_localize(NY, ambiguous="NaT", nonexistent="NaT")
    df = df[df.index.notna()]
    df = df[~df.index.duplicated()].sort_index().drop(columns=["ts", "v"])
    sch = nyse_schedule(a.start, a.end)
    out = []
    for d, row in sch.iterrows():
        g = df.loc[row["open"]:row["close"] - pd.Timedelta(minutes=1)]
        if len(g):
            g = g.copy()
            g["day"] = d
            out.append(g)
    res = pd.concat(out)
    res.to_pickle(RAW / "histdata_session_1m.pkl")
    print(f"histdata sessions: {res['day'].nunique()} of {len(sch)}, rows {len(res)}")


# ------------------------------------------------------------------ parse ---
def read_day(d: date, side: str) -> pd.DataFrame:
    path = DUKA_RAW / rel_path(d, side)
    cols = ["o", "h", "l", "c", "v"]
    if not path.exists():
        return pd.DataFrame(columns=cols, index=pd.DatetimeIndex([], tz="UTC"))
    body = path.read_bytes()
    if not body:
        return pd.DataFrame(columns=cols, index=pd.DatetimeIndex([], tz="UTC"))
    a = np.frombuffer(lzma.decompress(body), dtype=">i4").reshape(-1, 6)
    base = pd.Timestamp(datetime(d.year, d.month, d.day, tzinfo=timezone.utc))
    # record: seconds-from-midnight-UTC, open, close, low, high (int, 1/1000 pt), volume (float32)
    df = pd.DataFrame(
        {"o": a[:, 1] / PRICE_SCALE, "h": a[:, 4] / PRICE_SCALE, "l": a[:, 3] / PRICE_SCALE,
         "c": a[:, 2] / PRICE_SCALE, "v": a[:, 5].copy().view(">f4").astype(float)},
        index=base + pd.to_timedelta(a[:, 0].astype(np.int64), unit="s"),
    )
    return df


def cmd_build(a) -> None:
    """Cut every NYSE session (open..close, NY time, DST-aware) out of the UTC day files.

    Keeps 1-minute bars labelled by their START time, from 09:30 to the last minute
    before the session close. Dukascopy writes a row for every minute of the day; rows
    with zero volume are flat filler (no ticks) and are flagged, not dropped, so the
    timeline stays regular.
    """
    sch = nyse_schedule(a.start, a.end)
    if a.years:  # e.g. "2015-2017,2023-2026": restrict to years whose files are complete
        keep = set()
        for part in a.years.split(","):
            lo_y, _, hi_y = part.partition("-")
            keep |= set(range(int(lo_y), int(hi_y or lo_y) + 1))
        sch = sch[sch.index.year.isin(sorted(keep))]
    frames, gaps, modelled = [], [], []
    # session median spread per day where the ASK file exists (used only to fill a missing ASK day)
    day_spread = {}
    for d, row in sch.iterrows():
        dd = d.date()
        if (DUKA_RAW / rel_path(dd, "ASK")).exists() and (DUKA_RAW / rel_path(dd, "BID")).exists():
            lo, hi = row["open"].tz_convert("UTC"), row["close"].tz_convert("UTC")
            b, k = read_day(dd, "BID"), read_day(dd, "ASK")
            j = b.join(k, lsuffix="_b", rsuffix="_a", how="inner")
            j = j[(j.index >= lo) & (j.index < hi)]
            if len(j):
                day_spread[d] = float((j.c_a - j.c_b).median())
    sp_idx = pd.Series(day_spread).sort_index()
    for d, row in sch.iterrows():
        dd = d.date()
        bid, ask = read_day(dd, "BID"), read_day(dd, "ASK")
        if bid.empty:
            gaps.append({"day": str(dd), "issue": "missing or empty BID file"})
            continue
        if ask.empty:
            if sp_idx.empty:
                gaps.append({"day": str(dd), "issue": "missing ASK file, no spread to model"})
                continue
            near = sp_idx.iloc[(abs((sp_idx.index - d).days)).argsort()[:4]]
            sp = float(near.median())  # median session spread of the 4 nearest days with a real ASK file
            ask = bid[["o", "h", "l", "c"]] + sp
            ask["v"] = bid["v"]
            modelled.append(str(dd))
        lo, hi = row["open"].tz_convert("UTC"), row["close"].tz_convert("UTC")
        b = bid[(bid.index >= lo) & (bid.index < hi)]
        k = ask[(ask.index >= lo) & (ask.index < hi)]
        expected = int((hi - lo).total_seconds() // 60)
        m = b.join(k, lsuffix="_b", rsuffix="_a", how="inner")
        active = int(((m.v_b > 0) | (m.v_a > 0)).sum())
        if len(m) < expected or active < 0.9 * expected:
            gaps.append({"day": str(dd), "issue": f"{len(m)} rows, {active} active of {expected} expected minutes"})
        m["day"] = d
        frames.append(m)
    df = pd.concat(frames)
    df.index = df.index.tz_convert(NY)
    df.index.name = "t"
    df.to_pickle(SESSION_PKL)
    info = {
        "source": "Dukascopy Bank historical datafeed (datafeed.dukascopy.com)",
        "symbol": SYMBOL, "sides": ["BID", "ASK"], "bar": "1 minute, labelled by start time",
        "years_filter": a.years or "all", "sessions_requested": len(sch), "sessions_built": int(df["day"].nunique()),
        "sessions_by_year": {str(y): int(n) for y, n in df.groupby(df["day"].dt.year)["day"].nunique().items()},
        "first_day": str(df["day"].min().date()), "last_day": str(df["day"].max().date()),
        "rows": len(df), "ask_modelled_days": len(modelled),
        "_ask_modelled_note": "days whose ASK file was not (yet) cached: ASK = BID + median session spread of the 4 nearest days with a real ASK file",
        "ask_modelled_list": modelled, "gaps": gaps,
    }
    GAPS_JSON.write_text(json.dumps(info, indent=1))
    print(json.dumps({k: v for k, v in info.items() if k not in ("gaps", "ask_modelled_list")}, indent=1))
    print(f"{len(gaps)} sessions flagged; details in {GAPS_JSON.relative_to(ROOT)}")


def main() -> None:
    ap = argparse.ArgumentParser()
    ap.add_argument("cmd", choices=["fetch", "build", "status", "histdata"])
    ap.add_argument("--start", default=DEFAULT_START)
    ap.add_argument("--end", default=DEFAULT_END)
    ap.add_argument("--threads", type=int, default=24)
    ap.add_argument("--reverse", action="store_true", help="download newest first (to run a 2nd worker)")
    ap.add_argument("--years", default="", help='build only: e.g. "2015-2017,2023-2026"')
    a = ap.parse_args()
    {"fetch": cmd_fetch, "build": cmd_build, "status": cmd_status, "histdata": cmd_histdata}[a.cmd](a)


if __name__ == "__main__":
    sys.exit(main())