# Install necessary libraries
!pip install pandas numpy matplotlib seaborn

print("Dependencies installed successfully.")

import logging
from collections import deque

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

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

logger.info("All libraries imported successfully.")

def create_roll_model_state(prices: list | np.ndarray) -> dict:
    """
    Initializes the state dictionary for Roll model calculations.

    Parameters
    ----------
    prices : list | np.ndarray
        A list or NumPy array of transaction prices.

    Returns
    -------
    dict
        An initial state dictionary containing the 'prices' and a status.
    """
    logger.info("Creating initial Roll model state...")
    state = {
        'prices': np.asarray(prices),
        'status': 'initialized'
    }
    logger.info(f"Initial state created with {len(prices)} prices.")
    return state

def calculate_price_changes(state: dict) -> dict:
    """
    Calculates the first-order price differences (delta P_t).

    Parameters
    ----------
    state : dict
        Current state dictionary containing 'prices'.

    Returns
    -------
    dict
        Updated state with 'price_changes' and a status.
    """
    logger.info("Calculating price changes...")
    if 'prices' not in state or len(state['prices']) < 2:
        logger.warning("Insufficient prices to calculate changes. At least two prices are needed.")
        state['status'] = 'error_insufficient_prices'
        state['price_changes'] = np.array([])
        return state

    price_changes = np.diff(state['prices'])
    state['price_changes'] = price_changes
    state['status'] = 'price_changes_calculated'
    logger.info(f"Calculated {len(price_changes)} price changes.")
    return state

def estimate_roll_spread(state: dict) -> dict:
    """
    Estimates the effective bid-ask spread using the Roll model.

    Parameters
    ----------
    state : dict
        Current state dictionary containing 'price_changes'.

    Returns
    -------
    dict
        Updated state with 'roll_spread' and a status.
    """
    logger.info("Estimating Roll spread...")
    if 'price_changes' not in state or len(state['price_changes']) < 2:
        logger.warning("Insufficient price changes to estimate spread. At least two changes are needed.")
        state['status'] = 'error_insufficient_price_changes'
        state['roll_spread'] = np.nan
        return state

    dP = state['price_changes']
    # Calculate covariance between dP_t and dP_{t-1}
    # Covariance matrix for [dP[:-1], dP[1:]]
    cov_matrix = np.cov(dP[:-1], dP[1:])
    cov_dp_lagged_dp = cov_matrix[0, 1]

    logger.info(f"Covariance of successive price changes: {cov_dp_lagged_dp:.6f}")

    if cov_dp_lagged_dp >= 0:
        logger.warning("Covariance of successive price changes is non-negative. This may indicate a lack of bid-ask bounce or data issues.")
        state['status'] = 'warning_non_negative_covariance'
        state['roll_spread'] = np.nan # Spread cannot be calculated from non-negative covariance
    else:
        roll_spread = np.sqrt(-2 * cov_dp_lagged_dp)
        state['roll_spread'] = roll_spread
        state['status'] = 'roll_spread_estimated'
        logger.info(f"Roll model spread estimated: {roll_spread:.6f}")

    return state

# --- Data Simulation ---

# Parameters for simulation
N_TRADES = 1000  # Number of simulated trades
TRUE_SPREAD = 0.05 # Actual bid-ask spread (e.g., 5 cents)
MID_PRICE_START = 100.0
VOLATILITY = 0.01 # Volatility of mid-price movements

np.random.seed(42) # for reproducibility

logger.info(f"Simulating {N_TRADES} transaction prices with true spread of {TRUE_SPREAD}")

# Simulate mid-price movements (e.g., a random walk)
mid_prices = MID_PRICE_START + np.cumsum(np.random.normal(0, VOLATILITY, N_TRADES))

# Simulate transactions occurring at bid or ask based on an unobserved order flow
# Assume 50% chance of hitting bid or ask, creating the bid-ask bounce
indicator = np.random.choice([-1, 1], N_TRADES) # -1 for bid, 1 for ask
transaction_prices = mid_prices + indicator * (TRUE_SPREAD / 2)

# Create a DataFrame for better data handling and visualization
df_trades = pd.DataFrame({
    'trade_id': range(N_TRADES),
    'mid_price': mid_prices,
    'indicator': indicator,
    'transaction_price': transaction_prices
})

logger.info("Simulated data preview:")
display(df_trades.head())

# --- Apply Roll Model Functions ---

logger.info("Applying Roll model functions to simulated data...")

# 1. Create initial state
roll_state = create_roll_model_state(df_trades['transaction_price'].values)
logger.info(f"State after initialization: {roll_state['status']}")

# 2. Calculate price changes
roll_state = calculate_price_changes(roll_state)
logger.info(f"State after calculating price changes: {roll_state['status']}")

# 3. Estimate Roll spread
roll_state = estimate_roll_spread(roll_state)
logger.info(f"State after estimating Roll spread: {roll_state['status']}")

estimated_roll_spread = roll_state.get('roll_spread')

logger.info(f"True Spread: {TRUE_SPREAD:.4f}")
logger.info(f"Estimated Roll Spread: {estimated_roll_spread:.4f}")

# --- Visualization ---

# Plot 1: Transaction Prices Over Time
fig1, ax1 = plt.subplots(figsize=(12, 6))
sns.lineplot(x='trade_id', y='transaction_price', data=df_trades, ax=ax1, label='Transaction Price', alpha=0.8)
sns.lineplot(x='trade_id', y='mid_price', data=df_trades, ax=ax1, label='Mid-Price', linestyle='--', color='red')
ax1.set_title('Simulated Transaction Prices and Mid-Prices Over Time')
ax1.set_xlabel('Trade ID')
ax1.set_ylabel('Price')
ax1.legend()
ax1.grid(True, linestyle='--', alpha=0.6)
plt.tight_layout()
plt.show()

# Plot 2: Distribution of Price Changes
fig2, ax2 = plt.subplots(figsize=(10, 6))
sns.histplot(roll_state['price_changes'], bins=50, kde=True, ax=ax2)
ax2.set_title('Distribution of Transaction Price Changes')
ax2.set_xlabel('Price Change (dP)')
ax2.set_ylabel('Frequency')
ax2.grid(True, linestyle='--', alpha=0.6)
plt.tight_layout()
plt.show()

# Plot 3: Autocorrelation of Price Changes (Lag 1)
# We'll manually calculate and plot for demonstration
dP_t = roll_state['price_changes'][1:]
dP_t_minus_1 = roll_state['price_changes'][:-1]

fig3, ax3 = plt.subplots(figsize=(10, 6))
sns.scatterplot(x=dP_t_minus_1, y=dP_t, ax=ax3, alpha=0.6)
ax3.set_title('Autocorrelation Plot of Price Changes (dP_t vs dP_{t-1})')
ax3.set_xlabel('Price Change at t-1 (dP_{t-1})')
ax3.set_ylabel('Price Change at t (dP_t)')
ax3.grid(True, linestyle='--', alpha=0.6)
plt.tight_layout()
plt.show()

# --- Summary Statistics ---

summary_data = {
    'Metric': ['True Spread', 'Estimated Roll Spread', 'Number of Trades', 'Cov(dP_t, dP_{t-1})'],
    'Value': [TRUE_SPREAD, estimated_roll_spread, N_TRADES, np.cov(roll_state['price_changes'][:-1], roll_state['price_changes'][1:])[0, 1]]
}
summary_df = pd.DataFrame(summary_data)

logger.info("Summary of Roll Model Estimation:")
display(summary_df)