ValueError: shapes (1,1) and (5,) not aligned: 1 (dim 1) != 5 (dim 0)
Traceback:
File "/usr/local/lib/python3.10/dist-packages/streamlit/runtime/scriptrunner/exec_code.py", line 88, in exec_func_with_error_handling
    result = func()
File "/usr/local/lib/python3.10/dist-packages/streamlit/runtime/scriptrunner/script_runner.py", line 579, in code_to_exec
    exec(code, module.__dict__)
File "/content/portfolio_optimizer.py", line 3558, in <module>
    display_backtest_results(backtest_results, comparison_type, portfolio1, portfolio2)
File "/content/portfolio_optimizer.py", line 496, in display_backtest_results
    display_portfolio_vs_portfolio(results, portfolio1, portfolio2)
File "/content/portfolio_optimizer.py", line 885, in display_portfolio_vs_portfolio
    risk_contrib_a = calculate_risk_contribution(returns_a, weights_a)
File "/content/portfolio_optimizer.py", line 704, in calculate_risk_contribution
    portfolio_variance = np.dot(weights, np.dot(cov_matrix, weights))

%%writefile portfolio_optimizer.py
import streamlit as st
import pandas as pd
import plotly.graph_objects as go
import numpy as np
import yfinance as yf
import statsmodels.api as sm
import datetime
from datetime import datetime, timedelta
from pandas_datareader import data as pdr
import plotly.express as px
import warnings, requests
from requests.adapters import HTTPAdapter
from typing import Union
from urllib3.util.retry import Retry
from scipy.stats import skew, kurtosis, norm
from joblib import Parallel, delayed
from statsmodels.stats.stattools import durbin_watson
from functools import wraps
from scipy.optimize import minimize
from requests.packages.urllib3.exceptions import InsecureRequestWarning
from io import BytesIO

warnings.filterwarnings('ignore')

# ----------------------------
# Frequency Mapping
# ----------------------------
frequency_mapping = {
    "Daily": "D",
    "Weekly": "W",
    "Monthly": "M",
    "Quarterly": "Q",
    "Yearly": "Y"
}

# Inverse mapping to get frequency name from code
rebalance_freq_inverse_mapping = {v: k for k, v in frequency_mapping.items()}

# ----------------------------
# Example Portfolios Configuration
# ----------------------------

def create_example_portfolios() -> list:
    """
    Creates a list of predefined example portfolios.
    
    Returns:
        list: A list of dictionaries, each representing a portfolio.
    """
    today = pd.to_datetime(datetime.today())
    example_definitions = [
        {
            'name': "Retirement Portfolio",
            'years': 10,
            'benchmark_symbol': '^GSPC',
            'rebalance_freq': 'M',
            'selected': ['AAPL', 'MSFT', 'GOOGL', 'JPM', 'XOM'],
            'allocations': [20.0, 20.0, 20.0, 20.0, 20.0]
        },
        {
            'name': "Tech-Heavy Portfolio",
            'years': 5,
            'benchmark_symbol': '^IXIC',
            'rebalance_freq': 'M',
            'selected': ['AAPL', 'MSFT', 'GOOGL', 'AMZN', 'META'],
            'allocations': [25.0, 25.0, 20.0, 15.0, 15.0]
        },
        {
            'name': "Balanced Risk-Return Portfolio",
            'years': 7,
            'benchmark_symbol': '^GSPC',
            'rebalance_freq': 'M',
            'selected': ['AAPL', 'JNJ', 'V', 'PG', 'XOM'],
            'allocations': [20.0, 20.0, 20.0, 20.0, 20.0]
        },
        {
            'name': "Income-Focused Portfolio",
            'years': 7,
            'benchmark_symbol': '^DJI',
            'rebalance_freq': 'Q',
            'selected': ['PG', 'KO', 'JNJ', 'T', 'PFE'],
            'allocations': [20.0, 20.0, 20.0, 20.0, 20.0]
        },
        {
            'name': "Growth-Oriented Portfolio",
            'years': 3,
            'benchmark_symbol': '^IXIC',
            'rebalance_freq': 'M',
            'selected': ['TSLA', 'NVDA', 'AMD', 'META', 'NFLX'],
            'allocations': [30.0, 20.0, 20.0, 15.0, 15.0]
        },
        {
            'name': "Value Portfolio",
            'years': 10,
            'benchmark_symbol': '^DJI',
            'rebalance_freq': 'Q',
            'selected': ['KO', 'PFE', 'XOM', 'WMT', 'CVX'],
            'allocations': [20.0, 20.0, 20.0, 20.0, 20.0]
        },
        {
            'name': "Dividend-Focused Portfolio",
            'years': 7,
            'benchmark_symbol': 'DVY',
            'rebalance_freq': 'Q',
            'selected': ['T', 'VZ', 'PFE', 'KO', 'IBM'],
            'allocations': [20.0, 20.0, 20.0, 20.0, 20.0]
        },
        {
            'name': "Global Diversification Portfolio",
            'years': 5,
            'benchmark_symbol': 'ACWI',
            'rebalance_freq': 'M',
            'selected': ['AAPL', 'TSM', 'BABA', 'SAP', 'UL'],
            'allocations': [20.0, 20.0, 20.0, 20.0, 20.0]
        },
        {
            'name': "High-Risk Growth Portfolio",
            'years': 3,
            'benchmark_symbol': '^IXIC',
            'rebalance_freq': 'M',
            'selected': ['TSLA', 'ARKK', 'NVDA', 'SHOP', 'CRWD'],
            'allocations': [30.0, 20.0, 20.0, 15.0, 15.0]
        },
        {
            'name': "Emerging Markets Portfolio",
            'years': 5,
            'benchmark_symbol': 'EEM',
            'rebalance_freq': 'M',
            'selected': ['BABA', 'TSM', 'PDD', 'INFY', 'VALE'],
            'allocations': [25.0, 25.0, 20.0, 15.0, 15.0]
        }
    ]
    
    portfolios = []
    for port_def in example_definitions:
        portfolio = {
            'name': port_def['name'],
            'start_date': today - timedelta(days=365 * port_def['years']),
            'end_date': today,
            'rf_rate': 0.02,
            'broker_fee': 0.0,
            'benchmark_symbol': port_def['benchmark_symbol'],
            'rebalance_freq': port_def['rebalance_freq'],
            'selected': port_def['selected'],
            'allocations': port_def['allocations']
        }
        portfolios.append(portfolio)
    
    return portfolios

# Initialize example portfolios if not already in session state
if 'portfolios' not in st.session_state or not st.session_state.portfolios:
    st.session_state.portfolios = create_example_portfolios()
# ----------------------------
# Global Settings and Configurations
# ----------------------------

# Suppress only the InsecureRequestWarning, if bypassing SSL
warnings.simplefilter('ignore', InsecureRequestWarning)
requests.packages.urllib3.disable_warnings()  # Use cautiously
# ----------------------------
# Global Theme Configuration
# ----------------------------
THEME = 'plotly_dark'  # Options: 'plotly_dark', 'plotly_white', 'seaborn', etc.
# ----------------------------
# Global Plot Layout Configuration
# ----------------------------
FIG_HEIGHT = 600
HOVERMODE = 'x unified'

# ----------------------------
# Session State Initialization
# ----------------------------

def initialize_session_state():
    """
    Initializes Streamlit session state with default values.
    """
    # Initialize example portfolios
    st.session_state.setdefault('portfolios', create_example_portfolios())
    
    # Other session state variables
    st.session_state.setdefault('backtest_results', {})
    st.session_state.setdefault('step', "Configure Portfolio")
    st.session_state.setdefault('edit_portfolio', None)
    st.session_state.setdefault('add_new_portfolio', False)
    st.session_state.setdefault('default_config', {
        'rf_rate': 0.02,
        'broker_fee': 0.0,
        'benchmark_symbol': '^GSPC'
    })
    st.session_state.setdefault('show_proceed', False)
    st.session_state.setdefault('selected_time_frames', ['Annually'])
    st.session_state.setdefault('selected_rolling_periods', [252])

# Initialize session state
initialize_session_state()
# ----------------------------
# Utility Decorators
# ----------------------------

def handle_exceptions(func):
    """
    Decorator to handle exceptions in Streamlit applications.
    Displays errors using Streamlit's error message system.
    """
    @wraps(func)
    def wrapper(*args, **kwargs):
        try:
            return func(*args, **kwargs)
        except Exception as e:
            st.error(f"Error in {func.__name__}: {e}")
            return None
    return wrapper

# ----------------------------
# Helper Functions
# ----------------------------

@handle_exceptions
def get_company_name(ticker_df: pd.DataFrame, ticker: str) -> str:
    """
    Retrieve the company name for a given ticker.

    Parameters:
    - ticker_df (pd.DataFrame): DataFrame containing 'Ticker' and 'Company Name' columns.
    - ticker (str): Stock ticker symbol.

    Returns:
    - str: Company name or "Unknown" if not found.
    """
    match = ticker_df[ticker_df['Ticker'] == ticker]
    if not match.empty:
        return match.iloc[0]['Company Name']
    return "Unknown"

def format_asset_option(ticker: str, company_name: str) -> str:
    """
    Format the asset option display string.

    Parameters:
    - ticker (str): Stock ticker symbol.
    - company_name (str): Name of the company.

    Returns:
    - str: Formatted string combining ticker and company name.
    """
    return f"{ticker} - {company_name}"

def handle_portfolio_data(portfolio, date_range=None):
    """
    Helper function to download data and handle missing tickers.

    Parameters:
    - portfolio (dict): Portfolio details.
    - date_range (tuple, optional): (start_date, end_date)

    Returns:
    - available_selected (list): Tickers with available data.
    - missing_selected (set): Tickers missing data.
    - price_data (pd.DataFrame): Adjusted close prices.
    """
    start_date = date_range[0] if date_range else portfolio['start_date']
    end_date = date_range[1] if date_range else portfolio['end_date']
    price_data = download_data(portfolio['selected'], start_date, end_date)
    if price_data.empty:
        return [], set(portfolio['selected']), price_data
    available_selected = [ticker for ticker in portfolio['selected'] if ticker in price_data.columns]
    missing_selected = set(portfolio['selected']) - set(available_selected)
    if missing_selected:
        st.warning(f"Excluded tickers with no data: {', '.join(missing_selected)}")
    return available_selected, missing_selected, price_data

@handle_exceptions
def run_portfolio_comparison(comparison_type, portfolio1, portfolio2=None):
        """
        Generalized function to run portfolio comparisons.

        Parameters:
        - comparison_type (str): "vs_benchmark" or "vs_portfolio"
        - portfolio1 (dict): Primary portfolio.
        - portfolio2 (dict, optional): Secondary portfolio for "vs_portfolio".

        Returns:
        - result (dict): Backtest results or error.
        """
        try:
            if comparison_type == "vs_benchmark":
                available, missing, price_data = handle_portfolio_data(portfolio1)
                if not available:
                    return {"error": "No valid tickers with available data."}
                allocations = adjust_allocations([portfolio1['allocations'][portfolio1['selected'].index(t)] for t in available])
                returns, cum_returns = backtest(
                    weights=np.array(allocations) / 100,
                    prices=price_data[available],
                    rebalance_freq=portfolio1['rebalance_freq'],
                    broker_fee=portfolio1['broker_fee']
                )
                if cum_returns.empty:
                    return {"error": "Cumulative returns are empty."}
                # Benchmark processing...
                benchmark_data = download_data(
                    [portfolio1['benchmark_symbol']],
                    portfolio1['start_date'].strftime('%Y-%m-%d'),
                    portfolio1['end_date'].strftime('%Y-%m-%d')
                )
                if benchmark_data.empty:
                    return {"error": "Benchmark data could not be downloaded."}
                
                benchmark_returns, benchmark_cum_returns, benchmark_metrics = process_benchmark(benchmark_data, returns, portfolio1)
                portfolio_metrics = calculate_metrics(returns, cum_returns, portfolio1['rf_rate'], benchmark_returns)
                
                # Compile backtest results
                backtest_results = {
                    'returns': returns,
                    'cum_returns': cum_returns,
                    'weights': allocations,
                    'price_data': price_data[available],
                    'metrics': portfolio_metrics,
                    'benchmark_returns': benchmark_returns,
                    'benchmark_cum_returns': benchmark_cum_returns,
                    'benchmark_metrics': benchmark_metrics
                }
                
                # **Added Return Statement**
                return {"backtest_results": backtest_results}
    
            elif comparison_type == "vs_portfolio":
                        overlap_start = max(portfolio1['start_date'], portfolio2['start_date'])
                        overlap_end = min(portfolio1['end_date'], portfolio2['end_date'])
                        if overlap_start >= overlap_end:
                            return {"error": "No overlapping date ranges."}
                        available_a, missing_a, price_data_a = handle_portfolio_data(portfolio1, (overlap_start, overlap_end))
                        available_b, missing_b, price_data_b = handle_portfolio_data(portfolio2, (overlap_start, overlap_end))
                        if not available_a or not available_b:
                            return {"error": "One of the portfolios has no valid tickers with available data."}
                        allocations_a = adjust_allocations([portfolio1['allocations'][portfolio1['selected'].index(t)] for t in available_a])
                        allocations_b = adjust_allocations([portfolio2['allocations'][portfolio2['selected'].index(t)] for t in available_b])
                        returns_a, cum_returns_a = backtest(
                            weights=np.array(allocations_a) / 100,
                            prices=price_data_a[available_a],
                            rebalance_freq=portfolio1['rebalance_freq'],
                            broker_fee=portfolio1['broker_fee']
                        )
                        returns_b, cum_returns_b = backtest(
                            weights=np.array(allocations_b) / 100,
                            prices=price_data_b[available_b],
                            rebalance_freq=portfolio2['rebalance_freq'],
                            broker_fee=portfolio2['broker_fee']
                        )
                        if cum_returns_a.empty or cum_returns_b.empty:
                            return {"error": "One of the cumulative returns is empty."}
                        common_index = cum_returns_a.index.intersection(cum_returns_b.index)
                        if common_index.empty:
                            return {"error": "No overlapping dates after backtesting."}
                        returns_a = returns_a.loc[common_index]
                        returns_b = returns_b.loc[common_index]
                        cum_returns_a = cum_returns_a.loc[common_index]
                        cum_returns_b = cum_returns_b.loc[common_index]
                        metrics_a = calculate_metrics(returns_a, cum_returns_a, portfolio1['rf_rate'], None)
                        metrics_b = calculate_metrics(returns_b, cum_returns_b, portfolio2['rf_rate'], None)
                        backtest_results = {
                            'returns_a': returns_a,
                            'cum_returns_a': cum_returns_a,
                            'weights_a': allocations_a,
                            'price_data_a': price_data_a[available_a],
                            'metrics_a': metrics_a,
                            'returns_b': returns_b,
                            'cum_returns_b': cum_returns_b,
                            'weights_b': allocations_b,
                            'price_data_b': price_data_b[available_b],
                            'metrics_b': metrics_b
                        }
                        return {"backtest_results": backtest_results}
        except Exception as e:
            return {"error": f"An unexpected error occurred: {e}"}

def adjust_allocations(allocations):
    """Adjust allocations to sum to 100% if necessary."""
    total = sum(allocations)
    if not np.isclose(total, 100.0, atol=1e-4):
        st.warning("Allocations do not sum to 100%. Adjusting allocations proportionally.")
        allocations = [a / total * 100 for a in allocations]
    return allocations

def process_portfolio_allocations(portfolio, price_data):
    """Processes portfolio allocations and handles missing tickers."""
    available = [ticker for ticker in portfolio['selected'] if ticker in price_data.columns]
    missing = set(portfolio['selected']) - set(available)
    if missing:
        st.warning(f"Excluded tickers from '{portfolio['name']}': {', '.join(missing)}")
    if not available:
        st.error(f"No valid tickers for portfolio '{portfolio['name']}' in the overlapping period.")
        st.stop()
    allocations = [portfolio['allocations'][i] for i, ticker in enumerate(portfolio['selected']) if ticker in available]
    allocations = adjust_allocations(allocations)
    return allocations, available, missing

def initialize_benchmark_metrics():
    """Initializes benchmark metrics when benchmark data is unavailable."""
    metrics = {metric: "N/A" for metric in [
        'Start Balance', 'End Balance', 'Annualized Return (CAGR)', 'Best Year', 'Worst Year',
        'Arithmetic Mean (Monthly)', 'Arithmetic Mean (Annualized)', 'Geometric Mean (Monthly)',
        'Geometric Mean (Annualized)', 'Standard Deviation (Monthly)', 'Standard Deviation (Annualized)',
        'Downside Deviation (Monthly)', 'Maximum Drawdown', 'Sharpe Ratio', 'Sortino Ratio',
        'Gain/Loss Ratio', 'Skewness', 'Excess Kurtosis', 'Safe Withdrawal Rate',
        'Perpetual Withdrawal Rate', 'Positive Periods', 'Benchmark Correlation', 'Beta',
        'Alpha (annualized)', 'R2', 'Treynor Ratio', 'Calmar Ratio',
        'Modigliani–Modigliani Measure', 'Information Ratio', 'Tracking Error',
        'Active Return', 'Upside Capture Ratio', 'Downside Capture Ratio'
    ]}
    return None, pd.Series(dtype=float), metrics

def process_benchmark(benchmark_data, returns, portfolio):
    """Processes benchmark data and calculates benchmark metrics."""
    benchmark_returns = benchmark_data[portfolio['benchmark_symbol']].pct_change().dropna()
    benchmark_returns = benchmark_returns.reindex(returns.index, method='ffill').dropna()
    common_index = returns.index.intersection(benchmark_returns.index)
    if common_index.empty:
        st.error("No overlapping dates between portfolio returns and benchmark returns after alignment.")
        st.stop()
    returns_aligned = returns.loc[common_index]
    cum_returns_aligned = (1 + benchmark_returns.loc[common_index]).cumprod()
    benchmark_metrics = calculate_metrics(
        benchmark_returns, cum_returns_aligned, portfolio['rf_rate'], None
    )
    return benchmark_returns, cum_returns_aligned, benchmark_metrics

def create_new_portfolio():
    """Creates a new empty portfolio template with default values."""
    return {
        'name': "",
        'start_date': pd.to_datetime(datetime.today() - timedelta(days=365)),
        'end_date': pd.to_datetime(datetime.today()),
        'rf_rate': st.session_state.default_config['rf_rate'],
        'broker_fee': st.session_state.default_config['broker_fee'],
        'benchmark_symbol': st.session_state.default_config['benchmark_symbol'],
        'rebalance_freq': 'M',
        'selected': [],
        'allocations': []
    }

def configure_portfolio(portfolio):
    """Handles the configuration of a portfolio."""
    col1, col2 = st.columns(2)
    
    with col1:
        portfolio['name'] = st.text_input("Portfolio Name", value=portfolio.get('name', ''))
        portfolio['start_date'] = st.date_input(
            "Start Date",
            value=portfolio.get('start_date', pd.to_datetime(datetime.today() - timedelta(days=365)))
        )
        portfolio['rf_rate'] = st.number_input(
            "Risk-Free Rate (%)",
            value=portfolio.get('rf_rate', 0.02) * 100,
            step=0.1
        ) / 100
        portfolio['rebalance_freq'] = st.selectbox(
            "Rebalancing Frequency",
            options=list(frequency_mapping.keys()),
            index=list(frequency_mapping.values()).index(portfolio.get('rebalance_freq', 'M'))
        )
    
    with col2:
        portfolio['benchmark_symbol'] = st.text_input(
            "Benchmark Symbol",
            value=portfolio.get('benchmark_symbol', '^GSPC')
        )
        portfolio['end_date'] = st.date_input(
            "End Date",
            value=portfolio.get('end_date', pd.to_datetime(datetime.today()))
        )
        portfolio['broker_fee'] = st.number_input(
            "Broker Fee (%)",
            value=portfolio.get('broker_fee', 0.0) * 100,
            step=0.01
        ) / 100
    
    return portfolio

def update_session_state(updates: dict):
    """
    Helper function to update Streamlit session state.

    Parameters:
    - updates (dict): Dictionary of key-value pairs to update in session state.
    """
    for key, value in updates.items():
        st.session_state[key] = value

def display_backtest_results(results, comparison_type, portfolio1, portfolio2):
    """Displays the backtest results in the Streamlit app."""
    if comparison_type == "Portfolio vs Benchmark":
        display_portfolio_vs_benchmark(results)
    else:
        display_portfolio_vs_portfolio(results, portfolio1, portfolio2)

