"""1-minute candles, BID and ASK: as trading sessions in New York time, or as one continuous series in UTC.

The candles come from the Dukascopy historical feed (tools/ib_data.py downloads and cleans the Nasdaq sessions,
tools/fx_data.py the gold minutes). A chart shows the BID; a buy fills at the ASK. Keeping both sides is how the
spread ends up in every result.

Indian markets (the Booming Bulls audit) come as NSE sessions, 09:15 to 15:29 India time: the index spot
(tools/sb_backtest.py builds it from the Kaggle and Hugging Face minute files) and single stocks with their
traded volume. Those feeds carry one price, so BID and ASK are the same candles and costs are charged apart.
"""

from __future__ import annotations

import math
from dataclasses import dataclass
from pathlib import Path

import numpy as np
import pandas as pd

NEW_YORK = "America/New_York"


@dataclass(frozen=True)
class Candles:
    """Open, high, low and close of every minute in a session (one array per field)."""

    open: np.ndarray
    high: np.ndarray
    low: np.ndarray
    close: np.ndarray


@dataclass(frozen=True)
class Session:
    """One trading day: Nasdaq 09:30 to the closing bell, gold 17:00 to 17:00 New York."""

    day: pd.Timestamp
    minute: np.ndarray  # timestamp of every candle, New York time
    clock: np.ndarray  # the same minutes as integers: 9:30 -> 930
    bid: Candles
    ask: Candles
    volume: np.ndarray | None = None  # shares traded in each minute (stocks only)

    def __len__(self) -> int:
        return len(self.minute)


def load_sessions(path: Path, *, zero_spread: bool = False) -> list[Session]:
    """Read a cleaned session file. zero_spread=True prices every order at the BID (the no-cost check)."""
    frame = pd.read_pickle(path)
    sessions = []
    for day, rows in frame.groupby("day", sort=True):
        rows = rows.sort_index()
        bid = Candles(*(rows[f"{field}_b"].to_numpy(float) for field in "ohlc"))
        ask = bid if zero_spread else Candles(*(rows[f"{field}_a"].to_numpy(float) for field in "ohlc"))
        clock = (rows.index.hour * 100 + rows.index.minute).to_numpy()
        sessions.append(Session(pd.Timestamp(day), rows.index.to_numpy(), clock, bid, ask))
    return sessions


def load_gold_sessions(path: Path, *, spread: float) -> list[Session]:
    """Read the gold 1-minute file (XAUUSD_m1.npz from tools/fx_data.py) as New York 17:00-to-17:00 days.

    Gold trades almost around the clock. Its day is the 17:00 New York roll-over that TradingView's daily candle
    uses on broker feeds, labelled by the date it ends on (Sunday evening belongs to Monday). The ASK side is
    priced a fixed `spread` above the BID: the audit uses the 0.28 shown on the guru's own screen, not the
    feed's own, wider ASK.
    """
    z = np.load(path)
    scale = float(z["scale"][0])
    decimals = int(round(math.log10(scale)))  # the feed stores integer thousandths of a dollar
    bid = np.round(z["bid"].astype(np.float64) / scale, decimals)
    ny = pd.to_datetime(z["t"], unit="s", utc=True).tz_convert(NEW_YORK).tz_localize(None)
    label = (ny + pd.Timedelta(hours=7)).normalize()
    clock = (ny.hour * 100 + ny.minute).to_numpy()
    minute = ny.to_numpy()
    starts = np.flatnonzero(np.diff(label.asi8)) + 1
    sessions = []
    for a, b in zip(np.concatenate([[0], starts]), np.concatenate([starts, [len(ny)]]), strict=True):
        side = Candles(*(bid[a:b, field] for field in range(4)))
        ask = Candles(side.open + spread, side.high + spread, side.low + spread, side.close + spread)
        sessions.append(Session(pd.Timestamp(label[a]), minute[a:b], clock[a:b], side, ask))
    return sessions


@dataclass(frozen=True)
class Minutes:
    """Every traded minute of a market as one continuous series on the UTC clock.

    For rules that read structure across days and weeks (the Top G audit). The ASK is the feed's own, so every
    fill pays the spread that was really there at that minute.
    """

    time: np.ndarray  # UTC epoch seconds at the start of each minute
    bid: Candles
    ask: Candles

    def __len__(self) -> int:
        return len(self.time)


