import numpy as np
import pandas as pd
import torch
import torch.nn as nn
from scipy.optimize import minimize
from sklearn.covariance import LedoitWolf
import warnings
warnings.filterwarnings('ignore')

class CovarianceEstimator:
    """Robust covariance matrix estimation using shrinkage."""
    
    def __init__(self, method='ledoit_wolf'):
        self.method = method
        self.covariance_ = None
        
    def fit(self, returns):
        if self.method == 'ledoit_wolf':
            lw = LedoitWolf().fit(returns)
            self.covariance_ = lw.covariance_
            self.shrinkage_alpha_ = lw.shrinkage_
        elif self.method == 'exponential':
            self.covariance_ = self._exponential_cov(returns, halflife=60)
        elif self.method == 'factor':
            self.covariance_ = self._factor_cov(returns, n_factors=5)
        return self
    
    def _exponential_cov(self, returns, halflife):
        n = len(returns)
        weights = np.exp(-np.log(2) * np.arange(n)[::-1] / halflife)
        weights /= weights.sum()
        
        centered = returns - returns.mean(axis=0)
        return centered.T @ np.diag(weights) @ centered
    
    def _factor_cov(self, returns, n_factors):
        from sklearn.decomposition import PCA
        pca = PCA(n_components=n_factors)
        factors = pca.fit_transform(returns)
        loadings = pca.components_.T
        
        factor_cov = np.cov(factors, rowvar=False)
        residual_var = np.var(returns - returns @ loadings @ np.linalg.inv(loadings.T @ loadings) @ loadings.T, axis=0)
        
        return loadings @ factor_cov @ loadings.T + np.diag(residual_var)

class MeanVarianceOptimizer:
    """Classical mean-variance optimization with constraints."""
    
    def __init__(self, risk_aversion=1.0):
        self.risk_aversion = risk_aversion
        self.weights = None
        
    def optimize(self, expected_returns, covariance, constraints=None, short_selling=False):
        n = len(expected_returns)
        
        if constraints is None:
            constraints = {'max_weight': 0.3, 'min_weight': 0.0}
        
        bounds = tuple(
            (constraints.get('min_weight', 0.0), constraints.get('max_weight', 0.3))
            for _ in range(n)
        )
        
        if short_selling:
            bounds = tuple((-0.3, 0.3) for _ in range(n))
        
        objective = lambda w: (
            self.risk_aversion * w @ covariance @ w - expected_returns @ w
        )
        
        constraints_list = [{'type': 'eq', 'fun': lambda w: np.sum(w) - 1}]
        
        w0 = np.ones(n) / n
        result = minimize(objective, w0, method='SLSQP', bounds=bounds, constraints=constraints_list)
        
        self.weights = result.x
        return self.weights
    
    def portfolio_stats(self, weights, expected_returns, covariance, risk_free_rate=0.02):
        ret = expected_returns @ weights
        vol = np.sqrt(weights @ covariance @ weights)
        sharpe = (ret - risk_free_rate) / vol
        
        return {
            'expected_return': ret,
            'volatility': vol,
            'sharpe_ratio': sharpe
        }

class BlackLittermanModel:
    """Black-Litterman model combining market equilibrium with investor views."""
    
    def __init__(self, risk_aversion=2.5, tau=0.025):
        self.risk_aversion = risk_aversion
        self.tau = tau
        
    def market_implied_returns(self, market_caps, covariance):
        """Calculate market equilibrium returns."""
        w_mkt = market_caps / market_caps.sum()
        pi = self.risk_aversion * covariance @ w_mkt
        return pi
    
    def incorporate_views(self, equilibrium_returns, covariance, P, Q, omega=None):
        """Incorporate investor views into equilibrium returns."""
        if omega is None:
            omega = np.diag(np.diag(self.tau * P @ covariance @ P.T))
        
        M = np.linalg.inv(self.tau * covariance) + P.T @ np.linalg.inv(omega) @ P
        b = np.linalg.inv(self.tau * covariance) @ equilibrium_returns + P.T @ np.linalg.inv(omega) @ Q
        
        posterior_returns = np.linalg.solve(M, b)
        posterior_cov = np.linalg.inv(M)
        
        return posterior_returns, posterior_cov
    
    def optimize(self, posterior_returns, posterior_cov, constraints=None):
        optimizer = MeanVarianceOptimizer(risk_aversion=self.risk_aversion)
        return optimizer.optimize(posterior_returns, posterior_cov, constraints)