def display_portfolio_vs_benchmark(results):
    """
    Enhanced display function for Portfolio vs Benchmark comparison with improved visualizations and UX.
    """
    # Extract data from results
    returns = results['returns']
    cum_returns = results['cum_returns']
    weights = results['weights']
    price_data = results['price_data']
    metrics = results['metrics']
    benchmark_metrics = results['benchmark_metrics']
    benchmark_returns = results['benchmark_returns']
    benchmark_cum_returns = results['benchmark_cum_returns']

    # Create tabs for better organization
    main_tab, metrics_tab, risk_tab, analysis_tab = st.tabs([
        "📈 Main Charts", "📊 Performance Metrics", "🎯 Risk Analysis", "🔍 Deep Analysis"
    ])

    with main_tab:
        col1, col2 = st.columns(2)
        with col1:
            st.subheader("Portfolio Composition")
            fig_composition = px.pie(
                values=weights,
                names=price_data.columns,
                title="Current Portfolio Allocation",
                template=THEME
            )
            st.plotly_chart(fig_composition, use_container_width=True)

        with col2:
            st.subheader("Key Performance Indicators")
            metrics_display = {
                "Total Return": f"{(cum_returns.iloc[-1] - 1) * 100:.2f}%",
                "Annualized Return": metrics.get('Annualized Return (CAGR)', 'N/A'),
                "Sharpe Ratio": metrics.get('Sharpe Ratio', 'N/A'),
                "Max Drawdown": metrics.get('Maximum Drawdown', 'N/A')
            }
            for metric, value in metrics_display.items():
                st.metric(metric, value)

        # Enhanced Cumulative Returns Plot
        st.subheader("Cumulative Returns Comparison")
        fig_returns = go.Figure()
        fig_returns.add_trace(go.Scatter(
            x=cum_returns.index,
            y=cum_returns,
            name="Portfolio",
            line=dict(color='cyan', width=2)
        ))
        fig_returns.add_trace(go.Scatter(
            x=benchmark_cum_returns.index,
            y=benchmark_cum_returns,
            name="Benchmark",
            line=dict(color='orange', width=2, dash='dash')
        ))
        fig_returns.update_layout(
            template=THEME,
            hovermode='x unified',
            height=500
        )
        st.plotly_chart(fig_returns, use_container_width=True)

    with metrics_tab:
        st.subheader("Detailed Performance Metrics")
        
        # Create comparison metrics table
        metrics_comparison = pd.DataFrame({
            'Metric': metrics.keys(),
            "Portfolio": metrics.values(),
            "Benchmark": benchmark_metrics.values()
        })
        
        # Style the metrics table
        def highlight_better(row):
            styles = [''] * len(row)
            if isinstance(row["Portfolio"], (int, float)) and isinstance(row["Benchmark"], (int, float)):
                if 'Drawdown' in row['Metric'] or 'Risk' in row['Metric']:
                    better = row["Portfolio"] < row["Benchmark"]
                else:
                    better = row["Portfolio"] > row["Benchmark"]
                styles[1] = 'background-color: lightgreen' if better else 'background-color: lightcoral'
            return styles
        
        styled_metrics = metrics_comparison.style.apply(highlight_better, axis=1)
        st.dataframe(styled_metrics, height=400)

    with risk_tab:
        col1, col2 = st.columns(2)
        
        with col1:
            # Rolling Volatility
            st.subheader("Rolling Volatility")
            rolling_vol = returns.rolling(window=30).std() * np.sqrt(252)
            rolling_vol_bench = benchmark_returns.rolling(window=30).std() * np.sqrt(252)
            
            fig_vol = go.Figure()
            fig_vol.add_trace(go.Scatter(
                x=rolling_vol.index,
                y=rolling_vol,
                name="Portfolio",
                line=dict(color='cyan')
            ))
            fig_vol.add_trace(go.Scatter(
                x=rolling_vol_bench.index,
                y=rolling_vol_bench,
                name="Benchmark",
                line=dict(color='orange', dash='dash')
            ))
            st.plotly_chart(fig_vol, use_container_width=True)

        with col2:
            # Rolling Beta
            st.subheader("Rolling Beta")
            rolling_covar = returns.rolling(window=30).cov(benchmark_returns)
            rolling_var = benchmark_returns.rolling(window=30).var()
            rolling_beta = rolling_covar / rolling_var
            
            fig_beta = px.line(
                rolling_beta,
                title="30-Day Rolling Beta",
                template=THEME
            )
            st.plotly_chart(fig_beta, use_container_width=True)

    with analysis_tab:
        col1, col2 = st.columns(2)
        
        with col1:
            # Monthly Returns Heatmap
            st.subheader("Monthly Returns Heatmap")
            monthly_returns = returns.resample('M').agg(lambda x: (1 + x).prod() - 1)
            monthly_returns_matrix = monthly_returns.groupby([monthly_returns.index.year, monthly_returns.index.month]).first().unstack()
            
            fig_heatmap = px.imshow(
                monthly_returns_matrix,
                labels=dict(x="Month", y="Year", color="Returns"),
                color_continuous_scale="RdYlGn",
                template=THEME
            )
            st.plotly_chart(fig_heatmap, use_container_width=True)

        with col2:
            # Rolling Correlation
            st.subheader("Rolling Correlation with Benchmark")
            rolling_corr = returns.rolling(window=30).corr(benchmark_returns)
            
            fig_corr = px.line(
                rolling_corr,
                title="30-Day Rolling Correlation",
                template=THEME
            )
            st.plotly_chart(fig_corr, use_container_width=True)

        # Risk Contribution Analysis
        st.subheader("Risk Contribution Analysis")
        risk_contrib = calculate_risk_contribution(returns, weights)
        fig_risk = px.pie(
            values=risk_contrib,
            names=price_data.columns,
            title="Risk Contribution by Asset",
            template=THEME
        )
        st.plotly_chart(fig_risk, use_container_width=True)

    # Download Section
    st.subheader("📥 Download Analysis Results")
    
    # Prepare Excel file with multiple sheets
    output = BytesIO()
    with pd.ExcelWriter(output, engine='xlsxwriter') as writer:
        metrics_comparison.to_excel(writer, sheet_name='Performance Metrics')
        pd.DataFrame({
            'Date': cum_returns.index,
            "Portfolio": cum_returns.values,
            "Benchmark": benchmark_cum_returns.values
        }).to_excel(writer, sheet_name='Cumulative Returns')
        
        monthly_returns_matrix.to_excel(writer, sheet_name='Monthly Returns')
    
    st.download_button(
        label="Download Complete Analysis",
        data=output.getvalue(),
        file_name="portfolio_analysis.xlsx",
        mime="application/vnd.openxmlformats-officedocument.spreadsheetml.sheet"
    )

def calculate_risk_contribution(returns, weights, cov_matrix=None):
    """
    Calculates the risk contribution of each asset in the portfolio.

    Parameters:
    - returns (pd.DataFrame or pd.Series): Asset returns.
    - weights (list or np.array): Asset weights in the portfolio.
    - cov_matrix (pd.DataFrame, optional): Covariance matrix. If None, it will be calculated.

    Returns:
    - pd.Series: Risk contribution of each asset.
    """
    if isinstance(returns, pd.Series):
        returns = returns.to_frame()
    
    if cov_matrix is None:
        cov_matrix = returns.cov()
    
    portfolio_variance = np.dot(weights, np.dot(cov_matrix, weights))
    marginal_contribution = np.dot(cov_matrix, weights)
    risk_contribution = marginal_contribution * weights / portfolio_variance
    return pd.Series(risk_contribution, index=returns.columns)

def display_portfolio_vs_portfolio(results, portfolio1, portfolio2):
    """
    Enhanced display function for Portfolio vs Portfolio comparison with improved visualizations and UX.
    """
    # Extract data from results
    returns_a = results['returns_a']
    cum_returns_a = results['cum_returns_a']
    weights_a = results['weights_a']
    price_data_a = results['price_data_a']
    metrics_a = results['metrics_a']

    returns_b = results['returns_b']
    cum_returns_b = results['cum_returns_b']
    weights_b = results['weights_b']
    price_data_b = results['price_data_b']
    metrics_b = results['metrics_b']

    # Create tabs for better organization
    overview_tab, comparison_tab, risk_tab, analysis_tab = st.tabs([
        "📊 Overview", "🔄 Comparison", "🎯 Risk Analysis", "🔍 Deep Analysis"
    ])

    with overview_tab:
        col1, col2 = st.columns(2)
        
        with col1:
            st.subheader(f"{portfolio1} Composition")
            fig_comp_a = px.pie(
                values=weights_a,
                names=price_data_a.columns,
                title=f"{portfolio1} Allocation",
                template=THEME
            )
            st.plotly_chart(fig_comp_a, use_container_width=True)

            # Key metrics for portfolio 1
            metrics_display_a = {
                "Total Return": f"{(cum_returns_a.iloc[-1] - 1) * 100:.2f}%",
                "Annualized Return": metrics_a.get('Annualized Return (CAGR)', 'N/A'),
                "Sharpe Ratio": metrics_a.get('Sharpe Ratio', 'N/A'),
                "Max Drawdown": metrics_a.get('Maximum Drawdown', 'N/A')
            }
            for metric, value in metrics_display_a.items():
                st.metric(f"{portfolio1} {metric}", value)

        with col2:
            st.subheader(f"{portfolio2} Composition")
            fig_comp_b = px.pie(
                values=weights_b,
                names=price_data_b.columns,
                title=f"{portfolio2} Allocation",
                template=THEME
            )
            st.plotly_chart(fig_comp_b, use_container_width=True)

            # Key metrics for portfolio 2
            metrics_display_b = {
                "Total Return": f"{(cum_returns_b.iloc[-1] - 1) * 100:.2f}%",
                "Annualized Return": metrics_b.get('Annualized Return (CAGR)', 'N/A'),
                "Sharpe Ratio": metrics_b.get('Sharpe Ratio', 'N/A'),
                "Max Drawdown": metrics_b.get('Maximum Drawdown', 'N/A')
            }
            for metric, value in metrics_display_b.items():
                st.metric(f"{portfolio2} {metric}", value)

        # Cumulative Returns Comparison
        st.subheader("Cumulative Returns Comparison")
        fig_returns = go.Figure()
        fig_returns.add_trace(go.Scatter(
            x=cum_returns_a.index,
            y=cum_returns_a,
            name=portfolio1,
            line=dict(color='cyan', width=2)
        ))
        fig_returns.add_trace(go.Scatter(
            x=cum_returns_b.index,
            y=cum_returns_b,
            name=portfolio2,
            line=dict(color='orange', width=2)
        ))
        fig_returns.update_layout(
            template=THEME,
            hovermode='x unified',
            height=500
        )
        st.plotly_chart(fig_returns, use_container_width=True)

    with comparison_tab:
        # Enhanced Performance Metrics Comparison
        st.subheader("Detailed Performance Comparison")
        
        detailed_metrics = [
            'Annualized Return (CAGR)',
            'Sharpe Ratio',
            'Sortino Ratio',
            'Maximum Drawdown',
            'Standard Deviation (Annualized)',
            'Alpha (annualized)',
            'Beta',
            'Information Ratio',
            'Tracking Error',
            'Calmar Ratio'
        ]

        comparison_df = pd.DataFrame({
            'Metric': detailed_metrics,
            portfolio1: [metrics_a.get(metric, "N/A") for metric in detailed_metrics],
            portfolio2: [metrics_b.get(metric, "N/A") for metric in detailed_metrics]
        }).set_index('Metric')

        # Style the comparison table
        def highlight_better_performance(s):
            is_better = []
            for v in s:
                if isinstance(v, (int, float)):
                    if 'Drawdown' in s.name or 'Risk' in s.name:
                        is_better.append('background-color: lightgreen' if v == min(s) else '')
                    else:
                        is_better.append('background-color: lightgreen' if v == max(s) else '')
                else:
                    is_better.append('')
            return is_better

        styled_comparison = comparison_df.style.apply(highlight_better_performance, axis=1)
        st.dataframe(styled_comparison, height=400)

        # Rolling Performance Difference
        st.subheader("Rolling Performance Difference")
        rolling_diff = (returns_a - returns_b).rolling(window=30).mean() * 252
        fig_diff = px.line(
            rolling_diff,
            title="30-Day Rolling Return Difference",
            template=THEME
        )
        st.plotly_chart(fig_diff, use_container_width=True)

    with risk_tab:
        col1, col2 = st.columns(2)
        
        with col1:
            # Rolling Volatility Comparison
            st.subheader("Rolling Volatility")
            rolling_vol_a = returns_a.rolling(window=30).std() * np.sqrt(252)
            rolling_vol_b = returns_b.rolling(window=30).std() * np.sqrt(252)
            
            fig_vol = go.Figure()
            fig_vol.add_trace(go.Scatter(
                x=rolling_vol_a.index,
                y=rolling_vol_a,
                name=portfolio1,
                line=dict(color='cyan')
            ))
            fig_vol.add_trace(go.Scatter(
                x=rolling_vol_b.index,
                y=rolling_vol_b,
                name=portfolio2,
                line=dict(color='orange')
            ))
            st.plotly_chart(fig_vol, use_container_width=True)

        with col2:
            # Rolling Correlation
            st.subheader("Rolling Correlation")
            rolling_corr = returns_a.rolling(window=30).corr(returns_b)
            fig_corr = px.line(
                rolling_corr,
                title="30-Day Rolling Correlation",
                template=THEME
            )
            st.plotly_chart(fig_corr, use_container_width=True)

        # Risk Contribution Analysis
        st.subheader("Risk Contribution Analysis")
        col3, col4 = st.columns(2)
        
        with col3:
            risk_contrib_a = calculate_risk_contribution(returns_a, weights_a)
            fig_risk_a = px.pie(
                values=risk_contrib_a,
                names=price_data_a.columns,
                title=f"{portfolio1} Risk Contribution",
                template=THEME
            )
            st.plotly_chart(fig_risk_a, use_container_width=True)
        
        with col4:
            risk_contrib_b = calculate_risk_contribution(returns_b, weights_b)
            fig_risk_b = px.pie(
                values=risk_contrib_b,
                names=price_data_b.columns,
                title=f"{portfolio2} Risk Contribution",
                template=THEME
            )
            st.plotly_chart(fig_risk_b, use_container_width=True)

    with analysis_tab:
        # Monthly Returns Heatmap
        st.subheader("Monthly Returns Comparison")
        col5, col6 = st.columns(2)
        
        with col5:
            monthly_returns_a = returns_a.resample('M').agg(lambda x: (1 + x).prod() - 1)
            monthly_matrix_a = monthly_returns_a.groupby([monthly_returns_a.index.year, monthly_returns_a.index.month]).first().unstack()
            fig_heat_a = px.imshow(
                monthly_matrix_a,
                title=f"{portfolio1} Monthly Returns",
                labels=dict(x="Month", y="Year", color="Returns"),
                color_continuous_scale="RdYlGn",
                template=THEME
            )
            st.plotly_chart(fig_heat_a, use_container_width=True)
        
        with col6:
            monthly_returns_b = returns_b.resample('M').agg(lambda x: (1 + x).prod() - 1)
            monthly_matrix_b = monthly_returns_b.groupby([monthly_returns_b.index.year, monthly_returns_b.index.month]).first().unstack()
            fig_heat_b = px.imshow(
                monthly_matrix_b,
                title=f"{portfolio2} Monthly Returns",
                labels=dict(x="Month", y="Year", color="Returns"),
                color_continuous_scale="RdYlGn",
                template=THEME
            )
            st.plotly_chart(fig_heat_b, use_container_width=True)

        # Risk-Return Scatter Plot
        st.subheader("Risk-Return Profile")
        risk_return_data = pd.DataFrame({
            'Portfolio': [portfolio1, portfolio2],
            'Risk': [returns_a.std() * np.sqrt(252), returns_b.std() * np.sqrt(252)],
            'Return': [returns_a.mean() * 252, returns_b.mean() * 252]
        })
        
        fig_risk_return = px.scatter(
            risk_return_data,
            x='Risk',
            y='Return',
            text='Portfolio',
            title="Risk-Return Profile",
            template=THEME
        )
        fig_risk_return.update_traces(textposition='top center')
        st.plotly_chart(fig_risk_return, use_container_width=True)

    # Download Section
    st.subheader("📥 Download Analysis Results")
    
    # Prepare Excel file with multiple sheets
    output = BytesIO()
    with pd.ExcelWriter(output, engine='xlsxwriter') as writer:
        comparison_df.to_excel(writer, sheet_name='Performance Metrics')
        pd.DataFrame({
            'Date': cum_returns_a.index,
            portfolio1: cum_returns_a.values,
            portfolio2: cum_returns_b.values
        }).to_excel(writer, sheet_name='Cumulative Returns')
        
        monthly_matrix_a.to_excel(writer, sheet_name=f'{portfolio1} Monthly Returns')
        monthly_matrix_b.to_excel(writer, sheet_name=f'{portfolio2} Monthly Returns')
    
    st.download_button(
        label="Download Complete Analysis",
        data=output.getvalue(),
        file_name=f"portfolio_comparison_{portfolio1}_vs_{portfolio2}.xlsx",
        mime="application/vnd.openxmlformats-officedocument.spreadsheetml.sheet"
    )

    # Final Summary
    st.subheader("📝 Summary and Recommendations")
    
    # Calculate overall scores
    score_a = calculate_final_score(metrics_a, metrics_b)
    score_b = calculate_final_score(metrics_b, metrics_a)
    
    col7, col8 = st.columns(2)
    with col7:
        st.metric(f"{portfolio1} Score", f"{score_a:.2f}/100")
    with col8:
        st.metric(f"{portfolio2} Score", f"{score_b:.2f}/100")
    
    # Generate and display recommendations
    recommendations = generate_portfolio_comparison_recommendations(
        metrics_a, metrics_b, portfolio1, portfolio2
    )
    for rec in recommendations:
        st.markdown(f"• {rec}")

    # Update session state
    update_session_state({
        'comparison_results': results,
        'show_proceed': True
    })
    st.success("Portfolio comparison analysis completed! 🎉")

def create_performance_stats_df(metrics, benchmark_metrics):
    """Creates a DataFrame for performance statistics."""
    return pd.DataFrame({
        'Metric': [
            'Start Balance', 'End Balance', 'Annualized Return (CAGR)',
            'Standard Deviation (Annualized)', 'Best Year', 'Worst Year',
            'Maximum Drawdown', 'Sharpe Ratio', 'Sortino Ratio', 'Benchmark Correlation'
        ],
        'Portfolio': [
            metrics.get('Start Balance', "N/A"),
            metrics.get('End Balance', "N/A"),
            metrics.get('Annualized Return (CAGR)', "N/A"),
            metrics.get('Standard Deviation (Annualized)', "N/A"),
            metrics.get('Best Year', "N/A"),
            metrics.get('Worst Year', "N/A"),
            metrics.get('Maximum Drawdown', "N/A"),
            metrics.get('Sharpe Ratio', "N/A"),
            metrics.get('Sortino Ratio', "N/A"),
            metrics.get('Benchmark Correlation', "N/A")
        ],
        'Benchmark': [
            benchmark_metrics.get('Start Balance', "N/A"),
            benchmark_metrics.get('End Balance', "N/A"),
            benchmark_metrics.get('Annualized Return (CAGR)', "N/A"),
            benchmark_metrics.get('Standard Deviation (Annualized)', "N/A"),
            benchmark_metrics.get('Best Year', "N/A"),
            benchmark_metrics.get('Worst Year', "N/A"),
            benchmark_metrics.get('Maximum Drawdown', "N/A"),
            benchmark_metrics.get('Sharpe Ratio', "N/A"),
            benchmark_metrics.get('Sortino Ratio', "N/A"),
            benchmark_metrics.get('Benchmark Correlation', "N/A")
        ]
    }).set_index('Metric')

def create_advanced_metrics_df(metrics, benchmark_metrics):
    """Creates a DataFrame for advanced metrics."""
    return pd.DataFrame({
        'Metric': list(metrics.keys()),
        'Portfolio': list(metrics.values()),
        'Benchmark': list(benchmark_metrics.values())
    }).set_index('Metric')

def create_performance_stats_df_portfolio_vs_portfolio(metrics_a, metrics_b, portfolio1, portfolio2):
    """Creates a refined DataFrame for Portfolio vs Portfolio performance statistics."""
    metrics = [
        'Annualized Return (CAGR)',
        'Sharpe Ratio',
        'Sortino Ratio',
        'Maximum Drawdown',
        'Beta',
        'Alpha (annualized)',
        'Information Ratio'
    ]

    data = {
        'Metric': metrics,
        portfolio1: [metrics_a.get(metric, "N/A") for metric in metrics],
        portfolio2: [metrics_b.get(metric, "N/A") for metric in metrics]
    }

    performance_df = pd.DataFrame(data)
    performance_df.set_index('Metric', inplace=True)
    return performance_df

def create_advanced_metrics_df_portfolio_vs_portfolio(metrics_a, metrics_b, portfolio1, portfolio2):
    """Creates a refined DataFrame for Portfolio vs Portfolio advanced metrics."""
    metrics = [
        'Beta',
        'Alpha (annualized)',
        'Information Ratio',
        'Modigliani–Modigliani Measure',
        'Tracking Error',
        'Upside Capture Ratio',
        'Downside Capture Ratio'
    ]

    data = {
        'Metric': metrics,
        portfolio1: [metrics_a.get(metric, "N/A") for metric in metrics],
        portfolio2: [metrics_b.get(metric, "N/A") for metric in metrics]
    }

    advanced_metrics_df = pd.DataFrame(data)
    advanced_metrics_df.set_index('Metric', inplace=True)
    return advanced_metrics_df

