"""Portfolio weighting: mean-variance, Black-Litterman, HRP, and managed portfolios.

Wraps PyPortfolioOpt for standard optimization methods and implements a
signal-orthogonalization managed portfolio with rolling Sharpe-maximizing
lambda selection (see ManagedPortfolio docstring for details).
"""

from __future__ import annotations

from typing import Any

import numpy as np
import pandas as pd
from pypfopt import (
    BlackLittermanModel,
    EfficientCDaR,
    EfficientCVaR,
    EfficientFrontier,
    HRPOpt,
    risk_models,
)
from pypfopt.risk_models import CovarianceShrinkage


class ManagedPortfolio:
    """Signal-orthogonalization managed portfolio with rolling Sharpe-maximizing lambda.

    Stage 1: Standardize the raw signal s_t to zero mean / unit std.
    Stage 2: Orthogonalize: s_adj = s_norm - beta_s * f_t (residual after
             regressing standardized signal on external factor returns).
    Stage 3: Select lambda via rolling cv_window to maximize in-sample Sharpe
             of lambda * s_adj.

    Note: This is NOT the Moreira-Muir (2017) variance-timing approach.
    MM2017 Eq. 3 scales factor returns by the inverse of lagged conditional
    variance: f^VM_t = (c / sigma^2_{t-1}) * f_t, with no regression step.
    The present implementation is a signal management framework closer to
    Baz et al. (2015), "Dissecting Investment Strategies in the Cross Section
    and Time Series," with Sharpe-optimal scalar selection.

    Parameters
    ----------
    cv_window : int
        Rolling window length (in months) for cross-validation.
    lambda_grid : list[float] | None
        Grid of lambda values to search over. Defaults to
        [0.0, 0.25, 0.5, 0.75, 1.0, 1.5, 2.0, 3.0, 5.0].
    """

    def __init__(
        self,
        cv_window: int = 60,
        lambda_grid: list[float] | None = None,
    ) -> None:
        self.cv_window = cv_window
        self.lambda_grid = lambda_grid or [0.0, 0.25, 0.5, 0.75, 1.0, 1.5, 2.0, 3.0, 5.0]
        # Fitted attributes
        self._beta: np.ndarray[Any, np.dtype[np.floating[Any]]] | None = None
        self._lambda_opt: float = 1.0
        self._lambda_path: pd.Series[float] | None = None

    def fit(
        self,
        signal: pd.Series[float],
        factor_returns: pd.DataFrame,
    ) -> None:
        """Fit the managed portfolio model.

        Stage 1: Standardize the signal to zero mean and unit variance.
        Stage 2: Regress signal on factor returns to get residual beta.
        Lambda selection: Rolling CV over lambda_grid, picking the lambda
        that maximizes in-sample Sharpe ratio.

        Parameters
        ----------
        signal : pd.Series
            Time series of signal values (e.g., portfolio return from sorting).
        factor_returns : pd.DataFrame
            Factor return matrix aligned with signal index.
        """
        # Align signal and factors on common index
        common_idx = signal.index.intersection(factor_returns.index)
        s = np.asarray(signal.loc[common_idx].values, dtype=float)
        f = np.asarray(factor_returns.loc[common_idx].values, dtype=float)

        # Stage 1: Standardize signal
        s_mean = float(np.mean(s))
        s_std = float(np.std(s, ddof=1))
        if s_std > 0:
            s_norm: np.ndarray[Any, np.dtype[np.floating[Any]]] = (s - s_mean) / s_std
        else:
            s_norm = s - s_mean

        # Stage 2: Regress signal on factors to get beta
        # s_t = alpha + beta' * f_t + epsilon_t
        ones = np.ones((len(s_norm), 1))
        x_mat = np.hstack([ones, f])
        beta_full = np.linalg.lstsq(x_mat, s_norm, rcond=None)[0]
        self._beta = beta_full[1:]  # factor loadings (exclude intercept)

        # Compute managed adjusted signal: s_t - beta' * f_t
        s_adj = s_norm - f @ self._beta

        # Lambda CV: rolling window to select optimal lambda
        n = len(common_idx)
        lambda_path_values: list[float] = []
        lambda_path_idx: list[object] = []

        for t in range(self.cv_window, n):
            # In-sample window: [t - cv_window, t)
            best_lambda = self.lambda_grid[0]
            best_sharpe = -np.inf

            for lam in self.lambda_grid:
                # Managed return: lam * s_adj over the in-sample window
                managed_r = lam * s_adj[t - self.cv_window : t]
                m_mean = float(np.mean(managed_r))
                m_std = float(np.std(managed_r, ddof=1))
                sharpe = m_mean / m_std if m_std > 0 else 0.0

                if sharpe > best_sharpe:
                    best_sharpe = sharpe
                    best_lambda = lam

            lambda_path_values.append(best_lambda)
            lambda_path_idx.append(common_idx[t])

        self._lambda_opt = float(np.median(lambda_path_values)) if lambda_path_values else 1.0
        self._lambda_path = pd.Series(lambda_path_values, index=lambda_path_idx, dtype=float)

    def predict(
        self,
        signal: pd.Series[float],
        factor_returns: pd.DataFrame,
    ) -> pd.Series[float]:
        """Produce managed portfolio weights from new signal and factors.

        Parameters
        ----------
        signal : pd.Series
            Out-of-sample signal time series.
        factor_returns : pd.DataFrame
            Out-of-sample factor returns.

        Returns
        -------
        pd.Series
            Managed portfolio weights (lambda * adjusted signal).
        """
        if self._beta is None:
            raise RuntimeError("Must call fit() before predict().")

        common_idx = signal.index.intersection(factor_returns.index)
        s = np.asarray(signal.loc[common_idx].values, dtype=float)
        f = np.asarray(factor_returns.loc[common_idx].values, dtype=float)

        # Standardize using the signal's own mean/std (out-of-sample)
        s_std = float(np.std(s, ddof=1))
        s_norm = (s - float(np.mean(s))) / s_std if s_std > 0 else s - float(np.mean(s))

        # Adjusted signal
        s_adj = s_norm - f @ self._beta

        # Apply optimal lambda
        managed = self._lambda_opt * s_adj

        return pd.Series(managed, index=common_idx, dtype=float)

    @property
    def lambda_opt(self) -> float:
        """Optimal lambda selected by cross-validation."""
        return self._lambda_opt

    @property
    def lambda_path(self) -> pd.Series[float]:
        """Time series of selected lambda values from rolling CV."""
        if self._lambda_path is None:
            return pd.Series(dtype=float)
        return self._lambda_path


