"""
Stock Market Time Series Dataset - Portfolio Optimization

Author: Molla Samser
Organization: RSK World
Designer & Tester: Rima Khatun
Website: https://rskworld.in/
Email: help@rskworld.in
Phone: +91 93305 39277
Address: Nutanhat, Mongolkote, Purba Burdwan, West Bengal, India, 713147
"""

import pandas as pd
import numpy as np
import matplotlib.pyplot as plt
import seaborn as sns
from pathlib import Path
from typing import List, Dict, Tuple, Optional
from scipy.optimize import minimize
import warnings
warnings.filterwarnings('ignore')


class PortfolioOptimizer:
    """
    Modern Portfolio Theory optimization tools.
    """
    
    def __init__(self, stock_paths: List[str]):
        """
        Initialize portfolio optimizer.
        
        Args:
            stock_paths: List of paths to stock CSV files
        """
        self.stock_paths = [Path(p) for p in stock_paths]
        self.symbols = [p.stem for p in self.stock_paths]
        self.df = self._load_all_stocks()
        self.returns = self.df.pct_change().dropna()
        
        sns.set_style('whitegrid')
        plt.rcParams['figure.figsize'] = (12, 6)
    
    def _load_all_stocks(self) -> pd.DataFrame:
        """Load all stocks and merge closing prices."""
        dfs = []
        for path in self.stock_paths:
            df = pd.read_csv(path, parse_dates=['Date'], index_col='Date')
            symbol = path.stem
            dfs.append(df[['Close']].rename(columns={'Close': symbol}))
        
        return pd.concat(dfs, axis=1)
    
    def calculate_portfolio_metrics(self, weights: np.ndarray) -> Tuple[float, float]:
        """
        Calculate portfolio return and volatility.
        
        Args:
            weights: Portfolio weights
            
        Returns:
            Tuple of (annual_return, annual_volatility)
        """
        # Annual return
        portfolio_return = np.sum(self.returns.mean() * weights) * 252
        
        # Annual volatility
        portfolio_variance = np.dot(weights.T, np.dot(self.returns.cov() * 252, weights))
        portfolio_volatility = np.sqrt(portfolio_variance)
        
        return portfolio_return, portfolio_volatility
    
    def calculate_sharpe_ratio(self, weights: np.ndarray, risk_free_rate: float = 0.02) -> float:
        """
        Calculate Sharpe ratio.
        
        Args:
            weights: Portfolio weights
            risk_free_rate: Risk-free rate (annual)
            
        Returns:
            Sharpe ratio
        """
        portfolio_return, portfolio_volatility = self.calculate_portfolio_metrics(weights)
        sharpe_ratio = (portfolio_return - risk_free_rate) / portfolio_volatility
        return sharpe_ratio
    
    def optimize_max_sharpe(self, risk_free_rate: float = 0.02) -> Dict:
        """
        Optimize portfolio for maximum Sharpe ratio.
        
        Args:
            risk_free_rate: Risk-free rate (annual)
            
        Returns:
            Dictionary with optimization results
        """
        n_assets = len(self.symbols)
        
        # Objective function (negative Sharpe ratio for minimization)
        def objective(weights):
            return -self.calculate_sharpe_ratio(weights, risk_free_rate)
        
        # Constraints
        constraints = (
            {'type': 'eq', 'fun': lambda x: np.sum(x) - 1}  # Sum of weights = 1
        )
        
        # Bounds (0 <= weight <= 1)
        bounds = tuple((0, 1) for _ in range(n_assets))
        
        # Initial guess (equal weights)
        initial_weights = np.array([1/n_assets] * n_assets)
        
        # Optimize
        result = minimize(
            objective,
            initial_weights,
            method='SLSQP',
            bounds=bounds,
            constraints=constraints
        )
        
        optimal_weights = result.x
        portfolio_return, portfolio_volatility = self.calculate_portfolio_metrics(optimal_weights)
        sharpe_ratio = self.calculate_sharpe_ratio(optimal_weights, risk_free_rate)
        
        return {
            'weights': dict(zip(self.symbols, optimal_weights)),
            'return': portfolio_return,
            'volatility': portfolio_volatility,
            'sharpe_ratio': sharpe_ratio
        }
    
    def optimize_min_volatility(self) -> Dict:
        """
        Optimize portfolio for minimum volatility.
        
        Returns:
            Dictionary with optimization results
        """
        n_assets = len(self.symbols)
        
        # Objective function
        def objective(weights):
            _, portfolio_volatility = self.calculate_portfolio_metrics(weights)
            return portfolio_volatility
        
        # Constraints
        constraints = (
            {'type': 'eq', 'fun': lambda x: np.sum(x) - 1}
        )
        
        # Bounds
        bounds = tuple((0, 1) for _ in range(n_assets))
        
        # Initial guess
        initial_weights = np.array([1/n_assets] * n_assets)
        
        # Optimize
        result = minimize(
            objective,
            initial_weights,
            method='SLSQP',
            bounds=bounds,
            constraints=constraints
        )
        
        optimal_weights = result.x
        portfolio_return, portfolio_volatility = self.calculate_portfolio_metrics(optimal_weights)
        sharpe_ratio = self.calculate_sharpe_ratio(optimal_weights)
        
        return {
            'weights': dict(zip(self.symbols, optimal_weights)),
            'return': portfolio_return,
            'volatility': portfolio_volatility,
            'sharpe_ratio': sharpe_ratio
        }
    
    def efficient_frontier(self, n_portfolios: int = 100) -> pd.DataFrame:
        """
        Generate efficient frontier.
        
        Args:
            n_portfolios: Number of portfolios to generate
            
        Returns:
            DataFrame with portfolio metrics
        """
        n_assets = len(self.symbols)
        results = []
        
        # Generate random portfolios
        for _ in range(n_portfolios):
            weights = np.random.random(n_assets)
            weights /= np.sum(weights)
            
            portfolio_return, portfolio_volatility = self.calculate_portfolio_metrics(weights)
            sharpe_ratio = self.calculate_sharpe_ratio(weights)
            
            results.append({
                'return': portfolio_return,
                'volatility': portfolio_volatility,
                'sharpe_ratio': sharpe_ratio
            })
        
        return pd.DataFrame(results)
    
    def plot_efficient_frontier(self, save_path: Optional[str] = None):
        """
        Plot efficient frontier with optimal portfolios.
        
        Args:
            save_path: Path to save the plot
        """
        print(f"\n{'='*70}")
        print("Generating Efficient Frontier")
        print(f"{'='*70}\n")
        
        # Generate portfolios
        portfolios = self.efficient_frontier(n_portfolios=5000)
        
        # Find optimal portfolios
        max_sharpe = self.optimize_max_sharpe()
        min_vol = self.optimize_min_volatility()
        
        # Plot
        plt.figure(figsize=(12, 7))
        
        # Scatter plot of all portfolios
        scatter = plt.scatter(
            portfolios['volatility'] * 100,
            portfolios['return'] * 100,
            c=portfolios['sharpe_ratio'],
            cmap='viridis',
            alpha=0.5,
            s=10
        )
        plt.colorbar(scatter, label='Sharpe Ratio')
        
        # Plot optimal portfolios
        plt.scatter(
            max_sharpe['volatility'] * 100,
            max_sharpe['return'] * 100,
            color='red',
            marker='*',
            s=500,
            label=f'Max Sharpe Ratio ({max_sharpe["sharpe_ratio"]:.2f})',
            edgecolors='black',
            linewidths=2
        )
        
        plt.scatter(
            min_vol['volatility'] * 100,
            min_vol['return'] * 100,
            color='blue',
            marker='*',
            s=500,
            label=f'Min Volatility ({min_vol["volatility"]*100:.2f}%)',
            edgecolors='black',
            linewidths=2
        )
        
        plt.title('Efficient Frontier', fontsize=16, fontweight='bold')
        plt.xlabel('Volatility (% Annual)', fontsize=12)
        plt.ylabel('Expected Return (% Annual)', fontsize=12)
        plt.legend(loc='best', fontsize=10)
        plt.grid(True, alpha=0.3)
        plt.tight_layout()
        
        if save_path:
            plt.savefig(save_path, dpi=300, bbox_inches='tight')
            print(f"Efficient frontier plot saved to: {save_path}")
        
        plt.show()
        
        # Print results
        print(f"\n{'='*70}")
        print("Optimal Portfolio - Maximum Sharpe Ratio")
        print(f"{'='*70}")
        for symbol, weight in max_sharpe['weights'].items():
            print(f"{symbol:>10}: {weight:>8.2%}")
        print(f"{'-'*70}")
        print(f"{'Expected Return:':<30} {max_sharpe['return']:>10.2%}")
        print(f"{'Volatility:':<30} {max_sharpe['volatility']:>10.2%}")
        print(f"{'Sharpe Ratio:':<30} {max_sharpe['sharpe_ratio']:>10.4f}")
        print(f"{'='*70}\n")
        
        print(f"{'='*70}")
        print("Optimal Portfolio - Minimum Volatility")
        print(f"{'='*70}")
        for symbol, weight in min_vol['weights'].items():
            print(f"{symbol:>10}: {weight:>8.2%}")
        print(f"{'-'*70}")
        print(f"{'Expected Return:':<30} {min_vol['return']:>10.2%}")
        print(f"{'Volatility:':<30} {min_vol['volatility']:>10.2%}")
        print(f"{'Sharpe Ratio:':<30} {min_vol['sharpe_ratio']:>10.4f}")
        print(f"{'='*70}\n")
    
    def plot_correlation_matrix(self, save_path: Optional[str] = None):
        """
        Plot correlation matrix of returns.
        
        Args:
            save_path: Path to save the plot
        """
        correlation = self.returns.corr()
        
        plt.figure(figsize=(10, 8))
        sns.heatmap(correlation, annot=True, fmt='.3f', cmap='coolwarm',
                   center=0, square=True, linewidths=1, cbar_kws={"shrink": 0.8})
        plt.title('Returns Correlation Matrix', fontsize=16, fontweight='bold')
        plt.tight_layout()
        
        if save_path:
            plt.savefig(save_path, dpi=300, bbox_inches='tight')
            print(f"Correlation matrix saved to: {save_path}")
        
        plt.show()
    
    def monte_carlo_simulation(self, weights: np.ndarray, days: int = 252,
                              simulations: int = 1000, initial_value: float = 100000) -> pd.DataFrame:
        """
        Run Monte Carlo simulation for portfolio.
        
        Args:
            weights: Portfolio weights
            days: Number of days to simulate
            simulations: Number of simulations
            initial_value: Initial portfolio value
            
        Returns:
            DataFrame with simulation results
        """
        # Calculate mean returns and covariance
        mean_returns = self.returns.mean()
        cov_matrix = self.returns.cov()
        
        # Run simulations
        portfolio_simulations = np.zeros((days, simulations))
        
        for i in range(simulations):
            # Generate random returns
            Z = np.random.normal(size=(days, len(self.symbols)))
            L = np.linalg.cholesky(cov_matrix)
            daily_returns = mean_returns.values + np.inner(L, Z).T
            
            # Calculate portfolio returns
            portfolio_returns = np.dot(daily_returns, weights)
            
            # Calculate portfolio value
            portfolio_value = initial_value * (1 + portfolio_returns).cumprod()
            portfolio_simulations[:, i] = portfolio_value
        
        return pd.DataFrame(portfolio_simulations)
    
    def plot_monte_carlo(self, weights: Dict[str, float], days: int = 252,
                        simulations: int = 1000, save_path: Optional[str] = None):
        """
        Plot Monte Carlo simulation results.
        
        Args:
            weights: Portfolio weights dictionary
            days: Number of days to simulate
            simulations: Number of simulations
            save_path: Path to save the plot
        """
        print(f"\n{'='*70}")
        print(f"Running Monte Carlo Simulation ({simulations} simulations, {days} days)")
        print(f"{'='*70}\n")
        
        # Convert weights dict to array
        weights_array = np.array([weights[symbol] for symbol in self.symbols])
        
        # Run simulation
        results = self.monte_carlo_simulation(weights_array, days, simulations)
        
        # Plot
        fig, axes = plt.subplots(2, 1, figsize=(14, 10))
        
        # All simulations
        axes[0].plot(results, alpha=0.1, color='blue', linewidth=0.5)
        axes[0].plot(results.mean(axis=1), color='red', linewidth=2, label='Mean')
        axes[0].set_title('Monte Carlo Portfolio Simulations', fontsize=14, fontweight='bold')
        axes[0].set_ylabel('Portfolio Value ($)', fontsize=12)
        axes[0].legend()
        axes[0].grid(True, alpha=0.3)
        
        # Distribution of final values
        final_values = results.iloc[-1]
        axes[1].hist(final_values, bins=50, alpha=0.7, color='blue', edgecolor='black')
        axes[1].axvline(final_values.mean(), color='red', linestyle='--', 
                       linewidth=2, label=f'Mean: ${final_values.mean():,.0f}')
        axes[1].axvline(final_values.quantile(0.05), color='orange', linestyle='--', 
                       linewidth=2, label=f'5th Percentile: ${final_values.quantile(0.05):,.0f}')
        axes[1].axvline(final_values.quantile(0.95), color='green', linestyle='--', 
                       linewidth=2, label=f'95th Percentile: ${final_values.quantile(0.95):,.0f}')
        axes[1].set_title('Distribution of Final Portfolio Values', fontsize=12, fontweight='bold')
        axes[1].set_xlabel('Final Portfolio Value ($)', fontsize=10)
        axes[1].set_ylabel('Frequency', fontsize=10)
        axes[1].legend()
        axes[1].grid(True, alpha=0.3)
        
        plt.tight_layout()
        
        if save_path:
            plt.savefig(save_path, dpi=300, bbox_inches='tight')
            print(f"Monte Carlo plot saved to: {save_path}")
        
        plt.show()
        
        # Print statistics
        print(f"\n{'='*70}")
        print("Monte Carlo Simulation Results")
        print(f"{'='*70}")
        print(f"Mean Final Value:          ${final_values.mean():,.2f}")
        print(f"Median Final Value:        ${final_values.median():,.2f}")
        print(f"5th Percentile (VaR 95%):  ${final_values.quantile(0.05):,.2f}")
        print(f"95th Percentile:           ${final_values.quantile(0.95):,.2f}")
        print(f"Best Case:                 ${final_values.max():,.2f}")
        print(f"Worst Case:                ${final_values.min():,.2f}")
        print(f"{'='*70}\n")