def display_drawdowns(cum_returns, benchmark_cum_returns):
    """Displays drawdown details for portfolio and benchmark."""
    st.markdown("### 📈 Drawdowns for Portfolio")
    portfolio_drawdowns = get_drawdown_details(cum_returns)
    if portfolio_drawdowns:
        st.table(pd.DataFrame(portfolio_drawdowns))
    else:
        st.write("No drawdowns detected for the portfolio.")

    if not benchmark_cum_returns.empty:
        st.markdown("### 📈 Drawdowns for Benchmark")
        benchmark_drawdowns = get_drawdown_details(benchmark_cum_returns)
        if benchmark_drawdowns:
            st.table(pd.DataFrame(benchmark_drawdowns))
        else:
            st.write("No drawdowns detected for the benchmark.")
    else:
        st.write("Benchmark data not available for drawdown analysis.")

def perform_risk_factor_attribution(returns):
    """Performs and displays risk factor attribution analysis."""
    start_date = returns.index.min().strftime('%Y-%m-%d')
    end_date = returns.index.max().strftime('%Y-%m-%d')
    factors = fetch_fama_french_factors(start_date, end_date)
    if factors.empty:
        st.write("Failed to retrieve factor data.")
        return
    factors = factors.asfreq(returns.index.freq, method='ffill')
    attribution, error = calculate_risk_factor_attribution(returns, factors)

    if error:
        st.error(error)
        st.write("Risk factor attribution analysis is unavailable.")
    elif attribution.empty:
        st.write("Risk factor attribution analysis is unavailable.")
    else:
        st.write("Risk Factor Attribution:")
        st.table(attribution)
        fig_attribution = px.bar(
            attribution,
            x='Factor',
            y='Contribution (%)',
            title='Risk Factor Attribution',
            labels={'Contribution (%)': 'Contribution (%)'},
            template=THEME
        )
        st.plotly_chart(fig_attribution, use_container_width=True)

def display_visualizations(cum_returns, benchmark_cum_returns, returns, price_data, weights, metrics, benchmark_metrics):
    """Displays an enhanced set of visualizations for Portfolio vs Benchmark with improved interactivity and aesthetics."""
    if not cum_returns.empty and not benchmark_cum_returns.empty:
        plot_growth_comparison(cum_returns, benchmark_cum_returns, portfolio_name="Portfolio", benchmark_name="Benchmark")
    else:
        st.warning("Insufficient data to display Growth Comparison. Ensure both portfolio and benchmark have data.")

    plot_drawdown_comparison(
        (cum_returns / cum_returns.expanding().max() - 1) * 100,
        (benchmark_cum_returns / benchmark_cum_returns.expanding().max() - 1) * 100,
        portfolio_name="Portfolio",
        benchmark_name="Benchmark"
    )

    st.subheader("📈 Compound Annual Growth Rate (CAGR) Over Time")
    if st.session_state.selected_time_frames:
        plot_cagr_over_time(cum_returns, time_frames=st.session_state.selected_time_frames)
    else:
        st.warning("Please select at least one time frame for CAGR.")

    st.subheader("📦 Returns Distribution Box Plot")
    plot_box(returns)

    st.subheader("🔥 Correlation Heatmap of Assets")
    plot_correlation_heatmap(price_data.pct_change().dropna())

    st.subheader("🔍 Risk-Return Attribution Analysis")
    plot_risk_return_attribution(returns, weights)
    
    # **Corrected Rolling Beta Calculation**
    st.subheader("📉 Rolling Beta Over Time")
    if not benchmark_cum_returns.empty:
        # Align benchmark returns with portfolio returns
        benchmark_returns = benchmark_cum_returns.pct_change().dropna().reindex(returns.index, method='ffill').dropna()
        
        # Calculate rolling covariance and variance
        rolling_cov = returns.rolling(window=252).cov(benchmark_returns)
        rolling_var = benchmark_returns.rolling(window=252).var()
        
        # Compute rolling beta
        rolling_beta = rolling_cov / rolling_var
        
        # Align rolling_beta with returns
        rolling_beta = rolling_beta.dropna()
        
        fig_rolling_beta = px.line(
            rolling_beta,
            title='Rolling Beta (1-Year Window)',
            labels={'index': 'Date', 'value': 'Beta'},
            template=THEME
        )
        st.plotly_chart(fig_rolling_beta, use_container_width=True)
    else:
        st.warning("Benchmark returns not available to plot Rolling Beta.")

    st.subheader("🔥 Cumulative Returns Heatmap")
    plot_cumulative_returns_heatmap(cum_returns)

    st.subheader("🥧 Portfolio Allocation Pie Chart")
    plot_allocation_pie(weights, price_data.columns, title="🥧 Portfolio Allocation", hover_info="percent+name")

    st.subheader("📉 Rolling Metrics")
    if st.session_state.selected_rolling_periods:
        plot_rolling_metrics(returns, windows=st.session_state.selected_rolling_periods)
    else:
        st.warning("Please select at least one rolling period to display metrics.")
    # **Added: Rolling Sharpe Ratio Graph**
    st.subheader("📈 Rolling Sharpe Ratio Over Time")
    if not benchmark_cum_returns.empty:
        # Calculate rolling Sharpe Ratio
        rolling_sharpe = (returns.rolling(window=252).mean() - st.session_state.default_config['rf_rate']) / returns.rolling(window=252).std()
        rolling_sharpe = rolling_sharpe.dropna()
        
        fig_rolling_sharpe = px.line(
            rolling_sharpe,
            title='Rolling Sharpe Ratio (1-Year Window)',
            labels={'index': 'Date', 'value': 'Sharpe Ratio'},
            template=THEME
        )
        st.plotly_chart(fig_rolling_sharpe, use_container_width=True)
    else:
        st.warning("Benchmark returns not available to plot Rolling Sharpe Ratio.")
    st.markdown("### 💡 Recommendations")
    recommendations = generate_recommendations(
        allocations=[alloc for alloc in weights if alloc > 0],
        available_selected=price_data.columns
    )
    if recommendations:
        for rec in recommendations:
            st.markdown(f"• {rec}")
    else:
        st.success("🎉 Your portfolio allocations are well-balanced!")

    st.markdown("### 📝 Final Score")
    portfolio_score = calculate_final_score(metrics, benchmark_metrics)
    st.metric("📊 Portfolio Score", f"{portfolio_score:.2f} / 100")
    # Add interpretation/comment similar to Risk Factor Attribution
    st.markdown("""
    **📖 Interpretation:**
    - The **Final Score** represents a comprehensive evaluation of your portfolio's performance based on selected metrics.
    - A higher score indicates a better balance between risk and return, aligning with your investment objectives.
    - Use this score to assess overall portfolio health and identify areas for improvement.
    """)

def display_visualizations_portfolio_vs_portfolio(
    cum_returns_a, cum_returns_b, returns_a, returns_b,
    price_data_a, price_data_b, weights_a, weights_b,
    portfolio1, portfolio2, metrics_a, metrics_b
):
    """Displays various visualizations for Portfolio vs Portfolio."""
    # Growth Comparison
    fig_growth = px.line(title='Growth Comparison')
    fig_growth.add_scatter(x=cum_returns_a.index, y=cum_returns_a, mode='lines', name=portfolio1)
    fig_growth.add_scatter(x=cum_returns_b.index, y=cum_returns_b, mode='lines', name=portfolio2)
    st.plotly_chart(fig_growth, use_container_width=True)

    # Drawdown Comparison
    plot_drawdown_comparison(
        (cum_returns_a / cum_returns_a.expanding().max() - 1) * 100,
        (cum_returns_b / cum_returns_b.expanding().max() - 1) * 100
    )

    # Box Plot for Returns Distribution
    st.subheader("📦 Returns Distribution Box Plot")
    # Create a DataFrame with both portfolios' returns
    returns_df = pd.DataFrame({
        portfolio1: returns_a,
        portfolio2: returns_b
    })
    plot_box(
        returns_df,
        title="Returns Distribution for Both Portfolios",
        multiple=True,
        portfolio_names=[portfolio1, portfolio2]
    )

    # Heatmap of Correlations Between Assets
    st.subheader("🔥 Correlation Heatmap of Assets")
    plot_correlation_heatmap(price_data_a.pct_change().dropna())
    plot_correlation_heatmap(price_data_b.pct_change().dropna())

    # Portfolio Allocation Pie Chart
    st.subheader("🥧 Portfolio Allocation Pie Chart")
    plot_allocation_pie(weights_a, price_data_a.columns, title=f"{portfolio1} Allocation", hover_info="percent+name")
    plot_allocation_pie(weights_b, price_data_b.columns, title=f"{portfolio2} Allocation", hover_info="percent+name")

    st.markdown("### 📝 Final Score")
    score_a = calculate_final_score(metrics_a, metrics_b)
    score_b = calculate_final_score(metrics_b, metrics_a)
    st.write(f"**{portfolio1}:** {score_a:.2f} / 100")
    st.write(f"**{portfolio2}:** {score_b:.2f} / 100")

def plot_box(
    returns: Union[pd.Series, pd.DataFrame],
    title: str = "Returns Distribution Box Plot",
    labels: dict = {"Return": "Returns (%)", "Asset": "Asset"},
    color_sequence: list = px.colors.qualitative.Bold,
    multiple: bool = False,
    portfolio_names: list = None
):
    """
    Generic function to plot box plots for portfolio returns.

    Parameters:
    - returns (pd.DataFrame or pd.Series): Returns data.
    - title (str): Title of the plot.
    - labels (dict): Axis labels.
    - color_sequence (list): Colors for the boxes.
    - multiple (bool): If True, handle multiple portfolios.
    - portfolio_names (list): Names of the portfolios for labeling.
    """
    if isinstance(returns, pd.DataFrame):
        if multiple and portfolio_names:
            for i, col in enumerate(returns.columns):
                fig = px.box(
                    returns[[col]].reset_index(drop=True),
                    y=col,
                    title=f"{title} for {portfolio_names[i]}",
                    labels=labels,
                    points='all',
                    template=THEME,
                    color_discrete_sequence=[color_sequence[i % len(color_sequence)]]
                )
                fig.update_traces(boxmean='sd')
                fig.update_layout(xaxis_title="Assets", yaxis_title=labels.get('Return', 'Returns (%)'))
                st.plotly_chart(fig, use_container_width=True)
        else:
            melted_returns = returns.reset_index().melt(id_vars=returns.index.name if returns.index.name else 'Date', var_name='Asset', value_name='Return')
            fig = px.box(
                melted_returns,
                x='Asset',
                y='Return',
                title=title,
                labels=labels,
                points='all',
                template=THEME,
                color='Asset',
                color_discrete_sequence=color_sequence
            )
    elif isinstance(returns, pd.Series):
        fig = px.box(
            returns.rename("Return"),
            y='Return',
            title=title,
            labels=labels,
            points='all',
            template=THEME,
            color_discrete_sequence=['cyan']
        )
    else:
        st.warning("Returns data is neither a DataFrame nor a Series.")
        return

    fig.update_traces(boxmean='sd')
    fig.update_layout(
        xaxis_title=labels.get("Asset", "Assets"),
        yaxis_title=labels.get("Return", "Returns (%)"),
        legend_title="Assets"
    )
    st.plotly_chart(fig, use_container_width=True)

    st.markdown("""
        **📖 Interpretation:**
        - **Boxes:** Represent the interquartile range (IQR) where the middle 50% of the data lies.
        - **Median Line:** Indicates the median return.
        - **Whiskers:** Extend to show the range of the data excluding outliers.
        - **Outliers:** Individual points outside the whiskers represent atypical returns.
        - **Box Mean (Shaded):** Shows the mean and standard deviation of returns.
    """)

def plot_cumulative_returns_heatmap(cum_returns):
    """Plots the cumulative returns heatmap."""
    try:
        cum_returns_normalized = cum_returns / cum_returns.max()
        fig_heatmap = px.imshow(
            cum_returns_normalized.to_frame().T,
            labels=dict(x="Date", y="Portfolio", color="Normalized Cumulative Return"),
            title="Cumulative Returns Heatmap",
            aspect="auto",
            color_continuous_scale='Viridis',
            template=THEME
        )
        st.plotly_chart(fig_heatmap, use_container_width=True)
        st.markdown(
            "**Interpretation:** The heatmap visualizes the normalized cumulative returns over time, "
            "allowing for an intuitive comparison of portfolio performance across different periods."
        )
    except Exception as e:
        st.error(f"Error plotting cumulative returns heatmap: {e}")

def plot_returns_correlation_heatmap(price_data):
    """Plots the correlation heatmap of asset returns."""
    correlation_data = price_data.pct_change().dropna()
    plot_correlation_heatmap(correlation_data)

def plot_correlation_heatmap(correlation_data):
    """Plots an enhanced correlation heatmap using Plotly with improved aesthetics."""
    if correlation_data.empty:
        st.warning("No data available to plot correlation heatmap.")
        return
    corr = correlation_data.corr()
    fig_heatmap = px.imshow(
        corr,
        title="📈 Correlation Heatmap of Assets",
        labels=dict(x="Asset", y="Asset", color="Correlation"),
        color_continuous_scale='RdBu',
        zmin=-1,
        zmax=1,
        text_auto=True,
        aspect="auto",
        template=THEME
    )
    fig_heatmap.update_layout(
        title_x=0.5,
        xaxis_title="Assets",
        yaxis_title="Assets"
    )
    fig_heatmap.update_traces(
        hovertemplate="Asset1: %{x}<br>Asset2: %{y}<br>Correlation: %{z:.2f}<extra></extra>"
    )
    st.plotly_chart(fig_heatmap, use_container_width=True)

def plot_growth_comparison(cum_returns, benchmark_cum_returns, portfolio_name='Portfolio', benchmark_name='Benchmark'):
    """Plots an enhanced growth comparison between portfolio and benchmark with improved interactivity."""
    fig_growth = px.line(
        title='📈 Growth Comparison',
        labels={'value': 'Cumulative Returns', 'index': 'Date'},
        template=THEME
    )
    
    fig_growth.add_scatter(
        x=cum_returns.index,
        y=cum_returns,
        mode='lines',
        name=portfolio_name,
        line=dict(color='cyan', width=2)
    )
    fig_growth.add_scatter(
        x=benchmark_cum_returns.index,
        y=benchmark_cum_returns,
        mode='lines',
        name=benchmark_name,
        line=dict(color='orange', width=2, dash='dash')
    )
    
    fig_growth.update_layout(
        hovermode='x unified',
        xaxis=dict(showgrid=True),
        yaxis=dict(showgrid=True),
        legend=dict(title="Legend", x=0.01, y=0.99),
        height=600  # Added height parameter
    )
    
    st.plotly_chart(fig_growth, use_container_width=True, height=600, config={'responsive': True})

def plot_drawdown_comparison(drawdown_portfolio: pd.Series, drawdown_benchmark: pd.Series, portfolio_name='Portfolio', benchmark_name='Benchmark') -> None:
    """Plots an enhanced drawdown comparison between portfolio and benchmark with improved visuals."""
    fig_drawdown = px.line(
        title='📉 Drawdown Comparison',
        labels={'value': 'Drawdown (%)', 'index': 'Date'},
        template=THEME
    )
    fig_drawdown.add_scatter(
        x=drawdown_portfolio.index,
        y=drawdown_portfolio,
        mode='lines',
        name=portfolio_name,
        line=dict(color='red', width=2)
    )
    if not drawdown_benchmark.empty:
        fig_drawdown.add_scatter(
            x=drawdown_benchmark.index,
            y=drawdown_benchmark,
            mode='lines',
            name=benchmark_name,
            line=dict(color='blue', width=2, dash='dash')
        )
    fig_drawdown.update_layout(
        hovermode=HOVERMODE,
        yaxis=dict(showgrid=True),
        xaxis=dict(showgrid=True),
        legend=dict(title="Legend", x=0.01, y=0.99),
        height=FIG_HEIGHT
    )
    fig_drawdown.update_traces(hovertemplate="Date: %{x}<br>Drawdown: %{y:.2f}%")
    st.plotly_chart(fig_drawdown, use_container_width=True, height=600, config={'responsive': True})

def plot_cagr_over_time(cum_returns, time_frames):
    """Plots CAGR over different time frames."""
    for frame in time_frames:
        plot_cagr(cum_returns, frame)

def plot_risk_return_attribution(returns, weights):
    """Plots the risk-return attribution analysis."""
    fig_rra = px.scatter(
        x=returns.mean() * 252,
        y=returns.std() * np.sqrt(252),
        size=weights*100,
        color=returns.mean() / returns.std(),
        hover_name=returns.index,
        title='Risk-Return Attribution',
        labels={'x': 'Annualized Return', 'y': 'Annualized Risk (Std Dev)'},
        template='plotly_dark',
        size_max=60
    )
    st.plotly_chart(fig_rra, use_container_width=True)

def plot_rolling_metrics(returns, windows):
    """Plots rolling metrics based on selected window periods."""
    for window in windows:
        rolling_mean = returns.rolling(window=window).mean()
        rolling_std = returns.rolling(window=window).std()
        fig = px.line(title=f'Rolling {window}-Day Mean and Std Dev')
        fig.add_scatter(x=rolling_mean.index, y=rolling_mean, mode='lines', name='Rolling Mean')
        fig.add_scatter(x=rolling_std.index, y=rolling_std, mode='lines', name='Rolling Std Dev')
        st.plotly_chart(fig, use_container_width=True)

@handle_exceptions
def calculate_risk_factor_attribution(returns: pd.Series, factors: pd.DataFrame, annualization_factor: int = 252) -> tuple:
    """
    Decompose portfolio returns by risk factors using linear regression and provide
    comments on the factor contributions.

    Parameters:
    - returns (pd.Series): Portfolio returns.
    - factors (pd.DataFrame): DataFrame where each column represents a risk factor.
    - annualization_factor (int): Annualization factor (252 for daily, 12 for monthly returns).

    Returns:
    - pd.DataFrame: Factor names, their contribution percentages, and explanatory comments.
    - None: If an exception occurs, handled by the decorator.
    """
    # Ensure returns has a name, or assign one if missing
    if returns.name is None:
        returns = returns.rename("Portfolio_Returns")

    # Convert returns to a DataFrame
    returns = returns.to_frame()

    # Ensure both returns and factors are in decimal format
    if returns['Portfolio_Returns'].max() > 1:
        returns /= 100

    if factors.max().max() > 1:
        factors /= 100

    # Adjust returns to excess returns by subtracting RF (if it exists in factors)
    if 'RF' in factors.columns:
        returns['Portfolio_Returns'] -= factors['RF']
        factors = factors.drop(columns=['RF'])

    # Align the data using an inner join and drop any remaining NaN values
    aligned_data = pd.merge(returns, factors, left_index=True, right_index=True, how='inner').dropna()

    # Check that we have data after merging
    if aligned_data.empty:
        raise ValueError("No overlapping dates between returns and factors after dropping NaNs.")

    # Prepare X and y for regression
    X = aligned_data[factors.columns]
    X = sm.add_constant(X)  # Add intercept
    y = aligned_data['Portfolio_Returns']

    # Run the regression model
    model = sm.OLS(y, X).fit()
    coefficients = model.params.drop('const', errors='ignore')

    # Calculate factor contributions (non-annualized)
    factor_means = X.mean().drop('const', errors='ignore')
    contributions = coefficients * factor_means

    # Calculate total return
    total_return = returns['Portfolio_Returns'].mean()

    # Check if total_return is zero to avoid division by zero
    if np.isclose(total_return, 0):
        st.warning("Total portfolio return is zero. Cannot compute contribution percentages.")
        attribution_df = pd.DataFrame({
            'Factor': contributions.index,
            'Contribution (%)': [0.0] * len(contributions)
        })
    else:
        # Calculate percentage contributions
        attribution_df = pd.DataFrame({
            'Factor': contributions.index,
            'Contribution (%)': (contributions / total_return * 100).round(2)
        })

    # Remove the constant term if it exists
    if 'const' in attribution_df['Factor'].values:
        attribution_df = attribution_df[attribution_df['Factor'] != 'const']

    # Use full names for factors
    factor_full_names = {
        'Mkt-RF': 'Market Risk Premium',
        'SMB': 'Small Minus Big',
        'HML': 'High Minus Low'
    }
    attribution_df['Factor'] = attribution_df['Factor'].replace(factor_full_names)

    # Generate detailed comments based on factor contributions
    comments = []
    for _, row in attribution_df.iterrows():
        factor_name = row['Factor']
        contribution = row['Contribution (%)']

        # Market Risk Premium (Mkt-RF) Analysis
        if factor_name == 'Market Risk Premium':
            if contribution > 75:
                comments.append(f"{factor_name} = {contribution}%: Extremely high sensitivity to market movements.")
            elif contribution > 50:
                comments.append(f"{factor_name} = {contribution}%: High sensitivity to overall market movements.")
            elif contribution > 30:
                comments.append(f"{factor_name} = {contribution}%: Notable sensitivity to market risk.")
            elif contribution > 10:
                comments.append(f"{factor_name} = {contribution}%: Moderate exposure to market risk.")
            else:
                comments.append(f"{factor_name} = {contribution}%: Low exposure to market risk.")

        # Small Minus Big (SMB) Analysis
        elif factor_name == 'Small Minus Big':
            if contribution > 5:
                comments.append(f"{factor_name} = {contribution}%: Strong tilt towards small-cap stocks.")
            elif contribution > 0:
                comments.append(f"{factor_name} = {contribution}%: Mild preference for small-cap stocks.")
            elif contribution > -5:
                comments.append(f"{factor_name} = {contribution}%: Slight tilt towards larger-cap stocks.")
            else:
                comments.append(f"{factor_name} = {contribution}%: Clear preference for large-cap stocks.")

        # High Minus Low (HML) Analysis
        elif factor_name == 'High Minus Low':
            if contribution > 5:
                comments.append(f"{factor_name} = {contribution}%: Strong tilt toward value stocks.")
            elif contribution > 0:
                comments.append(f"{factor_name} = {contribution}%: Slight positive tilt toward value stocks.")
            elif contribution > -5:
                comments.append(f"{factor_name} = {contribution}%: Minor negative tilt toward growth stocks.")
            else:
                comments.append(f"{factor_name} = {contribution}%: Clear preference for growth stocks.")

    # Add comments to the DataFrame
    attribution_df['Comment'] = comments

    return attribution_df, None

