"""Parametric portfolio policy (Brandt-Santa-Clara-Valkanov 2009).

Implements the BSV parametric portfolio approach where weights are linear
functions of demeaned, standardized firm characteristics. Theta parameters
are estimated by maximizing expected CRRA utility.

This module covers optimization-based portfolio construction. For sorting-based
long-short portfolios, see eapctf.sorting.portfolio.
"""

from __future__ import annotations

import numpy as np
import pandas as pd
from scipy.optimize import minimize


class ParametricPolicy:
    """Brandt-Santa-Clara-Valkanov (2009) parametric portfolio policy.

    The weight of asset i at time t is:
        w_it = w_bench_it + (1/N_t) * sum_k theta_k * x_kit

    where x_kit are demeaned characteristics and theta is the parameter vector
    estimated by maximizing expected CRRA utility.

    Parameters
    ----------
    char_cols : list[str]
        Characteristic column names used in the weight function.
    gamma : float
        Risk aversion parameter for CRRA utility.
    benchmark_weights : str
        Benchmark weight scheme: "equal" or "vw".

    References
    ----------
    Brandt, M. W., Santa-Clara, P., & Valkanov, R. (2009). Parametric
        portfolio policies: Exploiting characteristics in the cross-section
        of equity returns. Review of Financial Studies, 22(9), 3411-3447.
        Eq. 1 (weight function), Eq. 5 (utility maximization objective).
    """

    def __init__(
        self,
        char_cols: list[str],
        gamma: float = 5.0,
        benchmark_weights: str = "equal",
    ) -> None:
        self.char_cols = char_cols
        self.gamma = gamma
        self.benchmark_weights = benchmark_weights
        # theta vector: one parameter per characteristic
        self.theta_: np.ndarray | None = None

    def fit(
        self,
        data: pd.DataFrame,
        ret_col: str = "ret",
        date_col: str = "date",
        entity_col: str = "permno",
    ) -> None:
        """Estimate theta by maximizing in-sample CRRA utility.

        The objective is: max_theta sum_t U(r_pt)
        where r_pt = sum_i w_it * r_{i,t+1}
        and U(r) = (1+r)^{1-gamma} / (1-gamma) for CRRA utility.

        Uses scipy.optimize.minimize with L-BFGS-B.

        Parameters
        ----------
        data : pd.DataFrame
            Panel data with returns, dates, entity ids, and characteristics.
        ret_col : str
            Return column name.
        date_col : str
            Date column name.
        entity_col : str
            Entity identifier column name.
        """
        if self.benchmark_weights == "vw":
            raise NotImplementedError(
                "Value-weighted benchmark is not yet implemented. "
                "Use benchmark_weights='equal'."
            )

        df = data.copy()
        k = len(self.char_cols)
        dates = sorted(df[date_col].unique())

        # Pre-process: standardize characteristics cross-sectionally each period
        # BSV (2009) Eq. 1 requires zero mean AND unit std for x~_kit
        for date_val in dates:
            mask = df[date_col] == date_val
            for col in self.char_cols:
                vals = df.loc[mask, col]
                std_val = float(vals.std())
                df.loc[mask, col] = (vals - vals.mean()) / (std_val if std_val > 1e-8 else 1.0)

        gamma = self.gamma

        def _neg_utility(theta: np.ndarray) -> float:
            """Negative average CRRA utility across all periods."""
            total_utility = 0.0
            n_dates = 0

            for date_val in dates:
                date_slice = df.loc[df[date_col] == date_val]
                n_t = len(date_slice)
                if n_t == 0:
                    continue

                # Benchmark weights: 1/N for equal weighting
                w_bench = np.ones(n_t) / n_t

                # Characteristic tilt: (1/N_t) * sum_k theta_k * x_kit
                char_matrix = date_slice[self.char_cols].values  # (N_t, K)
                tilt = (1.0 / n_t) * (char_matrix @ theta)

                # Total weights
                w = w_bench + tilt

                # Portfolio return for this period
                r_port = float(np.dot(w, date_slice[ret_col].values))

                # CRRA utility: (1+r)^{1-gamma} / (1-gamma)
                wealth = 1.0 + r_port
                if wealth <= 0:
                    # Penalize negative wealth heavily
                    total_utility -= 1e6
                else:
                    if abs(gamma - 1.0) < 1e-10:
                        total_utility += np.log(wealth)
                    else:
                        total_utility += (wealth ** (1.0 - gamma)) / (1.0 - gamma)

                n_dates += 1

            # Return negative because we minimize
            if n_dates == 0:
                return 0.0
            return -total_utility / n_dates

        # Initial theta = zeros
        theta0 = np.zeros(k)

        # Optimize using L-BFGS-B
        result = minimize(
            _neg_utility,
            theta0,
            method="L-BFGS-B",
            options={"maxiter": 1000, "ftol": 1e-10},
        )

        self.theta_ = result.x

    def weights(
        self,
        data: pd.DataFrame,
        date_col: str = "date",
        entity_col: str = "permno",
    ) -> pd.DataFrame:
        """Apply fitted theta to compute portfolio weights.

        Parameters
        ----------
        data : pd.DataFrame
            Panel data with characteristics. Must contain char_cols.
        date_col : str
            Date column name.
        entity_col : str
            Entity identifier column name.

        Returns
        -------
        pd.DataFrame
            Weight matrix indexed by (date, entity).
        """
        if self.theta_ is None:
            raise RuntimeError("Must call fit() before weights().")

        df = data.copy()
        dates = sorted(df[date_col].unique())
        theta = self.theta_

        # Standardize characteristics cross-sectionally (BSV 2009: zero mean, unit std)
        for date_val in dates:
            mask = df[date_col] == date_val
            for col in self.char_cols:
                vals = df.loc[mask, col]
                std_val = float(vals.std())
                df.loc[mask, col] = (vals - vals.mean()) / (std_val if std_val > 1e-8 else 1.0)

        # Compute weights for each date
        weight_list: list[pd.DataFrame] = []
        for date_val in dates:
            date_slice = df.loc[df[date_col] == date_val]
            n_t = len(date_slice)
            if n_t == 0:
                continue

            w_bench = np.ones(n_t) / n_t
            char_matrix = date_slice[self.char_cols].values
            tilt = (1.0 / n_t) * (char_matrix @ theta)
            w = w_bench + tilt

            w_df = pd.DataFrame(
                {date_col: date_val, entity_col: date_slice[entity_col].values, "weight": w}
            )
            weight_list.append(w_df)

        result = pd.concat(weight_list, ignore_index=True)
        return result
