# Install necessary libraries
!pip install nolds numpy pandas matplotlib seaborn

import logging
import time
import random
from collections import deque

import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
import seaborn as sns
import nolds

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

import inspect
print(inspect.signature(nolds.hurst_rs))

def create_state(initial_data: np.ndarray = None) -> dict:
    """
    Initializes the state dictionary for the Hurst exponent analysis.

    Parameters
    ----------
    initial_data : np.ndarray, optional
        An optional initial time series data to include in the state.
        If None, an empty list is initialized.

    Returns
    -------
    dict
        An initialized state dictionary with keys for data, parameters, and logging.

    Examples
    --------
    >>> state = create_state()
    >>> 'data' in state and 'parameters' in state
    True

    >>> sample_data = np.array([1, 2, 3, 4, 5])
    >>> state_with_data = create_state(initial_data=sample_data)
    >>> np.array_equal(state_with_data['data'], sample_data)
    True
    """
    logger = logging.getLogger(__name__)
    state = {
        'data': initial_data if initial_data is not None else np.array([]),
        'parameters': {
            'min_window_size': 10,
            'max_window_size': None, # Will be set based on data length
            'n_windows': 100, # Number of window sizes to test for R/S analysis
            'log_data_transformed': False # Flag to indicate if data has been transformed (e.g., log-returns)
        },
        'results': {},
        'history': deque(maxlen=100) # For tracking recent operations or data points
    }
    logger.info("State initialized.")
    return state

def generate_synthetic_series(state: dict, series_type: str, length: int, volatility: float,
                              mean_reversion_strength: float = 0.1, trend_strength: float = 0.001,
                              seed: int = None) -> np.ndarray:
    """
    Generates synthetic time series data (prices) based on specified characteristics.

    Parameters
    ----------
    state : dict
        The current state dictionary. (Not used for modification here, but included for consistent signature).
    series_type : str
        Type of series to generate: 'random_walk', 'mean_reverting', 'trending'.
    length : int
        The number of data points in the series.
    volatility : float
        The standard deviation of the daily returns for the series.
    mean_reversion_strength : float, optional
        For 'mean_reverting' series, controls how quickly the series reverts to its mean.
        Defaults to 0.1.
    trend_strength : float, optional
        For 'trending' series, controls the magnitude of the underlying trend.
        Defaults to 0.001.
    seed : int, optional
        Random seed for reproducibility. Defaults to None.

    Returns
    -------
    np.ndarray
        A NumPy array representing the generated synthetic time series.

    Raises
    ------
    ValueError
        If an unknown series_type is provided.

    Examples
    --------
    >>> state = create_state()
    >>> rw = generate_synthetic_series(state, 'random_walk', 100, 0.01, seed=42)
    >>> len(rw) == 100
    True

    >>> mr = generate_synthetic_series(state, 'mean_reverting', 100, 0.01, mean_reversion_strength=0.05, seed=42)
    >>> len(mr) == 100
    True
    """
    logger = logging.getLogger(__name__)
    if seed is not None:
        np.random.seed(seed)
        random.seed(seed)

    prices = np.zeros(length)
    prices[0] = 100.0 # Starting price

    if series_type == 'random_walk':
        returns = np.random.normal(0, volatility, length - 1)
        prices[1:] = prices[0] * np.exp(np.cumsum(returns))
        logger.info(f"Generated a random walk series of length {length}.")

    elif series_type == 'mean_reverting':
        # Ornstein-Uhlenbeck process approximation for mean reversion
        mean_level = prices[0]
        dt = 1.0 # Time step
        for t in range(1, length):
            # dx = alpha * (mu - x) * dt + sigma * dW
            # Here, alpha is mean_reversion_strength, mu is mean_level, sigma is volatility
            drift = mean_reversion_strength * (mean_level - prices[t-1]) * dt
            shock = volatility * np.random.normal(0, 1) * np.sqrt(dt)
            prices[t] = prices[t-1] + drift + shock
        logger.info(f"Generated a mean-reverting series of length {length} with strength {mean_reversion_strength}.")

    elif series_type == 'trending':
        # Geometric Brownian Motion with a drift
        returns = np.random.normal(0, volatility, length - 1) + trend_strength
        prices[1:] = prices[0] * np.exp(np.cumsum(returns))
        logger.info(f"Generated a trending series of length {length} with strength {trend_strength}.")

    else:
        logger.error(f"Unknown series type: {series_type}")
        raise ValueError(f"Unknown series type: {series_type}")

    return prices

