"""Order-Flow Imbalance features (gold layer) — faithful Python port of the Rust ml-features
`ofi_calculator.rs` stateless functions (Blocks F/G/H/CC of the 134-dim alpha vector). The Rust is the
oracle; these match it bit-for-bit.

OFI (Cont 2010) is the canonical "follow the money" signal: it traces where book volume is being added
vs pulled across consecutive snapshots — every level's signed contribution to net order flow.
"""
from __future__ import annotations

import math

from fxhnt.domain.marketdata import Mbp10Snapshot

_CLIP = 1e6


def _safe_clip(value: float, lo: float = -_CLIP, hi: float = _CLIP) -> float:
    if not math.isfinite(value):
        return 0.0
    return min(max(value, lo), hi)


def ofi_per_level(current: Mbp10Snapshot, previous: Mbp10Snapshot) -> list[float]:
    """Block F — Cont (2010) per-level OFI: bid_contrib − ask_contrib across the top-5 levels.
    bid_contrib = +V_b if P_b rose-or-held else −V_b^{prev};  ask_contrib = +V_a if P_a fell-or-held
    else −V_a^{prev}.  (Integer price compares, exactly like the Rust.)"""
    n = min(len(current.levels), len(previous.levels), 5)
    out = [0.0] * 5
    for i in range(n):
        c, p = current.levels[i], previous.levels[i]
        bid_contrib = float(c.bid_sz) if c.bid_px >= p.bid_px else -float(p.bid_sz)
        ask_contrib = float(c.ask_sz) if c.ask_px <= p.ask_px else -float(p.ask_sz)
        out[i] = _safe_clip(bid_contrib - ask_contrib)
    return out


def log_gofi(ofi: list[float]) -> list[float]:
    """Block G — signed log compression: sign(x)·ln(1+|x|). (x=0 → 0, since ln(1)=0.)"""
    return [math.copysign(1.0, x) * math.log1p(abs(x)) for x in ofi]


def book_deltas_per_level(current: Mbp10Snapshot, previous: Mbp10Snapshot) -> tuple[list[float], list[float]]:
    """Block H — raw per-level size changes (bid, ask) across the top-5 levels."""
    n = min(len(current.levels), len(previous.levels), 5)
    bid, ask = [0.0] * 5, [0.0] * 5
    for i in range(n):
        c, p = current.levels[i], previous.levels[i]
        bid[i] = _safe_clip(float(c.bid_sz) - float(p.bid_sz))
        ask[i] = _safe_clip(float(c.ask_sz) - float(p.ask_sz))
    return bid, ask


def slope_and_convexity(snap: Mbp10Snapshot) -> list[float]:
    """Block CC — shape of the cumulative-depth curve: OLS slope + discrete-2nd-difference convexity for
    bid and ask. Returns [bid_slope, bid_convexity, ask_slope, ask_convexity]; needs ≥3 levels."""
    n = min(len(snap.levels), 5)
    if n < 3:
        return [0.0, 0.0, 0.0, 0.0]
    cum_bid, cum_ask = [0.0] * 5, [0.0] * 5
    bid_acc = ask_acc = 0.0
    for i in range(n):
        bid_acc += float(snap.levels[i].bid_sz)
        ask_acc += float(snap.levels[i].ask_sz)
        cum_bid[i], cum_ask[i] = bid_acc, ask_acc

    nf = float(n)
    mean_x = sum(range(n)) / nf
    mean_bid = sum(cum_bid[:n]) / nf
    mean_ask = sum(cum_ask[:n]) / nf
    num_bid = num_ask = den = 0.0
    for i in range(n):
        x_dev = i - mean_x
        num_bid += x_dev * (cum_bid[i] - mean_bid)
        num_ask += x_dev * (cum_ask[i] - mean_ask)
        den += x_dev * x_dev
    bid_slope, ask_slope = (num_bid / den, num_ask / den) if den > 1e-12 else (0.0, 0.0)

    mid = n // 2
    half_span = (n - 1) / 2.0
    denom = max(half_span * half_span, 1e-12)
    bid_convex = (cum_bid[n - 1] - 2.0 * cum_bid[mid] + cum_bid[0]) / denom
    ask_convex = (cum_ask[n - 1] - 2.0 * cum_ask[mid] + cum_ask[0]) / denom
    return [bid_slope, bid_convex, ask_slope, ask_convex]


_LEVELS = (1, 2, 3, 4, 5)
FEATURE_NAMES = (
    *(f"ofi_l{i}" for i in _LEVELS),
    *(f"log_gofi_l{i}" for i in _LEVELS),
    *(f"bid_delta_l{i}" for i in _LEVELS),
    *(f"ask_delta_l{i}" for i in _LEVELS),
    "cum_bid_slope", "cum_bid_convexity", "cum_ask_slope", "cum_ask_convexity",
)


class OfiFeatures:
    """Stateful Block F+G+H+CC aggregator (24 features) — holds the previous snapshot for the
    consecutive-snapshot OFI/delta computations; slope/convexity are pure functions of the current book."""
    name = "ofi"
    feature_names = FEATURE_NAMES

    def __init__(self) -> None:
        self._prev: Mbp10Snapshot | None = None

    def update(self, snap: Mbp10Snapshot) -> dict[str, float]:
        sc = slope_and_convexity(snap)                       # current-only, valid even at cold start
        if self._prev is None:
            ofi = gofi = bid_d = ask_d = [0.0] * 5           # no prior snapshot -> no flow yet
        else:
            ofi = ofi_per_level(snap, self._prev)
            gofi = log_gofi(ofi)
            bid_d, ask_d = book_deltas_per_level(snap, self._prev)
        self._prev = snap
        return dict(zip(FEATURE_NAMES, [*ofi, *gofi, *bid_d, *ask_d, *sc]))

    def reset(self) -> None:
        self._prev = None