def cov_matrix(
    returns: pd.DataFrame,
    method: str = "ledoit_wolf",
    frequency: int = 12,
    **kwargs: Any,
) -> pd.DataFrame:
    """Estimate covariance matrix using various methods.

    Parameters
    ----------
    returns : pd.DataFrame
        Asset return matrix (T x N), each column is an asset.
    method : str
        Estimation method. One of:
        - "sample": sample covariance
        - "exp_cov": exponentially weighted covariance
        - "ledoit_wolf": Ledoit-Wolf shrinkage
        - "oas": Oracle Approximating Shrinkage
        - "min_cov_det": Minimum Covariance Determinant (robust)
    frequency : int
        Annualization factor (12 for monthly, 252 for daily).

    Returns
    -------
    pd.DataFrame
        N x N covariance matrix with asset names as index and columns.

    References
    ----------
    Ledoit, O., & Wolf, M. (2004). Honey, I shrunk the sample covariance
        matrix. Journal of Portfolio Management, 30(4), 110-119.
    """
    result: pd.DataFrame
    if method == "sample":
        result = risk_models.sample_cov(returns, frequency=frequency, **kwargs)
    elif method == "exp_cov":
        result = risk_models.exp_cov(returns, frequency=frequency, **kwargs)
    elif method == "ledoit_wolf":
        cs = CovarianceShrinkage(returns, frequency=frequency)
        result = cs.ledoit_wolf(**kwargs)
    elif method == "oas":
        cs = CovarianceShrinkage(returns, frequency=frequency)
        result = cs.oracle_approximating(**kwargs)
    elif method == "min_cov_det":
        result = risk_models.min_cov_determinant(returns, frequency=frequency, **kwargs)
    else:
        raise ValueError(
            f"Unknown method '{method}'. "
            "Use: sample, exp_cov, ledoit_wolf, oas, min_cov_det."
        )
    return result


