import sys

# Install yfinance for fetching stock data
!{sys.executable} -m pip install pandas numpy yfinance matplotlib seaborn

import pandas as pd
import numpy as np
import yfinance as yf
import logging
from datetime import datetime, timedelta
import matplotlib.pyplot as plt
import seaborn as sns
import random
import time

# Configure logging
logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')

def create_amihud_state() -> dict:
    """
    Initializes the state dictionary for Amihud ratio calculations.

    Returns
    -------
    dict
        An initialized state dictionary with keys for data storage,
        ratio storage, and metrics summary.

    Examples
    --------
    >>> state = create_amihud_state()
    >>> isinstance(state, dict)
    True
    >>> 'data_store' in state
    True
    """
    logging.info("Initializing Amihud state.")
    state = {
        'data_store': pd.DataFrame(),
        'ratios_store': pd.DataFrame(),
        'metrics_summary': {}
    }
    return state

def add_random_jitter(value: float, scale: float = 0.01) -> float:
    """
    Adds a small, normally distributed random jitter to a numerical value.

    Parameters
    ----------
    value : float
        The numerical value to which jitter will be added.
    scale : float, optional
        The standard deviation of the normal distribution for the jitter,
        by default 0.01.

    Returns
    -------
    float
        The original value with added random jitter.

    Examples
    --------
    >>> round(add_random_jitter(10.0, scale=0.001), 1)
    10.0
    """
    return value + np.random.normal(0, scale)

def fetch_stock_data(state: dict, ticker: str, start_date: str, end_date: str) -> dict:
    """
    Fetches historical stock data for a given ticker and date range.
    Includes retry logic with exponential backoff and random jitter.

    Parameters
    ----------
    state : dict
        The current state dictionary.
    ticker : str
        The stock ticker symbol (e.g., 'AAPL').
    start_date : str
        The start date for data fetching in 'YYYY-MM-DD' format.
    end_date : str
        The end date for data fetching in 'YYYY-MM-DD' format.

    Returns
    -------
    dict
        The updated state dictionary with fetched stock data.

    Examples
    --------
    >>> initial_state = create_amihud_state()
    >>> updated_state = fetch_stock_data(initial_state, 'AAPL', '2023-01-01', '2023-01-10')
    >>> 'AAPL' in updated_state['data_store'] # doctest: +SKIP
    True
    """
    logging.info(f"Fetching data for {ticker} from {start_date} to {end_date}")
    max_retries = 3
    base_delay = 1  # seconds

    for attempt in range(max_retries):
        try:
            df = yf.download(ticker, start=start_date, end=end_date, progress=False)
            if df.empty:
                logging.warning(f"No data fetched for {ticker}. It might be a delisted stock or incorrect ticker.")
                state['data_store'] = pd.DataFrame() # Store empty DF if no data
                return state
            df.index = pd.to_datetime(df.index)
            state['data_store'] = df
            logging.info(f"Successfully fetched data for {ticker}.")
            return state
        except Exception as e:
            logging.warning(f"Attempt {attempt + 1} to fetch data for {ticker} failed: {e}")
            if attempt < max_retries - 1:
                delay = add_random_jitter(base_delay * (2 ** attempt))
                logging.info(f"Retrying in {delay:.2f} seconds...")
                time.sleep(delay)
            else:
                logging.error(f"Failed to fetch data for {ticker} after {max_retries} attempts.")
                state['data_store'] = pd.DataFrame() # Store empty DF on failure
                return state
    return state # Should not be reached but for safety

def calculate_amihud_ratio(state: dict, ticker: str) -> dict:
    """
    Calculates the Amihud illiquidity ratio for the provided stock data.

    Parameters
    ----------
    state : dict
        The current state dictionary containing the stock data.
    ticker : str
        The stock ticker symbol.

    Returns
    -------
    dict
        The updated state dictionary with Amihud ratios stored.

    Examples
    --------
    >>> initial_state = create_amihud_state()
    >>> # Assume data is fetched into initial_state['data_store'] as a DataFrame
    >>> # Example DataFrame structure:
    >>> sample_data = pd.DataFrame({
    ...     'Close': [100, 101, 99, 102, 103],
    ...     'Volume': [1000000, 1200000, 900000, 1500000, 1100000]
    ... }, index=pd.to_datetime(['2023-01-01', '2023-01-02', '2023-01-03', '2023-01-04', '2023-01-05']))
    >>> initial_state['data_store'] = sample_data.copy() # Assign copy to avoid modifying original
    >>> updated_state = calculate_amihud_ratio(initial_state, 'TEST')
    >>> 'Amihud_Ratio' in updated_state['ratios_store'].columns # doctest: +SKIP
    True
    """
    logging.info(f"Calculating Amihud ratio for {ticker}.")
    data_df = state['data_store'].copy()

    if data_df.empty:
        logging.warning(f"No data available for {ticker} to calculate Amihud ratio.")
        state['ratios_store'] = pd.DataFrame(columns=['Return', 'Abs_Return', 'Dollar_Volume', 'Amihud_Ratio'])
        return state

    # Calculate daily returns
    data_df['Return'] = data_df['Close'].pct_change()

    # Calculate absolute daily returns
    data_df['Abs_Return'] = data_df['Return'].abs()

    # Calculate daily dollar volume
    data_df['Dollar_Volume'] = data_df['Volume'] * data_df['Close']

    # Calculate Amihud Ratio
    # Handle division by zero: if Dollar_Volume is 0, Amihud_Ratio is NaN
    data_df['Amihud_Ratio'] = data_df['Abs_Return'] / data_df['Dollar_Volume'].replace(0, np.nan)

    state['ratios_store'] = data_df
    logging.info(f"Amihud ratio calculation complete for {ticker}.")
    return state

