import numpy as np
import pandas as pd
import cvxpy as cp
from config import logger, Color
from constraints import make_nearest_psd

def exact_risk_parity_allocation(cov_matrix: pd.DataFrame, silent: bool = False) -> pd.Series:
    """
    Computes the Exact True Risk Parity (Equal Risk Contribution - ERC) portfolio.
    
    Instead of using heuristic or graph-theoretic approaches (like HRP), this uses
    the strictly convex Spinu logarithmic barrier formulation to guarantee mathematically 
    exact equal marginal risk contributions.
    
    The optimization problem is:
    min 0.5 * x^T * Sigma * x - c * sum(ln(x))
    subject to x >= 0
    
    The optimal portfolio weights are then w = x / sum(x).
    
    Args:
        cov_matrix: The asset covariance matrix (N x N)
        silent: If True, suppresses console logs
        
    Returns:
        pd.Series containing the normalized optimal Risk Parity weights.
    """
    if not silent:
        print(f" {Color.DIM}ℹ Solving Exact True Risk Parity (ERC) formulation via CVXPY...{Color.RESET}", end="", flush=True)
        
    tickers = list(cov_matrix.columns)
    n = len(tickers)
    
    # Ensure strict positive semi-definiteness for CVXPY convex solver
    sigma_vals = make_nearest_psd(cov_matrix.values)
    
    # Define variables
    x = cp.Variable(n, nonneg=True)
    
    # Spinu Objective: 0.5 * x^T * Sigma * x - sum(ln(x))
    # We use cp.sum(cp.log(x)) as the barrier. Note that cp.log enforces strict positivity internally.
    risk = 0.5 * cp.quad_form(x, cp.psd_wrap(sigma_vals))
    barrier = cp.sum(cp.log(x))
    
    objective = cp.Minimize(risk - barrier)
    
    # We do not constrain sum(x) == 1 here. We solve unconstrained (besides positivity)
    # and then normalize the raw allocations x into portfolio weights w.
    prob = cp.Problem(objective)
    
    try:
        # Solve using ECOS or SCS since it involves logs
        prob.solve(solver=cp.SCS, max_iters=5000, eps=1e-5)
    except Exception as e:
        logger.error(f"SCS Failed for Exact Risk Parity: {e}. Retrying with ECOS...")
        try:
            prob.solve(solver=cp.ECOS)
        except Exception as e2:
            raise RuntimeError(f"Exact Risk Parity optimization failed: {e2}")
            
    if prob.status not in [cp.OPTIMAL, cp.OPTIMAL_INACCURATE] or x.value is None:
        raise RuntimeError(f"Exact Risk Parity did not converge. Status: {prob.status}")
        
    # Normalize to get valid portfolio weights summing to 1
    raw_x = x.value
    weights = raw_x / np.sum(raw_x)
    
    # Clean up near-zero rounding errors
    weights[weights < 1e-6] = 0.0
    weights = weights / np.sum(weights)
    
    if not silent:
        print(f" {Color.GREEN}done.{Color.RESET}")
        
    return pd.Series(weights, index=tickers)