def mean_variance_weights(
    expected_returns: pd.Series[float],
    cov_matrix_input: pd.DataFrame | None = None,
    returns_data: pd.DataFrame | None = None,
    risk_aversion: float = 1.0,
    weight_bounds: tuple[float, float] = (0.0, 1.0),
    method: str = "max_sharpe",
    target: float | None = None,
    risk_free_rate: float = 0.0,
    market_neutral: bool = False,
) -> pd.Series[float]:
    """Mean-variance portfolio optimization via PyPortfolioOpt.

    Parameters
    ----------
    expected_returns : pd.Series
        Expected return for each asset.
    cov_matrix_input : pd.DataFrame | None
        Pre-computed covariance matrix. If None, estimated from returns_data.
    returns_data : pd.DataFrame | None
        Historical returns used to estimate covariance if cov_matrix_input is None.
    risk_aversion : float
        Risk aversion parameter for max_quadratic_utility.
    weight_bounds : tuple[float, float]
        Min and max weight per asset.
    method : str
        Optimization method:
        - "max_sharpe": maximize Sharpe ratio
        - "min_vol": minimize portfolio volatility
        - "max_quadratic_utility": maximize mu'w - (delta/2)*w'Sigma*w
        - "efficient_risk": minimize risk for target return
        - "efficient_return": maximize return for target risk
        - "cvar": minimize Conditional Value at Risk
        - "cdar": minimize Conditional Drawdown at Risk
    target : float | None
        Target volatility (efficient_risk) or target return (efficient_return),
        or target CVaR/CDaR.
    risk_free_rate : float
        Risk-free rate for Sharpe ratio computation.
    market_neutral : bool
        If True, weights sum to 0 instead of 1.

    Returns
    -------
    pd.Series
        Portfolio weights indexed by asset names. Sums to ~1.0 (or ~0.0).

    References
    ----------
    Markowitz, H. (1952). Portfolio selection. Journal of Finance, 7(1), 77-91.
    """
    # Estimate covariance if not provided
    if cov_matrix_input is None:
        if returns_data is None:
            raise ValueError("Either cov_matrix_input or returns_data must be provided.")
        cov = cov_matrix(returns_data, method="ledoit_wolf")
    else:
        cov = cov_matrix_input

    # CVaR and CDaR methods need returns_data and use different optimizers
    if method == "cvar":
        if returns_data is None:
            raise ValueError("returns_data required for CVaR optimization.")
        ef_cvar = EfficientCVaR(expected_returns, returns_data, weight_bounds=weight_bounds)
        if target is not None:
            # Maximize return subject to a CVaR target
            ef_cvar.efficient_risk(target)
        else:
            ef_cvar.min_cvar()
        weights_dict: dict[str, float] = dict(ef_cvar.clean_weights())
        return pd.Series(weights_dict, dtype=float)

    if method == "cdar":
        if returns_data is None:
            raise ValueError("returns_data required for CDaR optimization.")
        ef_cdar = EfficientCDaR(expected_returns, returns_data, weight_bounds=weight_bounds)
        if target is not None:
            # Maximize return subject to a CDaR target
            ef_cdar.efficient_risk(target)
        else:
            ef_cdar.min_cdar()
        weights_dict = dict(ef_cdar.clean_weights())
        return pd.Series(weights_dict, dtype=float)

    # Standard EfficientFrontier methods
    ef_obj = EfficientFrontier(
        expected_returns,
        cov,
        weight_bounds=weight_bounds,
    )

    if method == "max_sharpe":
        ef_obj.max_sharpe(risk_free_rate=risk_free_rate)
    elif method == "min_vol":
        ef_obj.min_volatility()
    elif method == "max_quadratic_utility":
        ef_obj.max_quadratic_utility(risk_aversion=risk_aversion, market_neutral=market_neutral)
    elif method == "efficient_risk":
        if target is None:
            raise ValueError("target (volatility) required for efficient_risk.")
        ef_obj.efficient_risk(target_volatility=target, market_neutral=market_neutral)
    elif method == "efficient_return":
        if target is None:
            raise ValueError("target (return) required for efficient_return.")
        ef_obj.efficient_return(target_return=target, market_neutral=market_neutral)
    else:
        raise ValueError(
            f"Unknown method '{method}'. Use: max_sharpe, min_vol, "
            "max_quadratic_utility, efficient_risk, efficient_return, cvar, cdar."
        )

    weights_dict = dict(ef_obj.clean_weights())
    return pd.Series(weights_dict, dtype=float)