@handle_exceptions
def fetch_fama_french_factors(start_date: str, end_date: str) -> pd.DataFrame:
    """
    Fetch Fama-French daily factor data between specified dates.

    Parameters:
    - start_date (str): Start date in 'YYYY-MM-DD' format.
    - end_date (str): End date in 'YYYY-MM-DD' format.

    Returns:
    - pd.DataFrame: DataFrame containing Fama-French factors.
    """
    # [Function Implementation Remains Unchanged]
    try:
        # Fetch daily Fama-French factors
        ff_data = pdr.get_data_famafrench('F-F_Research_Data_Factors_daily', start=start_date, end=end_date)
        if not ff_data:
            st.error("Fama-French data not found.")
            return pd.DataFrame()

        factors = ff_data[0]
        factors = factors.rename(columns=lambda x: x.strip())  # Remove any leading/trailing spaces
        factors /= 100  # Convert percentage returns to decimal format

        # Fill missing data by forward and backward filling
        factors = factors.fillna(method='ffill').fillna(method='bfill')

        # Check for necessary columns
        required_columns = {'Mkt-RF', 'SMB', 'HML', 'RF'}
        if not required_columns.issubset(factors.columns):
            missing = required_columns - set(factors.columns)
            st.error(f"Missing factor columns: {', '.join(missing)}")
            return pd.DataFrame()

        return factors
    except Exception as e:
        st.error(f"Error fetching factor data: {e}")
        return pd.DataFrame()

@handle_exceptions
def calculate_final_score(primary_metrics: dict, comparison_metrics: dict) -> float:
    metrics_to_compare = [
        'Annualized Return (CAGR)',
        'Sharpe Ratio',
        'Sortino Ratio',
        'Treynor Ratio',
        'Calmar Ratio',
        'Alpha (annualized)',
        'Information Ratio',
        'Modigliani–Modigliani Measure',
        'Upside Capture Ratio',
        'Gain/Loss Ratio'
    ]
    metric_weights = {
        'Annualized Return (CAGR)': 15,
        'Sharpe Ratio': 15,
        'Sortino Ratio': 10,
        'Treynor Ratio': 10,
        'Calmar Ratio': 10,
        'Alpha (annualized)': 10,
        'Information Ratio': 10,
        'Modigliani–Modigliani Measure': 10,
        'Upside Capture Ratio': 5,
        'Gain/Loss Ratio': 5
    }

    score = 0
    total = 0

    for metric in metrics_to_compare:
        primary = primary_metrics.get(metric, "N/A")
        comparison = comparison_metrics.get(metric, "N/A")
        if primary != "N/A" and comparison != "N/A":
            try:
                primary_val = float(primary.strip('%')) if '%' in primary else float(primary)
                comparison_val = float(comparison.strip('%')) if '%' in comparison else float(comparison)
                weight = metric_weights.get(metric, 1)
                if comparison_val == 0:
                    st.warning(f"Comparison metric '{metric}' has a value of zero. Skipping this metric to avoid division by zero.")
                    continue
                if primary_val > comparison_val:
                    score += weight * (primary_val / comparison_val)
                else:
                    score += weight * (primary_val / comparison_val) * 0.5  # Partial credit
                total += weight
            except ValueError:
                st.warning(f"Invalid numeric conversion for metric '{metric}'. Skipping.")
                continue
    return (score / total) * 100 if total > 0 else 0
    
@handle_exceptions
def plot_cagr_over_time(cum_returns: pd.Series, time_frames: list = ['Weekly', 'Monthly', 'Quarterly', 'Annually']) -> None:
    """
    Plot CAGR over multiple time frames.

    Parameters:
    - cum_returns (pd.Series): Cumulative returns of the portfolio.
    - time_frames (list): List of time frames to calculate CAGR.
    """
    frequency_map = {
        'Weekly': 'W',
        'Monthly': 'M',
        'Quarterly': 'Q',
        'Annually': 'Y'
    }
    fig = go.Figure()
    for tf in time_frames:
        freq = frequency_map.get(tf)
        if not freq:
            continue
        rolled = cum_returns.resample(freq).last()
        if rolled.empty:
            st.warning(f"No data available for {tf} CAGR calculation.")
            continue
        years = (rolled.index[-1] - rolled.index[0]).days / 365.25
        if years <= 0:
            st.warning(f"Not enough data to calculate {tf} CAGR.")
            continue
        cagr = (rolled / rolled.iloc[0]) ** (1/years) - 1
        fig.add_trace(go.Scatter(x=rolled.index, y=cagr, mode='lines', name=f'{tf} CAGR'))
    fig.update_layout(
        title='CAGR Over Multiple Time Frames',
        xaxis_title='Date',
        yaxis_title='CAGR',
        hovermode=HOVERMODE
    )
    st.plotly_chart(fig, use_container_width=True)

@handle_exceptions
def backtest(weights: list, prices: pd.DataFrame, rebalance_freq: str = 'M', broker_fee: float = 0.0, debug: bool = False) -> tuple:
    """
    Backtest the portfolio based on weights and price data.

    Parameters:
    - weights (list): Allocation weights for each asset.
    - prices (pd.DataFrame): Adjusted closing prices of assets.
    - rebalance_freq (str): Rebalancing frequency (e.g., 'D', 'W', 'M').
    - broker_fee (float): Broker fee as a percentage.
    - debug (bool): If True, prints debug information.

    Returns:
    - tuple: (portfolio_returns: pd.Series, cum_returns: pd.Series)
    """
    try:
        # Calculate returns and drop any NaN values
        returns = prices.pct_change().dropna()
        if returns.empty:
            st.error("Returns data is empty after calculating percentage changes. Please ensure your price data is accurate and spans multiple periods.")
            return pd.Series(dtype=float), pd.Series(dtype=float)

        # Calculate portfolio returns
        portfolio_returns = returns.dot(weights)
        if not isinstance(portfolio_returns, pd.Series):
            portfolio_returns = portfolio_returns.squeeze()

        # Identify rebalancing dates
        if rebalance_freq.upper() == 'D':
            rebalance_dates = returns.index
        else:
            rebalance_dates = returns.resample(rebalance_freq.upper()).last().dropna().index

        # Debug: Show rebalancing dates
        if debug:
            st.write(f"Rebalancing Dates: {rebalance_dates.tolist()}")

        # Ensure rebalance_dates are in the portfolio_returns index
        valid_rebalance_dates = rebalance_dates.intersection(portfolio_returns.index)
        if debug:
            st.write(f"Valid Rebalancing Dates: {valid_rebalance_dates.tolist()}")

        # Apply broker fees on rebalancing dates
        if not valid_rebalance_dates.empty:
            portfolio_returns.loc[valid_rebalance_dates] -= broker_fee / 100  # Convert to decimal
        else:
            st.warning("⚠️ No valid rebalancing dates found within the returns data.")

        # Calculate cumulative returns
        cum_returns = (1 + portfolio_returns).cumprod()
        if cum_returns.empty:
            st.error("Cumulative returns are empty. Check the data and allocations.")
            return pd.Series(dtype=float), pd.Series(dtype=float)

        # Debug: Show cumulative returns
        if debug:
            st.write("Cumulative Returns:")
            st.write(cum_returns)

        return portfolio_returns, cum_returns
    except Exception as e:
        st.error(f"Error during backtesting: {e}")
        return pd.Series(dtype=float), pd.Series(dtype=float)

def calculate_sharpe_ratio(returns: pd.Series, rf: float = 0.02) -> float:
    """
    Calculate the Sharpe Ratio for a given set of returns.

    Parameters:
    - returns (pd.Series): Daily returns of the portfolio.
    - rf (float): Risk-free rate (default is 2%).

    Returns:
    - float: Sharpe Ratio.
    """
    excess_return = returns.mean() * 252 - rf
    std_dev = returns.std() * np.sqrt(252)
    return excess_return / std_dev if std_dev != 0 else np.nan

def calculate_sortino_ratio(returns: pd.Series, rf: float = 0.02) -> float:
    """
    Calculate the Sortino Ratio for a given set of returns.

    Parameters:
    - returns (pd.Series): Daily returns of the portfolio.
    - rf (float): Risk-free rate (default is 2%).

    Returns:
    - float: Sortino Ratio.
    """
    excess_return = returns.mean() * 252 - rf
    downside_std = returns[returns < 0].std() * np.sqrt(252)
    return excess_return / downside_std if downside_std != 0 else np.nan

def calculate_treynor_ratio(returns: pd.Series, benchmark_returns: pd.Series, rf: float = 0.02) -> float:
    """
    Calculate the Treynor Ratio for a given set of returns.

    Parameters:
    - returns (pd.Series): Daily returns of the portfolio.
    - benchmark_returns (pd.Series): Daily returns of the benchmark.
    - rf (float): Risk-free rate (default is 2%).

    Returns:
    - float: Treynor Ratio.
    """
    beta = calculate_beta(returns, benchmark_returns)
    excess_return = returns.mean() * 252 - rf
    return excess_return / beta if beta != 0 else np.nan

def calculate_calmar_ratio(returns: pd.Series, cum_returns: pd.Series) -> float:
    """
    Calculate the Calmar Ratio for a given set of returns.

    Parameters:
    - returns (pd.Series): Daily returns of the portfolio.
    - cum_returns (pd.Series): Cumulative returns of the portfolio.

    Returns:
    - float: Calmar Ratio.
    """
    annual_return = returns.mean() * 252
    max_dd = drawdown(cum_returns)
    return annual_return / abs(max_dd) if max_dd != 0 else np.nan

def calculate_beta(returns: pd.Series, benchmark_returns: pd.Series) -> float:
    """
    Calculate the Beta of the portfolio relative to the benchmark.

    Parameters:
    - returns (pd.Series): Daily returns of the portfolio.
    - benchmark_returns (pd.Series): Daily returns of the benchmark.

    Returns:
    - float: Beta value.
    """
    if returns.empty or benchmark_returns.empty:
        return np.nan
    covariance_matrix = np.cov(returns, benchmark_returns)
    covariance = covariance_matrix[0, 1]
    benchmark_variance = covariance_matrix[1, 1]
    return covariance / benchmark_variance if benchmark_variance != 0 else np.nan

def calculate_alpha(returns: pd.Series, benchmark_returns: pd.Series, rf: float = 0.02) -> float:
    """
    Calculate the Alpha of the portfolio relative to the benchmark.

    Parameters:
    - returns (pd.Series): Daily returns of the portfolio.
    - benchmark_returns (pd.Series): Daily returns of the benchmark.
    - rf (float): Risk-free rate (default is 2%).

    Returns:
    - float: Alpha value.
    """
    beta = calculate_beta(returns, benchmark_returns)
    portfolio_return = returns.mean() * 252
    benchmark_return = benchmark_returns.mean() * 252
    return portfolio_return - (rf + beta * (benchmark_return - rf)) if not np.isnan(beta) else np.nan

def calculate_r_squared(returns: pd.Series, benchmark_returns: pd.Series) -> float:
    """
    Calculate the R-squared of the portfolio relative to the benchmark.

    Parameters:
    - returns (pd.Series): Daily returns of the portfolio.
    - benchmark_returns (pd.Series): Daily returns of the benchmark.

    Returns:
    - float: R-squared value.
    """
    if returns.empty or benchmark_returns.empty:
        return np.nan
    covariance = np.cov(returns, benchmark_returns)
    var_port = covariance[0,0]
    var_bench = covariance[1,1]
    cov = covariance[0,1]
    return (cov ** 2) / (var_port * var_bench) if var_port !=0 and var_bench !=0 else np.nan

def calculate_information_ratio(returns: pd.Series, benchmark_returns: pd.Series) -> float:
    """
    Calculate the Information Ratio of the portfolio relative to the benchmark.

    Parameters:
    - returns (pd.Series): Daily returns of the portfolio.
    - benchmark_returns (pd.Series): Daily returns of the benchmark.

    Returns:
    - float: Information Ratio.
    """
    if returns.empty or benchmark_returns.empty:
        return np.nan
    active_return = (returns.mean() - benchmark_returns.mean()) * 252
    tracking_error = calculate_tracking_error(returns, benchmark_returns)
    return active_return / tracking_error if tracking_error != 0 else np.nan

def generate_recommendations(allocations: list, available_selected: list) -> list:
    """
    Generate basic portfolio recommendations based on current allocations.

    Parameters:
    - allocations (list of float): Current allocation percentages.
    - available_selected (list of str): List of selected asset tickers.

    Returns:
    - list of str: Recommendations.
    """
    recommendations = []
    max_alloc = max(allocations)
    min_alloc = min(allocations)
    total_alloc = sum(allocations)

    if max_alloc > 40.0:
        idx = allocations.index(max_alloc)
        ticker = available_selected[idx]
        recommendations.append(f"🔹 Consider reducing the allocation to **{ticker}** since it constitutes **{max_alloc:.2f}%** of your portfolio.")

    if min_alloc < 5.0:
        idx = allocations.index(min_alloc)
        ticker = available_selected[idx]
        recommendations.append(f"🔸 Consider increasing the allocation to **{ticker}** to at least **5.00%** for better diversification.")

    if total_alloc < 100.0:
        recommendations.append("🔹 Consider allocating the remaining funds to additional assets to reach a total of **100%**.")

    return recommendations

def get_common_benchmarks(selected_tickers: list) -> dict:
    """
    Define benchmark suggestions based on asset sectors or indices.

    Parameters:
    - selected_tickers (list of str): List of selected asset tickers.

    Returns:
    - dict: Dictionary of benchmark names and their corresponding tickers.
    """
    # Example benchmark sets
    sp500 = {"AAPL", "MSFT", "GOOGL", "AMZN", "META", "TSLA"}  # Example S&P 500 tech companies
    nasdaq_tech = {"AAPL", "MSFT", "GOOGL", "AMZN", "META", "TSLA", "NVDA", "AMD"}

    selected_set = set(selected_tickers)

    if selected_set.issubset(sp500):
        return {
            "S&P 500": "^GSPC",
            "Dow Jones Industrial Average": "^DJI",
            "Russell 2000": "^RUT",
            "Custom": "CUSTOM"
        }
    elif selected_set.issubset(nasdaq_tech):
        return {
            "NASDAQ Composite": "^IXIC",
            "QQQ (Invesco QQQ ETF)": "QQQ",
            "Custom": "CUSTOM"
        }
    else:
        return {
            "S&P 500": "^GSPC",
            "NASDAQ Composite": "^IXIC",
            "Dow Jones Industrial Average": "^DJI",
            "Russell 2000": "^RUT",
            "Custom": "CUSTOM"
        }

def calculate_tracking_error(returns: pd.Series, benchmark_returns: pd.Series) -> float:
    """
    Calculate the Tracking Error of the portfolio relative to the benchmark.

    Parameters:
    - returns (pd.Series): Daily returns of the portfolio.
    - benchmark_returns (pd.Series): Daily returns of the benchmark.

    Returns:
    - float: Tracking Error.
    """
    if returns.empty or benchmark_returns.empty:
        return np.nan
    return np.std((returns - benchmark_returns)) * np.sqrt(252)

@handle_exceptions
def calculate_performance_attribution(returns: pd.Series, weights: list) -> pd.DataFrame:
    """
    Calculate performance attribution based on asset contributions.

    Parameters:
    - returns (pd.Series or pd.DataFrame): Returns of assets.
    - weights (list): Allocation weights for each asset.

    Returns:
    - pd.DataFrame: DataFrame containing asset names and their contribution percentages.
    """
    if isinstance(returns, pd.Series):
        returns = returns.to_frame('Asset')  # Convert Series to DataFrame with a default column name

    if len(weights) != len(returns.columns):
        st.error(f"Number of weights ({len(weights)}) does not match number of assets ({len(returns.columns)}).")
        return pd.DataFrame()

    annual_returns = returns.mean() * 252
    contributions = annual_returns * weights
    attribution_df = pd.DataFrame({
        'Asset': returns.columns,
        'Contribution (%)': (contributions / contributions.sum() * 100).round(2)
    }).sort_values(by='Contribution (%)', ascending=False)
    return attribution_df

def calculate_active_return(returns: pd.Series, benchmark_returns: pd.Series) -> float:
    """
    Calculate the Active Return of the portfolio relative to the benchmark.

    Parameters:
    - returns (pd.Series): Daily returns of the portfolio.
    - benchmark_returns (pd.Series): Daily returns of the benchmark.

    Returns:
    - float: Active Return percentage.
    """
    if returns.empty or benchmark_returns.empty:
        return np.nan
    return (returns.mean() - benchmark_returns.mean()) * 252 * 100

def calculate_gain_loss_ratio(returns: pd.Series) -> float:
    """
    Calculate the Gain/Loss Ratio of the portfolio.

    Parameters:
    - returns (pd.Series): Daily returns of the portfolio.

    Returns:
    - float: Gain/Loss Ratio.
    """
    gains = returns[returns > 0].sum()
    losses = -returns[returns < 0].sum()
    return gains / losses if losses != 0 else np.nan

def drawdown(cum_returns: pd.Series) -> float:
    """
    Calculate the Maximum Drawdown of the portfolio.

    Parameters:
    - cum_returns (pd.Series): Cumulative returns of the portfolio.

    Returns:
    - float: Maximum Drawdown value.
    """
    if cum_returns.empty:
        return np.nan
    peak = cum_returns.expanding(min_periods=1).max()
    dd = (cum_returns / peak) - 1
    return dd.min()

def calculate_capture_ratio(returns: pd.Series, benchmark_returns: pd.Series, upside: bool = True) -> float:
    """
    Calculate the Upside or Downside Capture Ratio.

    Parameters:
    - returns (pd.Series): Daily returns of the portfolio.
    - benchmark_returns (pd.Series): Daily returns of the benchmark.
    - upside (bool): If True, calculates Upside Capture Ratio; otherwise, Downside.

    Returns:
    - float: Capture Ratio percentage.
    """
    if returns.empty or benchmark_returns.empty:
        return np.nan
    mask = benchmark_returns > 0 if upside else benchmark_returns < 0
    if mask.sum() == 0:
        return np.nan
    portfolio = returns[mask]
    benchmark = benchmark_returns[mask]
    return (portfolio.sum() / benchmark.sum()) * 100 if benchmark.sum() != 0 else np.nan

def calculate_safe_withdrawal_rate(returns: pd.Series) -> float:
    """
    Calculate the Safe Withdrawal Rate based on portfolio returns.

    Parameters:
    - returns (pd.Series): Daily returns of the portfolio.

    Returns:
    - float: Safe Withdrawal Rate percentage.
    """
    if returns.empty:
        return np.nan
    return (returns.mean() / returns.std()) * 100

def calculate_perpetual_withdrawal_rate(returns: pd.Series) -> float:
    """
    Calculate the Perpetual Withdrawal Rate based on portfolio returns.

    Parameters:
    - returns (pd.Series): Daily returns of the portfolio.

    Returns:
    - float: Perpetual Withdrawal Rate percentage.
    """
    if returns.empty:
        return np.nan
    return (returns.mean() / returns.std()) * 100

def calculate_positive_periods(returns: pd.Series) -> str:
    """
    Calculate the number and percentage of positive return periods.

    Parameters:
    - returns (pd.Series): Daily returns of the portfolio.

    Returns:
    - str: String describing positive periods.
    """
    if returns.empty:
        return "N/A"
    positive = returns > 0
    return f"{positive.sum()} out of {len(returns)} ({(positive.sum()/len(returns))*100:.2f}%)"

def calculate_modigliani_miller(returns: pd.Series, benchmark_returns: pd.Series, rf: float = 0.02) -> float:
    """
    Calculate the Modigliani–Modigliani Measure (M2).

    Parameters:
    - returns (pd.Series): Daily returns of the portfolio.
    - benchmark_returns (pd.Series): Daily returns of the benchmark.
    - rf (float): Risk-free rate (default is 2%).

    Returns:
    - float: M2 value.
    """
    sharpe = calculate_sharpe_ratio(returns, rf)
    alpha = calculate_alpha(returns, benchmark_returns, rf)
    return alpha / sharpe if sharpe != 0 and not np.isnan(alpha) else np.nan

