pip install pandas numpy scipy matplotlib seaborn cvxpy yfinance

import logging
import random
import time
from collections import deque
from typing import Any, Dict, List, Tuple

import cvxpy as cp
import matplotlib.pyplot as plt
import numpy as np
import pandas as pd
import seaborn as sns
import yfinance as yf
from scipy.optimize import minimize

def create_logger(name: str, level: int = logging.INFO) -> logging.Logger:
    """
    Initializes and returns a custom logger instance.

    Parameters
    ----------
    name : str
        The name of the logger.
    level : int, optional
        The logging level (e.g., logging.INFO, logging.DEBUG), defaults to logging.INFO.

    Returns
    -------
    logging.Logger
        An initialized logger instance.

    Examples
    --------
    >>> log = create_logger('my_module')
    >>> log.info('This is an info message.')
    """
    logger = logging.getLogger(name)
    logger.setLevel(level)

    if not logger.handlers:
        handler = logging.StreamHandler()
        formatter = logging.Formatter('%(asctime)s - %(name)s - %(levelname)s - %(message)s')
        handler.setFormatter(formatter)
        logger.addHandler(handler)
    logger.info(f"Logger '{name}' initialized with level {logging.getLevelName(level)}.")
    return logger

# Initialize a global logger for the notebook
log = create_logger('CapitalAllocationNotebook')

def create_initial_state(
    num_assets: int = 5,
    investment_horizon: int = 252,
    random_seed: int = 42,
) -> Dict:
    """
    Initializes the primary state dictionary for the capital allocation process.

    Parameters
    ----------
    num_assets : int, optional
        The number of assets or strategies to consider, defaults to 5.
    investment_horizon : int, optional
        The number of periods for simulation/analysis, defaults to 252.
    random_seed : int, optional
        A seed for the random number generator, defaults to 42.

    Returns
    -------
    Dict
        An initial state dictionary containing configuration parameters.

    Examples
    --------
    >>> initial_state = create_initial_state(num_assets=3, investment_horizon=100)
    >>> print(initial_state['num_assets'])
    3
    """
    log.info("Creating initial state dictionary.")
    state = {
        'num_assets': num_assets,
        'investment_horizon': investment_horizon,
        'random_seed': random_seed,
        'asset_names': [f'Asset_{i+1}' for i in range(num_assets)],
        'current_weights': np.array([1/num_assets] * num_assets), # Start with equal weights
        'portfolio_returns': pd.Series(dtype=float),
        'portfolio_value': pd.Series(dtype=float),
        'simulation_history': [],
        'backoff_attempts': 0
    }
    log.debug(f"Initial state created: {state}")
    return state

def generate_asset_returns(
    state: Dict,
    num_periods: int,
    mean_returns: np.ndarray = None,
    stds: np.ndarray = None,
    corr_matrix: np.ndarray = None,
    add_jitter: bool = False
) -> Dict:
    """
    Generates simulated daily asset returns using a multivariate normal distribution.

    Parameters
    ----------
    state : Dict
        The current state dictionary.
    num_periods : int
        The number of periods (e.g., days) for which to generate returns.
    mean_returns : np.ndarray, optional
        Array of mean daily returns for each asset. If None, defaults to small positive values.
    stds : np.ndarray, optional
        Array of daily standard deviations for each asset. If None, defaults to small positive values.
    corr_matrix : np.ndarray, optional
        Correlation matrix between assets. If None, an identity matrix (no correlation) is used.
    add_jitter : bool, optional
        If True, adds a small random jitter to the generated returns, defaults to False.

    Returns
    -------
    Dict
        The updated state dictionary, including a DataFrame of simulated returns.

    Examples
    --------
    >>> state = create_initial_state(num_assets=2)
    >>> state = generate_asset_returns(state, num_periods=10)
    >>> print(state['asset_returns'].shape)
    (10, 2)
    """
    log.info(f"Generating {num_periods} periods of asset returns for {state['num_assets']} assets.")
    np.random.seed(state['random_seed'])

    num_assets = state['num_assets']

    if mean_returns is None:
        mean_returns = np.random.rand(num_assets) * 0.0005  # Daily mean returns
    if stds is None:
        stds = np.random.rand(num_assets) * 0.01  # Daily standard deviations

    if corr_matrix is None:
        corr_matrix = np.eye(num_assets)
    else:
        if corr_matrix.shape != (num_assets, num_assets):
            log.error("Correlation matrix dimensions do not match the number of assets.")
            raise ValueError("Correlation matrix dimensions must match num_assets.")

    # Convert correlation matrix to covariance matrix
    cov_matrix = np.diag(stds) @ corr_matrix @ np.diag(stds)

    # Generate returns
    daily_returns = np.random.multivariate_normal(mean_returns, cov_matrix, num_periods)

    if add_jitter:
        jitter = np.random.normal(0, 0.0001, daily_returns.shape)
        daily_returns += jitter
        log.debug("Added random jitter to asset returns.")

    state['asset_returns'] = pd.DataFrame(daily_returns, columns=state['asset_names'])
    log.info("Asset returns generated successfully.")
    log.debug(f"Generated returns head:\n{state['asset_returns'].head()}")
    return state