class RiskParityOptimizer:
    """Risk parity portfolio construction."""
    
    def __init__(self):
        self.weights = None
        
    def optimize(self, covariance):
        n = covariance.shape[0]
        
        def risk_contribution_objective(w):
            port_vol = np.sqrt(w @ covariance @ w)
            marginal_contrib = covariance @ w
            risk_contrib = w * marginal_contrib / port_vol
            target_contrib = port_vol / n
            
            return np.sum((risk_contrib - target_contrib)**2)
        
        constraints = [{'type': 'eq', 'fun': lambda w: np.sum(w) - 1}]
        bounds = [(0.01, 0.5)] * n
        w0 = np.ones(n) / n
        
        result = minimize(risk_contribution_objective, w0, method='SLSQP',
                         bounds=bounds, constraints=constraints)
        
        self.weights = result.x
        return self.weights
    
    def risk_decomposition(self, weights, covariance):
        port_vol = np.sqrt(weights @ covariance @ weights)
        marginal_contrib = covariance @ weights
        risk_contrib = weights * marginal_contrib / port_vol
        
        return {
            'total_risk': port_vol,
            'risk_contribution': risk_contrib,
            'pct_contribution': risk_contrib / port_vol
        }

class DeepPortfolioOptimizer(nn.Module):
    """Neural network for learning portfolio optimization."""
    
    def __init__(self, input_dim=100, hidden_dim=64):
        super().__init__()
        self.network = nn.Sequential(
            nn.Linear(input_dim, hidden_dim),
            nn.ReLU(),
            nn.BatchNorm1d(hidden_dim),
            nn.Dropout(0.2),
            nn.Linear(hidden_dim, hidden_dim),
            nn.ReLU(),
            nn.Linear(hidden_dim, 1)
        )
        self.weight_layer = nn.Sequential(
            nn.Linear(1, hidden_dim),
            nn.ReLU(),
            nn.Linear(hidden_dim, input_dim),
            nn.Softmax(dim=-1)
        )
        
    def forward(self, features):
        risk_pref = self.network(features)
        weights = self.weight_layer(risk_pref)
        return weights, risk_pref

class EfficientFrontierGenerator:
    """Generate the efficient frontier by solving for multiple risk levels."""
    
    def __init__(self, n_points=50):
        self.n_points = n_points
        
    def generate(self, expected_returns, covariance, constraints=None):
        optimizer = MeanVarianceOptimizer()
        
        min_ret = np.min(expected_returns)
        max_ret = np.max(expected_returns)
        target_returns = np.linspace(min_ret, max_ret, self.n_points)
        
        frontier = []
        for target in target_returns:
            try:
                cons = [
                    {'type': 'eq', 'fun': lambda w: np.sum(w) - 1},
                    {'type': 'eq', 'fun': lambda w, t=target: expected_returns @ w - t}
                ]
                
                n = len(expected_returns)
                bounds = tuple((0, 0.3) for _ in range(n))
                w0 = np.ones(n) / n
                
                result = minimize(
                    lambda w: w @ covariance @ w,
                    w0, method='SLSQP', bounds=bounds, constraints=cons
                )
                
                if result.success:
                    w = result.x
                    ret = expected_returns @ w
                    vol = np.sqrt(w @ covariance @ w)
                    sharpe = (ret - 0.02) / vol
                    
                    frontier.append({
                        'return': ret,
                        'volatility': vol,
                        'sharpe': sharpe,
                        'weights': w
                    })
            except:
                continue
        
        return pd.DataFrame(frontier)

def generate_synthetic_data(n_assets=20, n_days=500):
    """Generate synthetic asset returns with realistic properties."""
    np.random.seed(42)
    
    mu = np.random.uniform(0.05, 0.15, n_assets) / 252
    
    A = np.random.randn(n_assets, n_assets * 2) * 0.1
    base_cov = A @ A.T / (n_assets * 2)
    sigma = np.sqrt(np.diag(base_cov))
    corr = base_cov / np.outer(sigma, sigma)
    np.fill_diagonal(corr, 1)
    
    cov = np.outer(sigma, sigma) * corr
    
    returns = np.random.multivariate_normal(mu, cov, n_days)
    
    asset_names = [f'Asset_{i}' for i in range(n_assets)]
    return pd.DataFrame(returns, columns=asset_names)