@handle_exceptions
def get_drawdown_details(cum_returns: pd.Series) -> list:
    """
    Generate detailed drawdown information.

    Parameters:
    - cum_returns (pd.Series): Cumulative returns of the portfolio.

    Returns:
    - list: List of dictionaries containing drawdown details.
    """
    drawdowns = []
    if cum_returns.empty:
        return drawdowns

    peak = cum_returns.iloc[0]
    peak_date = cum_returns.index[0]
    trough = cum_returns.iloc[0]
    trough_date = cum_returns.index[0]

    for date, value in cum_returns.items():
        if value > peak:
            if trough < peak:
                drawdowns.append({
                    'Start': peak_date.strftime('%b %Y'),
                    'End': trough_date.strftime('%b %Y'),
                    'Drawdown': f"{((trough / peak) - 1) * 100:.2f}%"
                })
            peak = value
            peak_date = date
            trough = value
            trough_date = date
        elif value < trough:
            trough = value
            trough_date = date

    # Final drawdown
    if trough < peak:
        drawdowns.append({
            'Start': peak_date.strftime('%b %Y'),
            'End': trough_date.strftime('%b %Y'),
            'Drawdown': f"{((trough / peak) - 1) * 100:.2f}%"
        })

    # Sort drawdowns by severity (most negative first)
    drawdowns_sorted = sorted(drawdowns, key=lambda x: float(x['Drawdown'].strip('%')), reverse=False)
    return drawdowns_sorted[:10]  # Return top 10 most severe drawdowns

@handle_exceptions
def optimize_portfolio(
    returns: pd.DataFrame,
    benchmark_returns: pd.Series = None,
    objectives: list = ['sharpe'],
    rf: float = 0.02,
    max_weight: float = 1.0,
    min_weight: float = 0.0,
    target_return: float = None
) -> np.ndarray:
    """
    Optimize a portfolio based on specified objectives.

    Parameters:
    - returns (pd.DataFrame): Historical returns of assets.
    - benchmark_returns (pd.Series, optional): Returns of the benchmark index.
    - objectives (list): Objectives to optimize.
        Options include 'sharpe', 'min_variance', 'max_return', 'min_drawdown', 'maximize_alpha', 'minimize_beta'.
    - rf (float): Risk-free rate for alpha calculation.
    - max_weight (float): Maximum weight per asset.
    - min_weight (float): Minimum weight per asset.
    - target_return (float, optional): Target return for the portfolio.

    Returns:
    - np.ndarray: Optimized asset weights or None if optimization fails.
    """
    from scipy.optimize import minimize

    def calculate_beta(portfolio_returns: pd.Series, benchmark_returns: pd.Series) -> float:
        covariance = np.cov(portfolio_returns, benchmark_returns)
        variance = covariance[1, 1]
        return covariance[0, 1] / variance if variance != 0 else 0

    def calculate_var(portfolio_returns: np.ndarray, confidence_level: float = 0.95) -> float:
        return np.percentile(portfolio_returns, (1 - confidence_level) * 100)

    def calculate_cvar(portfolio_returns: np.ndarray, confidence_level: float = 0.95) -> float:
        var = calculate_var(portfolio_returns, confidence_level)
        return portfolio_returns[portfolio_returns <= var].mean()

    # Define individual objective functions
    def sharpe_ratio(weights: np.ndarray) -> float:
        portfolio_return = np.dot(returns.mean(), weights) * 252
        portfolio_volatility = np.sqrt(np.dot(weights.T, np.dot(returns.cov() * 252, weights)))
        return -portfolio_return / portfolio_volatility if portfolio_volatility > 1e-6 else 0  # Negative for maximization

    def min_variance(weights: np.ndarray) -> float:
        return np.dot(weights.T, np.dot(returns.cov() * 252, weights))

    def max_return(weights: np.ndarray) -> float:
        return -np.dot(returns.mean(), weights) * 252  # Negative for maximization

    def min_drawdown(weights: np.ndarray) -> float:
        portfolio_returns = returns.dot(weights)
        cvar = calculate_cvar(portfolio_returns, 0.95)
        return cvar  # Minimizing CVaR as a proxy for drawdown

    # Initialize list of objective functions
    objective_functions = []

    # Map objectives to functions
    for obj in objectives:
        if obj == 'sharpe':
            objective_functions.append(sharpe_ratio)
        elif obj == 'min_variance':
            objective_functions.append(min_variance)
        elif obj == 'max_return':
            objective_functions.append(max_return)
        elif obj == 'min_drawdown':
            objective_functions.append(min_drawdown)
        elif obj == 'maximize_alpha':
            if benchmark_returns is None:
                st.error("benchmark_returns must be provided for 'maximize_alpha' objective.")
                return None

            benchmark_return = benchmark_returns.mean() * 252  # Assuming daily returns annualized

            def maximize_alpha(weights: np.ndarray) -> float:
                portfolio_return = np.dot(returns.mean(), weights) * 252
                portfolio_volatility = np.sqrt(np.dot(weights.T, np.dot(returns.cov() * 252, weights)))
                portfolio_returns = returns.dot(weights)
                beta = calculate_beta(portfolio_returns, benchmark_returns)
                alpha = portfolio_return - (rf + beta * (benchmark_return - rf))
                return -alpha  # Negative for maximization

            objective_functions.append(maximize_alpha)
        elif obj == 'minimize_beta':
            if benchmark_returns is None:
                st.error("benchmark_returns must be provided for 'minimize_beta' objective.")
                return None

            def minimize_beta(weights: np.ndarray) -> float:
                portfolio_returns = returns.dot(weights)
                beta = calculate_beta(portfolio_returns, benchmark_returns)
                return beta

            objective_functions.append(minimize_beta)
        else:
            st.error(f"Invalid optimization objective: {obj}")
            return None

    # Composite objective function: weighted sum of individual objectives
    def composite_objective(weights: np.ndarray) -> float:
        return sum(fn(weights) for fn in objective_functions)

    # Define constraints
    constraints = [
        {'type': 'eq', 'fun': lambda x: np.sum(x) - 1},
    ]

    if target_return is not None:
        def target_return_constraint(x: np.ndarray) -> float:
            return np.dot(x, returns.mean()) * 252 - target_return
        constraints.append({
            'type': 'ineq',
            'fun': target_return_constraint
        })

    # Define bounds for weights
    bounds = tuple((min_weight, max_weight) for _ in range(returns.shape[1]))

    # Initial guess (equally distributed weights)
    initial_guess = np.array([1.0 / returns.shape[1]] * returns.shape[1])

    # Perform optimization
    result = minimize(
        composite_objective,
        initial_guess,
        method='SLSQP',
        bounds=bounds,
        constraints=constraints
    )

    if result.success:
        optimized_weights = result.x
        portfolio_volatility = np.sqrt(np.dot(optimized_weights.T, np.dot(returns.cov() * 252, optimized_weights)))
        portfolio_return = np.dot(returns.mean(), optimized_weights) * 252
        sharpe_ratio_val = (portfolio_return - rf) / portfolio_volatility if portfolio_volatility > 0 else np.nan

        if portfolio_volatility <= 1e-6:
            st.error("Optimized portfolio has near-zero volatility. Optimization constraints may be too restrictive.")
            return None

        # Store optimized results in session state
        st.session_state.backtest_results['optimized_weights'] = optimized_weights
        st.session_state.backtest_results['optimized_return'] = portfolio_return
        st.session_state.backtest_results['optimized_volatility'] = portfolio_volatility
        st.session_state.backtest_results['optimized_sharpe'] = sharpe_ratio_val

        return optimized_weights
    else:
        st.error("Optimization failed. Try adjusting your constraints or target return.")
        return None

def monte_carlo_simulation(
    returns: pd.DataFrame,
    num_simulations: int = 1000,
    periods: int = 252,
    mean_returns: pd.Series = None,
    cov_matrix: pd.DataFrame = None,
    mean_reversion: bool = False,
    mean_reversion_speed: float = 0.1,
    long_term_mean: pd.Series = None,
    time_varying_vol: bool = False,
    vol_change_rate: float = 0.0,
    stress_shocks: dict = None,
    stress_period: int = None,
    return_distribution: str = 'normal'
) -> np.ndarray:
    """
    Perform Monte Carlo simulations for portfolio returns.

    Parameters:
    - returns (pd.DataFrame): Historical returns of assets.
    - num_simulations (int): Number of simulation paths.
    - periods (int): Number of periods to simulate.
    - mean_returns (pd.Series, optional): Mean returns of assets.
    - cov_matrix (pd.DataFrame, optional): Covariance matrix of asset returns.
    - mean_reversion (bool): If True, applies mean reversion.
    - mean_reversion_speed (float): Speed of mean reversion.
    - long_term_mean (pd.Series, optional): Long-term mean for mean reversion.
    - time_varying_vol (bool): If True, allows volatility to change over time.
    - vol_change_rate (float): Rate at which volatility changes.
    - stress_shocks (dict, optional): Dict of shocks to apply at specified period.
    - stress_period (int, optional): Period at which to apply stress shocks.
    - return_distribution (str): Distribution type ('normal' or 'log-normal').

    Returns:
    - np.ndarray: Array of simulated cumulative returns.
    """
    if mean_returns is None:
        mean_returns = returns.mean()
    if cov_matrix is None:
        cov_matrix = returns.cov()

    num_assets = returns.shape[1]
    available_selected = returns.columns.tolist()

    def simulate():
        weights = np.random.random(num_assets)
        weights /= np.sum(weights)
        simulated_prices = [1]  # Initial price
        current_mean = mean_returns.copy()
        current_vol = np.sqrt(np.diag(cov_matrix))
        for t in range(periods):
            if mean_reversion and long_term_mean is not None:
                current_mean += mean_reversion_speed * (long_term_mean - current_mean)
            if time_varying_vol:
                current_vol += vol_change_rate
                current_vol = np.clip(current_vol, 0, None)  # Ensure volatility doesn't go negative

            adjusted_cov_matrix = np.outer(current_vol, current_vol) * np.corrcoef(returns.T)
            if return_distribution == 'log-normal':
                simulated_returns = np.random.lognormal(mean=np.log(1 + current_mean), sigma=current_vol) - 1
            else:
                simulated_returns = np.random.multivariate_normal(current_mean, adjusted_cov_matrix)

            # Apply stress shocks if applicable
            if stress_shocks and t == stress_period:
                for idx, ticker in enumerate(available_selected):
                    simulated_returns[idx] += stress_shocks.get(ticker, 0)

            portfolio_return = np.dot(simulated_returns, weights)
            simulated_prices.append(simulated_prices[-1] * (1 + portfolio_return))
        cumulative_return = simulated_prices[-1] - 1
        return cumulative_return

    portfolio_returns = Parallel(n_jobs=-1)(
        delayed(simulate)() for _ in range(num_simulations)
    )

    return np.array(portfolio_returns)

@handle_exceptions
def plot_efficient_frontier(returns: pd.DataFrame, num_portfolios: int = 1000, rf: float = 0.0, portfolio_name='Portfolio') -> None:
    """
    Plot an enhanced Efficient Frontier with optimized portfolio marked.

    Parameters:
    - returns (pd.DataFrame): Historical returns of assets.
    - num_portfolios (int): Number of portfolios to simulate.
    - rf (float): Risk-free rate for Sharpe Ratio calculation.
    - portfolio_name (str): Name of the portfolio being optimized.
    """
    from scipy.optimize import minimize

    num_assets = returns.shape[1]

    def generate_portfolio_metrics(weights: np.ndarray) -> tuple:
        portfolio_return = np.dot(returns.mean(), weights) * 252
        portfolio_volatility = np.sqrt(np.dot(weights.T, np.dot(returns.cov() * 252, weights)))
        sharpe_ratio = (portfolio_return - rf) / portfolio_volatility if portfolio_volatility > 0 else 0
        return portfolio_volatility, portfolio_return, sharpe_ratio

    # Generate random portfolios
    weights_list = np.random.dirichlet(np.ones(num_assets), num_portfolios)
    portfolio_metrics = np.array([generate_portfolio_metrics(w) for w in weights_list])

    # Filter out invalid portfolios
    portfolio_metrics = portfolio_metrics[~np.isnan(portfolio_metrics).any(axis=1)]

    if portfolio_metrics.size == 0:
        st.warning("No valid portfolios to plot on the Efficient Frontier.")
        return

    ef_df = pd.DataFrame({
        'Std Dev': portfolio_metrics[:, 0],
        'Return': portfolio_metrics[:, 1],
        'Sharpe Ratio': portfolio_metrics[:, 2]
    })
    fig = px.scatter(
        ef_df,
        x='Std Dev',
        y='Return',
        color='Sharpe Ratio',
        color_continuous_scale='Turbo',
        title='📈 Efficient Frontier',
        labels={
            'Std Dev': 'Annualized Volatility (Std Dev)',
            'Return': 'Annualized Return (%)',
            'Sharpe Ratio': 'Sharpe Ratio'
        },
        hover_data={
            'Sharpe Ratio': ':.2f',
            'Std Dev': ':.2f',
            'Return': ':.2f'
        },
        template=THEME
    )

    # Highlight the portfolio with the maximum Sharpe ratio
    max_sharpe_idx = ef_df['Sharpe Ratio'].idxmax()
    max_sharpe = ef_df.loc[max_sharpe_idx]
    fig.add_trace(go.Scatter(
        x=[max_sharpe['Std Dev']],
        y=[max_sharpe['Return']],
        mode='markers+text',
        marker=dict(color='gold', size=12, symbol='star'),
        name='Max Sharpe Ratio',
        text=["Max Sharpe"],
        textposition="top center",
        hoverinfo='text'
    ))
    fig.add_annotation(
        x=max_sharpe['Std Dev'],
        y=max_sharpe['Return'],
        text="Max Sharpe",
        showarrow=True,
        arrowhead=1,
        ax=0,
        ay=-40
    )

    # Add optimized portfolio if available
    optimized_weights = st.session_state.backtest_results.get('optimized_weights', None)
    optimized_return = st.session_state.backtest_results.get('optimized_return', None)
    optimized_volatility = st.session_state.backtest_results.get('optimized_volatility', None)
    optimized_sharpe = st.session_state.backtest_results.get('optimized_sharpe', None)

    if all(v is not None for v in [optimized_weights, optimized_return, optimized_volatility, optimized_sharpe]):
        if optimized_volatility > 1e-6 and np.isfinite(optimized_sharpe):
            fig.add_trace(go.Scatter(
                x=[optimized_volatility],
                y=[optimized_return],
                mode='markers+text',
                marker=dict(color='red', size=16, symbol='diamond'),
                name=f'{portfolio_name} Optimized',
                text=["Optimized"],
                textposition="top center",
                hoverinfo='text'
            ))
            fig.add_annotation(
                x=optimized_volatility,
                y=optimized_return,
                text="Optimized",
                showarrow=True,
                arrowhead=2,
                ax=0,
                ay=-50
            )

    fig.update_layout(
        hovermode=HOVERMODE,
        xaxis=dict(showgrid=True, title=dict(text='Annualized Volatility (Std Dev)')),
        yaxis=dict(showgrid=True, title=dict(text='Annualized Return (%)')),
        legend=dict(title="Portfolio", x=0.01, y=0.99)
    )
    fig.update_traces(marker=dict(line=dict(width=1, color='DarkSlateGrey')), selector=dict(mode='markers'))
    st.plotly_chart(fig, use_container_width=True, config={'responsive': True, 'scrollZoom': True})

@handle_exceptions
def generate_portfolio(returns: pd.DataFrame, rf: float = 0.0) -> tuple:
    """
    Generate a random portfolio and calculate its metrics.

    Parameters:
    - returns (pd.DataFrame): Historical returns of assets.
    - rf (float): Risk-free rate for Sharpe Ratio calculation.

    Returns:
    - tuple: (volatility: float, return: float, sharpe_ratio: float)
    """
    try:
        weights = np.random.dirichlet(np.ones(returns.shape[1]))
        portfolio_return = np.dot(returns.mean(), weights) * 252
        portfolio_volatility = np.sqrt(np.dot(weights.T, np.dot(returns.cov() * 252, weights)))
        sharpe_ratio = (portfolio_return - rf) / portfolio_volatility if portfolio_volatility > 0 else 0
        return portfolio_volatility, portfolio_return, sharpe_ratio
    except Exception as e:
        st.error(f"Error generating portfolio: {e}")
        return np.nan, np.nan, np.nan  # Return NaNs to indicate failure

def plot_allocation_pie(weights: list, assets: list, title: str = '🥧 Portfolio Allocation', hover_info: str = "percent+name") -> None:
    """
    Plots the portfolio allocation pie chart.
    
    Parameters:
    - weights (list): Allocation weights for each asset.
    - assets (list): List of asset ticker symbols.
    - title (str): Title of the pie chart.
    - hover_info (str): Information to display on hover ('percent+name', etc.).
    
    Returns:
    - None
    """
    allocation_df = pd.DataFrame({
        'Asset': assets,
        'Weight': weights
    })
    
    # Determine hover data based on hover_info
    if hover_info == "percent+name":
        hover_data = {'Weight': ':.2f%'}
    elif hover_info == "name":
        hover_data = {}
    else:
        hover_data = {'Weight': ':.2f%'}
    
    fig = px.pie(
        allocation_df,
        names='Asset',
        values='Weight',
        title=title,
        hover_data=hover_data,
        hole=0.3  # Donut chart appearance
    )
    
    fig.update_traces(textposition='inside', textinfo='percent+label')
    fig.update_layout(
        uniformtext_minsize=12,
        uniformtext_mode='hide',
        legend_title="Assets",
        showlegend=True,
        height=FIG_HEIGHT
    )
    
    st.plotly_chart(fig, use_container_width=True, height=600, config={'responsive': True})

def plot_rolling_cagr(cum_returns: pd.Series, window: int = 252) -> None:
    """
    Plot Rolling CAGR over a specified window.

    Parameters:
    - cum_returns (pd.Series): Cumulative returns of the portfolio.
    - window (int): Window size for rolling CAGR calculation.
    """
    try:
        rolling_years = window / 252  # Assuming daily data
        rolling_cagr = (cum_returns / cum_returns.shift(window)) ** (1 / rolling_years) - 1
        fig = px.line(
            rolling_cagr.dropna(),
            title=f'Rolling {int(rolling_years)}-Year CAGR',
            labels={'index': 'Date', 'value': 'Rolling CAGR'},
            template=THEME
        )
        st.plotly_chart(fig, use_container_width=True)
    except Exception as e:
        st.error(f"Error plotting rolling CAGR: {e}")

@handle_exceptions
def plot_rolling_metrics(returns: pd.Series, windows: list = [30, 90, 180, 252]) -> None:
    """
    Plot rolling metrics (Volatility and Sharpe Ratio) over multiple windows.

    Parameters:
    - returns (pd.Series): Daily returns of the portfolio.
    - windows (list): List of window sizes for rolling calculations.
    """
    metrics = {}
    for window in windows:
        rolling_return = returns.rolling(window).mean() * 252
        rolling_vol = returns.rolling(window).std() * np.sqrt(252)
        rolling_sharpe = (rolling_return - 0.02) / rolling_vol
        metrics[f'Rolling {window}-Day Volatility'] = rolling_vol
        metrics[f'Rolling {window}-Day Sharpe Ratio'] = rolling_sharpe
    fig = go.Figure()
    for metric_name, metric_series in metrics.items():
        fig.add_trace(go.Scatter(x=metric_series.index, y=metric_series, mode='lines', name=metric_name))
    fig.update_layout(
        title='Rolling Metrics Over User-Defined Periods',
        xaxis_title='Date',
        yaxis_title='Value',
        hovermode='x unified',
        template=THEME
    )
    st.plotly_chart(fig, use_container_width=True)

@handle_exceptions
def plot_risk_return_attribution(returns: pd.Series, weights: list) -> None:
    """
    Plot Risk-Return Attribution for the portfolio.

    Parameters:
    - returns (pd.Series or pd.DataFrame): Returns of assets.
    - weights (list): Allocation weights for each asset.
    """
    try:
        if isinstance(returns, pd.Series):
            asset_returns = returns.mean() * 252
            asset_volatility = returns.std() * np.sqrt(252)
            asset_contribution = weights[0] * asset_volatility if len(weights) > 0 else 0
            hover_names = [returns.name] if returns.name else ['Asset']
            x = [asset_volatility]
            y = [asset_returns]
            size = [asset_contribution]
            color = [weights[0]]
        else:
            asset_returns = returns.mean() * 252
            asset_volatility = returns.std() * np.sqrt(252)
            asset_contribution = np.array(weights) * asset_volatility
            hover_names = returns.columns.tolist()
            x = asset_volatility.values
            y = asset_returns.values
            size = asset_contribution
            color = weights

        fig = px.scatter(
            x=x,
            y=y,
            size=size,
            color=color,
            hover_name=hover_names,
            title='Risk-Return Attribution',
            labels={'x': 'Annualized Volatility', 'y': 'Annualized Return', 'color': 'Weight (%)'},
            size_max=60,
            color_continuous_scale='Viridis'
        )
        fig.update_layout(template='plotly_dark')
        st.plotly_chart(fig, use_container_width=True)
    except Exception as e:
        st.error(f"Error plotting risk-return attribution: {e}")