def get_tracking_metrics(
    state: Dict,
    returns: pd.Series,
    annualization_factor: int = 252,
    risk_free_rate: float = 0.01
) -> Dict:
    """
    Calculates and summarizes key performance metrics for a given portfolio return series.

    Parameters
    ----------
    state : Dict
        The current state dictionary.
    returns : pd.Series
        A pandas Series of daily portfolio returns.
    annualization_factor : int, optional
        The factor to annualize daily metrics, defaults to 252.
    risk_free_rate : float, optional
        The risk-free rate for Sharpe ratio calculation, defaults to 0.01.

    Returns
    -------
    Dict
        A dictionary containing the calculated performance metrics.

    Examples
    --------
    >>> state = create_initial_state()
    >>> synthetic_returns = pd.Series(np.random.normal(0.0005, 0.01, 252))
    >>> metrics = get_tracking_metrics(state, synthetic_returns)
    >>> print(metrics['annualized_return'])
    """
    if returns.empty:
        log.warning("Returns series is empty, cannot calculate metrics.")
        return {}

    log.info("Calculating performance metrics.")

    # Cumulative returns
    cumulative_returns = (1 + returns).prod() - 1

    # Annualized return
    annualized_return = (1 + returns).prod()**(annualization_factor / len(returns)) - 1

    # Annualized volatility
    annualized_volatility = returns.std() * np.sqrt(annualization_factor)

    # Sharpe Ratio
    # Ensure risk_free_rate is on a daily basis for daily returns calculation
    daily_risk_free_rate = (1 + risk_free_rate)**(1/annualization_factor) - 1
    sharpe_ratio = (returns.mean() - daily_risk_free_rate) / returns.std() * np.sqrt(annualization_factor)

    # Max Drawdown
    cumulative_wealth = (1 + returns).cumprod()
    peak = cumulative_wealth.expanding(min_periods=1).max()
    drawdown = (cumulative_wealth / peak) - 1
    max_drawdown = drawdown.min()

    metrics = {
        'total_return': cumulative_returns,
        'annualized_return': annualized_return,
        'annualized_volatility': annualized_volatility,
        'sharpe_ratio': sharpe_ratio,
        'max_drawdown': max_drawdown
    }
    log.debug(f"Calculated metrics: {metrics}")
    return metrics

def calculate_portfolio_returns(
    state: Dict,
    asset_returns: pd.DataFrame,
    weights: np.ndarray
) -> pd.Series:
    """
    Computes the portfolio returns for each period given individual asset returns and weights.

    Parameters
    ----------
    state : Dict
        The current state dictionary.
    asset_returns : pd.DataFrame
        A DataFrame where each column represents an asset's returns and each row is a period.
    weights : np.ndarray
        A 1D NumPy array representing the weights allocated to each asset.

    Returns
    -------
    pd.Series
        A pandas Series containing the calculated portfolio returns for each period.

    Examples
    --------
    >>> state = create_initial_state(num_assets=2)
    >>> state['asset_returns'] = pd.DataFrame({'Asset_1': [0.01, 0.005], 'Asset_2': [0.002, 0.015]})
    >>> weights = np.array([0.5, 0.5])
    >>> portfolio_rets = calculate_portfolio_returns(state, state['asset_returns'], weights)
    >>> print(portfolio_rets.iloc[0])
    0.006
    """
    log.info("Calculating portfolio returns.")
    if not np.isclose(np.sum(weights), 1.0):
        log.warning(f"Weights do not sum to 1.0. Sum is {np.sum(weights)}. Normalizing weights.")
        weights = weights / np.sum(weights)

    portfolio_returns = asset_returns.dot(weights)
    log.debug(f"Portfolio returns calculated. First 5 values:\n{portfolio_returns.head()}")
    return portfolio_returns

def simulate_portfolio_performance(
    state: Dict,
    weights: np.ndarray,
    initial_capital: float = 1000.0
) -> Dict:
    """
    Simulates the performance of a portfolio over time with fixed weights.

    Parameters
    ----------
    state : Dict
        The current state dictionary, must contain 'asset_returns'.
    weights : np.ndarray
        A 1D NumPy array representing the constant weights allocated to each asset.
    initial_capital : float, optional
        The starting capital for the simulation, defaults to 1000.0.

    Returns
    -------
    Dict
        The updated state dictionary, including 'portfolio_returns' and 'portfolio_value'.

    Examples
    --------
    >>> state = create_initial_state(num_assets=2)
    >>> state = generate_asset_returns(state, num_periods=10)
    >>> weights = np.array([0.5, 0.5])
    >>> updated_state = simulate_portfolio_performance(state, weights)
    >>> print(updated_state['portfolio_value'].iloc[-1])
    # Expected output: (approximately 1000 * (1 + sum(daily_portfolio_returns))) after 10 days
    """
    log.info("Simulating portfolio performance with fixed weights.")
    if 'asset_returns' not in state:
        log.error("asset_returns not found in state. Please generate asset returns first.")
        raise ValueError("asset_returns not found in state.")

    portfolio_returns = calculate_portfolio_returns(state, state['asset_returns'], weights)

    cumulative_returns = (1 + portfolio_returns).cumprod()
    portfolio_value = initial_capital * cumulative_returns

    state['portfolio_returns'] = portfolio_returns
    state['portfolio_value'] = portfolio_value

    log.info("Portfolio simulation complete.")
    log.debug(f"Final portfolio value: {portfolio_value.iloc[-1]:.2f}")
    return state

def allocate_equal_weight(
    state: Dict
) -> Dict:
    """
    Allocates capital equally among all assets or strategies.

    Parameters
    ----------
    state : Dict
        The current state dictionary, must contain 'num_assets'.

    Returns
    -------
    Dict
        The updated state dictionary with 'current_weights' set to equal weights.

    Examples
    --------
    >>> state = create_initial_state(num_assets=3)
    >>> state = allocate_equal_weight(state)
    >>> print(state['current_weights'])
    [0.33333333 0.33333333 0.33333333]
    """
    log.info("Allocating capital using an equal-weight strategy.")
    num_assets = state['num_assets']
    weights = np.array([1.0 / num_assets] * num_assets)
    state['current_weights'] = weights
    log.debug(f"Equal weights assigned: {weights}")
    return state

