"""Classical (non-LLM) trading policies. Each is a small, documented heuristic with explicit
parameters; none sees the latent fundamental, only the Observation.

Price arithmetic is in ticks; quantities in lots. Every policy returns at most a handful of intents.
"""
from __future__ import annotations

from dataclasses import dataclass

import numpy as np

from ..exchange.orders import Action, CancelIntent, OrderIntent, OrderType, Side
from .base import Observation, clip_qty


def _limit_or_market(rng: np.random.Generator, p_market: float) -> OrderType:
    return OrderType.MARKET if rng.random() < p_market else OrderType.LIMIT


def _cancel_all(obs: Observation, asset: int | None = None) -> list[CancelIntent]:
    return [CancelIntent(o.order_id) for o in obs.open_orders if asset is None or o.asset == asset]


def _affordable(obs: Observation, price_ticks: float, tick_cents: int) -> int:
    return int(max(0, obs.free_cash // max(1, int(price_ticks * tick_cents * 1.001))))


@dataclass
class Fundamentalist:
    """Value trader. Posterior value = public PV adjusted by the private signal (if informed):
        V_hat = public_value + w * (signal - x_hat_public) / (1+r-rho)  with w = prec/(1+prec).
    Trades toward the value when the relative mispricing exceeds `threshold`; order size scales with
    the mispricing (capped) and with available cash/shares. The signal adjustment is relative to the
    public posterior mean x_hat (shown in the Observation), scaled by 1/(1+r-rho)."""
    name: str = "fundamentalist"
    threshold: float = 0.005
    aggressiveness: float = 20.0     # lots per 1% mispricing
    p_market: float = 0.3
    max_qty: int = 50
    improve_ticks: int = 1           # limit orders are placed inside the spread by this many ticks
    cross_mult: float = 2.0          # beyond cross_mult*threshold the limit order is marketable (at the far quote)

    def decide(self, obs: Observation, rng: np.random.Generator) -> Action:
        act = Action()
        for a, book in obs.books.items():
            if book.halted:
                continue
            tc = obs.tick_cents[a]
            v_pub = obs.public_value[a] / tc          # ticks
            v_hat = v_pub
            if a in obs.private_signals:
                prec = obs.signal_precision.get(a, 0.0)
                w = prec / (1.0 + prec)
                rho = 0.98
                # posterior mean of x_n combines the public prior (in v_pub) with the private signal, weight w
                v_hat = v_pub + w * (obs.private_signals[a] - obs.x_hat_public.get(a, 0.0)) / (1.0 + obs.r_period - rho) / tc
            mid = book.mid
            if mid <= 0:
                continue
            m = (v_hat - mid) / mid
            if abs(m) < self.threshold:
                continue
            q = clip_qty(abs(m) * 100.0 * self.aggressiveness, min(self.max_qty, obs.max_order_qty))
            if q == 0:
                continue
            act.cancels += _cancel_all(obs, a)
            otype = _limit_or_market(rng, self.p_market)
            if m > 0:
                q = min(q, _affordable(obs, book.best_ask or mid, tc))
                if q <= 0:
                    continue
                if otype is OrderType.MARKET:
                    price = None
                elif abs(m) > self.cross_mult * self.threshold and book.best_ask is not None:
                    price = int(min(v_hat, book.best_ask))
                else:
                    price = int(min(v_hat, (book.best_bid or mid) + self.improve_ticks))
                act.orders.append(OrderIntent(a, Side.BUY, otype, q, price, tag=self.name))
            else:
                held = obs.positions.get(a, 0)
                q = min(q, max(held, 0)) if not obs.short_allowed else q
                if q <= 0:
                    continue
                if otype is OrderType.MARKET:
                    price = None
                elif abs(m) > self.cross_mult * self.threshold and book.best_bid is not None:
                    price = int(max(v_hat, book.best_bid))
                else:
                    price = int(max(v_hat, (book.best_ask or mid) - self.improve_ticks))
                act.orders.append(OrderIntent(a, Side.SELL, otype, q, price, tag=self.name))
        return act


@dataclass
class TrendFollower:
    """Momentum: signal = mean(mid, last 5) / mean(mid, last 30) - 1 (simple moving averages); buys when positive."""
    name: str = "trend"
    short_window: int = 5
    long_window: int = 30
    threshold: float = 0.002
    size: int = 20
    p_market: float = 0.6
    max_qty: int = 50

    def decide(self, obs: Observation, rng: np.random.Generator) -> Action:
        act = Action()
        for a, book in obs.books.items():
            if book.halted:
                continue
            h = obs.mid_history[a]
            if len(h) < self.long_window:
                continue
            s = h[-self.short_window:].mean() / h[-self.long_window:].mean() - 1.0
            if abs(s) < self.threshold:
                continue
            tc = obs.tick_cents[a]
            q = clip_qty(self.size * min(abs(s) / self.threshold, 3.0), min(self.max_qty, obs.max_order_qty))
            act.cancels += _cancel_all(obs, a)
            otype = _limit_or_market(rng, self.p_market)
            if s > 0:
                q = min(q, _affordable(obs, book.best_ask or book.mid, tc))
                if q <= 0:
                    continue
                price = None if otype is OrderType.MARKET else int(book.best_ask or book.mid)
                act.orders.append(OrderIntent(a, Side.BUY, otype, q, price, tag=self.name))
            else:
                held = obs.positions.get(a, 0)
                q = q if obs.short_allowed else min(q, max(held, 0))
                if q <= 0:
                    continue
                price = None if otype is OrderType.MARKET else int(book.best_bid or book.mid)
                act.orders.append(OrderIntent(a, Side.SELL, otype, q, price, tag=self.name))
        return act


@dataclass
class MeanReverter:
    """Contrarian: z = (mid - MA_w)/sd_w; sells when z > thr, buys when z < -thr."""
    name: str = "meanrev"
    window: int = 40
    threshold: float = 1.5
    size: int = 15
    p_market: float = 0.2
    max_qty: int = 50

    def decide(self, obs: Observation, rng: np.random.Generator) -> Action:
        act = Action()
        for a, book in obs.books.items():
            if book.halted:
                continue
            h = obs.mid_history[a]
            if len(h) < self.window:
                continue
            w = h[-self.window:]
            sd = w.std()
            if sd <= 0:
                continue
            z = (book.mid - w.mean()) / sd
            if abs(z) < self.threshold:
                continue
            tc = obs.tick_cents[a]
            q = clip_qty(self.size * min(abs(z) / self.threshold, 3.0), min(self.max_qty, obs.max_order_qty))
            act.cancels += _cancel_all(obs, a)
            otype = _limit_or_market(rng, self.p_market)
            if z < 0:
                q = min(q, _affordable(obs, book.best_ask or book.mid, tc))
                if q <= 0:
                    continue
                price = None if otype is OrderType.MARKET else int((book.best_bid or book.mid) + 1)
                act.orders.append(OrderIntent(a, Side.BUY, otype, q, price, tag=self.name))
            else:
                held = obs.positions.get(a, 0)
                q = q if obs.short_allowed else min(q, max(held, 0))
                if q <= 0:
                    continue
                price = None if otype is OrderType.MARKET else int((book.best_ask or book.mid) - 1)
                act.orders.append(OrderIntent(a, Side.SELL, otype, q, price, tag=self.name))
        return act


@dataclass
class NoiseTrader:
    """Zero-intelligence-plus liquidity demander: random side, size ~ 1+Geometric, limit price at
    mid ± |N(0, spread_sd)| ticks on the passive side (or a market order with prob p_market)."""
    name: str = "noise"
    p_market: float = 0.3
    mean_size: float = 8.0
    price_sd_frac: float = 0.001     # sd of the limit-price offset as a fraction of mid (0.1%)
    max_qty: int = 40

    def decide(self, obs: Observation, rng: np.random.Generator) -> Action:
        act = Action()
        for a, book in obs.books.items():
            if book.halted:
                continue
            side = Side.BUY if rng.random() < 0.5 else Side.SELL
            q = clip_qty(1 + rng.geometric(1.0 / self.mean_size), min(self.max_qty, obs.max_order_qty))
            tc = obs.tick_cents[a]
            if side is Side.BUY:
                q = min(q, _affordable(obs, book.best_ask or book.mid, tc))
            elif not obs.short_allowed:
                q = min(q, max(obs.positions.get(a, 0), 0))
            if q <= 0:
                continue
            if len(obs.open_orders) > 3:
                act.cancels += _cancel_all(obs, a)
            otype = _limit_or_market(rng, self.p_market)
            if otype is OrderType.MARKET:
                act.orders.append(OrderIntent(a, side, otype, q, None, tag=self.name))
            else:
                off = abs(rng.standard_normal()) * self.price_sd_frac * book.mid
                ref = (book.best_bid if side is Side.BUY else book.best_ask) or book.mid
                price = int(round(ref - off)) if side is Side.BUY else int(round(ref + off))
                act.orders.append(OrderIntent(a, side, otype, q, max(price, 1), tag=self.name))
        return act


@dataclass
class MarketMaker:
    """Inventory-aware two-sided quoter (bounded Avellaneda–Stoikov-style rule).

    Reservation price  r = mid * (1 - skew * inv / max_inventory)          (inventory skew, fraction of mid)
    Half-spread        h = clip(base_half + vol_mult * sigma_ticks, min_half, max_half)
    with sigma_ticks the std of recent tick-to-tick mid changes. The original A–S formulas assume
    Poisson fills and a terminal date; here the skew and spread are bounded fractions of the price so
    that quotes always stay inside the exchange price band. Re-quotes every activation."""
    name: str = "market_maker"
    skew: float = 0.02               # 2% of mid at full inventory deviation
    base_half: float = 8.0           # ticks (8 cents on a 200.00 stock = 4 bp half-spread floor)
    vol_mult: float = 1.0
    quote_qty: int = 20
    max_inventory: int = 200         # inventory deviation (lots) at which the skew is maximal
    min_half_spread: int = 1
    max_half_spread: int = 60
    vol_window: int = 30
    withdrawn: bool = False          # liquidity-shock treatment: stop quoting
    _inv0: dict | None = None        # inventory target = endowment (set at first decision)

    def decide(self, obs: Observation, rng: np.random.Generator) -> Action:
        act = Action()
        act.cancels += _cancel_all(obs)
        if self.withdrawn:
            return act
        if self._inv0 is None:
            self._inv0 = dict(obs.positions)
        for a, book in obs.books.items():
            if book.halted:
                continue
            h = obs.mid_history[a]
            sigma = float(np.std(np.diff(h[-self.vol_window:]))) if len(h) > 5 else 2.0
            inv = obs.positions.get(a, 0) - self._inv0.get(a, 0)
            res = book.mid * (1.0 - self.skew * inv / self.max_inventory)
            half = float(np.clip(self.base_half + self.vol_mult * sigma, self.min_half_spread, self.max_half_spread))
            bid = int(np.floor(res - half))
            ask = int(np.ceil(res + half))
            if ask - bid < 2 * self.min_half_spread:
                ask = bid + 2 * self.min_half_spread
            tc = obs.tick_cents[a]
            qb = self.quote_qty if inv < self.max_inventory else 0
            qa = self.quote_qty if inv > -self.max_inventory else 0
            qb = min(qb, _affordable(obs, bid, tc))
            if not obs.short_allowed:
                qa = min(qa, max(obs.positions.get(a, 0), 0))
            if qb > 0:
                act.orders.append(OrderIntent(a, Side.BUY, OrderType.LIMIT, qb, max(bid, 1), tag=self.name))
            if qa > 0:
                act.orders.append(OrderIntent(a, Side.SELL, OrderType.LIMIT, qa, ask, tag=self.name))
        return act


@dataclass
class StatArb:
    """Cross-asset arbitrageur: trades the residual of a rolling regression of asset j's mid on asset
    i's mid (pairs trading on the two first assets). Single-asset runs: inactive."""
    name: str = "statarb"
    window: int = 60
    threshold: float = 1.5
    size: int = 10
    max_qty: int = 40

    def decide(self, obs: Observation, rng: np.random.Generator) -> Action:
        act = Action()
        assets = sorted(obs.books)
        if len(assets) < 2:
            return act
        i, j = assets[0], assets[1]
        hi, hj = obs.mid_history[i], obs.mid_history[j]
        n = min(len(hi), len(hj))
        if n < self.window:
            return act
        x, y = np.log(hi[-self.window:]), np.log(hj[-self.window:])
        beta = np.cov(x, y)[0, 1] / max(np.var(x), 1e-12)
        resid = y - beta * x
        z = (resid[-1] - resid.mean()) / max(resid.std(), 1e-9)
        if abs(z) < self.threshold:
            return act
        act.cancels += _cancel_all(obs)
        q = clip_qty(self.size, min(self.max_qty, obs.max_order_qty))
        bj, bi = obs.books[j], obs.books[i]
        if bj.halted or bi.halted:
            return act
        # z>0: j rich relative to i -> sell j, buy i
        sell_a, buy_a = (j, i) if z > 0 else (i, j)
        held = obs.positions.get(sell_a, 0)
        qs = q if obs.short_allowed else min(q, max(held, 0))
        qb = min(q, _affordable(obs, obs.books[buy_a].best_ask or obs.books[buy_a].mid, obs.tick_cents[buy_a]))
        if qs > 0:
            act.orders.append(OrderIntent(sell_a, Side.SELL, OrderType.LIMIT, qs, int(obs.books[sell_a].best_ask or obs.books[sell_a].mid) - 1, tag=self.name))
        if qb > 0:
            act.orders.append(OrderIntent(buy_a, Side.BUY, OrderType.LIMIT, qb, int(obs.books[buy_a].best_bid or obs.books[buy_a].mid) + 1, tag=self.name))
        return act


@dataclass
class Leveraged:
    """A trend follower that may borrow (leverage_cap > 1 on the account). The policy is identical to
    `TrendFollower`; leverage, margin calls and forced liquidation are enforced by the exchange, so
    this class exists to make the treatment explicit in configs."""
    name: str = "leveraged"
    inner: TrendFollower | None = None

    def __post_init__(self) -> None:
        if self.inner is None:
            self.inner = TrendFollower(name="leveraged", size=30, p_market=0.7)

    def decide(self, obs: Observation, rng: np.random.Generator) -> Action:
        return self.inner.decide(obs, rng)  # type: ignore[union-attr]


CLASSICAL_POLICIES = {
    "fundamentalist": Fundamentalist,
    "trend": TrendFollower,
    "meanrev": MeanReverter,
    "noise": NoiseTrader,
    "market_maker": MarketMaker,
    "statarb": StatArb,
    "leveraged": Leveraged,
}