"""Position sizing functions for use inside ``generate_signal()``.

Each function returns a weight (float) or a dict of weights.
All thresholds and parameters are configurable with sensible defaults.
"""

from __future__ import annotations

from typing import Any, Dict, Optional

import numpy as np

from cpz.indicators._helpers import EPSILON, TRADING_DAYS_PER_YEAR


def vol_target(
    bars: Any,
    risk: float = 0.02,
    *,
    period: int = 20,
    annualize: bool = False,
    trading_days: int = TRADING_DAYS_PER_YEAR,
    col_close: str = "close",
) -> float:
    """Size position by inverse realised volatility.

    The returned weight targets *risk* units of daily volatility
    (e.g. ``risk=0.02`` means the position contributes ~2 % daily vol).

    Args:
        bars: DataFrame or array of prices.
        risk: Target volatility contribution per position.
        period: Rolling window for volatility estimation.
        annualize: If ``True``, interpret *risk* as annualised vol target.
        trading_days: Trading days per year (for annualisation).
        col_close: Column name for close prices.
    """
    from cpz.indicators._helpers import ensure_array

    close = ensure_array(bars, col_close)
    if len(close) < period + 1:
        return 0.0

    log_ret = np.diff(np.log(np.maximum(close[-(period + 1) :], EPSILON)))
    vol = float(np.std(log_ret, ddof=1))

    if annualize:
        vol *= np.sqrt(trading_days)

    return risk / vol if vol > EPSILON else 0.0


def atr_size(
    bars: Any,
    risk_per_trade: float = 0.01,
    *,
    atr_period: int = 14,
    smoothing: str = "wilder",
    col_high: str = "high",
    col_low: str = "low",
    col_close: str = "close",
) -> float:
    """Size position so that one ATR move equals *risk_per_trade* of portfolio.

    Args:
        bars: DataFrame with OHLC data.
        risk_per_trade: Fraction of portfolio to risk per ATR unit.
        atr_period: ATR look-back.
        smoothing: ATR smoothing (``"wilder"``, ``"ema"``, ``"sma"``).
        col_*: Column name overrides.
    """
    from cpz.indicators import atr as compute_atr
    from cpz.indicators._helpers import ensure_array

    atr_val = compute_atr(bars, atr_period, smoothing=smoothing,
                          col_high=col_high, col_low=col_low, col_close=col_close)
    close = ensure_array(bars, col_close)
    price = float(close[-1]) if len(close) > 0 else 0.0

    if price < EPSILON or atr_val < EPSILON:
        return 0.0
    return risk_per_trade / (atr_val / price)


def kelly(
    win_rate: float,
    avg_win: float,
    avg_loss: float,
    *,
    fraction: float = 0.25,
    max_size: float = 1.0,
) -> float:
    """Kelly criterion with fractional scaling.

    Args:
        win_rate: Probability of a winning trade (0-1).
        avg_win: Average gain on winning trades.
        avg_loss: Average loss on losing trades (positive number).
        fraction: Fraction of full Kelly to use (default 0.25 = quarter-Kelly).
        max_size: Maximum position size cap.
    """
    if avg_loss < EPSILON:
        return 0.0
    edge = win_rate - (1.0 - win_rate) / (avg_win / avg_loss)
    sized = max(0.0, edge * fraction)
    return min(sized, max_size)


def equal_weight(n_assets: int) -> float:
    """Equal-weight allocation: ``1 / N``."""
    return 1.0 / n_assets if n_assets > 0 else 0.0


def risk_parity_weights(
    bars_dict: Dict[str, Any],
    *,
    period: int = 60,
    col_close: str = "close",
) -> Dict[str, float]:
    """Inverse-variance risk parity across symbols.

    Each asset's weight is proportional to ``1 / variance``, normalised
    to sum to 1.

    Args:
        bars_dict: ``{symbol: DataFrame}`` mapping.
        period: Rolling window for variance estimation.
        col_close: Close price column name.
    """
    from cpz.indicators._helpers import ensure_array

    inv_var: Dict[str, float] = {}
    for sym, bars in bars_dict.items():
        close = ensure_array(bars, col_close)
        if len(close) < period + 1:
            continue
        log_ret = np.diff(np.log(np.maximum(close[-(period + 1) :], EPSILON)))
        var = float(np.var(log_ret, ddof=1))
        if var > EPSILON:
            inv_var[sym] = 1.0 / var

    total = sum(inv_var.values())
    if total < EPSILON:
        n = len(bars_dict)
        return {s: 1.0 / n for s in bars_dict} if n > 0 else {}
    return {s: v / total for s, v in inv_var.items()}


def inverse_vol_weights(
    bars_dict: Dict[str, Any],
    *,
    period: int = 20,
    annualize: bool = False,
    trading_days: int = TRADING_DAYS_PER_YEAR,
    col_close: str = "close",
) -> Dict[str, float]:
    """Inverse-volatility weights: weight ∝ 1 / vol.

    Args:
        bars_dict: ``{symbol: DataFrame}`` mapping.
        period: Rolling window for volatility.
        annualize: Annualise the volatility before inverting.
        trading_days: Trading days per year.
        col_close: Close price column name.
    """
    from cpz.indicators._helpers import ensure_array

    inv_vol: Dict[str, float] = {}
    for sym, bars in bars_dict.items():
        close = ensure_array(bars, col_close)
        if len(close) < period + 1:
            continue
        log_ret = np.diff(np.log(np.maximum(close[-(period + 1) :], EPSILON)))
        vol = float(np.std(log_ret, ddof=1))
        if annualize:
            vol *= np.sqrt(trading_days)
        if vol > EPSILON:
            inv_vol[sym] = 1.0 / vol

    total = sum(inv_vol.values())
    if total < EPSILON:
        n = len(bars_dict)
        return {s: 1.0 / n for s in bars_dict} if n > 0 else {}
    return {s: v / total for s, v in inv_vol.items()}