def calculate_portfolio_statistics(
    state: Dict,
    annualization_factor: int = 252
) -> Dict:
    """
    Calculates the expected annual returns, annual standard deviations, and
    the annual covariance matrix of asset returns.

    Parameters
    ----------
    state : Dict
        The current state dictionary, must contain 'asset_returns'.
    annualization_factor : int, optional
        The factor to annualize daily statistics, defaults to 252.

    Returns
    -------
    Dict
        The updated state dictionary, including 'expected_annual_returns',
        'annual_stds', and 'annual_cov_matrix'.

    Examples
    --------
    >>> state = create_initial_state(num_assets=2)
    >>> state = generate_asset_returns(state, num_periods=252)
    >>> updated_state = calculate_portfolio_statistics(state)
    >>> print(updated_state['annual_cov_matrix'].shape)
    (2, 2)
    """
    log.info("Calculating portfolio statistics (expected returns, stds, covariance matrix).")
    if 'asset_returns' not in state or state['asset_returns'].empty:
        log.error("asset_returns not found or is empty in state.")
        raise ValueError("asset_returns must be present and non-empty in state.")

    asset_returns = state['asset_returns']

    # Expected daily returns and daily standard deviations
    daily_mean_returns = asset_returns.mean()
    daily_stds = asset_returns.std()
    daily_cov_matrix = asset_returns.cov()

    # Annualize
    expected_annual_returns = daily_mean_returns * annualization_factor
    annual_stds = daily_stds * np.sqrt(annualization_factor)
    annual_cov_matrix = daily_cov_matrix * annualization_factor

    state['expected_annual_returns'] = expected_annual_returns.values
    state['annual_stds'] = annual_stds.values
    state['annual_cov_matrix'] = annual_cov_matrix.values

    log.info("Portfolio statistics calculated.")
    log.debug(f"Expected Annual Returns: {expected_annual_returns}")
    log.debug(f"Annual Covariance Matrix head:\n{annual_cov_matrix.head()}")
    return state

def allocate_mean_variance_optimization(
    state: Dict,
    target_return: float = None,
    max_sharpe: bool = True
) -> Dict:
    """
    Allocates capital using Mean-Variance Optimization (MVO).

    Parameters
    ----------
    state : Dict
        The current state dictionary, must contain 'expected_annual_returns'
        and 'annual_cov_matrix'.
    target_return : float, optional
        The target annual return for the portfolio. If None, optimizes for max Sharpe.
    max_sharpe : bool, optional
        If True, optimizes for the maximum Sharpe ratio. If False and target_return is None,
        it minimizes portfolio variance. Defaults to True.

    Returns
    -------
    Dict
        The updated state dictionary with 'current_weights' set to MVO-optimized weights.

    Examples
    --------
    >>> state = create_initial_state(num_assets=3)
    >>> state = generate_asset_returns(state, num_periods=252)
    >>> state = calculate_portfolio_statistics(state)
    >>> state = allocate_mean_variance_optimization(state, max_sharpe=True)
    >>> print(state['current_weights'].sum())
    1.0
    """
    log.info("Performing Mean-Variance Optimization (MVO).")
    if 'expected_annual_returns' not in state or 'annual_cov_matrix' not in state:
        log.error("Missing expected_annual_returns or annual_cov_matrix in state.")
        raise ValueError("State must contain expected_annual_returns and annual_cov_matrix.")

    num_assets = state['num_assets']
    expected_returns = state['expected_annual_returns']
    cov_matrix = state['annual_cov_matrix']

    # Define variables
    weights = cp.Variable(num_assets)

    # Define objective
    portfolio_return = expected_returns @ weights
    portfolio_variance = cp.quad_form(weights, cov_matrix)

    # Define constraints
    constraints = [cp.sum(weights) == 1, weights >= 0] # Sum of weights = 1, no short-selling

    if max_sharpe:
        log.info("Optimizing for maximum Sharpe Ratio.")
        # To maximize Sharpe ratio, we maximize (portfolio_return - risk_free_rate) / sqrt(portfolio_variance)
        # This is a fractional programming problem. A common approach is to transform it.
        # This transformation is valid when (return - rf) is positive.
        # Let y = weights / (portfolio_return - risk_free_rate)
        # Maximize (portfolio_return - risk_free_rate) / sqrt(portfolio_variance)
        # is equivalent to minimizing portfolio_variance for a given (portfolio_return - risk_free_rate) = 1 (or any positive constant)
        # and then scale weights.
        # However, a simpler approach is to optimize for a target volatility and then check the return, or vice versa, to build the efficient frontier.
        # For directly maximizing Sharpe, we can use the method described by Robert Merton (1972) or a numerical solver. cvxpy can handle this directly with quadratic forms.

        # Standard approach for max Sharpe when risk-free rate is zero or incorporated:
        # Maximize objective: portfolio_return / sqrt(portfolio_variance)
        # This is equivalent to minimizing variance for target_return = R, and then scaling.
        # A common way with cvxpy is to solve a series of variance minimization problems to find the efficient frontier
        # and then pick the one with the highest sharpe.

        # More direct way: find weights that maximize (portfolio_return - rf) / sqrt(portfolio_variance)
        # If we assume risk-free rate is 0 for simplicity, then maximize (portfolio_return / sqrt(portfolio_variance))
        # This can be formulated by introducing an auxiliary variable k = 1 / (return_portfolio - rf)
        # and then maximizing (return_portfolio - rf) * k subject to (weights.T @ cov @ weights) * k**2 <= 1.
        # A simpler form is to solve: max (r_p - rf) subject to var_p = X, and repeat for X, or directly find tangent portfolio.

        # For simplicity and direct CVXPY usage, we'll maximize return for a given level of risk (or minimize risk for a given return) and select from efficient frontier.
        # The problem formulation directly maximizes the quadratic utility function `return - lambda * variance` where lambda is risk aversion.
        # For max Sharpe, we are looking for the tangent portfolio. We can achieve this by solving a problem with a 'risk aversion' parameter.
        # This formulation is standard in many MVO libraries.
        risk_free_rate = 0.01 # Assuming an annualized risk-free rate
        objective = cp.Maximize(portfolio_return - risk_free_rate - 0.5 * portfolio_variance)
        problem = cp.Problem(objective, constraints)
        problem.solve()

    elif target_return is not None:
        log.info(f"Optimizing for target return: {target_return:.4f}.")
        constraints.append(portfolio_return >= target_return)
        objective = cp.Minimize(portfolio_variance)
        problem = cp.Problem(objective, constraints)
        problem.solve()

    else:
        log.info("Optimizing for minimum variance.")
        objective = cp.Minimize(portfolio_variance)
        problem = cp.Problem(objective, constraints)
        problem.solve()

    if problem.status == cp.OPTIMAL or problem.status == cp.OPTIMAL_INACCURATE:
        state['current_weights'] = weights.value
        log.debug(f"MVO optimized weights: {weights.value}")
    else:
        log.warning("MVO problem did not converge to an optimal solution. Keeping previous weights.")
        log.warning(f"Problem status: {problem.status}")
        if 'current_weights' not in state:
            state['current_weights'] = np.array([1.0 / num_assets] * num_assets)

    # Ensure weights are non-negative and sum to 1 due to potential floating point inaccuracies
    state['current_weights'][state['current_weights'] < 0] = 0
    state['current_weights'] = state['current_weights'] / np.sum(state['current_weights'])

    return state

