import numpy as np
import scipy.optimize as opt
import cvxpy as cp


def max_sharpe_ratio(returns):
    """
    Optimize portfolio weights to maximize Sharpe ratio
    Input: returns - numpy array with shape (N, T) where N is number of assets, T is time periods
    Output: optimal weights vector
    """
    # Calculate expected returns and covariance matrix
    mu = np.mean(returns, axis=1)
    Sigma = np.cov(returns)

    # Define optimization problem
    n = len(mu)

    def negative_sharpe(w):
        port_return = np.dot(w, mu)
        port_vol = np.sqrt(np.dot(w, np.dot(Sigma, w)))
        return -port_return / port_vol

    # Constraints
    constraints = {"type": "eq", "fun": lambda w: np.sum(w) - 1}
    bounds = tuple((0, 1) for _ in range(n))  # Long-only constraints

    # Initial guess
    w0 = np.ones(n) / n

    # Optimize
    result = opt.minimize(
        negative_sharpe, w0, method="SLSQP", bounds=bounds, constraints=constraints
    )

    return result.x


def min_cvar(returns, alpha=0.05):
    """
    Minimize Conditional Value-at-Risk (CVaR)
    Input:
        returns - numpy array with shape (N, T)
        alpha - confidence level (e.g., 0.05 for 95% CVaR)
    Output: optimal weights vector
    """
    T = returns.shape[1]
    n = returns.shape[0]
    mu = np.mean(returns, axis=1)

    # CVaR optimization variables
    w = cp.Variable(n)
    zeta = cp.Variable()
    s = cp.Variable(T)

    # Objective: minimize CVaR
    objective = cp.Minimize(zeta + (1 / (alpha * T)) * cp.sum(s))

    # Constraints
    constraints = [
        cp.sum(w) == 1,
        w >= 0,  # Long-only
        s >= 0,
        -returns.T @ w - zeta <= s,
    ]

    # Solve
    prob = cp.Problem(objective, constraints)
    prob.solve()

    return w.value


def max_sharpe_ratio_l1(returns, l1_lambda=0.1):
    """
    Max Sharpe Ratio with L1 regularization (sparse)
    """
    mu = np.mean(returns, axis=1)
    Sigma = np.cov(returns)
    n = len(mu)

    def negative_sharpe_l1(w):
        port_return = np.dot(w, mu)
        port_vol = np.sqrt(np.dot(w, np.dot(Sigma, w)))
        sharpe = port_return / port_vol
        l1_penalty = l1_lambda * np.sum(np.abs(w))
        return -sharpe + l1_penalty

    constraints = {"type": "eq", "fun": lambda w: np.sum(w) - 1}
    bounds = tuple((0, 1) for _ in range(n))  # Long-only
    w0 = np.ones(n) / n

    result = opt.minimize(
        negative_sharpe_l1, w0, method="SLSQP", bounds=bounds, constraints=constraints
    )

    return result.x


def min_cvar_l1(returns, alpha=0.05, l1_lambda=0.1):
    """
    Minimize CVaR with L1 regularization (sparse)
    """
    T = returns.shape[1]
    n = returns.shape[0]

    w = cp.Variable(n)
    zeta = cp.Variable()
    s = cp.Variable(T)

    cvar_term = zeta + (1 / (alpha * T)) * cp.sum(s)
    l1_penalty = l1_lambda * cp.norm1(w)
    objective = cp.Minimize(cvar_term + l1_penalty)

    constraints = [cp.sum(w) == 1, w >= 0, s >= 0, -returns.T @ w - zeta <= s]

    prob = cp.Problem(objective, constraints)
    prob.solve()

    return w.value


def mean_variance(returns, risk_aversion=1.0):
    """
    Mean-variance optimization (maximize return - risk_aversion * variance)
    Input: 
        returns - numpy array with shape (N, T)
        risk_aversion - risk aversion parameter (higher means more risk averse)
    Output: optimal weights vector
    """
    mu = np.mean(returns, axis=1)
    Sigma = np.cov(returns)
    n = len(mu)

    # Define the optimization problem
    w = cp.Variable(n)

    # Objective function
    objective = cp.Maximize(w @ mu - risk_aversion * cp.quad_form(w, Sigma))

    # Constraints
    constraints = [cp.sum(w) == 1, w >= 0]  # Long-only

    # Solve
    prob = cp.Problem(objective, constraints)
    prob.solve()

    return w.value


def max_sharpe_ratio_smooth(returns, gamma=0.1):
    """
    Optimize portfolio weights to maximize Sharpe ratio with L2 regularization
    Input: 
        returns - numpy array with shape (N, T)
        gamma - regularization strength (higher = more diversification)
    Output: optimal weights vector
    """
    mu = np.mean(returns, axis=1)
    Sigma = np.cov(returns)
    n = len(mu)

    def negative_sharpe(w):
        port_return = np.dot(w, mu)
        port_vol = np.sqrt(np.dot(w, np.dot(Sigma, w)))
        # Add L2 regularization to encourage diversification
        regularization = gamma * np.sum(w ** 2)
        return -(port_return / port_vol) + regularization

    constraints = {"type": "eq", "fun": lambda w: np.sum(w) - 1}
    bounds = tuple((0, 1) for _ in range(n))
    w0 = np.ones(n) / n

    result = opt.minimize(
        negative_sharpe, w0, method="SLSQP", bounds=bounds, constraints=constraints
    )
    return result.x / np.sum(result.x)  # Ensure weights sum to 1


def mean_variance_smooth(returns, risk_aversion=1.0, gamma=0.1):
    """
    Mean-variance optimization with entropy regularization
    Input: 
        returns - numpy array with shape (N, T)
        risk_aversion - risk aversion parameter
        gamma - entropy regularization strength
    Output: optimal weights vector
    """
    mu = np.mean(returns, axis=1)
    Sigma = np.cov(returns)
    n = len(mu)

    w = cp.Variable(n)

    # Entropy regularization (maximizing entropy encourages more uniform weights)
    entropy = -cp.sum(cp.entr(w))

    objective = cp.Maximize(
        w @ mu - risk_aversion * cp.quad_form(w, Sigma) + gamma * entropy
    )

    constraints = [
        cp.sum(w) == 1,
        w >= 0.01,  # Minimum weight of 1% to prevent zero weights
        w <= 0.5,  # Maximum weight of 50% to prevent concentration
    ]

    prob = cp.Problem(objective, constraints)
    prob.solve()

    return w.value


def min_cvar_smooth(returns, alpha=0.05, gamma=0.1):
    """
    Minimize CVaR with L2 regularization
    Input:
        returns - numpy array with shape (N, T)
        alpha - confidence level
        gamma - regularization strength
    Output: optimal weights vector
    """
    T = returns.shape[1]
    n = returns.shape[0]
    mu = np.mean(returns, axis=1)

    w = cp.Variable(n)
    zeta = cp.Variable()
    s = cp.Variable(T)

    # L2 regularization term
    reg_term = gamma * cp.sum_squares(w)

    objective = cp.Minimize(zeta + (1 / (alpha * T)) * cp.sum(s) + reg_term)

    constraints = [
        cp.sum(w) == 1,
        w >= 0.05,  # Minimum 5% allocation to each asset
        w <= 0.3,  # Maximum 30% allocation to any single asset
        s >= 0,
        -returns.T @ w - zeta <= s,
    ]

    prob = cp.Problem(objective, constraints)
    prob.solve()

    return w.value