def plot_var_cvar_distribution(portfolio_returns: np.ndarray, var: float, cvar: float) -> None:
    """
    Plot the distribution of portfolio returns with VaR and CVaR lines.

    Parameters:
    - portfolio_returns (np.ndarray): Simulated portfolio returns.
    - var (float): Value at Risk.
    - cvar (float): Conditional Value at Risk.
    """
    fig = px.histogram(portfolio_returns, nbins=50, title='Returns Distribution with VaR and CVaR',
                       labels={'value': 'Returns', 'count': 'Frequency'})
    fig.add_vline(x=var, line_dash="dash", line_color="red", annotation_text=f"VaR: {var:.2f}", annotation_position="top left")
    fig.add_vline(x=cvar, line_dash="dash", line_color="blue", annotation_text=f"CVaR: {cvar:.2f}", annotation_position="top left")
    st.plotly_chart(fig, use_container_width=True)

def plot_var_cvar_over_time(portfolio_returns: pd.Series, var: float, cvar: float) -> None:
    """
    Plot cumulative returns with VaR and CVaR over time.

    Parameters:
    - portfolio_returns (pd.Series): Daily returns of the portfolio.
    - var (float): Value at Risk.
    - cvar (float): Conditional Value at Risk.
    """
    cumulative_returns = (1 + portfolio_returns).cumprod()
    fig = px.line(cumulative_returns, title='Cumulative Returns with VaR and CVaR Over Time', labels={'value': 'Cumulative Returns', 'index': 'Date'})
    fig.add_hline(y=var, line_dash="dash", line_color="red", annotation_text=f"VaR: {var:.2f}", annotation_position="bottom right")
    fig.add_hline(y=cvar, line_dash="dash", line_color="blue", annotation_text=f"CVaR: {cvar:.2f}", annotation_position="bottom right")
    st.plotly_chart(fig, use_container_width=True)

def calculate_var(returns: np.ndarray, confidence_level: float = 0.95) -> float:
    """
    Calculate Value at Risk (VaR) for a given set of returns.

    Parameters:
    - returns (np.ndarray): Portfolio returns.
    - confidence_level (float): Confidence level for VaR.

    Returns:
    - float: VaR value.
    """
    return np.percentile(returns, 100 * (1 - confidence_level))

def calculate_cvar(returns: np.ndarray, confidence_level: float = 0.95) -> float:
    """
    Calculate Conditional Value at Risk (CVaR) for a given set of returns.

    Parameters:
    - returns (np.ndarray): Portfolio returns.
    - confidence_level (float): Confidence level for CVaR.

    Returns:
    - float: CVaR value.
    """
    var = calculate_var(returns, confidence_level)
    return returns[returns <= var].mean()

# ----------------------------
# Caching Functions
# ----------------------------

@st.cache_data(show_spinner=False, persist=True)
def get_tickers() -> pd.DataFrame:
    """
    Fetch and combine S&P 500 and NASDAQ-100 tickers from Wikipedia.

    Returns:
    - pd.DataFrame: DataFrame containing combined tickers and company names.
    """
    try:
        # Fetch S&P 500 companies
        sp500_url = 'https://en.wikipedia.org/wiki/List_of_S%26P_500_companies'
        sp500_response = requests.get(sp500_url, verify=False)
        sp500_table = pd.read_html(sp500_response.text)[0]
        sp500 = sp500_table[['Symbol', 'Security']].rename(columns={'Symbol': 'Ticker', 'Security': 'Company Name'})
        sp500['Ticker'] = sp500['Ticker'].str.replace('.', '-', regex=False)

        # Fetch NASDAQ-100 companies
        nasdaq100_url = 'https://en.wikipedia.org/wiki/NASDAQ-100'
        nasdaq100_response = requests.get(nasdaq100_url, verify=False)
        nasdaq100_tables = pd.read_html(nasdaq100_response.text)

        nasdaq100 = pd.DataFrame()
        for table in nasdaq100_tables:
            if 'Ticker' in table.columns and 'Company' in table.columns:
                nasdaq100 = table[['Ticker', 'Company']].rename(columns={'Ticker': 'Ticker', 'Company': 'Company Name'})
                nasdaq100['Ticker'] = nasdaq100['Ticker'].str.replace('.', '-', regex=False)
                break

        combined = pd.concat([sp500, nasdaq100], ignore_index=True)
        combined = combined.drop_duplicates(subset=['Ticker'])
        return combined.sort_values('Ticker').reset_index(drop=True)
    except Exception as e:
        st.error(f"Error fetching tickers: {e}")
        return pd.DataFrame(columns=['Ticker', 'Company Name'])

@st.cache_data(show_spinner=False, persist=True)
def download_data(tickers: list, start: str, end: str, retries: int = 3, backoff_factor: float = 0.3) -> pd.DataFrame:
    """
    Download historical adjusted closing prices for specified tickers.

    Parameters:
    - tickers (list): List of ticker symbols.
    - start (str): Start date in 'YYYY-MM-DD' format.
    - end (str): End date in 'YYYY-MM-DD' format.
    - retries (int): Number of retries for failed requests.
    - backoff_factor (float): Backoff factor for retries.

    Returns:
    - pd.DataFrame: DataFrame containing adjusted closing prices.
    """
    try:
        # Configure retry strategy for requests (used by yfinance internally)
        session = requests.Session()
        retry = Retry(
            total=retries,
            read=retries,
            connect=retries,
            backoff_factor=backoff_factor,
            status_forcelist=(500, 502, 504),
        )
        adapter = HTTPAdapter(max_retries=retry)
        session.mount('http://', adapter)
        session.mount('https://', adapter)

        # Attempt to download data
        data = yf.download(tickers, start=start, end=end, progress=False, session=session)['Adj Close']

        # Handle potential empty data
        if isinstance(data, pd.Series):
            data = data.to_frame()
        if data.empty:
            st.warning("No price data available for the selected portfolio. Please check the ticker symbols and date range.")
            return pd.DataFrame()

        if data.index.tz is not None:
            data.index = data.index.tz_localize(None)

        # Fill missing data
        data_filled = data.fillna(method='ffill').fillna(method='bfill')

        # Check for remaining missing values
        if data_filled.isnull().values.any():
            missing_count = data_filled.isnull().sum().sum()
            st.warning(f"Data contains {missing_count} missing values after filling. Some calculations may be affected.")
            # Optionally, you can drop remaining missing values or handle them as needed
            data_filled = data_filled.dropna()

            if data_filled.empty:
                st.error("All data was missing after filling. Please revise ticker selections or date range.")
                return pd.DataFrame()

        return data_filled
    except Exception as e:
        st.error(f"Error downloading data: {e}")
        return pd.DataFrame()

def calculate_metrics(
    returns: pd.Series, 
    cum_returns: pd.Series, 
    rf: float = 0.02, 
    benchmark_returns: pd.Series = None
) -> dict:
    """
    Calculate comprehensive portfolio metrics, optionally relative to a benchmark.

    Parameters:
    - returns (pd.Series): Portfolio returns.
    - cum_returns (pd.Series): Cumulative returns of the portfolio.
    - rf (float): Risk-free rate.
    - benchmark_returns (pd.Series, optional): Benchmark returns.

    Returns:
    - dict: Dictionary of calculated metrics.
    """
    metrics = {}
    metrics['Start Balance'] = "$10,000.00"
    if not cum_returns.empty and not returns.empty:
        end_balance = 10000 * cum_returns.iloc[-1]
        years = (cum_returns.index[-1] - cum_returns.index[0]).days / 365.25
        try:
            cagr = (end_balance / 10000) ** (1 / years) - 1 if years > 0 else np.nan
        except:
            cagr = np.nan
        metrics.update({
            'End Balance': f"${end_balance:,.2f}",
            'Annualized Return (CAGR)': f"{cagr * 100:.2f}%" if not np.isnan(cagr) else "N/A",
            'Best Year': f"{returns.resample('Y').apply(lambda x: (1 + x).prod() - 1).max() * 100:.2f}%" if not returns.empty else "N/A",
            'Worst Year': f"{returns.resample('Y').apply(lambda x: (1 + x).prod() - 1).min() * 100:.2f}%" if not returns.empty else "N/A",
            'Arithmetic Mean (Monthly)': f"{returns.mean() * 100:.2f}%",
            'Arithmetic Mean (Annualized)': f"{returns.mean() * 252 * 100:.2f}%",
            'Geometric Mean (Monthly)': f"{(np.exp(np.log1p(returns).mean()) - 1) * 100:.2f}%" if not returns.empty else "N/A",
            'Geometric Mean (Annualized)': f"{(np.exp(np.log1p(returns).mean() * 252) - 1) * 100:.2f}%" if not returns.empty else "N/A",
            'Standard Deviation (Monthly)': f"{returns.std() * 100:.2f}%",
            'Standard Deviation (Annualized)': f"{returns.std() * np.sqrt(252) * 100:.2f}%",
            'Downside Deviation (Monthly)': f"{returns[returns < 0].std() * 100:.2f}%" if not returns.empty else "N/A",
            'Maximum Drawdown': f"{drawdown(cum_returns) * 100:.2f}%" if not cum_returns.empty else "N/A",
            'Sharpe Ratio': f"{calculate_sharpe_ratio(returns, rf):.2f}" if not cum_returns.empty else "N/A",
            'Sortino Ratio': f"{calculate_sortino_ratio(returns, rf):.2f}" if not cum_returns.empty else "N/A",
            'Gain/Loss Ratio': f"{calculate_gain_loss_ratio(returns):.2f}" if not cum_returns.empty else "N/A",
            'Skewness': f"{skew(returns):.2f}" if not returns.empty else "N/A",
            'Excess Kurtosis': f"{kurtosis(returns):.2f}" if not returns.empty else "N/A",
            'Safe Withdrawal Rate': f"{calculate_safe_withdrawal_rate(returns):.6f}%" if not returns.empty else "N/A",
            'Perpetual Withdrawal Rate': f"{calculate_perpetual_withdrawal_rate(returns):.6f}%" if not returns.empty else "N/A",
            'Positive Periods': calculate_positive_periods(returns) if not returns.empty else "N/A"
        })
    else:
        # Assign "N/A" to all metrics if cum_returns or returns are empty
        metrics.update({
            'End Balance': "N/A",
            'Annualized Return (CAGR)': "N/A",
            'Best Year': "N/A",
            'Worst Year': "N/A",
            'Arithmetic Mean (Monthly)': "N/A",
            'Arithmetic Mean (Annualized)': "N/A",
            'Geometric Mean (Monthly)': "N/A",
            'Geometric Mean (Annualized)': "N/A",
            'Standard Deviation (Monthly)': "N/A",
            'Standard Deviation (Annualized)': "N/A",
            'Downside Deviation (Monthly)': "N/A",
            'Maximum Drawdown': "N/A",
            'Sharpe Ratio': "N/A",
            'Sortino Ratio': "N/A",
            'Gain/Loss Ratio': "N/A",
            'Skewness': "N/A",
            'Excess Kurtosis': "N/A",
            'Safe Withdrawal Rate': "N/A",
            'Perpetual Withdrawal Rate': "N/A",
            'Positive Periods': "N/A"
        })
    
    if benchmark_returns is not None and not benchmark_returns.empty:
        correlation = returns.corr(benchmark_returns)
        beta = calculate_beta(returns, benchmark_returns)
        alpha = calculate_alpha(returns, benchmark_returns, rf)
        r2 = calculate_r_squared(returns, benchmark_returns)
        treynor = calculate_treynor_ratio(returns, benchmark_returns, rf)
        calmar = calculate_calmar_ratio(returns, cum_returns)
        m2 = calculate_modigliani_miller(returns, benchmark_returns, rf)
        info_ratio = calculate_information_ratio(returns, benchmark_returns)
        tracking_error = calculate_tracking_error(returns, benchmark_returns)
        active_return = calculate_active_return(returns, benchmark_returns)
        upside_capture = calculate_capture_ratio(returns, benchmark_returns, upside=True)
        downside_capture = calculate_capture_ratio(returns, benchmark_returns, upside=False)
        
        metrics.update({
            'Benchmark Correlation': f"{correlation:.2f}" if not np.isnan(correlation) else "N/A",
            'Beta': f"{beta:.2f}" if not np.isnan(beta) else "N/A",
            'Alpha (annualized)': f"{alpha * 100:.2f}%" if not np.isnan(alpha) else "N/A",
            'R2': f"{r2 * 100:.2f}%" if not np.isnan(r2) else "N/A",
            'Treynor Ratio': f"{treynor:.2f}" if not np.isnan(treynor) else "N/A",
            'Calmar Ratio': f"{calmar:.2f}" if not np.isnan(calmar) else "N/A",
            'Modigliani–Modigliani Measure': f"{m2 * 100:.2f}%" if not np.isnan(m2) else "N/A",
            'Information Ratio': f"{info_ratio:.2f}" if not np.isnan(info_ratio) else "N/A",
            'Tracking Error': f"{tracking_error * 100:.2f}%" if not np.isnan(tracking_error) else "N/A",
            'Active Return': f"{active_return:.2f}%" if not np.isnan(active_return) else "N/A",
            'Upside Capture Ratio': f"{upside_capture:.2f}%" if not np.isnan(upside_capture) else "N/A",
            'Downside Capture Ratio': f"{downside_capture:.2f}%" if not np.isnan(downside_capture) else "N/A"
        })
    else:
        # Assign "N/A" to all benchmark metrics
        metrics.update({
            'Benchmark Correlation': "N/A",
            'Beta': "N/A",
            'Alpha (annualized)': "N/A",
            'R2': "N/A",
            'Treynor Ratio': "N/A",
            'Calmar Ratio': "N/A",
            'Modigliani–Modigliani Measure': "N/A",
            'Information Ratio': "N/A",
            'Tracking Error': "N/A",
            'Active Return': "N/A",
            'Upside Capture Ratio': "N/A",
            'Downside Capture Ratio': "N/A"
        })
    
    return metrics

# ----------------------------
# Additional Helper Functions
# ----------------------------

@handle_exceptions
def plot_efficient_frontier_comparison(
    simulated_metrics: pd.DataFrame,
    optimized_metrics: tuple
) -> None:
    """
    Plot Efficient Frontier and mark the optimized portfolio.

    Parameters:
    - simulated_metrics (pd.DataFrame): Simulated portfolio metrics.
    - optimized_metrics (tuple): Metrics of the optimized portfolio.
    """
    if simulated_metrics.empty:
        st.warning("No simulated metrics available to plot.")
        return

    fig = px.scatter(
        simulated_metrics,
        x='Std Dev',
        y='Return',
        color='Sharpe Ratio',
        color_continuous_scale='Viridis',
        title='Efficient Frontier',
        labels={
            'Std Dev': 'Annualized Volatility (Std Dev)',
            'Return': 'Annualized Return',
            'Sharpe Ratio': 'Sharpe Ratio'
        },
        hover_data={
            'Sharpe Ratio': ':.2f',
            'Std Dev': ':.2f',
            'Return': ':.2f'
        },
        template=THEME
    )

    # Add optimized portfolio if available
    if optimized_metrics:
        vol, ret, sharpe = optimized_metrics
        if vol > 0 and np.isfinite(sharpe):
            fig.add_trace(go.Scatter(
                x=[vol],
                y=[ret],
                mode='markers+text',
                marker=dict(color='red', size=14, symbol='star'),
                name='Optimized Portfolio',
                text=["Optimized"],
                textposition="top center",
                hoverinfo='text'
            ))
            fig.add_annotation(
                x=vol,
                y=ret,
                text="Optimized",
                showarrow=True,
                arrowhead=2,
                ax=0,
                ay=-40
            )

    fig.update_layout(
        hovermode='x unified',
        xaxis=dict(rangeslider=dict(visible=True), type='linear'),
    )
    st.plotly_chart(fig, use_container_width=True, config={
        'responsive': True,
        'scrollZoom': True,
        'displayModeBar': True
    })

# ----------------------------
# End of Helper Functions
# ----------------------------

# ----------------------------
# Streamlit Layout
# ----------------------------
st.set_page_config(page_title="🎯 Portfolio Optimizer", layout="wide")

# Display loading spinner
with st.spinner("🎯 Portfolio Optimizer is loading, please be patient..."):
    # Simulate loading delay if necessary
    pass

nav_steps = [
    {"name": "Configure Portfolio", "icon": "📁"},
    {"name": "Run Backtest", "icon": "📊"},
    {"name": "Optimize Portfolio", "icon": "🔧"},
    {"name": "Monte Carlo Simulations", "icon": "📈"},
    {"name": "Risk Analysis", "icon": "⚠️"}
]

with st.sidebar:
    st.header("📂 Navigation")
    # Create a list of (label, value) tuples for radio buttons
    nav_options = [(f"{step['icon']} {step['name']}", step['name']) for step in nav_steps]
    labels, values = zip(*nav_options)
    
    # Use radio buttons for direct navigation with a callback
    def set_step():
        selected_label = st.session_state.nav_radio  # Access the selected label from session_state
        selected_idx = labels.index(selected_label)
        st.session_state.step = values[selected_idx]
    
    selected_label = st.radio("Select a Step:", options=labels, key="nav_radio", on_change=set_step)

current_step = st.session_state.step

tickers = get_tickers()
# Define asset_options by formatting ticker symbols with company names
asset_options = tickers.apply(
    lambda row: format_asset_option(row['Ticker'], row['Company Name']), 
    axis=1
).tolist()