import typing

def exponential_backoff_retry(
    state: typing.Dict,
    func: callable,
    *args,
    max_retries: int = 5,
    base_delay: float = 1.0,
    max_delay: float = 60.0,
    add_jitter: bool = True,
    **kwargs
) -> typing.Any:
    """
    Implements an exponential backoff retry mechanism for a given function.

    Parameters
    ----------
    state : typing.Dict
        The current state dictionary.
    func : callable
        The function to be executed and retried.
    *args
        Positional arguments to pass to `func`.
    max_retries : int, optional
        The maximum number of retry attempts, defaults to 5.
    base_delay : float, optional
        The initial delay in seconds before the first retry, defaults to 1.0.
    max_delay : float, optional
        The maximum delay between retries, defaults to 60.0.
    add_jitter : bool, optional
        If True, adds random jitter to the delay, defaults to True.
    **kwargs
        Keyword arguments to pass to `func`.

    Returns
    -------
    typing.Any
        The result of the `func` if successful.

    Raises
    -------
    Exception
        If the function fails after `max_retries`.

    Examples
    --------
    >>> state = create_initial_state()
    >>> def flaky_function(x):
    ...     if random.random() < 0.7:  # 70% chance of failure
    ...         raise ValueError("Simulated transient error")
    ...     return x * 2
    >>> # result = exponential_backoff_retry(state, flaky_function, 10, max_retries=3)
    >>> # print(result)
    """
    for i in range(max_retries):
        try:
            log.debug(f"Attempt {i + 1}/{max_retries} for function '{func.__name__}'.")
            return func(*args, **kwargs)
        except Exception as e:
            log.warning(f"Attempt {i + 1} failed for '{func.__name__}': {e}")
            if i == max_retries - 1:
                log.error(f"Max retries ({max_retries}) exceeded for '{func.__name__}'.")
                raise

            delay = min(base_delay * (2 ** i), max_delay)
            if add_jitter:
                jitter = random.uniform(0, delay * 0.1) # Add up to 10% jitter
                delay += jitter
                log.debug(f"Adding {jitter:.2f}s jitter.")

            log.info(f"Retrying '{func.__name__}' in {delay:.2f} seconds...")
            time.sleep(delay)
            state['backoff_attempts'] += 1
    return None # Should not be reached due to raise above

def calculate_risk_contributions(
    state: Dict,
    weights: np.ndarray,
    cov_matrix: np.ndarray
) -> np.ndarray:
    """
    Calculates the risk contribution of each asset to the total portfolio volatility.

    Parameters
    ----------
    state : Dict
        The current state dictionary.
    weights : np.ndarray
        A 1D NumPy array of asset weights.
    cov_matrix : np.ndarray
        The covariance matrix of asset returns.

    Returns
    -------
    np.ndarray
        A 1D NumPy array representing the risk contribution of each asset.

    Examples
    --------
    >>> state = create_initial_state(num_assets=2)
    >>> weights = np.array([0.5, 0.5])
    >>> cov_matrix = np.array([[0.01, 0.005], [0.005, 0.02]])
    >>> risk_contribs = calculate_risk_contributions(state, weights, cov_matrix)
    >>> print(risk_contribs.sum())
    0.15811388300841897
    """
    log.debug("Calculating risk contributions.")
    portfolio_variance = weights.T @ cov_matrix @ weights
    portfolio_std = np.sqrt(portfolio_variance)

    # Marginal Contribution to Risk (MCR)
    marginal_contribution_to_risk = (cov_matrix @ weights) / portfolio_std

    # Component Contribution to Risk (CCR) or Risk Contribution (RC)
    risk_contributions = weights * marginal_contribution_to_risk
    log.debug(f"Risk contributions calculated: {risk_contributions}")
    return risk_contributions