if __name__ == "__main__":
    print("=" * 70)
    print("Stock Market Time Series Dataset - Portfolio Optimization")
    print("=" * 70)
    print("\nAuthor: Molla Samser | RSK World")
    print("Website: https://rskworld.in/")
    print("=" * 70)
    
    # Initialize optimizer with all stocks
    stock_paths = [
        'data/AAPL.csv',
        'data/GOOGL.csv',
        'data/MSFT.csv',
        'data/AMZN.csv',
        'data/TSLA.csv'
    ]
    
    optimizer = PortfolioOptimizer(stock_paths)
    
    # Plot correlation matrix
    print("\n1. Analyzing correlations...")
    optimizer.plot_correlation_matrix()
    
    # Generate and plot efficient frontier
    print("\n2. Generating efficient frontier...")
    optimizer.plot_efficient_frontier()
    
    # Run Monte Carlo simulation for max Sharpe portfolio
    print("\n3. Running Monte Carlo simulation for optimal portfolio...")
    max_sharpe = optimizer.optimize_max_sharpe()
    optimizer.plot_monte_carlo(max_sharpe['weights'], days=252, simulations=1000)
    
    print("\n" + "=" * 70)
    print("Visit https://rskworld.in/ for more datasets and tools!")
    print("=" * 70)