"""Capital allocation engine for fund strategies."""

from __future__ import annotations

import logging
import math
from collections import deque
from typing import Any

from horizon.fund._types import AllocationMethod, FundConfig

logger = logging.getLogger("horizon.fund.allocator")


class CapitalAllocator:
    """Allocates fund capital across strategies.

    Supports equal-weight, risk-parity, performance-based, and custom
    allocation methods. Returns a mapping of strategy_name -> capital.
    """

    def __init__(self, config: FundConfig) -> None:
        self._config = config
        self._custom_weights = dict(config.custom_weights)
        # Rolling window of daily PnL per strategy for improved vol estimation
        self._daily_pnl_history: dict[str, deque[float]] = {}

    def allocate(
        self,
        strategy_names: list[str],
        engines: dict[str, Any] | None = None,
    ) -> dict[str, float]:
        """Compute capital allocation for each strategy.

        Args:
            strategy_names: Active strategy names.
            engines: Optional mapping of name -> Engine (needed for
                     performance/risk-parity methods).

        Returns:
            Mapping of strategy_name -> allocated capital.
        """
        if not strategy_names:
            return {}

        total = self._config.total_capital
        method = self._config.allocation_method

        if method == AllocationMethod.EQUAL:
            return self._equal(strategy_names, total)
        elif method == AllocationMethod.RISK_PARITY:
            return self._risk_parity(strategy_names, total, engines)
        elif method == AllocationMethod.PERFORMANCE:
            return self._performance(strategy_names, total, engines)
        elif method == AllocationMethod.CUSTOM:
            return self._custom(strategy_names, total)
        else:
            return self._equal(strategy_names, total)

    def set_custom_weight(self, name: str, weight: float) -> None:
        """Set or update a custom weight for a strategy."""
        self._custom_weights[name] = weight

    def remove_weight(self, name: str) -> None:
        """Remove a custom weight."""
        self._custom_weights.pop(name, None)

    # ------------------------------------------------------------------
    # Allocation methods
    # ------------------------------------------------------------------

    @staticmethod
    def _equal(names: list[str], total: float) -> dict[str, float]:
        per_strategy = total / len(names)
        return {n: per_strategy for n in names}

    def _risk_parity(
        self,
        names: list[str],
        total: float,
        engines: dict[str, Any] | None,
    ) -> dict[str, float]:
        """Allocate inversely proportional to realized volatility.

        Uses a rolling window of daily PnL values to compute standard
        deviation as the volatility proxy. Falls back to a simpler
        proxy (daily_pnl / total_pnl) when history is insufficient,
        and to equal-weight if no engine data is available.
        """
        if not engines:
            return self._equal(names, total)

        inv_vols: dict[str, float] = {}
        for name in names:
            engine = engines.get(name)
            if engine is None:
                inv_vols[name] = 1.0
                continue
            try:
                status = engine.status()
                # Track daily PnL history for rolling stdev
                if name not in self._daily_pnl_history:
                    self._daily_pnl_history[name] = deque(maxlen=60)
                self._daily_pnl_history[name].append(status.daily_pnl)

                history = self._daily_pnl_history[name]
                if len(history) >= 5:
                    # Compute stdev of daily PnL as volatility proxy
                    mean_pnl = sum(history) / len(history)
                    variance = sum((x - mean_pnl) ** 2 for x in history) / len(history)
                    vol = max(math.sqrt(variance), 0.001)
                else:
                    # Fallback: simple proxy when not enough history
                    pnl = abs(status.total_realized_pnl + status.total_unrealized_pnl)
                    vol = max(abs(status.daily_pnl) / max(pnl + 1.0, 1.0), 0.001)
                inv_vols[name] = 1.0 / vol
            except Exception:
                inv_vols[name] = 1.0

        total_inv = sum(inv_vols.values())
        if total_inv <= 0:
            return self._equal(names, total)

        return {n: (inv_vols[n] / total_inv) * total for n in names}

    def _performance(
        self,
        names: list[str],
        total: float,
        engines: dict[str, Any] | None,
    ) -> dict[str, float]:
        """Allocate proportional to cumulative P&L (with floor).

        Strategies with negative P&L get a minimum allocation.
        Falls back to equal if no engine data available.
        """
        if not engines:
            return self._equal(names, total)

        pnls: dict[str, float] = {}
        for name in names:
            engine = engines.get(name)
            if engine is None:
                pnls[name] = 0.0
                continue
            try:
                status = engine.status()
                pnls[name] = status.total_realized_pnl + status.total_unrealized_pnl
            except Exception:
                pnls[name] = 0.0

        # Shift all P&L values so the minimum is at least 1.0
        # This gives losing strategies a small but nonzero allocation
        min_pnl = min(pnls.values()) if pnls else 0.0
        shift = abs(min_pnl) + 1.0 if min_pnl < 1.0 else 0.0
        shifted = {n: pnls[n] + shift for n in names}

        total_shifted = sum(shifted.values())
        if total_shifted <= 0:
            return self._equal(names, total)

        return {n: (shifted[n] / total_shifted) * total for n in names}

    def _custom(self, names: list[str], total: float) -> dict[str, float]:
        """Allocate based on user-defined weights.

        Missing strategies get equal share of remaining capital.
        """
        result: dict[str, float] = {}
        weighted_total = 0.0
        unweighted: list[str] = []

        for name in names:
            w = self._custom_weights.get(name)
            if w is not None:
                result[name] = w * total
                weighted_total += w
            else:
                unweighted.append(name)

        # Distribute remaining capital equally among unweighted strategies
        remaining = max(0.0, (1.0 - weighted_total) * total)
        if unweighted:
            per_unweighted = remaining / len(unweighted)
            for name in unweighted:
                result[name] = per_unweighted

        return result