def allocate_risk_parity(
    state: Dict
) -> Dict:
    """
    Allocates capital using a Risk Parity strategy.
    The goal is for each asset to contribute equally to the total portfolio risk.

    Parameters
    ----------
    state : Dict
        The current state dictionary, must contain 'annual_cov_matrix'.

    Returns
    -------
    Dict
        The updated state dictionary with 'current_weights' set to Risk Parity optimized weights.

    Examples
    --------
    >>> state = create_initial_state(num_assets=3)
    >>> state = generate_asset_returns(state, num_periods=252)
    >>> state = calculate_portfolio_statistics(state)
    >>> state = allocate_risk_parity(state)
    >>> print(state['current_weights'].sum())
    1.0
    """
    log.info("Performing Risk Parity allocation.")
    if 'annual_cov_matrix' not in state:
        log.error("Missing annual_cov_matrix in state.")
        raise ValueError("State must contain annual_cov_matrix.")

    num_assets = state['num_assets']
    cov_matrix = state['annual_cov_matrix']

    # Objective function for risk parity: minimize the sum of squared differences
    # between each asset's risk contribution and the average risk contribution.
    def risk_parity_objective(weights):
        weights = np.array(weights)
        if not np.isclose(weights.sum(), 1.0) or np.any(weights < 0):
            return 1e10 # Penalize invalid weights heavily
        risk_contributions = calculate_risk_contributions(state, weights, cov_matrix)
        target_risk_contribution = np.sum(risk_contributions) / num_assets
        # Minimize sum of squared differences from target risk contribution
        return np.sum((risk_contributions - target_risk_contribution)**2)

    # Constraints: weights must sum to 1, and be non-negative
    constraints = [{'type': 'eq', 'fun': lambda w: np.sum(w) - 1}]
    bounds = tuple((0.0, 1.0) for _ in range(num_assets))

    # Initial guess: equal weights
    initial_weights = np.array([1.0 / num_assets] * num_assets)

    result = minimize(risk_parity_objective, initial_weights, method='SLSQP', bounds=bounds, constraints=constraints)

    if result.success:
        state['current_weights'] = result.x
        log.debug(f"Risk Parity optimized weights: {result.x}")
    else:
        log.warning("Risk Parity optimization failed. Keeping previous weights.")
        log.warning(f"Optimization status: {result.message}")
        if 'current_weights' not in state:
            state['current_weights'] = np.array([1.0 / num_assets] * num_assets)

    # Ensure weights are non-negative and sum to 1 due to potential floating point inaccuracies
    state['current_weights'][state['current_weights'] < 0] = 0
    state['current_weights'] = state['current_weights'] / np.sum(state['current_weights'])

    return state

import typing

def create_investment_portfolio(
    state: typing.Dict,
    tickers: typing.List[str],
    start_date: str,
    end_date: str
) -> typing.Dict:
    """
    Retrieves historical stock price data and calculates daily returns for a given set of tickers.

    Parameters
    ----------
    state : typing.Dict
        The current state dictionary.
    tickers : typing.List[str]
        A list of stock ticker symbols.
    start_date : str
        The start date for fetching data in 'YYYY-MM-DD' format.
    end_date : str
        The end date for fetching data in 'YYYY-MM-DD' format.

    Returns
    -------
    typing.Dict
        The updated state dictionary, including 'asset_returns' from real stock data,
        and updated 'asset_names' and 'num_assets'.

    Examples
    --------
    >>> state = create_initial_state()
    >>> tickers = ['AAPL', 'MSFT']
    >>> state = create_investment_portfolio(state, tickers, '2020-01-01', '2021-01-01')
    >>> print(state['asset_returns'].head())
    """
    log.info(f"Fetching historical data for tickers: {tickers} from {start_date} to {end_date}.")
    try:
        # Using exponential backoff for yfinance download to handle potential rate limits or transient errors
        data = exponential_backoff_retry(
            state,
            yf.download,
            tickers=tickers,
            start=start_date,
            end=end_date,
            interval='1d',
            progress=False,
            max_retries=5,
            base_delay=5.0
        )

        if data.empty:
            log.warning(f"No data fetched for tickers: {tickers}. Proceeding to generate mock data.")
            # Trigger the except block to handle mock data generation
            raise ValueError("yf.download returned empty data.")

        # If only one ticker, yfinance returns a Series, convert to DataFrame
        if len(tickers) == 1:
            adj_close_prices = pd.DataFrame(data['Adj Close'])
        else:
            adj_close_prices = data['Adj Close']

        # Calculate daily returns
        asset_returns = adj_close_prices.pct_change().dropna()
        asset_returns.columns = tickers # Ensure columns are ticker names

        state['asset_returns'] = asset_returns
        state['asset_names'] = list(tickers)
        state['num_assets'] = len(tickers)

        log.info(f"Successfully fetched and processed data for {len(tickers)} assets.")
        log.debug(f"Real asset returns head:\n{state['asset_returns'].head()}")

    except Exception as e:
        log.error(f"Failed to fetch real stock data for tickers {tickers}: {e}. Generating mock data instead.")
        # Calculate approximate num_periods from start_date and end_date
        start_dt = pd.to_datetime(start_date)
        end_dt = pd.to_datetime(end_date)
        num_periods = (end_dt - start_dt).days
        if num_periods <= 0:
            num_periods = 252 # Default to one trading year if date range is invalid or too small

        # Update state with mock asset names and count for generate_asset_returns
        state['num_assets'] = len(tickers)
        state['asset_names'] = [f'Mock_{ticker}' for ticker in tickers]

        # Generate mock asset returns
        state = generate_asset_returns(
            state,
            num_periods=num_periods,
            mean_returns=np.random.rand(len(tickers)) * 0.0005, # Small positive mean returns
            stds=np.random.rand(len(tickers)) * 0.01, # Reasonable daily stds
            add_jitter=True
        )

        log.warning("Using mock data for real-world demonstration due to previous data fetching issues.")

    return state