def load_minutes(path: Path, *, zero_spread: bool = False) -> Minutes:
    """Read a 1-minute file built by tools/fx_data.py (e.g. XAUUSD_m1.npz). zero_spread=True prices every order
    at the BID (the no-cost check)."""
    z = np.load(path)
    scale = float(z["scale"][0])
    decimals = int(round(math.log10(scale)))  # the feed stores integer thousandths of a dollar

    def side(name: str) -> Candles:
        prices = np.round(z[name].astype(np.float64) / scale, decimals)
        return Candles(*(np.ascontiguousarray(prices[:, field]) for field in range(4)))

    bid = side("bid")
    return Minutes(z["t"], bid, bid if zero_spread else side("ask"))


def gold_days(time: np.ndarray) -> np.ndarray:
    """The gold day of each UTC minute, as days since 1970: New York 17:00 to 17:00, labelled by the date it
    ends on (the same day as load_gold_sessions)."""
    new_york = pd.to_datetime(time, unit="s", utc=True).tz_convert(NEW_YORK).tz_localize(None)
    return (new_york + pd.Timedelta(hours=7)).normalize().to_numpy().astype("datetime64[D]").astype(np.int64)


@dataclass(frozen=True)
class Bars:
    """Candles of a longer timeframe built from the minutes. Bar k holds minutes first[k] to end[k] - 1."""

    time: np.ndarray  # UTC epoch seconds at the start of each bar
    first: np.ndarray  # index of the bar's first minute
    end: np.ndarray  # index one past its last minute
    bid: Candles

    def __len__(self) -> int:
        return len(self.time)

    def closed_before(self, minute: int) -> int:
        """The last bar that had closed when `minute` began (-1: none)."""
        return int(np.searchsorted(self.end, minute, side="right")) - 1

    def holding(self, minute: int) -> int:
        """The bar that contains `minute`."""
        return self.closed_before(minute) + 1


def resample(m: Minutes, minutes: int) -> Bars:
    """BID candles of `minutes` minutes on the UTC clock, the way a broker chart draws them: a period with no
    trading has no candle, and a candle opens and closes on its first and last traded minute."""
    period = 60 * minutes
    bucket = m.time // period
    first = np.concatenate([[0], np.flatnonzero(np.diff(bucket)) + 1])
    end = np.append(first[1:], len(m))
    bid = Candles(
        m.bid.open[first],
        np.maximum.reduceat(m.bid.high, first),
        np.minimum.reduceat(m.bid.low, first),
        m.bid.close[end - 1],
    )
    return Bars(bucket[first] * period, first, end, bid)


# ---------------------------------------------------------------------------------------------------------------
# NSE (India): sessions 09:15-15:29 India time. The feeds' timestamps are India wall-clock time.

NSE_OPEN, NSE_LAST = 9 * 60 + 15, 15 * 60 + 29  # minutes after midnight
NSE_MIN_MINUTES = 300  # a session with fewer minutes (a special or broken day) is left out


def _nse_sessions(
    seconds: np.ndarray, o: np.ndarray, h: np.ndarray, lo: np.ndarray, c: np.ndarray, volume: np.ndarray | None
) -> list[Session]:
    """Cut India wall-clock epoch seconds into regular sessions, in time order."""
    day, of_day = seconds // 86400, (seconds % 86400) // 60
    regular = (of_day >= NSE_OPEN) & (of_day <= NSE_LAST)
    columns = [day, of_day, o, h, lo, c] + ([volume] if volume is not None else [])
    columns = [x[regular] for x in columns]
    order = np.lexsort((columns[1], columns[0]))
    day, of_day, o, h, lo, c, *rest = (x[order] for x in columns)
    volume = rest[0] if rest else None
    starts = np.concatenate([[0], np.flatnonzero(np.diff(day)) + 1, [len(day)]])
    sessions = []
    for a, b in zip(starts[:-1], starts[1:], strict=True):
        if b - a < NSE_MIN_MINUTES:
            continue
        minute = (seconds[regular][order][a:b]).astype("datetime64[s]").astype("datetime64[m]")
        clock = of_day[a:b] // 60 * 100 + of_day[a:b] % 60
        candles = Candles(o[a:b], h[a:b], lo[a:b], c[a:b])
        session_volume = volume[a:b] if volume is not None else None
        day_label = pd.Timestamp(np.datetime64(int(day[a]), "D"))
        sessions.append(Session(day_label, minute, clock.astype(np.int64), candles, candles, session_volume))
    return sessions