def calculate_hurst_exponent(state: dict, time_series: np.ndarray, fit_method: str = 'RANSAC',
                             min_window_size: int = None, max_window_size: int = None,
                             n_windows: int = None, use_log_returns: bool = True) -> dict:
    """
    Calculates the Hurst exponent of a time series using R/S analysis.

    Parameters
    ----------
    state : dict
        The current state dictionary, which will be updated with the calculated
        Hurst exponent and potentially the log_data_transformed flag.
    time_series : np.ndarray
        The input time series data (e.g., price series).
    fit_method : str, optional
        The method used by `nolds.hurst_rs` to fit the line for R/S analysis.
        Defaults to 'RANSAC' (can also be 'poly').
    min_window_size : int, optional
        The minimum window size for R/S analysis. If not provided, it's taken
        from `state['parameters']`. Defaults to 10.
    max_window_size : int, optional
        The maximum window size for R/S analysis. If not provided, it's taken
        from `state['parameters']` or defaults to half the series length.
    n_windows : int, optional
        The number of different window sizes to use for R/S analysis. If not provided,
        it's taken from `state['parameters']`. Defaults to 100.
    use_log_returns : bool, optional
        If `True`, the function will calculate the Hurst exponent on the log-returns
        of the time series instead of the raw values. Defaults to True.

    Returns
    -------
    dict
        The updated state dictionary including the calculated Hurst exponent in
        `state['results']['hurst_exponent']`.

    Examples
    --------
    >>> state = create_state()
    >>> prices = np.array([100, 101, 100, 102, 103, 102.5, 104, 103.5, 105, 106])
    >>> updated_state = calculate_hurst_exponent(state, prices)
    >>> 'hurst_exponent' in updated_state['results']
    True

    >>> # Example with a known random walk (H=0.5)
    >>> np.random.seed(42)
    >>> rw_prices = np.cumsum(np.random.randn(1000)) + 100
    >>> state_rw = create_state()
    >>> updated_state_rw = calculate_hurst_exponent(state_rw, rw_prices)
    >>> hurst_rw = updated_state_rw['results']['hurst_exponent']
    >>> 0.4 < hurst_rw < 0.6 # Hurst for random walk should be around 0.5
    True
    """
    logger = logging.getLogger(__name__)

    if len(time_series) < state['parameters']['min_window_size'] * 2:
        logger.warning(f"Time series length ({len(time_series)}) is too short for Hurst exponent calculation. Min required: {state['parameters']['min_window_size'] * 2}.")
        state['results']['hurst_exponent'] = np.nan
        state['results']['hurst_calculation_error'] = "Series too short"
        return state

    data_for_hurst = time_series
    if use_log_returns:
        # Calculate log returns: log(P_t / P_{t-1})
        # Handle potential zero or negative prices by adding a small epsilon or filtering
        if np.any(time_series <= 0):
            logger.warning("Time series contains non-positive values. Cannot calculate log returns. Using raw series.")
            use_log_returns = False
        else:
            data_for_hurst = np.diff(np.log(time_series))
            state['parameters']['log_data_transformed'] = True
            logger.info("Calculating Hurst exponent on log returns.")

    if len(data_for_hurst) == 0:
        logger.warning("Data for Hurst calculation is empty after transformation. Cannot calculate.")
        state['results']['hurst_exponent'] = np.nan
        state['results']['hurst_calculation_error'] = "Empty data after transformation"
        return state

    # Determine window parameters, preferring function args, then state, then defaults
    min_win = min_window_size if min_window_size is not None else state['parameters'].get('min_window_size', 10)

    # Ensure max_win is an integer. Default to half the series length if not provided in args or state.
    default_max_win = len(data_for_hurst) // 2
    max_win = max_window_size if max_window_size is not None else state['parameters'].get('max_window_size', default_max_win)
    # If max_win is still None (e.g., if state['parameters']['max_window_size'] was None), use default_max_win
    if max_win is None:
        max_win = default_max_win

    # Clamp min_win and max_win to valid ranges
    min_win = max(10, min_win)
    max_win = min(max_win, len(data_for_hurst) // 2)

    # Ensure max_win is greater than min_win for logspace to work and nolds requirements
    if max_win <= min_win:
        logger.warning(f"Adjusting max_window_size. Original: {max_win}, Adjusted: {min_win + 1}")
        max_win = min_win + 1
        if max_win > len(data_for_hurst):
            max_win = len(data_for_hurst) # Ensure it doesn't exceed data length

    n_win = n_windows if n_windows is not None else state['parameters'].get('n_windows', 20) # Use 20 as suggested for nvals

    # Generate nvals for nolds.hurst_rs
    nvals = np.unique(
        np.logspace(
            np.log10(min_win),
            np.log10(max_win),
            num=n_win
        ).astype(int)
    )

    # Exponential backoff for retries in case of nolds internal issues (e.g. singular matrix)
    max_retries = 3
    for attempt in range(max_retries):
        try:
            hurst_exponent = nolds.hurst_rs(
                                                  data_for_hurst,
                                                  nvals=nvals,
                                                  fit=fit_method,
                                                  debug_plot=False) # Keep debug_plot False for non-interactive environments
            state['results']['hurst_exponent'] = hurst_exponent
            logger.info(f"Successfully calculated Hurst exponent: {hurst_exponent:.4f} for series length {len(time_series)}.")
            break # Exit loop on success
        except Exception as e:
            logger.warning(f"Attempt {attempt + 1} failed to calculate Hurst exponent: {e}")
            if attempt < max_retries - 1:
                sleep_time = 2 ** attempt + random.uniform(0, 1) # Exponential backoff with jitter
                logger.info(f"Retrying in {sleep_time:.2f} seconds...")
                time.sleep(sleep_time)
            else:
                logger.error(f"Failed to calculate Hurst exponent after {max_retries} attempts.")
                state['results']['hurst_exponent'] = np.nan
                state['results']['hurst_calculation_error'] = str(e)

    return state

def track_metrics(state: dict, metric_name: str, value: any, series_label: str = 'default') -> dict:
    """
    Tracks and stores key metrics in the state dictionary.

    This function is designed to collect various metrics (e.g., Hurst exponent)
    along with their associated labels (e.g., series type, parameters) into a
    structured format within the state dictionary. It initializes a DataFrame
    for tracking if it doesn't exist and appends new metric entries.

    Parameters
    ----------
    state : dict
        The current state dictionary containing analysis results.
    metric_name : str
        A descriptive name for the metric being tracked (e.g., 'hurst_exponent').
    value : Any
        The value of the metric to track.
    series_label : str, optional
        A label to identify the series from which the metric was derived
        (e.g., 'Random Walk', 'Mean Reverting'). Defaults to 'default'.

    Returns
    -------
    dict
        The updated state dictionary with the new metric appended to
        `state['results']['tracked_metrics']`.

    Examples
    --------
    >>> state = create_state()
    >>> state = track_metrics(state, 'hurst_exponent', 0.52, 'Random Walk')
    >>> 'tracked_metrics' in state['results']
    True
    >>> state['results']['tracked_metrics'].shape[0] == 1
    True

    >>> state = track_metrics(state, 'hurst_exponent', 0.25, 'Mean Reverting')
    >>> state['results']['tracked_metrics'].shape[0] == 2
    True
    """
    logger = logging.getLogger(__name__)

    if 'tracked_metrics' not in state['results'] or not isinstance(state['results']['tracked_metrics'], pd.DataFrame):
        state['results']['tracked_metrics'] = pd.DataFrame(columns=['Series Type', 'Metric Name', 'Value'])
        logger.debug("Initialized 'tracked_metrics' DataFrame.")

    new_metric = pd.DataFrame([{'Series Type': series_label, 'Metric Name': metric_name, 'Value': value}])
    state['results']['tracked_metrics'] = pd.concat([state['results']['tracked_metrics'], new_metric], ignore_index=True)
    logger.info(f"Tracked metric: {metric_name} = {value} for series type: {series_label}.")

    return state

def calculate_rolling_hurst(state: dict, time_series: np.ndarray, window_size: int, step: int = 1,
                            use_log_returns: bool = True) -> dict:
    """
    Calculates the Hurst exponent over a rolling window of the time series.

    Parameters
    ----------
    state : dict
        The current state dictionary. This will be updated with the rolling Hurst exponents.
    time_series : np.ndarray
        The input time series data.
    window_size : int
        The size of the rolling window.
    step : int, optional
        The number of steps to advance the window after each calculation. Defaults to 1.
    use_log_returns : bool, optional
        If `True`, Hurst exponent is calculated on log-returns. Defaults to True.

    Returns
    -------
    dict
        The updated state dictionary with a new entry `state['results']['rolling_hurst']`
        containing a Pandas DataFrame of rolling Hurst exponents.

    Examples
    --------
    >>> state = create_state()
    >>> np.random.seed(42)
    >>> data = np.cumsum(np.random.randn(500)) + 100
    >>> updated_state = calculate_rolling_hurst(state, data, window_size=100, step=10)
    >>> 'rolling_hurst' in updated_state['results']
    True
    >>> isinstance(updated_state['results']['rolling_hurst'], pd.DataFrame)
    True
    >>> updated_state['results']['rolling_hurst'].shape[0] > 0
    True
    """
    logger = logging.getLogger(__name__)

    if len(time_series) < window_size:
        logger.warning(f"Time series length ({len(time_series)}) is less than window size ({window_size}). Cannot calculate rolling Hurst.")
        state['results']['rolling_hurst'] = pd.DataFrame(columns=['Index', 'Hurst Exponent'])
        return state
    if window_size < state['parameters']['min_window_size'] * 2:
        logger.warning(f"Window size ({window_size}) is too small for Hurst exponent calculation. Min required: {state['parameters']['min_window_size'] * 2}.")
        state['results']['rolling_hurst'] = pd.DataFrame(columns=['Index', 'Hurst Exponent'])
        return state

    rolling_hurst_results = []
    data_deque = deque(maxlen=window_size)

    for i in range(len(time_series)):
        data_deque.append(time_series[i])

        if len(data_deque) == window_size and (i - window_size + 1) % step == 0:
            current_window = np.array(data_deque)
            temp_state = state.copy() # Use a temporary state for calculation to avoid modifying main state excessively
            temp_state = calculate_hurst_exponent(temp_state, current_window, use_log_returns=use_log_returns)
            hurst_val = temp_state['results'].get('hurst_exponent')

            if not np.isnan(hurst_val):
                rolling_hurst_results.append({'Index': i - window_size + 1, 'Hurst Exponent': hurst_val})
                logger.debug(f"Calculated rolling Hurst at index {i - window_size + 1}: {hurst_val:.4f}")
            else:
                logger.warning(f"Could not calculate Hurst for window ending at index {i}. Error: {temp_state['results'].get('hurst_calculation_error', 'Unknown error')}")

    state['results']['rolling_hurst'] = pd.DataFrame(rolling_hurst_results)
    logger.info(f"Calculated rolling Hurst exponent for {len(rolling_hurst_results)} windows.")

    return state

# Initialize state
initial_state = create_state()

# --- Generate Synthetic Series ---
series_length = 1000
series_volatility = 0.01
seed = 42

# Random Walk (H ~ 0.5)
rw_series = generate_synthetic_series(initial_state, 'random_walk', series_length, series_volatility, seed=seed)

# Mean Reverting (H < 0.5)
mr_series = generate_synthetic_series(initial_state, 'mean_reverting', series_length, series_volatility, mean_reversion_strength=0.05, seed=seed)

# Trending (H > 0.5)
trend_series = generate_synthetic_series(initial_state, 'trending', series_length, series_volatility, trend_strength=0.0005, seed=seed)

# Store series in a dictionary for easy access and plotting
synthetic_series = {
    'Random Walk': rw_series,
    'Mean Reverting': mr_series,
    'Trending': trend_series
}

# --- Visualize Synthetic Series ---
plt.figure(figsize=(15, 6))
sns.set_style('whitegrid')

for name, series_data in synthetic_series.items():
    plt.plot(series_data, label=name, alpha=0.8)

plt.title('Synthetic Time Series: Random Walk, Mean-Reverting, and Trending')
plt.xlabel('Time Step')
plt.ylabel('Price')
plt.legend()
plt.grid(True)
plt.show()

logging.info("Generated and visualized synthetic series.")

# Initialize a fresh state for Hurst calculations
hurst_state = create_state()

# Calculate Hurst for each synthetic series
for name, series_data in synthetic_series.items():
    logging.info(f"Calculating Hurst for {name} series...")
    # For Mean Reverting, consider using raw series as suggested, otherwise use log returns
    if name == 'Mean Reverting':
        hurst_state = calculate_hurst_exponent(hurst_state, series_data, use_log_returns=False)
    else:
        hurst_state = calculate_hurst_exponent(hurst_state, series_data, use_log_returns=True)

    hurst_value = hurst_state['results'].get('hurst_exponent', np.nan)
    hurst_state = track_metrics(hurst_state, 'hurst_exponent', hurst_value, name)
    logging.info(f"Hurst Exponent for {name}: {hurst_value:.4f}")

# Display summary of Hurst exponents
print("\n--- Summary of Hurst Exponents ---")
display(hurst_state['results']['tracked_metrics'])

logging.info("Calculated and summarized Hurst exponents for synthetic series.")

# Choose one series for rolling analysis, e.g., the Mean Reverting series
series_to_analyze = synthetic_series['Mean Reverting']
series_label = 'Mean Reverting'

# Parameters for rolling Hurst calculation
rolling_window_size = 250
rolling_step = 50

# Initialize a fresh state for rolling Hurst calculation
rolling_hurst_state = create_state()

logging.info(f"Calculating rolling Hurst for {series_label} series with window size {rolling_window_size} and step {rolling_step}...")
# Calculate rolling Hurst on raw series for Mean Reverting as suggested
rolling_hurst_state = calculate_rolling_hurst(rolling_hurst_state, series_to_analyze, rolling_window_size, rolling_step, use_log_returns=False)

# Get rolling Hurst results
rolling_hurst_df = rolling_hurst_state['results']['rolling_hurst']

if not rolling_hurst_df.empty:
    # --- Visualize Rolling Hurst Exponent ---
    plt.figure(figsize=(15, 7))
    sns.set_style('whitegrid')

    # Subplot 1: Original Series
    plt.subplot(2, 1, 1)
    plt.plot(series_to_analyze, label=f'{series_label} Series', color='cornflowerblue')
    plt.title(f'Original {series_label} Series')
    plt.xlabel('Time Step')
    plt.ylabel('Price')
    plt.legend()
    plt.grid(True)

    # Subplot 2: Rolling Hurst Exponent
    plt.subplot(2, 1, 2)
    plt.plot(rolling_hurst_df['Index'], rolling_hurst_df['Hurst Exponent'], color='firebrick', marker='o', markersize=4, linestyle='-')
    plt.axhline(0.5, color='gray', linestyle='--', label='H = 0.5 (Random Walk)')
    plt.fill_between(rolling_hurst_df['Index'], 0, 0.5, color='green', alpha=0.1, label='Mean Reverting (H < 0.5)')
    plt.fill_between(rolling_hurst_df['Index'], 0.5, 1, color='red', alpha=0.1, label='Trending (H > 0.5)')
    plt.title(f'Rolling Hurst Exponent for {series_label} (Window Size: {rolling_window_size}, Step: {rolling_step})')
    plt.xlabel('Starting Index of Window')
    plt.ylabel('Hurst Exponent')
    plt.ylim(0, 1)
    plt.legend()
    plt.grid(True)

    plt.tight_layout()
    plt.show()
    logging.info("Visualized rolling Hurst exponent.")
else:
    logging.warning(f"No rolling Hurst data to plot for {series_label}.")