def plot_portfolio_performance(
    state: Dict,
    title: str = 'Portfolio Performance'
) -> None:
    """
    Visualizes the portfolio's value and returns over time.

    Parameters
    ----------
    state : Dict
        The current state dictionary, must contain 'portfolio_value' and 'portfolio_returns'.
    title : str, optional
        A title for the plot, defaults to 'Portfolio Performance'.

    Returns
    -------
    None
        Displays the plot directly.

    Examples
    --------
    >>> state = create_initial_state()
    >>> state = generate_asset_returns(state, num_periods=252)
    >>> weights = np.array([0.2, 0.2, 0.2, 0.2, 0.2])
    >>> state = simulate_portfolio_performance(state, weights)
    >>> plot_portfolio_performance(state, title='Equal Weight Portfolio Performance')
    """
    log.info(f"Plotting portfolio performance: {title}.")
    if 'portfolio_value' not in state or 'portfolio_returns' not in state:
        log.error("Missing 'portfolio_value' or 'portfolio_returns' in state for plotting.")
        return

    fig, (ax1, ax2) = plt.subplots(2, 1, figsize=(12, 10), sharex=True)

    # Plot Portfolio Value
    ax1.plot(state['portfolio_value'], label='Portfolio Value', color='blue')
    ax1.set_title(f'{title} - Portfolio Value Over Time')
    ax1.set_ylabel('Portfolio Value')
    ax1.legend()
    ax1.grid(True)

    # Plot Portfolio Returns
    ax2.plot(state['portfolio_returns'], label='Daily Returns', color='green', alpha=0.7)
    ax2.set_title(f'{title} - Daily Portfolio Returns')
    ax2.set_xlabel('Time (Periods)')
    ax2.set_ylabel('Daily Return')
    ax2.legend()
    ax2.grid(True)

    plt.tight_layout()
    plt.show()
    log.info("Portfolio performance plot displayed.")

def plot_asset_returns_distribution(
    state: Dict,
    title: str = 'Asset Returns Distribution'
) -> None:
    """
    Visualizes the distribution of individual asset returns.

    Parameters
    ----------
    state : Dict
        The current state dictionary, must contain 'asset_returns'.
    title : str, optional
        A title for the plot, defaults to 'Asset Returns Distribution'.

    Returns
    -------
    None
        Displays the plot directly.

    Examples
    --------
    >>> state = create_initial_state()
    >>> state = generate_asset_returns(state, num_periods=252)
    >>> plot_asset_returns_distribution(state)
    """
    log.info(f"Plotting asset returns distribution: {title}.")
    if 'asset_returns' not in state or state['asset_returns'].empty:
        log.error("Missing or empty 'asset_returns' in state for plotting.")
        return

    asset_returns = state['asset_returns']

    fig, axes = plt.subplots(len(asset_returns.columns), 1, figsize=(10, 4 * len(asset_returns.columns)), sharex=True)
    if len(asset_returns.columns) == 1: # Handle single subplot case
        axes = [axes]

    for i, col in enumerate(asset_returns.columns):
        sns.histplot(asset_returns[col], kde=True, ax=axes[i], bins=50, color='skyblue')
        axes[i].set_title(f'Distribution of {col} Returns')
        axes[i].set_xlabel('Daily Return')
        axes[i].set_ylabel('Frequency')
        axes[i].grid(True)

    plt.suptitle(title, y=1.02, fontsize=16)
    plt.tight_layout()
    plt.show()

    # Also plot box plot for comparison
    plt.figure(figsize=(12, 6))
    sns.boxplot(data=asset_returns, palette='viridis')
    plt.title(f'{title} - Asset Return Comparison')
    plt.xlabel('Asset')
    plt.ylabel('Daily Return')
    plt.grid(True)
    plt.tight_layout()
    plt.show()
    log.info("Asset returns distribution plots displayed.")

