"""Constraint-aware portfolio optimizer replacing simple CapitalAllocator."""

from __future__ import annotations

import logging
import math
import threading
import time
from collections import deque
from dataclasses import dataclass, field
from typing import Any

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


# ---------------------------------------------------------------------------
# Dataclasses
# ---------------------------------------------------------------------------

@dataclass(frozen=True)
class OptimizationConstraints:
    """Hard constraints for portfolio optimization."""
    max_per_market: float = 0.10
    max_per_exchange: float = 0.40
    max_per_sector: float = 0.30
    max_turnover: float = 0.20
    max_positions: int = 50


@dataclass(frozen=True)
class RebalanceTrigger:
    """Reason a rebalance was triggered."""
    trigger_type: str  # "drift" / "var_breach" / "correlation_change"
    details: str
    timestamp: float


@dataclass(frozen=True)
class OptimizationResult:
    """Output of portfolio optimization."""
    target_weights: dict[str, float]
    expected_return: float
    expected_risk: float
    turnover: float
    violations: list[str]
    rebalance_triggers: list[RebalanceTrigger]


# ---------------------------------------------------------------------------
# Constants
# ---------------------------------------------------------------------------

_MIN_HISTORY = 10
_GRADIENT_ITERS = 10
_DRIFT_THRESHOLD = 0.05
_EPSILON = 1e-12


# ---------------------------------------------------------------------------
# PortfolioOptimizer
# ---------------------------------------------------------------------------

