"""Tick-level trade-size distribution factors.

Methodology:

  Each trade is bucketed by size:
    retail   < 100 shares
    small    100-1,000
    medium   1,000-10,000
    block    >= 10,000

  Each trade is signed (buy vs sell) via the Lee-Ready tick rule:
    if price > mid:        buy
    if price < mid:        sell
    if price == mid:       use prior tick's sign (default to 0)

  Two factors are produced:

    block_share      = block_vol / total_vol         (range 0..1)
    block_aggression = (block_buys - block_sells) / block_vol  (range -1..+1)

  ``block_share`` is z-scored cross-sectionally; ``block_aggression`` is
  used directly (already bounded in [-1, +1]).

  ``buy_ratio`` is the all-size signed volume ratio; included here for
  convenience so the scoring code can pick it up alongside the other
  intraday metrics.
"""

from __future__ import annotations

from typing import Optional

import numpy as np
import pandas as pd

from .factor_sources import get_data_source


# Bucket thresholds (shares)
RETAIL_MAX = 100
SMALL_MAX = 1_000
MEDIUM_MAX = 10_000


def _bucket(size: int) -> str:
    if size < RETAIL_MAX:
        return "retail"
    if size < SMALL_MAX:
        return "small"
    if size < MEDIUM_MAX:
        return "medium"
    return "block"


def _sign_trades(ticks: pd.DataFrame) -> pd.Series:
    """Lee-Ready tick rule: sign each trade vs the prevailing mid."""
    mid = (ticks["bid"] + ticks["ask"]) / 2.0
    sign = pd.Series(0, index=ticks.index, dtype=int)
    sign[ticks["price"] > mid] = 1
    sign[ticks["price"] < mid] = -1

    # Trades at the mid: carry forward the prior sign
    at_mid = ticks["price"] == mid
    if at_mid.any():
        prior = sign.replace(0, np.nan).ffill().fillna(0).astype(int)
        sign[at_mid] = prior[at_mid]
    return sign


def compute_tick_factors(
    ticker: str,
    source=None,
    date: Optional[str] = None,
) -> dict[str, float]:
    """Return ``{block_share, block_aggression, buy_ratio}`` for ``ticker``."""
    empty = {"block_share": 0.0, "block_aggression": 0.0, "buy_ratio": 0.5}
    if source is None:
        source = get_data_source()
    ticks = source.get_ticks(ticker, date=date)
    if ticks is None or ticks.empty:
        return empty

    if "bid" not in ticks.columns or "ask" not in ticks.columns:
        # Fall back to rolling mid from price
        ticks = ticks.copy()
        ticks["mid"] = ticks["price"].rolling(20, min_periods=1).mean()
        ticks["bid"] = ticks["mid"] - ticks["mid"] * 0.0003
        ticks["ask"] = ticks["mid"] + ticks["mid"] * 0.0003

    sign = _sign_trades(ticks)
    ticks = ticks.assign(sign=sign, bucket=ticks["size"].apply(_bucket))

    total_vol = int(ticks["size"].sum())
    if total_vol <= 0:
        return empty

    # Block bucket
    block = ticks[ticks["bucket"] == "block"]
    block_vol = int(block["size"].sum())
    block_buys = int(block.loc[block["sign"] == 1, "size"].sum())
    block_sells = int(block.loc[block["sign"] == -1, "size"].sum())

    block_share = block_vol / total_vol
    block_aggression = (
        (block_buys - block_sells) / block_vol if block_vol > 0 else 0.0
    )
    buy_ratio = float((sign == 1).sum()) / max(1, len(sign))

    return {
        "block_share": float(block_share),
        "block_aggression": float(block_aggression),
        "buy_ratio": float(buy_ratio),
    }


def compute_tick_factors_batch(
    tickers: list[str],
    source=None,
) -> pd.DataFrame:
    """Return a DataFrame indexed by ticker with the three tick factors."""
    if source is None:
        source = get_data_source()
    rows = []
    for t in tickers:
        f = compute_tick_factors(t, source=source)
        f["ticker"] = t
        rows.append(f)
    return pd.DataFrame(rows).set_index("ticker")