if current_step == "Configure Portfolio":
    st.title("🎯 Portfolio Optimizer")
    # Display success message after deletion
    if 'delete_success' in st.session_state:
        st.success(st.session_state.delete_success)
        if st.session_state.get('show_balloons', False):
            st.balloons()
            st.session_state.show_balloons = False  # Reset balloon flag
        del st.session_state.delete_success
    # Display success message after editing
    if 'edit_success' in st.session_state:
        st.success(st.session_state.edit_success)
        del st.session_state.edit_success
    
    # Initialize show_proceed in session state
    if 'show_proceed' not in st.session_state:
        st.session_state.show_proceed = False
        
    # Define callback function to hide welcome message
    def hide_welcome():
        st.session_state.show_proceed = True
        st.session_state.selected_subtab = "➕ Create New Portfolio"  # Set default tab

    # Display onboarding message and button if not proceeded
    if not st.session_state.show_proceed:
        st.info("Welcome! Let's get you started with your first portfolio.")
        st.write("""
            1. **Configure your portfolio** by entering a name, start date, end date, and other details.
            2. **Select the assets** you want to include and set their allocations.
            3. **Run a backtest** to see how your portfolio performs.
            4. **Compare your portfolio** with a benchmark or another portfolio.
            5. **Explore the performance metrics** and visualizations.
        """)
        st.button("Got U, Let's Explore", key="explore_button", on_click=hide_welcome)
    # ----------------------------
    # Configure Portfolio
    # ----------------------------

    if st.session_state.add_new_portfolio:
        if st.session_state.add_new_portfolio:
            if st.session_state.edit_portfolio:
                st.header(f"✏️ Editing {st.session_state.edit_portfolio} Portfolio")
            else:
                st.header("➕ Create New Portfolio")
        
        # Initialize or fetch the portfolio being edited
        if st.session_state.edit_portfolio:
            # Editing existing portfolio
            portfolio_to_edit = next((p for p in st.session_state.portfolios if p['name'] == st.session_state.edit_portfolio), None)
            if portfolio_to_edit:
                new_portfolio = portfolio_to_edit.copy()
                st.text_input("📛 Portfolio Name", value=new_portfolio['name'], key="portfolio_name")
                start_date_input = st.date_input(
                    "📅 Start Date:",
                    value=new_portfolio['start_date'],
                    min_value=datetime(1900, 1, 1),
                    help="Choose the start date for backtesting."
                )
                end_date_input = st.date_input(
                    "📅 End Date:",
                    value=new_portfolio['end_date'],
                    min_value=datetime(1900, 1, 1),
                    help="Choose the end date for backtesting."
                )
                rf_rate = st.number_input(
                    "📈 Risk-Free Rate (%)",
                    min_value=0.0,
                    max_value=10.0,
                    value=new_portfolio['rf_rate'] * 100,
                    format="%.2f",
                    help="Enter the risk-free rate as a percentage."
                ) / 100
                broker_fee = st.number_input(
                    "💸 Broker Fee (%)",
                    min_value=0.0,
                    max_value=10.0,
                    value=new_portfolio['broker_fee'] * 100,
                    step=0.0001,
                    format="%.4f",
                    help="Enter the broker fee as a percentage per transaction."
                ) / 100
                rebalance_freq = st.selectbox(
                    "🔄 Rebalance Frequency",
                    options=list(frequency_mapping.keys()),
                    index=list(frequency_mapping.values()).index(new_portfolio['rebalance_freq']),
                    help="Choose how often the portfolio should be rebalanced."
                )
                benchmark_symbol = st.text_input(
                    "🏦 Benchmark Symbol",
                    value=new_portfolio['benchmark_symbol'],
                    help="Enter the benchmark symbol (e.g., ^GSPC for S&P 500)."
                )
                selected_assets_new = st.multiselect(
                    "🗂️ Select Assets:",
                    options=asset_options,
                    format_func=lambda x: f"{x} - {get_company_name(tickers, x)}",
                    default=[f"{ticker} - {get_company_name(tickers, ticker)}" for ticker in new_portfolio['selected']],
                    help="Choose the assets you want to include in your portfolio."
                )
            else:
                st.error("❌ Selected portfolio to edit was not found.")
                st.session_state.add_new_portfolio = False
                st.session_state.edit_portfolio = None
                st.stop()
        else:
            # Creating a new portfolio
            st.text_input("📛 Portfolio Name", key="portfolio_name")
            start_date_input = st.date_input(
                "📅 Start Date:",
                value=datetime.today() - timedelta(days=365 * 15),
                min_value=datetime(1900, 1, 1),
                help="Choose the start date for backtesting."
            )
            end_date_input = st.date_input(
                "📅 End Date:",
                value=datetime.today(),
                min_value=datetime(1900, 1, 1),
                help="Choose the end date for backtesting."
            )
            rf_rate = st.number_input(
                "📈 Risk-Free Rate (%)",
                min_value=0.0,
                max_value=10.0,
                value=st.session_state.default_config['rf_rate'] * 100,
                format="%.2f",
                help="Enter the risk-free rate as a percentage."
            ) / 100
            broker_fee = st.number_input(
                "💸 Broker Fee (%)",
                min_value=0.0,
                max_value=10.0,
                value=st.session_state.default_config['broker_fee'] * 100,
                step=0.0001,
                format="%.4f",
                help="Enter the broker fee as a percentage per transaction."
            ) / 100
            rebalance_freq = st.selectbox(
                "🔄 Rebalance Frequency",
                options=list(frequency_mapping.keys()),
                index=list(frequency_mapping.keys()).index('Monthly'),
                help="Choose how often the portfolio should be rebalanced."
            )
            benchmark_symbol = st.text_input(
                "🏦 Benchmark Symbol",
                value=st.session_state.default_config['benchmark_symbol'],
                help="Enter the benchmark symbol (e.g., ^GSPC for S&P 500)."
            )
            selected_assets_new = st.multiselect(
                "🗂️ Select Assets:",
                options=asset_options,
                format_func=lambda x: f"{x} - {get_company_name(tickers, x)}",
                default=[],
                help="Choose the assets you want to include in your portfolio."
            )
        
        if selected_assets_new:
            # Extract tickers from selected assets
            selected_tickers_new = [option.split(' - ')[0] for option in selected_assets_new if ' - ' in option]
            # Asset Allocation
            st.subheader("💰 Asset Allocation")
            allocation_df = pd.DataFrame({
                'Ticker': selected_tickers_new,
                'Allocation (%)': [100.0 / len(selected_tickers_new)] * len(selected_tickers_new)
            })
            allocations = []
            cols = st.columns(len(selected_tickers_new))
            for idx, ticker in enumerate(selected_tickers_new):
                with cols[idx]:
                    allocation = st.number_input(
                        f"{ticker} (%)",
                        min_value=0.0,
                        max_value=100.0,
                        value=100.0 / len(selected_tickers_new),
                        step=0.1,
                        key=f"alloc_{ticker}",
                        help=f"Set the allocation percentage for {ticker}."
                    )
                    allocations.append(allocation)
                    allocation_df.at[idx, 'Allocation (%)'] = allocation
        
            total_allocation_new = sum(allocations)
            st.markdown(f"**🧮 Total Allocation:** {total_allocation_new:.2f}%")
        
            if not np.isclose(total_allocation_new, 100.0, atol=1e-2):
                st.warning("⚠️ Allocations must sum to 100%. Please adjust the allocations.")
        
        else:
            st.info("🧐 Select at least one asset to allocate your portfolio.")
        
        # Submit Button
        submit_button_label = "✅ Update Portfolio" if st.session_state.edit_portfolio else "✅ Add Portfolio"
        if st.button(submit_button_label):
            # Validation
            if not st.session_state.portfolio_name:
                st.error("❌ Please provide a portfolio name.")
            elif not selected_assets_new:
                st.error("❌ Please select at least one asset.")
            elif benchmark_symbol == "" and benchmark_symbol != "CUSTOM":
                st.error("❌ Please enter a benchmark symbol.")
            elif not np.isclose(total_allocation_new, 100.0, atol=1e-2):
                st.error(f"❌ Allocations must sum to 100%. Currently sum to {total_allocation_new:.2f}%.")
            else:
                # Prepare the new or updated portfolio
                new_portfolio = {
                    'name': st.session_state.portfolio_name,
                    'start_date': pd.to_datetime(start_date_input),
                    'end_date': pd.to_datetime(end_date_input),
                    'rf_rate': rf_rate,
                    'broker_fee': broker_fee,
                    'benchmark_symbol': benchmark_symbol,
                    'rebalance_freq': frequency_mapping.get(rebalance_freq, 'M'),
                    'selected': selected_tickers_new,
                    'allocations': allocations
                }
        
                if st.session_state.edit_portfolio:
                    # Update existing portfolio
                    for idx, p in enumerate(st.session_state.portfolios):
                        if p['name'] == st.session_state.edit_portfolio:
                            st.session_state.portfolios[idx] = new_portfolio
                            break
                    st.success(f"🎉 Portfolio '{new_portfolio['name']}' updated successfully!")
                    st.session_state.edit_portfolio = None  # Reset edit mode
                else:
                    # Add as a new portfolio
                    if new_portfolio['name'] in [p['name'] for p in st.session_state.portfolios]:
                        st.error("❌ Portfolio name already exists. Please choose a unique name.")
                    else:
                        st.session_state.portfolios.append(new_portfolio)
                        st.success(f"🎉 Portfolio '{new_portfolio['name']}' added successfully!")
        
                st.balloons()
                st.session_state.add_new_portfolio = False  # Hide the form
                st.rerun()
        
        if st.button("❌ Cancel"):
            st.session_state.add_new_portfolio = False
            st.session_state.edit_portfolio = None
            st.rerun()
    
    # Show portfolio management interface if proceeded
    elif st.session_state.show_proceed:
        # Show existing portfolios if any
        if st.session_state.portfolios:
            st.subheader("📁 Manage Existing Portfolios")
            portfolio_names = [p['name'] for p in st.session_state.portfolios]
            selected_portfolio = st.selectbox("🔍 Select a Portfolio", portfolio_names)
            
            if selected_portfolio:
                portfolio = next(p for p in st.session_state.portfolios if p['name'] == selected_portfolio)
                
                with st.expander(f"ℹ️ Details for **{selected_portfolio}**"):
                    st.write(f"**Start Date:** {portfolio['start_date'].strftime('%Y-%m-%d')}")
                    st.write(f"**End Date:** {portfolio['end_date'].strftime('%Y-%m-%d')}")
                    st.write(f"**Benchmark:** {portfolio['benchmark_symbol']}")
                    st.write(f"**Rebalance Frequency:** {rebalance_freq_inverse_mapping.get(portfolio['rebalance_freq'], portfolio['rebalance_freq'])}")
                    st.write("**Allocations:**")
                    allocation_df = pd.DataFrame({
                        'Ticker': portfolio['selected'],
                        'Allocation (%)': portfolio['allocations']
                    })
                    st.table(allocation_df)
                
                # Action Buttons with Icons
                action_col0, action_col1, action_col2, action_col3 = st.columns(4)

                with action_col0:
                    if st.button("➕ Create New Portfolio", key="create_new_portfolio"):
                        st.session_state.add_new_portfolio = True
                        st.rerun()

                with action_col1:
                    if st.button(f"🗑️ Delete {selected_portfolio}", key=f"delete_{selected_portfolio}"):
                        st.session_state.portfolios = [p for p in st.session_state.portfolios if p['name'] != selected_portfolio]
                        st.session_state.delete_success = f"✅ Portfolio '{selected_portfolio}' deleted successfully."
                        st.session_state.show_balloons = True  # Optional: To control balloon display
                        st.rerun()

                with action_col2:
                    if st.button(f"✏️ Edit {selected_portfolio}", key=f"edit_{selected_portfolio}"):
                        st.session_state.edit_portfolio = selected_portfolio
                        st.session_state.add_new_portfolio = True  # Ensure the form is visible
                        st.rerun()

                with action_col3:
                    if st.button(f"📄 Duplicate {selected_portfolio}", key=f"duplicate_{selected_portfolio}"):
                        duplicated_portfolio = portfolio.copy()
                        duplicated_portfolio['name'] = f"{selected_portfolio}_Copy"
                        # Ensure the new name is unique
                        if duplicated_portfolio['name'] in portfolio_names:
                            counter = 1
                            while f"{selected_portfolio}_Copy{counter}" in portfolio_names:
                                counter += 1
                            duplicated_portfolio['name'] = f"{selected_portfolio}_Copy{counter}"
                        st.session_state.portfolios.append(duplicated_portfolio)
                        st.success(f"🎉 Portfolio '{duplicated_portfolio['name']}' duplicated successfully!")
                        st.rerun()
        
        st.markdown("---")  # Separator line for clarity

    # ----------------------------
    # Documentation
    # ----------------------------
    with st.expander("📖 Documentation"):  
        st.markdown("""
        ### Documentation

        **Portfolio Configuration:**
        - **Portfolio Name:** Enter a unique name for your portfolio.
        - **Risk-Free Rate:** The theoretical return of an investment with zero risk, often based on government bonds.
        - **Sharpe Ratio:** Measures the performance of an investment compared to a risk-free asset, after adjusting for its risk.
            *Formula:* $$\\frac{R_p - R_f}{\\sigma_p}$$
        - **VaR (Value at Risk):** Estimates the maximum potential loss over a specific time frame at a given confidence level.
        - **CVaR (Conditional Value at Risk):** The expected loss exceeding the VaR, providing insight into tail risk.
        - **Sortino Ratio:** Similar to the Sharpe Ratio but only penalizes downside volatility.
            *Formula:* $$\\frac{R_p - R_f}{\\sigma_d}$$
        - **Maximum Drawdown:** The largest peak-to-trough decline in the portfolio's value over a specific period.
        - **Alpha:** Measures the active return on an investment compared to a market index.
            *Formula:* $$R_p - [R_f + \\beta (R_m - R_f)]$$
        - **Beta:** Indicates the volatility of an investment relative to the market.
            *Formula:* $$\\beta = \\frac{Cov(R_p, R_m)}{Var(R_m)}$$
        - **Information Ratio:** Measures portfolio returns beyond the returns of a benchmark, adjusted for the volatility of those returns.
            *Formula:* $$\\frac{R_p - R_b}{TE}$$
        - **Gain/Loss Ratio:** The ratio of total gains to total losses in the portfolio.
        - **Modigliani–Modigliani Measure (M²):** Adjusts the portfolio return to the risk of a benchmark, allowing for comparison.
            *Formula:* $$M² = \\text{Sharpe Ratio}_{\\text{Portfolio}} \\times \\sigma_{\\text{Benchmark}} + R_f$$
        - **Tracking Error:** Measures the standard deviation of the difference between portfolio returns and benchmark returns.
        - **Upside/Downside Capture Ratio:** Measures how well the portfolio captures the benchmark's positive and negative movements respectively.
        - **Risk Factor Attribution:** Decomposes portfolio returns based on exposure to different risk factors like Momentum, Value, and Size.
        - **Performance Attribution:** Breaks down portfolio performance by individual assets, showing each asset's contribution to overall returns.
        - **Recommendation Engine:** Provides suggestions for portfolio adjustments to optimize performance based on current allocations and objectives.
        """)

if current_step == "Run Backtest":
        st.title("📊 Run Backtest")
        if not st.session_state.portfolios:
            st.warning("Please add at least one portfolio to run backtest.")
        else:
            st.subheader("🔍 Select Portfolios to Compare")
            comparison_type = st.selectbox(
                "Comparison Type",
                ["Portfolio vs Benchmark", "Portfolio vs Portfolio"],
                key="comparison_type"
            )

            portfolio_names = [p['name'] for p in st.session_state.portfolios]

            # Dynamic labels based on comparison_type
            portfolio1_label = "Select Portfolio" if comparison_type == "Portfolio vs Benchmark" else "Select First Portfolio"
            portfolio2_disabled = comparison_type == "Portfolio vs Benchmark"

            portfolio1 = st.selectbox(
                portfolio1_label,
                portfolio_names,
                key="p1"
            )

            portfolio2 = st.selectbox(
                "Select Second Portfolio",
                portfolio_names,
                key="p2",
                disabled=portfolio2_disabled
            )

            st.subheader("⚙️ Visualization Settings")
            selected_time_frames = st.multiselect(
                "Select Time Frames for CAGR",
                options=['Weekly', 'Monthly', 'Quarterly', 'Annually'],
                default=['Annually'],
                key="selected_time_frames",
                help="Choose one or more time frames to view CAGR over different horizons."
            )
            selected_rolling_periods = st.multiselect(
                "Select Rolling Periods (Days)",
                options=[30, 90, 180, 252],
                default=[252],
                key="selected_rolling_periods",
                help="Choose one or more periods to view rolling metrics."
            )

            run = st.button("Run Backtest")

            if run:
                with st.spinner("Running backtest..."):
                    if comparison_type == "Portfolio vs Benchmark":
                        portfolio = next(p for p in st.session_state.portfolios if p['name'] == portfolio1)
                        result = run_portfolio_comparison("vs_benchmark", portfolio)
                    else:
                        if portfolio1 == portfolio2:
                            st.error("Please select two different portfolios for comparison.")
                            st.stop()
                        portfolio_a = next(p for p in st.session_state.portfolios if p['name'] == portfolio1)
                        portfolio_b = next(p for p in st.session_state.portfolios if p['name'] == portfolio2)
                        result = run_portfolio_comparison("vs_portfolio", portfolio_a, portfolio_b)

                    if result.get("error"):
                        st.error(result["error"])
                    else:
                        backtest_results = result["backtest_results"]
                        display_backtest_results(backtest_results, comparison_type, portfolio1, portfolio2)

if current_step == "Optimize Portfolio":
    st.title("🔧 Optimize Portfolio")
    
    if not st.session_state.portfolios:
        st.warning("Please add at least one portfolio to optimize.")
    else:
        portfolio_names = [p['name'] for p in st.session_state.portfolios]
        selected_portfolio_name = st.selectbox("📁 Select Portfolio to Optimize", portfolio_names)
        portfolio = next(p for p in st.session_state.portfolios if p['name'] == selected_portfolio_name)
        
        # Download and prepare data
        price_data = download_data(
            portfolio['selected'],
            portfolio['start_date'],
            portfolio['end_date']
        )
        
        if price_data.empty:
            st.error("No price data available for the selected portfolio.")
        else:
            available_selected = [ticker for ticker in portfolio['selected'] if ticker in price_data.columns]
            missing_selected = list(set(portfolio['selected']) - set(available_selected))
            if missing_selected:
                st.warning(f"Excluded tickers with no data: {', '.join(missing_selected)}")
            
            if not available_selected:
                st.error("No selected tickers have available data for the chosen date range.")
            else:
                returns = price_data[available_selected].pct_change().dropna()
                
                st.header("🚀 Portfolio Optimization Settings")
                
                # Constraints Section
                with st.expander("⚙️ Constraints", expanded=True):
                    constraints_col1, constraints_col2 = st.columns(2)
                    
                    with constraints_col1:
                        max_weight = st.slider(
                            "📈 Maximum Weight per Asset (%)",
                            min_value=0.0,
                            max_value=100.0,
                            value=100.0,
                            step=0.1,
                            help="Set the maximum allocation percentage for any single asset."
                        ) / 100  # Convert to decimal
                
                    with constraints_col2:
                        min_weight = st.slider(
                            "📉 Minimum Weight per Asset (%)",
                            min_value=0.0,
                            max_value=100.0,
                            value=0.0,
                            step=0.1,
                            help="Set the minimum allocation percentage for any single asset."
                        ) / 100  # Convert to decimal
                
                    # Validation: Ensure min_weight does not exceed max_weight
                    if min_weight > max_weight:
                        st.error("⚠️ **Minimum weight cannot exceed Maximum weight.** Please adjust your settings.")
                
                # Objectives Section
                with st.expander("🎯 Objectives", expanded=True):
                    objective_options = {
                        "Maximize Sharpe Ratio": "sharpe",
                        "Minimize Variance": "min_variance",
                        "Maximize Return": "max_return",
                        "Minimize Drawdown": "min_drawdown",
                        "Maximize Alpha": "maximize_alpha",
                        "Minimize Beta": "minimize_beta"
                    }
                    selected_objectives = st.multiselect(
                        "Select Optimization Objectives",
                        options=list(objective_options.keys()),
                        default=["Maximize Sharpe Ratio"],
                        help="Choose one or more objectives to optimize your portfolio."
                    )
                    
                    if not selected_objectives:
                        st.warning("⚠️ Please select at least one optimization objective.")
                    
                # Target Return Section
                with st.expander("🎯 Target Return", expanded=True):
                    target_return = st.number_input(
                        "📈 Target Annual Return (%)",
                        min_value=0.0,
                        max_value=100.0,
                        value=10.0,
                        step=0.1,
                        help="Set the target annual return for portfolio optimization."
                    ) / 100  # Convert to decimal
                
                # Optimization Trigger
                if st.button("✅ Run Optimization"):
                    if not selected_objectives:
                        st.error("❌ Please select at least one optimization objective.")
                    else:
                        # Map objectives to their internal codes
                        objective_map = {v: k for k, v in objective_options.items()}
                        objectives_selected = [objective_options[obj] for obj in selected_objectives]
                        
                        with st.spinner("🧮 Optimizing portfolio..."):
                            # Fetch benchmark returns
                            benchmark_symbol = portfolio['benchmark_symbol']
                            benchmark_data = download_data([benchmark_symbol], portfolio['start_date'], portfolio['end_date'])
                            
                            if benchmark_data.empty:
                                st.error(f"Benchmark symbol '{benchmark_symbol}' does not have available data for the chosen period.")
                            else:
                                benchmark_returns = benchmark_data[benchmark_symbol].pct_change().dropna()
                                benchmark_returns = benchmark_returns.reindex(returns.index, method='ffill').dropna()
    
                                # Run optimization
                                optimized_weights = optimize_portfolio(
                                    returns,
                                    benchmark_returns=benchmark_returns,
                                    objectives=objectives_selected,
                                    rf=portfolio['rf_rate'],
                                    max_weight=max_weight,
                                    min_weight=min_weight,
                                    target_return=target_return
                                )
    
                                if optimized_weights is not None:
                                    # Ensure allocations sum to 100%
                                    total_allocation = sum(optimized_weights)
                                    if not np.isclose(total_allocation, 1.0, atol=1e-4):
                                        st.warning(f"🔄 Allocations sum to {total_allocation*100:.2f}%. Adjusting proportionally.")
                                        optimized_weights = [w / total_allocation for w in optimized_weights]
                                    
                                    # Display Optimized Weights
                                    st.subheader("📊 Optimized Portfolio Allocation")
                                    weights_df = pd.DataFrame({
                                        'Ticker': available_selected,
                                        'Allocation (%)': [f"{w*100:.2f}" for w in optimized_weights]
                                    })
                                    st.table(weights_df)
    
                                    # Plot Allocation Pie Chart
                                    plot_allocation_pie(
                                        optimized_weights,
                                        available_selected,
                                        title="🥧 Optimized Allocation",
                                        hover_info="percent+name"
                                    )
    
                                    # Display Performance Metrics
                                    st.subheader("📈 Optimized Portfolio Performance Metrics")
                                    optimized_return = np.dot(optimized_weights, returns.mean()) * 252
                                    optimized_volatility = np.sqrt(np.dot(optimized_weights, np.dot(returns.cov() * 252, optimized_weights)))
                                    optimized_sharpe = (optimized_return - portfolio['rf_rate']) / optimized_volatility if optimized_volatility > 0 else np.nan
    
                                    optimized_metrics = {
                                        'Expected Annual Return (%)': f"{optimized_return * 100:.2f}%",
                                        'Annualized Volatility (%)': f"{optimized_volatility * 100:.2f}%",
                                        'Sharpe Ratio': f"{optimized_sharpe:.2f}"
                                    }
    
                                    optimized_metrics_df = pd.DataFrame(list(optimized_metrics.items()), columns=['Metric', 'Value'])
                                    st.table(optimized_metrics_df)
    
                                    # Store optimized weights and metrics in session state
                                    st.session_state.backtest_results.update({
                                        'optimized_weights': optimized_weights,
                                        'optimized_return': optimized_return,
                                        'optimized_volatility': optimized_volatility,
                                        'optimized_sharpe': optimized_sharpe
                                    })
    
                                    st.success("🎉 Portfolio optimization completed successfully!")
                                else:
                                    st.error("❌ Optimization did not return any weights.")
    #else:
        #st.warning("Please add at least one portfolio to optimize.")