def plot_efficient_frontier(
    state: Dict,
    num_portfolios: int = 10000,
    risk_free_rate: float = 0.01
) -> None:
    """
    Visualizes the Efficient Frontier, plotting optimal portfolios based on risk and return.

    Parameters
    ----------
    state : Dict
        The current state dictionary, must contain 'expected_annual_returns',
        'annual_cov_matrix', and 'annual_stds'.
    num_portfolios : int, optional
        The number of random portfolios to generate for visualization, defaults to 10000.
    risk_free_rate : float, optional
        The risk-free rate used for Sharpe ratio calculation, defaults to 0.01.

    Returns
    -------
    None
        Displays the plot directly.

    Examples
    --------
    >>> state = create_initial_state(num_assets=5)
    >>> state = generate_asset_returns(state, num_periods=252)
    >>> state = calculate_portfolio_statistics(state)
    >>> plot_efficient_frontier(state)
    """
    log.info("Plotting Efficient Frontier.")
    if 'expected_annual_returns' not in state or 'annual_cov_matrix' not in state or 'annual_stds' not in state:
        log.error("Missing required statistics in state for plotting Efficient Frontier.")
        return

    expected_returns = state['expected_annual_returns']
    cov_matrix = state['annual_cov_matrix']
    num_assets = state['num_assets']

    portfolio_returns = []
    portfolio_volatility = []
    portfolio_weights = []

    for _ in range(num_portfolios):
        weights = np.random.random(num_assets)
        weights /= np.sum(weights)
        portfolio_weights.append(weights)

        ret = np.sum(weights * expected_returns)
        vol = np.sqrt(np.dot(weights.T, np.dot(cov_matrix, weights)))

        portfolio_returns.append(ret)
        portfolio_volatility.append(vol)

    # Convert to numpy arrays for easier manipulation
    portfolio_returns = np.array(portfolio_returns)
    portfolio_volatility = np.array(portfolio_volatility)

    # Calculate Sharpe Ratios
    sharpe_ratios = (portfolio_returns - risk_free_rate) / portfolio_volatility

    # Find the portfolio with the maximum Sharpe Ratio
    max_sharpe_idx = np.argmax(sharpe_ratios)
    max_sharpe_return = portfolio_returns[max_sharpe_idx]
    max_sharpe_volatility = portfolio_volatility[max_sharpe_idx]

    # Find the portfolio with the minimum volatility
    min_vol_idx = np.argmin(portfolio_volatility)
    min_vol_return = portfolio_returns[min_vol_idx]
    min_vol_volatility = portfolio_volatility[min_vol_idx]

    plt.figure(figsize=(12, 8))
    plt.scatter(portfolio_volatility, portfolio_returns, c=sharpe_ratios, cmap='viridis', marker='o', alpha=0.5)
    plt.title('Efficient Frontier with Simulated Portfolios')
    plt.xlabel('Annualized Volatility (Standard Deviation)')
    plt.ylabel('Annualized Return')
    plt.colorbar(label='Sharpe Ratio')

    # Plot individual assets
    asset_volatilities = state['annual_stds']
    plt.scatter(asset_volatilities, expected_returns, marker='X', color='red', s=100, label='Individual Assets')
    for i, txt in enumerate(state['asset_names']):
        plt.annotate(txt, (asset_volatilities[i], expected_returns[i]), xytext=(5, -5), textcoords='offset points')

    # Highlight Max Sharpe and Min Volatility Portfolios
    plt.scatter(max_sharpe_volatility, max_sharpe_return, marker='*', color='gold', s=300, label='Max Sharpe Ratio Portfolio')
    plt.scatter(min_vol_volatility, min_vol_return, marker='P', color='blue', s=300, label='Minimum Volatility Portfolio')

    plt.grid(True)
    plt.legend()
    plt.tight_layout()
    plt.show()
    log.info("Efficient Frontier plot displayed.")

def plot_risk_contributions(
    state: Dict,
    title: str = 'Asset Risk Contributions'
) -> None:
    """
    Visualizes the risk contribution of each asset to the total portfolio risk.

    Parameters
    ----------
    state : Dict
        The current state dictionary, must contain 'annual_cov_matrix', 'current_weights',
        and 'asset_names'.
    title : str, optional
        A title for the plot, defaults to 'Asset Risk Contributions'.

    Returns
    -------
    None
        Displays the plot directly.

    Examples
    --------
    >>> state = create_initial_state(num_assets=3)
    >>> state = generate_asset_returns(state, num_periods=252)
    >>> state = calculate_portfolio_statistics(state)
    >>> state = allocate_risk_parity(state) # Assuming Risk Parity weights are calculated
    >>> plot_risk_contributions(state)
    """
    log.info(f"Plotting risk contributions: {title}.")
    if 'annual_cov_matrix' not in state or 'current_weights' not in state or 'asset_names' not in state:
        log.error("Missing required data in state for plotting risk contributions.")
        return

    weights = state['current_weights']
    cov_matrix = state['annual_cov_matrix']
    asset_names = state['asset_names']

    if np.sum(weights) == 0: # Avoid division by zero if all weights are zero
        log.warning("All weights are zero, cannot calculate risk contributions. Displaying empty plot.")
        risk_contributions_abs = np.zeros(len(asset_names))
    else:
        risk_contributions_abs = calculate_risk_contributions(state, weights, cov_matrix)

    # Normalize to percentage for easier interpretation
    if np.sum(risk_contributions_abs) != 0:
        risk_contributions_percent = (risk_contributions_abs / np.sum(risk_contributions_abs)) * 100
    else:
        risk_contributions_percent = np.zeros(len(asset_names))

    plt.figure(figsize=(10, 6))
    sns.barplot(x=asset_names, y=risk_contributions_percent, palette='viridis')
    plt.title(title)
    plt.xlabel('Asset')
    plt.ylabel('Risk Contribution (%)')
    plt.grid(axis='y')
    plt.tight_layout()
    plt.show()
    log.info("Risk contributions plot displayed.")

def plot_asset_correlation_heatmap(
    state: Dict,
    title: str = 'Asset Correlation Heatmap'
) -> None:
    """
    Visualizes the correlation matrix of asset returns using a heatmap.

    Parameters
    ----------
    state : Dict
        The current state dictionary, must contain 'asset_returns'.
    title : str, optional
        A title for the plot, defaults to 'Asset Correlation Heatmap'.

    Returns
    -------
    None
        Displays the plot directly.

    Examples
    --------
    >>> state = create_initial_state(num_assets=5)
    >>> state = generate_asset_returns(state, num_periods=252)
    >>> plot_asset_correlation_heatmap(state)
    """
    log.info(f"Plotting asset correlation heatmap: {title}.")
    if 'asset_returns' not in state or state['asset_returns'].empty:
        log.error("Missing or empty 'asset_returns' in state for plotting correlation heatmap.")
        return

    asset_returns = state['asset_returns']
    correlation_matrix = asset_returns.corr()

    plt.figure(figsize=(10, 8))
    sns.heatmap(
        correlation_matrix,
        annot=True,
        cmap='coolwarm',
        fmt=".2f",
        linewidths=.5,
        cbar_kws={'label': 'Correlation Coefficient'}
    )
    plt.title(title)
    plt.xlabel('Assets')
    plt.ylabel('Assets')
    plt.tight_layout()
    plt.show()
    log.info("Asset correlation heatmap displayed.")