def load_nse_index(path: Path) -> list[Session]:
    """An index's 1-minute spot candles (NIFTY_m1.npz, BANKNIFTY_m1.npz: UTC epoch seconds) as NSE sessions."""
    z = np.load(path)
    india = z["t"].astype(np.int64) + 5 * 3600 + 30 * 60
    return _nse_sessions(india, *(z[field].astype(float) for field in "ohlc"), None)


def load_nse_stock(parts: list[tuple[Path, float, str, str]]) -> list[Session]:
    """A stock's 1-minute candles with volume, stitched from (csv file, price scale, first day, last day) parts.

    The scale undoes a back-adjustment, so every part is on one price scale. A later part never overwrites a
    minute an earlier part already has.
    """
    frames = []
    for path, scale, first, last in parts:
        frame = pd.read_csv(path)
        when = pd.to_datetime(frame["date" if "date" in frame.columns else "timestamp"].astype(str).str.slice(0, 19))
        frame = frame.assign(t=when)
        frame = frame[(frame.t >= pd.Timestamp(first)) & (frame.t < pd.Timestamp(last) + pd.Timedelta(days=1))]
        frame = frame.dropna(subset=["open", "high", "low", "close"]).drop_duplicates("t")
        for field in ("open", "high", "low", "close"):
            frame[field] = frame[field].astype(float) * scale
        frames.append(frame)
    frame = pd.concat(frames).drop_duplicates("t").sort_values("t")
    seconds = frame.t.to_numpy("datetime64[s]").astype(np.int64)
    prices = (frame[field].to_numpy() for field in ("open", "high", "low", "close"))
    return _nse_sessions(seconds, *prices, frame.volume.to_numpy(float))


@dataclass(frozen=True)
class SessionBars:
    """Candles of a longer timeframe cut from each session, aligned to 09:15 and never spanning two sessions.

    Bar k belongs to sessions[session[k]] and holds its minutes first[k] to last[k]. Times are minutes after
    midnight: a 5-minute bar opens_at 555 (09:15) and closes_at 560 (09:20).
    """

    session: np.ndarray
    opens_at: np.ndarray
    closes_at: np.ndarray
    first: np.ndarray
    last: np.ndarray
    candles: Candles
    volume: np.ndarray | None

    def __len__(self) -> int:
        return len(self.session)


def session_bars(sessions: list[Session], minutes: int) -> SessionBars:
    """BID candles of `minutes` minutes, the way an Indian broker chart draws them."""
    rows: dict[str, list] = {key: [] for key in ("session", "opens_at", "first", "last", "o", "h", "l", "c", "v")}
    for q, s in enumerate(sessions):
        of_day = s.clock // 100 * 60 + s.clock % 100
        bucket = (of_day - NSE_OPEN) // minutes
        cuts = np.concatenate([[0], np.flatnonzero(np.diff(bucket)) + 1, [len(bucket)]])
        for a, b in zip(cuts[:-1], cuts[1:], strict=True):
            rows["session"].append(q)
            rows["opens_at"].append(NSE_OPEN + int(bucket[a]) * minutes)
            rows["first"].append(a)
            rows["last"].append(b - 1)
            rows["o"].append(s.bid.open[a])
            rows["h"].append(s.bid.high[a:b].max())
            rows["l"].append(s.bid.low[a:b].min())
            rows["c"].append(s.bid.close[b - 1])
            rows["v"].append(float(s.volume[a:b].sum()) if s.volume is not None else 0.0)
    opens_at = np.asarray(rows["opens_at"], dtype=np.int64)
    candles = Candles(*(np.asarray(rows[key], dtype=float) for key in "ohlc"))
    volume = np.asarray(rows["v"], dtype=float) if sessions and sessions[0].volume is not None else None
    return SessionBars(
        np.asarray(rows["session"], dtype=np.int64),
        opens_at,
        opens_at + minutes,
        np.asarray(rows["first"], dtype=np.int64),
        np.asarray(rows["last"], dtype=np.int64),
        candles,
        volume,
    )