class PortfolioOptimizer:
    """Constraint-aware portfolio optimizer with Kelly sizing and
    Ledoit-Wolf covariance shrinkage.

    Replaces the simpler :class:`CapitalAllocator` with projected
    gradient descent that respects per-market, per-exchange, turnover,
    and position-count constraints.
    """

    def __init__(
        self,
        constraints: OptimizationConstraints | None = None,
        total_capital: float = 100_000,
    ) -> None:
        self._constraints = constraints or OptimizationConstraints()
        self._total_capital = total_capital
        self._current_weights: dict[str, float] = {}
        self._return_history: dict[str, deque[float]] = {}
        self._strategy_exchanges: dict[str, str] = {}
        self._lock = threading.Lock()

    # ------------------------------------------------------------------
    # Public helpers
    # ------------------------------------------------------------------

    def update_returns(self, strategy_name: str, ret: float) -> None:
        """Append a return observation for *strategy_name*."""
        if not math.isfinite(ret):
            return  # silently ignore non-finite returns
        with self._lock:
            if strategy_name not in self._return_history:
                self._return_history[strategy_name] = deque(maxlen=100)
            self._return_history[strategy_name].append(ret)

    def set_exchange(self, strategy_name: str, exchange: str) -> None:
        """Register the exchange a strategy trades on."""
        with self._lock:
            self._strategy_exchanges[strategy_name] = exchange

    def set_current_weights(self, weights: dict[str, float]) -> None:
        """Update the current portfolio weights snapshot."""
        with self._lock:
            self._current_weights = dict(weights)

    # ------------------------------------------------------------------
    # Core optimization
    # ------------------------------------------------------------------

    def optimize(
        self,
        strategy_names: list[str],
        engines: dict[str, Any] | None = None,
        probs: list[float] | None = None,
        prices: list[float] | None = None,
    ) -> OptimizationResult:
        """Run constrained portfolio optimization.

        Args:
            strategy_names: Active strategy names.
            engines: Optional mapping of name -> Engine (for return history).
            probs: Outcome probabilities per strategy (for Kelly sizing).
            prices: Market prices per strategy (for Kelly sizing).

        Returns:
            :class:`OptimizationResult` with target weights, risk metrics,
            constraint violations, and rebalance triggers.
        """
        if not strategy_names:
            return OptimizationResult(
                target_weights={},
                expected_return=0.0,
                expected_risk=0.0,
                turnover=0.0,
                violations=[],
                rebalance_triggers=[],
            )

        n = len(strategy_names)
        violations: list[str] = []

        # Enforce max positions
        if n > self._constraints.max_positions:
            violations.append(
                f"Position count {n} exceeds max {self._constraints.max_positions}"
            )

        # --- Step 1: Starting weights from Kelly or equal ----------------
        raw_weights = self._initial_weights(strategy_names, probs, prices)

        # --- Step 2: Covariance (Ledoit-Wolf if enough history) ----------
        cov_matrix: list[list[float]] | None = None
        with self._lock:
            has_history = all(
                strategy_name in self._return_history
                and len(self._return_history[strategy_name]) >= _MIN_HISTORY
                for strategy_name in strategy_names
            )

        if has_history:
            cov_matrix = self._build_covariance(strategy_names)

        # --- Step 3: Projected gradient descent --------------------------
        weights = list(raw_weights)
        with self._lock:
            current = [self._current_weights.get(name, 0.0) for name in strategy_names]
            exchanges = [
                self._strategy_exchanges.get(name, "unknown")
                for name in strategy_names
            ]
            constraints = self._constraints

        for _ in range(_GRADIENT_ITERS):
            # (a) Clip each weight to max_per_market
            weights = [min(w, constraints.max_per_market) for w in weights]
            weights = [max(w, 0.0) for w in weights]

            # (b) Normalize to sum <= 1
            total_w = sum(weights)
            if total_w > 1.0 + _EPSILON:
                weights = [w / total_w for w in weights]

            # (c) Exchange concentration: clip per-exchange group
            weights = self._clip_exchange_concentration(
                weights, exchanges, constraints.max_per_exchange
            )

            # (d) Turnover constraint
            turnover = sum(abs(weights[i] - current[i]) for i in range(n))
            if turnover > constraints.max_turnover and turnover > _EPSILON:
                # Interpolate toward current weights to reduce turnover
                alpha = constraints.max_turnover / turnover
                weights = [
                    current[i] + alpha * (weights[i] - current[i])
                    for i in range(n)
                ]

        # Final normalization (ensure non-negative, sum <= 1)
        weights = [max(w, 0.0) for w in weights]
        total_w = sum(weights)
        if total_w > 1.0 + _EPSILON:
            weights = [w / total_w for w in weights]

        # --- Step 4: Risk/return metrics ---------------------------------
        expected_return = sum(weights)  # proxy: sum of weights
        expected_risk = self._portfolio_risk(weights, cov_matrix)
        turnover = sum(abs(weights[i] - current[i]) for i in range(n))

        target_weights = {
            name: weights[i] * self._total_capital
            for i, name in enumerate(strategy_names)
        }

        # --- Step 5: Check rebalance triggers ----------------------------
        with self._lock:
            cw = dict(self._current_weights)
        tw_pct = {name: weights[i] for i, name in enumerate(strategy_names)}
        triggers = self.check_rebalance(cw, tw_pct)

        return OptimizationResult(
            target_weights=target_weights,
            expected_return=expected_return,
            expected_risk=expected_risk,
            turnover=turnover,
            violations=violations,
            rebalance_triggers=triggers,
        )

    # ------------------------------------------------------------------
    # Rebalance triggers
    # ------------------------------------------------------------------

    def check_rebalance(
        self,
        current_weights: dict[str, float],
        target_weights: dict[str, float],
    ) -> list[RebalanceTrigger]:
        """Check whether rebalancing is warranted.

        Args:
            current_weights: Current portfolio weight per strategy (0-1).
            target_weights: Target portfolio weight per strategy (0-1).

        Returns:
            List of triggers that fired.
        """
        triggers: list[RebalanceTrigger] = []
        now = time.time()

        all_keys = set(current_weights) | set(target_weights)
        max_drift = 0.0
        drift_name = ""
        for key in all_keys:
            drift = abs(current_weights.get(key, 0.0) - target_weights.get(key, 0.0))
            if drift > max_drift:
                max_drift = drift
                drift_name = key

        if max_drift > _DRIFT_THRESHOLD:
            triggers.append(RebalanceTrigger(
                trigger_type="drift",
                details=(
                    f"Strategy '{drift_name}' drifted {max_drift:.2%} "
                    f"from target (threshold: {_DRIFT_THRESHOLD:.0%})"
                ),
                timestamp=now,
            ))

        return triggers

    # ------------------------------------------------------------------
    # Simplified interface
    # ------------------------------------------------------------------

    def allocate(
        self,
        strategy_names: list[str],
        engines: dict[str, Any] | None = None,
    ) -> dict[str, float]:
        """Simplified interface matching :meth:`CapitalAllocator.allocate`.

        Calls :meth:`optimize` and returns ``{name: capital_amount}``.
        """
        result = self.optimize(strategy_names, engines=engines)
        return result.target_weights

    # ------------------------------------------------------------------
    # Internal helpers
    # ------------------------------------------------------------------

    def _initial_weights(
        self,
        strategy_names: list[str],
        probs: list[float] | None,
        prices: list[float] | None,
    ) -> list[float]:
        """Compute starting weights from Kelly fractions or equal weight."""
        n = len(strategy_names)
        if probs is not None and prices is not None and len(probs) == n and len(prices) == n:
            try:
                from horizon._horizon import multi_kelly

                fractions = multi_kelly(probs, prices, 1.0)
                total = sum(fractions)
                if total > _EPSILON:
                    return [f / total for f in fractions]
            except Exception:
                logger.debug("multi_kelly failed, falling back to equal weights")

        # Equal weights
        w = 1.0 / n if n > 0 else 0.0
        return [w] * n

    def _build_covariance(
        self, strategy_names: list[str]
    ) -> list[list[float]] | None:
        """Build Ledoit-Wolf shrunk covariance matrix from return history."""
        with self._lock:
            # Build T x N returns matrix
            histories = [
                list(self._return_history[name]) for name in strategy_names
            ]

        # Align to shortest common length
        min_len = min(len(h) for h in histories)
        if min_len < _MIN_HISTORY:
            return None

        # Each row is one time step with N asset returns
        returns_matrix: list[list[float]] = []
        for t in range(min_len):
            row = [histories[i][t] for i in range(len(strategy_names))]
            returns_matrix.append(row)

        try:
            from horizon._horizon import ledoit_wolf_shrinkage

            cov, _shrinkage_intensity = ledoit_wolf_shrinkage(returns_matrix)
            return cov
        except Exception:
            logger.debug("ledoit_wolf_shrinkage failed, skipping covariance")
            return None

    @staticmethod
    def _clip_exchange_concentration(
        weights: list[float],
        exchanges: list[str],
        max_per_exchange: float,
    ) -> list[float]:
        """Clip weights so no single exchange exceeds *max_per_exchange*."""
        # Group indices by exchange
        exchange_groups: dict[str, list[int]] = {}
        for i, ex in enumerate(exchanges):
            exchange_groups.setdefault(ex, []).append(i)

        result = list(weights)
        for _ex, indices in exchange_groups.items():
            group_total = sum(result[i] for i in indices)
            if group_total > max_per_exchange and group_total > _EPSILON:
                scale = max_per_exchange / group_total
                for i in indices:
                    result[i] *= scale

        return result

    @staticmethod
    def _portfolio_risk(
        weights: list[float],
        cov_matrix: list[list[float]] | None,
    ) -> float:
        """Compute portfolio risk (std dev) from weights and covariance.

        Falls back to sum-of-squared-weights proxy when covariance is
        unavailable.
        """
        if cov_matrix is not None:
            n = len(weights)
            variance = 0.0
            for i in range(n):
                for j in range(n):
                    variance += weights[i] * weights[j] * cov_matrix[i][j]
            return math.sqrt(max(variance, 0.0))

        # Proxy: sqrt of sum of squared weights (assumes unit variance, zero corr)
        return math.sqrt(sum(w * w for w in weights))
