"""Transaction cost models: Almgren-Chriss, linear/sqrt impact, spread, turnover analysis.

All functions are pure: data in, results out.
"""

from __future__ import annotations

from dataclasses import dataclass, field
from typing import Any, Dict, List, Optional

import numpy as np

EPSILON: float = 1e-15


@dataclass
class CostEstimate:
    """Breakdown of transaction cost for a single trade."""
    permanent_cost_bps: float = 0.0
    temporary_cost_bps: float = 0.0
    spread_cost_bps: float = 0.0
    commission_bps: float = 0.0
    total_cost_bps: float = 0.0


@dataclass
class TurnoverResult:
    """Portfolio turnover analysis."""
    one_way_turnover: float = 0.0
    two_way_turnover: float = 0.0
    estimated_cost_bps: float = 0.0
    per_asset_turnover: Dict[str, float] = field(default_factory=dict)
    per_asset_cost_bps: Dict[str, float] = field(default_factory=dict)


@dataclass
class CostModel:
    """Pluggable cost model for portfolio optimisation.

    All parameters are configurable so users can calibrate to their
    broker and market conditions.
    """
    spread_bps: float = 5.0
    impact_coefficient: float = 0.1
    impact_exponent: float = 0.5
    impact_model: str = "sqrt"
    commission_bps: float = 0.0
    permanent_fraction: float = 0.5

    def estimate(
        self,
        notional: float,
        adv: float = 0.0,
        volatility: float = 0.0,
    ) -> float:
        """Estimate round-trip cost in basis points.

        Args:
            notional: Trade notional value.
            adv: Average daily volume (notional).
            volatility: Daily volatility of the asset.
        """
        spread = self.spread_bps / 2.0

        if self.impact_model == "sqrt":
            impact = sqrt_impact(notional, adv, coefficient=self.impact_coefficient,
                                 volatility=volatility, exponent=self.impact_exponent)
        elif self.impact_model == "linear":
            impact = linear_impact(notional, adv, coefficient=self.impact_coefficient)
        elif self.impact_model == "almgren_chriss":
            result = almgren_chriss(notional, adv, volatility)
            impact = result.permanent_cost_bps + result.temporary_cost_bps
        else:
            impact = 0.0

        return spread + impact + self.commission_bps


def almgren_chriss(
    notional: float,
    adv: float,
    volatility: float,
    *,
    gamma: float = 0.314,
    eta: float = 0.142,
    alpha: float = 0.6,
    beta: float = 0.6,
    participation_rate: float = 0.1,
) -> CostEstimate:
    """Almgren-Chriss (2001) market impact model.

    permanent_impact = gamma * sigma * (Q/V)^alpha
    temporary_impact = eta * sigma * (Q/(V*T))^beta

    Args:
        notional: Trade notional.
        adv: Average daily volume (notional).
        volatility: Daily price volatility.
        gamma: Permanent impact coefficient.
        eta: Temporary impact coefficient.
        alpha: Permanent impact exponent.
        beta: Temporary impact exponent.
        participation_rate: Fraction of ADV consumed per time unit.
    """
    if adv < EPSILON or volatility < EPSILON:
        return CostEstimate()

    q_over_v = notional / adv
    perm = gamma * volatility * (q_over_v ** alpha) * 10000.0

    execution_time = q_over_v / max(participation_rate, EPSILON)
    temp = eta * volatility * ((q_over_v / max(execution_time, EPSILON)) ** beta) * 10000.0

    return CostEstimate(
        permanent_cost_bps=round(perm, 4),
        temporary_cost_bps=round(temp, 4),
        total_cost_bps=round(perm + temp, 4),
    )


def linear_impact(
    notional: float,
    adv: float,
    *,
    coefficient: float = 0.1,
) -> float:
    """Linear market impact in basis points.

    ``cost_bps = coefficient * (notional / ADV) * 10000``

    Args:
        notional: Trade notional.
        adv: Average daily volume.
        coefficient: Impact coefficient (calibrate to your market).
    """
    if adv < EPSILON:
        return 0.0
    return coefficient * (notional / adv) * 10000.0


def sqrt_impact(
    notional: float,
    adv: float,
    *,
    coefficient: float = 0.1,
    volatility: float = 0.0,
    exponent: float = 0.5,
) -> float:
    """Square-root market impact (Torre 1997, Barra).

    ``cost_bps = coefficient * sigma * (notional / ADV)^exponent * 10000``

    Args:
        notional: Trade notional.
        adv: Average daily volume.
        coefficient: Impact coefficient.
        volatility: Daily price volatility (if 0, uses coefficient alone).
        exponent: Impact exponent (default 0.5 = square root).
    """
    if adv < EPSILON:
        return 0.0
    sigma = volatility if volatility > EPSILON else 1.0
    return coefficient * sigma * ((notional / adv) ** exponent) * 10000.0


