"""Robust Portfolio Optimization pipeline functions.

Pipeline functions:
    robust_allocator - Robust portfolio allocation with uncertainty
"""

from __future__ import annotations

from typing import Any, Callable

from horizon.context import Context


def robust_allocator(
    markets: list[str],
    uncertainty_budget: float = 0.1,
    max_weight: float = 0.3,
    min_weight: float = 0.0,
) -> Callable[[Context], dict[str, Any]]:
    """Create a robust portfolio allocation pipeline function.

    Computes worst-case optimal portfolio weights accounting for
    estimation uncertainty in expected returns.

    Args:
        markets: List of market/feed names.
        uncertainty_budget: Size of uncertainty set.
        max_weight: Maximum weight per asset.
        min_weight: Minimum weight per asset.

    Returns:
        Pipeline function: (Context) -> dict with keys:
            weights, worst_case_return, nominal_return
    """
    from horizon._horizon import auth_require_pro
    auth_require_pro()

    from collections import deque
    price_histories: dict[str, deque] = {}

    def _robust_allocator(ctx: Context) -> dict[str, Any]:
        for market in markets:
            if market not in price_histories:
                price_histories[market] = deque(maxlen=200)
            fd = ctx.feeds.get(market)
            if fd is not None and fd.price > 0:
                price_histories[market].append(fd.price)

        result = {
            "weights": [1.0 / len(markets)] * len(markets),
            "worst_case_return": 0.0,
            "nominal_return": 0.0,
        }

        min_len = min(
            (len(price_histories[m]) for m in markets if m in price_histories),
            default=0,
        )

        if min_len < 10:
            return result

        try:
            from horizon._horizon import robust_optimize
            import math

            n = len(markets)

            # Compute returns
            returns_data = []
            for market in markets:
                prices = list(price_histories[market])[-min_len:]
                rets = [
                    (prices[i] - prices[i - 1]) / prices[i - 1]
                    for i in range(1, len(prices))
                    if prices[i - 1] > 0
                ]
                returns_data.append(rets)

            m = min(len(r) for r in returns_data)
            if m < 5:
                return result

            # Expected returns
            expected = [sum(r[:m]) / m for r in returns_data]

            # Covariance
            cov = [[0.0] * n for _ in range(n)]
            for i in range(n):
                for j in range(n):
                    mean_i = expected[i]
                    mean_j = expected[j]
                    cov[i][j] = sum(
                        (returns_data[i][k] - mean_i) * (returns_data[j][k] - mean_j)
                        for k in range(m)
                    ) / m

            # Uncertainty = std of returns
            uncertainty = [math.sqrt(max(cov[i][i], 1e-15)) for i in range(n)]

            opt = robust_optimize(
                expected, cov, uncertainty, uncertainty_budget,
                max_weight, min_weight,
            )
            result["weights"] = list(opt.weights)
            result["worst_case_return"] = opt.worst_case_return
            result["nominal_return"] = opt.nominal_return

        except (ValueError, RuntimeError, ImportError):
            pass

        return result

    _robust_allocator.__name__ = "robust_allocator"
    return _robust_allocator