if current_step == "Monte Carlo Simulations":
    st.title("📈 Monte Carlo Simulations")
    
    if st.session_state.portfolios:
        portfolio_names = [p['name'] for p in st.session_state.portfolios]
        selected_portfolio = st.selectbox("📁 Select Portfolio for Simulation", portfolio_names)
        portfolio = next(p for p in st.session_state.portfolios if p['name'] == selected_portfolio)
        
        # Download and prepare data
        price_data = download_data(
            portfolio['selected'],
            portfolio['start_date'],
            portfolio['end_date']
        )
        
        if price_data.empty:
            st.error("No price data available for the selected portfolio.")
        else:
            available_selected = [ticker for ticker in portfolio['selected'] if ticker in price_data.columns]
            missing_selected = list(set(portfolio['selected']) - set(available_selected))
            if missing_selected:
                st.warning(f"Excluded tickers with no data: {', '.join(missing_selected)}")
            
            if not available_selected:
                st.error("No selected tickers have available data for the chosen date range.")
            else:
                returns = price_data[available_selected].pct_change().dropna()
                
                # Organize simulation settings and results into tabs
                simulation_tabs = st.tabs(["🔧 Settings", "📊 Results"])
                
                # Settings Tab
                with simulation_tabs[0]:
                    st.header("🔧 Simulation Settings")
                    
                    # Simulation Parameters
                    sim_col1, sim_col2 = st.columns(2)
                    
                    with sim_col1:
                        num_simulations = st.number_input(
                            "🔢 Number of Simulations",
                            min_value=100,
                            max_value=10000,
                            value=1000,
                            step=100,
                            help="Define how many simulation paths to generate."
                        )
                        periods = st.number_input(
                            "📆 Number of Periods",
                            min_value=1,
                            max_value=252,
                            value=252,
                            step=1,
                            help="Set the number of periods (e.g., days) for each simulation."
                        )
                    
                    with sim_col2:
                        return_distribution = st.selectbox(
                            "📊 Return Distribution",
                            options=["Normal", "Log-Normal"],
                            index=0,
                            help="Choose the statistical distribution for asset returns."
                        )
                        enable_mean_reversion = st.checkbox("🔄 Enable Mean Reversion", help="Implement mean reversion in simulated returns.")
                    
                    # Mean Reversion Parameters
                    if enable_mean_reversion:
                        mr_col1, mr_col2 = st.columns(2)
                        with mr_col1:
                            mean_reversion_speed = st.slider(
                                "⚡ Mean Reversion Speed",
                                min_value=0.0,
                                max_value=1.0,
                                value=0.1,
                                step=0.01,
                                help="Adjust the speed at which returns revert to the mean."
                            )
                        with mr_col2:
                            long_term_mean_input = st.number_input(
                                "📈 Long-term Mean Return (%)",
                                value=0.0,
                                step=0.1,
                                help="Set the target mean return for mean reversion."
                            ) / 100  # Convert to decimal
                
                    # Time-Varying Volatility
                    enable_time_varying_vol = st.checkbox("📉 Enable Time-Varying Volatility", help="Allow volatility to change over time in simulations.")
                    if enable_time_varying_vol:
                        vol_change_rate = st.slider(
                            "📈 Volatility Change Rate",
                            min_value=-0.5,
                            max_value=0.5,
                            value=0.0,
                            step=0.01,
                            help="Set the rate at which volatility changes each period."
                        )
                
                    # Stress Testing
                    st.subheader("⚠️ Stress Testing Parameters")
                    enable_stress_testing = st.checkbox("⚠️ Enable Stress Testing")
                    stress_shocks = {}
                    if enable_stress_testing:
                        selected_stress_assets = st.multiselect(
                            "🔍 Select Assets to Apply Stress Shocks",
                            options=available_selected,
                            help="Choose assets to apply specific stress shocks."
                        )
                        for ticker in selected_stress_assets:
                            shock = st.number_input(
                                f"💥 Stress Shock to {ticker} (%)",
                                min_value=-100.0,
                                max_value=100.0,
                                value=0.0,
                                step=0.1,
                                help=f"Define the percentage shock for {ticker}."
                            ) / 100  # Convert to decimal
                            stress_shocks[ticker] = shock
                
                    # Sensitivity Analysis (Optional)
                    st.subheader("🔄 Sensitivity Analysis Parameters")
                    enable_sensitivity_analysis = st.checkbox("🔎 Enable Sensitivity Analysis", help="Assess how changes in market factors affect portfolio risk.")
                    sensitivity_adjustments = {}
                    if enable_sensitivity_analysis:
                        sensitivity_col1, sensitivity_col2 = st.columns(2)
                        with sensitivity_col1:
                            interest_rate_change = st.number_input(
                                "💹 Change in Interest Rates (bps)",
                                value=0.0,
                                step=0.1,
                                help="Specify the change in interest rates in basis points."
                            ) / 10000  # Convert to decimal
                        with sensitivity_col2:
                            inflation_rate_change = st.number_input(
                                "📈 Change in Inflation Rates (bps)",
                                value=0.0,
                                step=0.1,
                                help="Specify the change in inflation rates in basis points."
                            ) / 10000  # Convert to decimal
                        
                        for ticker in available_selected:
                            col_a, col_b = st.columns(2)
                            with col_a:
                                interest_sens = st.number_input(
                                    f"📊 Interest Rate Sensitivity for {ticker}",
                                    value=1.0,
                                    step=0.1,
                                    help=f"Set sensitivity of {ticker} to interest rate changes."
                                )
                            with col_b:
                                inflation_sens = st.number_input(
                                    f"📊 Inflation Rate Sensitivity for {ticker}",
                                    value=1.0,
                                    step=0.1,
                                    help=f"Set sensitivity of {ticker} to inflation rate changes."
                                )
                            adjustment = interest_sens * interest_rate_change + inflation_sens * inflation_rate_change
                            sensitivity_adjustments[ticker] = adjustment
                
                    # Simulation Trigger
                    if st.button("✅ Run Simulations"):
                        with st.spinner("🧮 Running Monte Carlo simulations..."):
                            simulated_returns = monte_carlo_simulation(
                                returns=returns, 
                                num_simulations=int(num_simulations), 
                                periods=int(periods),
                                mean_returns=returns.mean(), 
                                cov_matrix=returns.cov(),
                                mean_reversion=enable_mean_reversion,
                                mean_reversion_speed=mean_reversion_speed if enable_mean_reversion else 0.0,
                                long_term_mean=np.full(len(available_selected), long_term_mean_input / 252) if enable_mean_reversion else None,
                                time_varying_vol=enable_time_varying_vol,
                                vol_change_rate=vol_change_rate if enable_time_varying_vol else 0.0,
                                stress_shocks=stress_shocks if enable_stress_testing else {},
                                return_distribution=return_distribution.lower()
                            )
                            st.session_state.simulated_returns = simulated_returns
                            st.success("🎉 Simulations completed successfully!")
                
                # Results Tab
                with simulation_tabs[1]:
                    if 'simulated_returns' in st.session_state:
                        simulated_returns = st.session_state.simulated_returns
                        st.header("📊 Simulation Results")
        
                        # Descriptive Statistics
                        st.subheader("📈 Descriptive Statistics")
                        st.write(pd.Series(simulated_returns).describe())
        
                        # Histogram
                        st.subheader("📊 Returns Distribution")
                        fig_hist = px.histogram(
                            simulated_returns,
                            nbins=50,
                            title='📊 Simulated Returns Distribution',
                            labels={'value': 'Simulated Returns', 'count': 'Frequency'},
                            template=THEME
                        )
                        st.plotly_chart(fig_hist, use_container_width=True)
        
                        # Cumulative Distribution Function (CDF)
                        st.subheader("📉 Cumulative Distribution Function (CDF)")
                        fig_cdf = px.ecdf(
                            simulated_returns,
                            title='📉 Cumulative Distribution of Simulated Returns',
                            labels={'value': 'Simulated Returns', 'cumcount': 'CDF'},
                            template=THEME
                        )
                        st.plotly_chart(fig_cdf, use_container_width=True)
        
                        # Box Plot
                        st.subheader("📦 Box Plot of Simulated Returns")
                        fig_box = px.box(
                            pd.DataFrame(simulated_returns, columns=['Returns']),
                            y='Returns',
                            title='📦 Box Plot of Simulated Returns',
                            template=THEME
                        )
                        st.plotly_chart(fig_box, use_container_width=True)
        
                        # Summary Statistics Table
                        st.subheader("📝 Summary Statistics")
                        summary_stats = pd.DataFrame({
                            'Metric': ['Mean', 'Median', 'Standard Deviation', 'Minimum', 'Maximum'],
                            'Value': [
                                f"{np.mean(simulated_returns):.4f}",
                                f"{np.median(simulated_returns):.4f}",
                                f"{np.std(simulated_returns):.4f}",
                                f"{np.min(simulated_returns):.4f}",
                                f"{np.max(simulated_returns):.4f}"
                            ]
                        })
                        st.table(summary_stats)
        
                        # Interpretation
                        st.markdown("""
                        **📖 Interpretation:**
                        - **Descriptive Statistics:** Provides an overview of the distribution of simulated portfolio returns.
                        - **Histogram:** Visualizes the frequency distribution of returns.
                        - **CDF:** Shows the probability that a return is less than or equal to a particular value.
                        - **Box Plot:** Highlights the median, quartiles, and potential outliers in the return distribution.
                        - **Summary Statistics:** Summarizes key metrics from the simulation runs.
                        """)
                    else:
                        st.info("🔍 Run simulations to view results here.")
    #else:
        #st.warning("Please add at least one portfolio to run Monte Carlo simulations.")

# Main Risk Analysis Section
if current_step == "Risk Analysis":
    st.title("⚠️ Risk Analysis")
    
    if st.session_state.portfolios:
        portfolio_names = [p['name'] for p in st.session_state.portfolios]
        selected_portfolio = st.selectbox("📁 Select Portfolio for Risk Analysis", portfolio_names)
        portfolio = next(p for p in st.session_state.portfolios if p['name'] == selected_portfolio)
        
        # Download and prepare data
        price_data = download_data(
            portfolio['selected'],
            portfolio['start_date'],
            portfolio['end_date']
        )
        
        if price_data.empty:
            st.error("No price data available for the selected portfolio.")
        else:
            available_selected = [ticker for ticker in portfolio['selected'] if ticker in price_data.columns]
            missing_selected = list(set(portfolio['selected']) - set(available_selected))
            if missing_selected:
                st.warning(f"Excluded tickers with no data: {', '.join(missing_selected)}")
            
            if not available_selected:
                st.error("No selected tickers have available data for the chosen date range.")
            else:
                returns = price_data[available_selected].pct_change().dropna()
                
                # Organize risk analysis into tabs
                risk_analysis_tabs = st.tabs(["📉 Risk Metrics", "📈 Visualizations", "🛠️ Sensitivity Analysis"])
                
                # Risk Metrics Tab
                with risk_analysis_tabs[0]:
                    st.header("📉 Risk Metrics")
                    
                    # Risk Parameters
                    risk_params_col1, risk_params_col2 = st.columns(2)
                    
                    with risk_params_col1:
                        confidence_level = st.number_input(
                            "🔍 Confidence Level for VaR/CVaR (%)",
                            min_value=90.0,
                            max_value=99.0,
                            value=95.0,
                            step=1.0,
                            help="Set the confidence level for risk metrics like VaR and CVaR."
                        ) / 100  # Convert to decimal
                    
                    with risk_params_col2:
                        enable_stress_testing = st.checkbox("⚠️ Enable Stress Testing", help="Apply specific shocks to assets to assess portfolio resilience.")
                    
                    if enable_stress_testing:
                        st.markdown("#### ⚠️ Stress Testing Shocks")
                        stress_shocks = {}
                        for ticker in available_selected:
                            shock = st.number_input(
                                f"💥 Stress Shock to {ticker} (%)",
                                min_value=-100.0,
                                max_value=100.0,
                                value=0.0,
                                step=0.1,
                                help=f"Define the percentage shock for {ticker}."
                            ) / 100  # Convert to decimal
                            stress_shocks[ticker] = shock
                        if not any(shock != 0 for shock in stress_shocks.values()):
                            st.warning("⚠️ No stress shocks applied. Consider adding at least one to perform stress testing.")
                    
                    # Sensitivity Analysis Inputs
                    st.subheader("🔎 Sensitivity Analysis Parameters")
                    enable_sensitivity_analysis = st.checkbox("🔄 Enable Sensitivity Analysis", help="Assess how changes in market factors affect portfolio risk.")
                    sensitivity_adjustments = {}
                    if enable_sensitivity_analysis:
                        sensitivity_col1, sensitivity_col2 = st.columns(2)
                        with sensitivity_col1:
                            interest_rate_change = st.number_input(
                                "💹 Change in Interest Rates (bps)",
                                value=0.0,
                                step=0.1,
                                help="Specify the change in interest rates in basis points."
                            ) / 10000  # Convert to decimal
                        with sensitivity_col2:
                            inflation_rate_change = st.number_input(
                                "📈 Change in Inflation Rates (bps)",
                                value=0.0,
                                step=0.1,
                                help="Specify the change in inflation rates in basis points."
                            ) / 10000  # Convert to decimal
                        
                        for ticker in available_selected:
                            col_a, col_b = st.columns(2)
                            with col_a:
                                interest_sens = st.number_input(
                                    f"📊 Interest Rate Sensitivity for {ticker}",
                                    value=1.0,
                                    step=0.1,
                                    help=f"Set sensitivity of {ticker} to interest rate changes."
                                )
                            with col_b:
                                inflation_sens = st.number_input(
                                    f"📊 Inflation Rate Sensitivity for {ticker}",
                                    value=1.0,
                                    step=0.1,
                                    help=f"Set sensitivity of {ticker} to inflation rate changes."
                                )
                            adjustment = interest_sens * interest_rate_change + inflation_sens * inflation_rate_change
                            sensitivity_adjustments[ticker] = adjustment
                
                    # Calculate Risk Metrics Button
                    if st.button("📊 Calculate Risk Metrics"):
                        with st.spinner("🧮 Calculating risk metrics..."):
                            # Apply stress shocks if enabled
                            portfolio_returns = returns.copy()
                            if enable_stress_testing:
                                for ticker, shock in stress_shocks.items():
                                    portfolio_returns[ticker] += shock
                            
                            # Apply sensitivity adjustments if enabled
                            if enable_sensitivity_analysis:
                                for ticker, adjustment in sensitivity_adjustments.items():
                                    portfolio_returns[ticker] += adjustment
                            
                            # Ensure no extreme negative returns
                            portfolio_returns = portfolio_returns.clip(lower=-1.0)
                            
                            # Calculate weighted portfolio returns
                            weights = np.array(portfolio['allocations']) / 100
                            portfolio_returns = portfolio_returns.dot(weights)
                            
                            # Store in session state
                            st.session_state.portfolio_returns = portfolio_returns
                            
                            # Calculate metrics
                            cum_returns = (1 + portfolio_returns).cumprod()
                            var = calculate_var(portfolio_returns, confidence_level)
                            cvar = calculate_cvar(portfolio_returns, confidence_level)
                            skewness = skew(portfolio_returns)
                            kurt = kurtosis(portfolio_returns)
                            annual_vol = portfolio_returns.std() * np.sqrt(252)
                            max_drawdown_val = drawdown(cum_returns) * 100
                            
                            # Store metrics
                            risk_metrics = {
                                'Value at Risk (VaR)': f"{var * 100:.2f}%",
                                'Conditional Value at Risk (CVaR)': f"{cvar * 100:.2f}%",
                                'Skewness': f"{skewness:.2f}",
                                'Kurtosis': f"{kurt:.2f}",
                                'Annualized Volatility (%)': f"{annual_vol * 100:.2f}%",
                                'Maximum Drawdown (%)': f"{max_drawdown_val:.2f}%"
                            }
                            
                            st.session_state.risk_metrics = risk_metrics
                            st.success("✅ Risk metrics calculated successfully!")
        
                # Visualizations Tab
                with risk_analysis_tabs[1]:
                    st.header("📈 Risk Visualizations")
                    
                    if 'risk_metrics' in st.session_state:
                        var = st.session_state.risk_metrics.get('Value at Risk (VaR)', None)
                        cvar = st.session_state.risk_metrics.get('Conditional Value at Risk (CVaR)', None)
                        portfolio_returns_local = st.session_state.portfolio_returns
                        cum_returns = (1 + portfolio_returns_local).cumprod()
                        
                        # VaR and CVaR Plot
                        st.subheader("📉 Cumulative Returns with VaR and CVaR")
                        fig_cum = px.line(
                            cum_returns,
                            x=cum_returns.index,
                            y=cum_returns,
                            title='📉 Cumulative Returns Over Time',
                            labels={'y': 'Cumulative Returns', 'x': 'Date'},
                            template=THEME
                        )
                        if var is not None and cvar is not None:
                            fig_cum.add_hline(
                                y=1 + var,
                                line_dash="dash",
                                line_color="red",
                                annotation_text=f"VaR ({confidence_level*100:.0f}%): {var*100:.2f}%",
                                annotation_position="bottom right"
                            )
                            fig_cum.add_hline(
                                y=1 + cvar,
                                line_dash="dash",
                                line_color="blue",
                                annotation_text=f"CVaR ({confidence_level*100:.0f}%): {cvar*100:.2f}%",
                                annotation_position="bottom right"
                            )
                        st.plotly_chart(fig_cum, use_container_width=True)
                        
                        # Returns Distribution with VaR and CVaR
                        st.subheader("📊 Returns Distribution with VaR and CVaR")
                        fig_dist = px.histogram(
                            portfolio_returns_local,
                            nbins=50,
                            title='📊 Simulated Returns Distribution',
                            labels={'value': 'Returns', 'count': 'Frequency'},
                            template=THEME
                        )
                        if var is not None and cvar is not None:
                            fig_dist.add_vline(
                                x=var,
                                line_dash="dash",
                                line_color="red",
                                annotation_text=f"VaR: {var*100:.2f}%",
                                annotation_position="top left"
                            )
                            fig_dist.add_vline(
                                x=cvar,
                                line_dash="dash",
                                line_color="blue",
                                annotation_text=f"CVaR: {cvar*100:.2f}%",
                                annotation_position="top left"
                            )
                        st.plotly_chart(fig_dist, use_container_width=True)
                        
                        # Box Plot
                        st.subheader("📦 Box Plot of Portfolio Returns")
                        fig_box = px.box(
                            pd.DataFrame(portfolio_returns_local, columns=['Returns']),
                            y='Returns',
                            title='📦 Box Plot of Portfolio Returns',
                            template=THEME
                        )
                        st.plotly_chart(fig_box, use_container_width=True)
                        
                        # Interpretation
                        st.markdown("""
                        **📖 Interpretation:**
                        - **VaR (Value at Risk):** Represents the maximum expected loss at the specified confidence level.
                        - **CVaR (Conditional Value at Risk):** Indicates the average loss exceeding the VaR.
                        - **Skewness:** Measures the asymmetry of the return distribution.
                        - **Kurtosis:** Indicates the "tailedness" of the return distribution.
                        - **Annualized Volatility:** Measures the dispersion of returns, indicating risk.
                        - **Maximum Drawdown:** Shows the largest peak-to-trough decline, reflecting potential risk exposure.
                        """)
                    else:
                        st.info("🔍 Calculate risk metrics to view visualizations.")
                
                # Sensitivity Analysis Tab
                with risk_analysis_tabs[2]:
                    st.header("🛠️ Sensitivity Analysis")
                    
                    if 'portfolio_returns' in st.session_state:
                        portfolio_returns_local = st.session_state.portfolio_returns
                        st.subheader("📈 Adjusted Returns Distribution")
                        fig_adj_returns = px.histogram(
                            portfolio_returns_local,
                            nbins=50,
                            title='📈 Adjusted Portfolio Returns Distribution',
                            labels={'value': 'Returns', 'count': 'Frequency'},
                            template=THEME
                        )
                        st.plotly_chart(fig_adj_returns, use_container_width=True)
                        
                        st.subheader("🔄 Risk-Return Scatter Plot")
                        fig_risk_return = px.scatter(
                            x=portfolio_returns_local.std() * np.sqrt(252),
                            y=portfolio_returns_local.mean() * 252,
                            size=[w * 100 for w in weights],
                            color=[w * 100 for w in weights],
                            hover_name=available_selected,
                            title='🔄 Risk vs Return Scatter Plot',
                            labels={'x': 'Annualized Volatility (Std Dev)', 'y': 'Annualized Return'},
                            template=THEME
                        )
                        st.plotly_chart(fig_risk_return, use_container_width=True)
                        
                        # Interpretation
                        st.markdown("""
                        **📖 Interpretation:**
                        - **Adjusted Returns Distribution:** Reflects how stress shocks and sensitivity adjustments affect portfolio returns.
                        - **Risk-Return Scatter Plot:** Visualizes the relationship between risk (volatility) and return for each asset in the portfolio.
                        - **Size & Color Indicators:** Represent the allocation percentage, highlighting each asset's influence on the overall portfolio.
                        """)
                    else:
                        st.info("🔍 Calculate risk metrics to perform sensitivity analysis.")
    #else:
        #st.warning("Please add at least one portfolio to run risk analysis.")