def backtest_portfolio(returns_df, weights, rebalance_freq=21):
    """Backtest a portfolio strategy with periodic rebalancing."""
    n_days, n_assets = returns_df.shape
    
    portfolio_returns = []
    current_weights = weights.copy()
    
    for t in range(n_days):
        daily_return = returns_df.iloc[t].values @ current_weights
        portfolio_returns.append(daily_return)
        
        if (t + 1) % rebalance_freq == 0:
            current_weights = weights.copy()
        else:
            asset_returns = returns_df.iloc[t].values
            current_weights = current_weights * (1 + asset_returns)
            current_weights /= current_weights.sum()
    
    portfolio_returns = np.array(portfolio_returns)
    
    cumulative = (1 + portfolio_returns).cumprod()
    total_return = cumulative[-1] - 1
    annual_return = (1 + total_return) ** (252 / n_days) - 1
    annual_vol = np.std(portfolio_returns) * np.sqrt(252)
    sharpe = annual_return / annual_vol
    
    running_max = np.maximum.accumulate(cumulative)
    drawdown = (cumulative - running_max) / running_max
    max_drawdown = np.min(drawdown)
    
    return {
        'total_return': total_return,
        'annual_return': annual_return,
        'annual_volatility': annual_vol,
        'sharpe_ratio': sharpe,
        'max_drawdown': max_drawdown,
        'cumulative_returns': cumulative
    }

# Example usage
if __name__ == "__main__":
    returns_df = generate_synthetic_data(n_assets=15, n_days=500)
    
    estimator = CovarianceEstimator(method='ledoit_wolf')
    cov_matrix = estimator.fit(returns_df.values).covariance_
    expected_returns = returns_df.mean().values * 252
    
    print("Covariance Shrinkage Intensity:", estimator.shrinkage_alpha_)
    
    mv_optimizer = MeanVarianceOptimizer(risk_aversion=2.0)
    mv_weights = mv_optimizer.optimize(expected_returns, cov_matrix)
    mv_stats = mv_optimizer.portfolio_stats(mv_weights, expected_returns, cov_matrix)
    print(f"\nMean-Variance Portfolio:")
    print(f"  Expected Return: {mv_stats['expected_return']:.4f}")
    print(f"  Volatility: {mv_stats['volatility']:.4f}")
    print(f"  Sharpe Ratio: {mv_stats['sharpe_ratio']:.4f}")
    
    market_caps = np.random.uniform(1e9, 100e9, len(returns_df.columns))
    bl_model = BlackLittermanModel()
    pi = bl_model.market_implied_returns(market_caps, cov_matrix)
    
    P = np.zeros((2, len(returns_df.columns)))
    P[0, 0] = 1
    P[1, 1] = -1
    Q = np.array([0.05, -0.02])
    
    posterior_ret, posterior_cov = bl_model.incorporate_views(pi, cov_matrix, P, Q)
    bl_weights = bl_model.optimize(posterior_ret, posterior_cov)
    print(f"\nBlack-Litterman weights: {np.round(bl_weights, 3)}")
    
    rp_optimizer = RiskParityOptimizer()
    rp_weights = rp_optimizer.optimize(cov_matrix)
    rp_decomp = rp_optimizer.risk_decomposition(rp_weights, cov_matrix)
    print(f"\nRisk Parity Portfolio:")
    print(f"  Total Risk: {rp_decomp['total_risk']:.4f}")
    print(f"  Risk Contributions: {np.round(rp_decomp['pct_contribution'], 3)}")
    
    frontier_gen = EfficientFrontierGenerator(n_points=20)
    frontier = frontier_gen.generate(expected_returns, cov_matrix)
    print(f"\nEfficient Frontier:")
    print(f"  Points: {len(frontier)}")
    print(f"  Max Sharpe: {frontier['sharpe'].max():.4f}")
    
    backtest_results = backtest_portfolio(returns_df, mv_weights)
    print(f"\nBacktest Results:")
    print(f"  Annual Return: {backtest_results['annual_return']:.4f}")
    print(f"  Annual Volatility: {backtest_results['annual_volatility']:.4f}")
    print(f"  Sharpe Ratio: {backtest_results['sharpe_ratio']:.4f}")
    print(f"  Max Drawdown: {backtest_results['max_drawdown']:.4f}")