def summarize_amihud_metrics(state: dict, ticker: str) -> dict:
    """
    Calculates summary statistics for the Amihud illiquidity ratios.

    Parameters
    ----------
    state : dict
        The current state dictionary containing the calculated Amihud ratios.
    ticker : str
        The stock ticker symbol.

    Returns
    -------
    dict
        The updated state dictionary with summary statistics.

    Examples
    --------
    >>> current_state = create_amihud_state()
    >>> # Assume current_state['ratios_store'] contains 'Amihud_Ratio' column
    >>> sample_ratios = pd.DataFrame({'Amihud_Ratio': [0.001, 0.002, 0.0015, 0.003, 0.0005]})
    >>> current_state['ratios_store'] = sample_ratios
    >>> updated_state = summarize_amihud_metrics(current_state, 'TEST')
    >>> 'TEST' in updated_state['metrics_summary'] # doctest: +SKIP
    True
    """
    logging.info(f"Summarizing Amihud metrics for {ticker}.")
    ratios_df = state['ratios_store']

    if 'Amihud_Ratio' not in ratios_df.columns or ratios_df['Amihud_Ratio'].isnull().all():
        logging.warning(f"No valid Amihud Ratios available for {ticker} to summarize.")
        state['metrics_summary'][ticker] = {
            'mean': np.nan,
            'median': np.nan,
            'std_dev': np.nan,
            'min': np.nan,
            'max': np.nan
        }
        return state

    summary = {
        'mean': ratios_df['Amihud_Ratio'].mean(),
        'median': ratios_df['Amihud_Ratio'].median(),
        'std_dev': ratios_df['Amihud_Ratio'].std(),
        'min': ratios_df['Amihud_Ratio'].min(),
        'max': ratios_df['Amihud_Ratio'].max()
    }
    state['metrics_summary'][ticker] = summary
    logging.info(f"Summary metrics for {ticker} generated: {summary}")
    return state

# 1. Initialize the state
state = create_amihud_state()

# 2. Define parameters for demonstration
ticker = 'GOOG'
end_date = datetime.now().strftime('%Y-%m-%d')
start_date = (datetime.now() - timedelta(days=365*2)).strftime('%Y-%m-%d') # 2 years of data

logging.info(f"Demonstration started for ticker: {ticker}")

# 3. Fetch stock data
state = fetch_stock_data(state, ticker, start_date, end_date)
df_data = state['data_store']

if df_data.empty:
    logging.error(f"Cannot proceed with demonstration as no data was fetched for {ticker}.")
else:
    print(f"\n--- Fetched Data for {ticker} (First 5 Rows) ---")
    display(df_data.head())
    display(df_data.tail())

    # 4. Calculate Amihud Ratio
    state = calculate_amihud_ratio(state, ticker)
    df_ratios = state['ratios_store']

    print(f"\n--- Amihud Ratio Calculation for {ticker} (First 5 Rows) ---")
    display(df_ratios[['Close', 'Volume', 'Return', 'Abs_Return', 'Dollar_Volume', 'Amihud_Ratio']].head())

    # 5. Summarize Amihud Metrics
    state = summarize_amihud_metrics(state, ticker)
    summary_metrics = state['metrics_summary'].get(ticker, {})

    print(f"\n--- Summary Statistics for Amihud Ratio ({ticker}) ---")
    summary_df = pd.DataFrame([summary_metrics])
    display(summary_df.round(8))

    # 6. Visualization: Amihud Ratio over Time
    if 'Amihud_Ratio' in df_ratios.columns and not df_ratios['Amihud_Ratio'].isnull().all():
        plt.figure(figsize=(14, 7))
        sns.lineplot(x=df_ratios.index, y='Amihud_Ratio', data=df_ratios)
        plt.title(f'Amihud Illiquidity Ratio for {ticker} Over Time', fontsize=16)
        plt.xlabel('Date', fontsize=12)
        plt.ylabel('Amihud Ratio', fontsize=12)
        plt.yscale('log') # Use log scale due to potential large fluctuations
        plt.grid(True, linestyle='--', alpha=0.7)
        plt.xticks(rotation=45)
        plt.tight_layout()
        plt.show()

        # Visualization: Distribution of Amihud Ratio
        plt.figure(figsize=(10, 6))
        sns.histplot(df_ratios['Amihud_Ratio'].dropna(), bins=50, kde=True)
        plt.title(f'Distribution of Amihud Illiquidity Ratio for {ticker}', fontsize=16)
        plt.xlabel('Amihud Ratio', fontsize=12)
        plt.ylabel('Frequency', fontsize=12)
        plt.yscale('log') # Log scale for y-axis to better show distribution
        plt.grid(True, linestyle='--', alpha=0.7)
        plt.tight_layout()
        plt.show()
    else:
        logging.warning("Skipping visualizations as Amihud Ratio data is not available or entirely null.")

logging.info("Demonstration complete.")