def black_litterman_weights(
    returns: pd.DataFrame,
    market_caps: pd.Series[float],
    views: dict[str, float],
    view_confidences: list[float],
    weight_bounds: tuple[float, float] = (0.0, 1.0),
    risk_free_rate: float = 0.0,
) -> pd.Series[float]:
    """Black-Litterman portfolio weights via pypfopt.BlackLittermanModel.

    Combines market equilibrium with investor views to produce posterior
    expected returns, then optimizes via mean-variance.

    Parameters
    ----------
    returns : pd.DataFrame
        Historical asset returns (T x N).
    market_caps : pd.Series
        Market capitalization for each asset.
    views : dict[str, float]
        Absolute views on expected returns (e.g., {"AAPL": 0.10}).
    view_confidences : list[float]
        Confidence in each view (0 to 1).
    weight_bounds : tuple[float, float]
        Min and max weight per asset.
    risk_free_rate : float
        Risk-free rate.

    Returns
    -------
    pd.Series
        Portfolio weights indexed by asset names.

    References
    ----------
    Black, F., & Litterman, R. (1992). Global portfolio optimization.
        Financial Analysts Journal, 48(5), 28-43.
    """
    # Estimate covariance matrix using Ledoit-Wolf
    cov = cov_matrix(returns, method="ledoit_wolf")

    # Build Black-Litterman model
    bl = BlackLittermanModel(
        cov,
        absolute_views=views,
        pi="market",
        market_caps=market_caps,
        risk_free_rate=risk_free_rate,
        omega="idzorek",
        view_confidences=view_confidences,
    )

    # Get posterior expected returns
    bl_returns = bl.bl_returns()

    # Optimize using mean-variance on posterior returns
    ef = EfficientFrontier(bl_returns, cov, weight_bounds=weight_bounds)
    ef.max_sharpe(risk_free_rate=risk_free_rate)
    weights_dict: dict[str, float] = dict(ef.clean_weights())
    return pd.Series(weights_dict, dtype=float)


def hrp_weights(
    returns: pd.DataFrame,
    linkage_method: str = "single",
) -> pd.Series[float]:
    """Hierarchical Risk Parity weights via pypfopt.HRPOpt.

    Uses only the covariance structure of returns (no expected returns needed).
    Assets are clustered by correlation similarity, then weights are allocated
    inversely proportional to cluster variance.

    Parameters
    ----------
    returns : pd.DataFrame
        Historical asset returns (T x N).
    linkage_method : str
        Hierarchical clustering linkage method (e.g., "single", "complete", "average").

    Returns
    -------
    pd.Series
        Portfolio weights indexed by asset names.

    References
    ----------
    Lopez de Prado, M. (2016). Building diversified portfolios that outperform
        out of sample. Journal of Portfolio Management, 42(4), 59-69.
    """
    hrp = HRPOpt(returns)
    hrp.optimize(linkage_method=linkage_method)
    weights_dict: dict[str, float] = dict(hrp.clean_weights())
    return pd.Series(weights_dict, dtype=float)