# 1. Initialize state and generate simulated asset returns
log.info("--- Demonstration: Initializing state and generating simulated asset returns ---")
state_simulated = create_initial_state(num_assets=5, investment_horizon=500, random_seed=42)

# Define custom mean returns and stds for more interesting simulation
sim_mean_returns = np.array([0.0003, 0.0005, 0.0004, 0.0002, 0.0006])
sim_stds = np.array([0.008, 0.012, 0.010, 0.007, 0.015])

# Create a more realistic correlation matrix
sim_corr_matrix = np.array([
    [1.0, 0.6, 0.4, 0.2, 0.1],
    [0.6, 1.0, 0.7, 0.3, 0.2],
    [0.4, 0.7, 1.0, 0.5, 0.3],
    [0.2, 0.3, 0.5, 1.0, 0.8],
    [0.1, 0.2, 0.3, 0.8, 1.0]
])

state_simulated = generate_asset_returns(
    state_simulated,
    num_periods=state_simulated['investment_horizon'],
    mean_returns=sim_mean_returns,
    stds=sim_stds,
    corr_matrix=sim_corr_matrix,
    add_jitter=True
)

# Display head of simulated asset returns
log.info("Simulated Asset Returns (first 5 periods):")
display(state_simulated['asset_returns'].head())

# Visualize distribution of simulated asset returns
plot_asset_returns_distribution(state_simulated, title='Simulated Asset Returns Distribution')

# Visualize correlation of simulated asset returns
plot_asset_correlation_heatmap(state_simulated, title='Simulated Asset Returns Correlation Heatmap')

# Calculate portfolio statistics needed for MVO and Risk Parity
state_simulated = calculate_portfolio_statistics(state_simulated)

log.info("Expected Annual Returns:")
display(pd.DataFrame({
    'Asset': state_simulated['asset_names'],
    'Expected Annual Return': state_simulated['expected_annual_returns']
}))

log.info("Annual Standard Deviations:")
display(pd.DataFrame({
    'Asset': state_simulated['asset_names'],
    'Annual Std Dev': state_simulated['annual_stds']
}))

log.info("Annual Covariance Matrix:")
display(pd.DataFrame(
    state_simulated['annual_cov_matrix'],
    index=state_simulated['asset_names'],
    columns=state_simulated['asset_names']
))

# 2. Apply Equal Weight Allocation
log.info("--- Demonstration: Equal Weight Allocation ---")
state_equal_weight = state_simulated.copy()
state_equal_weight = allocate_equal_weight(state_equal_weight)

# Simulate performance with equal weights
state_equal_weight = simulate_portfolio_performance(state_equal_weight, state_equal_weight['current_weights'])

# 3. Calculate and display performance metrics
equal_weight_metrics = get_tracking_metrics(state_equal_weight, state_equal_weight['portfolio_returns'])
log.info("Equal Weight Portfolio Performance Metrics:")
display(pd.DataFrame([equal_weight_metrics], index=['Equal Weight']))

# 4. Plot performance
plot_portfolio_performance(state_equal_weight, title='Equal Weight Portfolio Performance')

# Plot risk contributions for equal weight (should be equal if risks are similar)
plot_risk_contributions(state_equal_weight, title='Equal Weight Risk Contributions')

# 5. Apply Mean-Variance Optimization (Max Sharpe)
log.info("--- Demonstration: Mean-Variance Optimization (Max Sharpe) ---")
state_mvo = state_simulated.copy()
# Ensure portfolio statistics are already calculated in state_simulated

state_mvo = allocate_mean_variance_optimization(state_mvo, max_sharpe=True)

# Simulate performance with MVO weights
state_mvo = simulate_portfolio_performance(state_mvo, state_mvo['current_weights'])

# 6. Calculate and display performance metrics
mvo_metrics = get_tracking_metrics(state_mvo, state_mvo['portfolio_returns'])
log.info("MVO (Max Sharpe) Portfolio Performance Metrics:")
display(pd.DataFrame([mvo_metrics], index=['MVO (Max Sharpe)']))

# 7. Plot performance
plot_portfolio_performance(state_mvo, title='MVO (Max Sharpe) Portfolio Performance')

# Plot Efficient Frontier
plot_efficient_frontier(state_simulated)

# Plot risk contributions for MVO
plot_risk_contributions(state_mvo, title='MVO (Max Sharpe) Risk Contributions')

# 8. Apply Risk Parity Allocation
log.info("--- Demonstration: Risk Parity Allocation ---")
state_risk_parity = state_simulated.copy()
# Ensure portfolio statistics are already calculated in state_simulated

state_risk_parity = allocate_risk_parity(state_risk_parity)

# Simulate performance with Risk Parity weights
state_risk_parity = simulate_portfolio_performance(state_risk_parity, state_risk_parity['current_weights'])

# 9. Calculate and display performance metrics
risk_parity_metrics = get_tracking_metrics(state_risk_parity, state_risk_parity['portfolio_returns'])
log.info("Risk Parity Portfolio Performance Metrics:")
display(pd.DataFrame([risk_parity_metrics], index=['Risk Parity']))

# 10. Plot performance
plot_portfolio_performance(state_risk_parity, title='Risk Parity Portfolio Performance')

# Plot risk contributions for Risk Parity (should be close to equal)
plot_risk_contributions(state_risk_parity, title='Risk Parity Risk Contributions')

# 11. Aggregate and display metrics for comparison
log.info("--- Comparison of Allocation Strategies ---")
all_metrics = pd.DataFrame({
    'Equal Weight': equal_weight_metrics,
    'MVO (Max Sharpe)': mvo_metrics,
    'Risk Parity': risk_parity_metrics
}).T

log.info("Summary of Performance Metrics:")
display(all_metrics)