def spread_cost(bid_ask_spread_bps: float) -> float:
    """Half-spread cost: ``cost = spread / 2``. The minimum cost of any trade."""
    return bid_ask_spread_bps / 2.0


def total_cost(
    notional: float,
    adv: float,
    volatility: float,
    *,
    bid_ask_bps: float = 5.0,
    commission_bps: float = 0.0,
    impact_model: str = "sqrt",
    impact_coefficient: float = 0.1,
    impact_exponent: float = 0.5,
) -> CostEstimate:
    """Combined transaction cost: spread + impact + commission.

    Args:
        notional: Trade notional.
        adv: Average daily volume.
        volatility: Daily price volatility.
        bid_ask_bps: Bid-ask spread in bps.
        commission_bps: Broker commission in bps.
        impact_model: ``"sqrt"`` (default), ``"linear"``, or ``"almgren_chriss"``.
        impact_coefficient: Impact coefficient.
        impact_exponent: Impact exponent (for sqrt model).
    """
    spread = spread_cost(bid_ask_bps)

    if impact_model == "sqrt":
        impact = sqrt_impact(notional, adv, coefficient=impact_coefficient,
                             volatility=volatility, exponent=impact_exponent)
        return CostEstimate(
            spread_cost_bps=round(spread, 4),
            temporary_cost_bps=round(impact, 4),
            commission_bps=commission_bps,
            total_cost_bps=round(spread + impact + commission_bps, 4),
        )
    elif impact_model == "linear":
        impact = linear_impact(notional, adv, coefficient=impact_coefficient)
        return CostEstimate(
            spread_cost_bps=round(spread, 4),
            temporary_cost_bps=round(impact, 4),
            commission_bps=commission_bps,
            total_cost_bps=round(spread + impact + commission_bps, 4),
        )
    elif impact_model == "almgren_chriss":
        ac = almgren_chriss(notional, adv, volatility)
        ac.spread_cost_bps = round(spread, 4)
        ac.commission_bps = commission_bps
        ac.total_cost_bps = round(
            ac.permanent_cost_bps + ac.temporary_cost_bps + spread + commission_bps, 4
        )
        return ac
    else:
        raise ValueError(f"Unknown impact_model '{impact_model}'. Use: sqrt, linear, almgren_chriss")


def turnover_analysis(
    current_weights: Dict[str, float],
    target_weights: Dict[str, float],
    *,
    cost_model: Optional[CostModel] = None,
    adv: Optional[Dict[str, float]] = None,
    volatility: Optional[Dict[str, float]] = None,
    portfolio_value: float = 1_000_000.0,
) -> TurnoverResult:
    """Analyse portfolio turnover and estimated rebalancing cost.

    Args:
        current_weights: ``{symbol: current_weight}``.
        target_weights: ``{symbol: target_weight}``.
        cost_model: Cost model for per-asset cost estimation.
        adv: ``{symbol: average_daily_volume}`` (optional).
        volatility: ``{symbol: daily_volatility}`` (optional).
        portfolio_value: Total portfolio value for notional calculation.
    """
    all_symbols = set(current_weights.keys()) | set(target_weights.keys())
    per_asset_turn: Dict[str, float] = {}
    per_asset_cost: Dict[str, float] = {}

    for sym in all_symbols:
        curr = current_weights.get(sym, 0.0)
        tgt = target_weights.get(sym, 0.0)
        turn = abs(tgt - curr)
        per_asset_turn[sym] = round(turn, 6)

        if cost_model and turn > EPSILON:
            notional = turn * portfolio_value
            sym_adv = (adv or {}).get(sym, portfolio_value * 0.01)
            sym_vol = (volatility or {}).get(sym, 0.02)
            per_asset_cost[sym] = round(cost_model.estimate(notional, sym_adv, sym_vol), 4)
        else:
            per_asset_cost[sym] = 0.0

    one_way = sum(per_asset_turn.values()) / 2.0
    two_way = sum(per_asset_turn.values())
    total_cost = sum(per_asset_cost.values())

    return TurnoverResult(
        one_way_turnover=round(one_way, 6),
        two_way_turnover=round(two_way, 6),
        estimated_cost_bps=round(total_cost, 4),
        per_asset_turnover=per_asset_turn,
        per_asset_cost_bps=per_asset_cost,
    )
