import numpy as np
import pandas as pd
import warnings
import warnings
warnings.filterwarnings('ignore')

def edge_rolling(data: pd.DataFrame, window: int, sign: bool = False, **kwargs) -> pd.Series:
    """
    Compute rolling bid-ask spread estimates from OHLC prices.

    This function uses a rolling window to estimate the bid-ask spread as described in
    Ardia, Guidotti, & Kroencke (2024). A value of 0.01 means a 1% spread.

    Parameters:
        data : pd.DataFrame
            DataFrame with columns 'open', 'high', 'low', and 'close' (case-insensitive).
        window : int or valid rolling window parameter
            Size of the moving window.
        sign : bool, default False
            If True, the returned spread keeps its sign.
        kwargs : dict
            Extra parameters for the pandas rolling function.

    Returns:
        pd.Series: Rolling spread estimates.
    """
    # Standardize column names and take log-prices.
    df = data.rename(columns=str.lower, inplace=False)
    log_open  = np.log(df['open'])
    log_high  = np.log(df['high'])
    log_low   = np.log(df['low'])
    log_close = np.log(df['close'])
    log_mid   = (log_high + log_low) / 2.0

    # Create lagged series for previous period prices.
    log_high_prev  = log_high.shift(1)
    log_low_prev   = log_low.shift(1)
    log_close_prev = log_close.shift(1)
    log_mid_prev   = log_mid.shift(1)

    # Compute various log-returns.
    r1 = log_mid - log_open        # mid - open
    r2 = log_open - log_mid_prev     # open - previous mid
    r3 = log_mid - log_close_prev    # mid - previous close
    r4 = log_close_prev - log_mid_prev  # previous close - previous mid
    r5 = log_open - log_close_prev   # open - previous close

    # Create an indicator for non-flat periods.
    tau = np.where(
        np.isnan(log_high) | np.isnan(log_low) | np.isnan(log_close_prev),
        np.nan,
        (log_high != log_low) | (log_low != log_close_prev)
    )
    # Indicators to check if open or previous close differ from high/low.
    ind_o_h = tau * np.where(np.isnan(log_open) | np.isnan(log_high), np.nan, log_open != log_high)
    ind_o_l = tau * np.where(np.isnan(log_open) | np.isnan(log_low), np.nan, log_open != log_low)
    ind_c_h = tau * np.where(np.isnan(log_close_prev) | np.isnan(log_high_prev), np.nan, log_close_prev != log_high_prev)
    ind_c_l = tau * np.where(np.isnan(log_close_prev) | np.isnan(log_low_prev), np.nan, log_close_prev != log_low_prev)

    # Calculate products of returns that will be used in moment conditions.
    prod_12 = r1 * r2
    prod_34 = r3 * r4
    prod_15 = r1 * r5
    prod_45 = r4 * r5
    tau_r1  = tau * r1
    tau_r2  = tau * r2
    tau_r4  = tau * r4
    tau_r5  = tau * r5

    # Collect all intermediate values in a DataFrame.
    vals = pd.DataFrame({
        'prod_12': prod_12,
        'prod_34': prod_34,
        'prod_15': prod_15,
        'prod_45': prod_45,
        'tau': tau,
        'r1': r1,
        'tau_r2': tau_r2,
        'r3': r3,
        'tau_r4': tau_r4,
        'r5': r5,
        'prod_12_sq': prod_12 ** 2,
        'prod_34_sq': prod_34 ** 2,
        'prod_15_sq': prod_15 ** 2,
        'prod_45_sq': prod_45 ** 2,
        'prod_12_34': prod_12 * prod_34,
        'prod_15_45': prod_15 * prod_45,
        'tau_r2_r2': tau_r2 * r2,
        'tau_r4_r4': tau_r4 * r4,
        'tau_r5_r5': tau_r5 * r5,
        'tau_r2_prod12': tau_r2 * prod_12,
        'tau_r4_prod34': tau_r4 * prod_34,
        'tau_r5_prod15': tau_r5 * prod_15,
        'tau_r4_prod45': tau_r4 * prod_45,
        'tau_r4_prod12': tau_r4 * prod_12,
        'tau_r2_prod34': tau_r2 * prod_34,
        'tau_r2_r4': tau_r2 * r4,
        'tau_r1_prod45': tau_r1 * prod_45,
        'tau_r5_prod45': tau_r5 * prod_45,
        'tau_r4_r5': tau_r4 * r5,
        'tau_r5_only': tau_r5,
        'ind_o_h': ind_o_h,
        'ind_o_l': ind_o_l,
        'ind_c_h': ind_c_h,
        'ind_c_l': ind_c_l
    }, index=df.index)

    # The first observation is not usable (due to shifting).
    vals.iloc[0] = np.nan

    # Adjust window length and min_periods to account for the lag.
    window_adj = window - 1 if isinstance(window, (int, np.integer)) else window
    if 'min_periods' in kwargs and isinstance(kwargs['min_periods'], (int, np.integer)):
        kwargs['min_periods'] = max(0, kwargs['min_periods'] - 1)

    # Compute rolling means for each column.
    roll_vals = vals.rolling(window=window_adj, **kwargs).mean()

    # Calculate probabilities needed for the estimator.
    p_tau = roll_vals['tau']
    p_open = roll_vals['ind_o_h'] + roll_vals['ind_o_l']
    p_close = roll_vals['ind_c_h'] + roll_vals['ind_c_l']

    # Count valid tau observations.
    count_tau = vals['tau'].rolling(window=window_adj, **kwargs).sum()
    # Mark window as missing if there are fewer than 2 valid periods or zero probabilities.
    roll_vals[(count_tau < 2) | (p_open == 0) | (p_close == 0)] = np.nan

    # Compute coefficients from the rolling means.
    a1 = -4.0 / p_open
    a2 = -4.0 / p_close
    a3 = roll_vals['r1'] / p_tau
    a4 = roll_vals['tau_r4'] / p_tau
    a5 = roll_vals['r3'] / p_tau
    a6 = roll_vals['r5'] / p_tau

    a12 = 2 * a1 * a2
    a11 = a1 ** 2
    a22 = a2 ** 2
    a33 = a3 ** 2
    a55 = a5 ** 2
    a66 = a6 ** 2

    # Compute expectations from moment conditions.
    E1 = a1 * (roll_vals['prod_12'] - a3 * roll_vals['tau_r2']) + \
         a2 * (roll_vals['prod_34'] - a4 * roll_vals['r3'])
    E2 = a1 * (roll_vals['prod_15'] - a3 * roll_vals['tau_r5_only']) + \
         a2 * (roll_vals['prod_45'] - a4 * roll_vals['r5'])

    # Compute variances from the moments.
    V1 = - E1**2 + (
        a11 * (roll_vals['prod_12_sq'] - 2 * a3 * roll_vals['tau_r2_prod12'] + a33 * roll_vals['tau_r2_r2']) +
        a22 * (roll_vals['prod_34_sq'] - 2 * a5 * roll_vals['tau_r4_prod34'] + a55 * roll_vals['tau_r4_r4']) +
        a12 * (roll_vals['prod_12_34'] - a3 * roll_vals['tau_r2_prod34'] - a5 * roll_vals['tau_r4_prod12'] + a3 * a5 * roll_vals['tau_r2_r4'])
    )
    V2 = - E2**2 + (
        a11 * (roll_vals['prod_15_sq'] - 2 * a3 * roll_vals['tau_r5_prod15'] + a33 * roll_vals['tau_r5_r5']) +
        a22 * (roll_vals['prod_45_sq'] - 2 * a6 * roll_vals['tau_r4_prod45'] + a66 * roll_vals['tau_r4_r4']) +
        a12 * (roll_vals['prod_15_45'] - a3 * roll_vals['tau_r5_prod45'] - a6 * roll_vals['tau_r4_r5'] + a3 * a6 * roll_vals['tau_r4_r5'])
    )

    tot_var = V1 + V2
    # If variance is positive, use a weighted combination; otherwise, take a simple average.
    s2 = np.where(tot_var > 0, (V2 * E1 + V1 * E2) / tot_var, (E1 + E2) / 2.0)
    spread = np.sqrt(np.abs(s2))
    if sign:
        spread *= np.sign(s2)

    return pd.Series(spread, index=df.index)

import yfinance as yf
import pandas as pd

ticker = "AAPL"
data = yf.download(ticker, start="2015-01-01", end="2025-12-31", auto_adjust=True)

#print("Downloaded data:")
print(data.head())
#print("Columns:", data.columns)

# Check if the columns are a MultiIndex
if isinstance(data.columns, pd.MultiIndex):
    #print("MultiIndex detected.")
    #print("Level 0 values:", data.columns.get_level_values(0).unique())
    #print("Level 1 values:", data.columns.get_level_values(1).unique())
    # In the new structure, level 0 contains the actual field names.
    data.columns = data.columns.get_level_values(0)
    #print("Columns after resetting:", data.columns)

# Extract the OHLC data.
ohlc = data[['Open', 'High', 'Low', 'Close']]
#print("OHLC data:")
ohlc

# Ensure we have the required data
ohlc = data[['Open', 'High', 'Low', 'Close']]
volume = data['Volume']

# Compute rolling spread with a 20-day window
rolling_spreads = edge_rolling(ohlc, window=20, min_periods=10, sign=False)

rolling_spreads

import matplotlib.pyplot as plt
import matplotlib.dates as mdates

# Compute daily returns and rolling volatility (20-day window)
returns = ohlc['Close'].pct_change()
rolling_vol = returns.rolling(window=20).std()

# Compute upper and lower spread bands around close price
upper_band = ohlc['Close'] * (1 + rolling_spreads / 2)
lower_band = ohlc['Close'] * (1 - rolling_spreads / 2)

# Use dark theme
plt.style.use('dark_background')

fig, ax1 = plt.subplots(figsize=(26, 7))

# --- 1) Volume (Left y-axis) ---
ax1.bar(ohlc.index, volume, color='gray', alpha=0.2, label='Volume')
ax1.set_ylabel("Volume", color='gray')
ax1.tick_params(axis='y', colors='gray')

# --- 2) Close Price (Right y-axis) ---
ax2 = ax1.twinx()
ax2.plot(ohlc.index, ohlc['Close'], color='lime', linewidth=2, label='Close Price')
ax2.set_ylabel("Close Price", color='lime')
ax2.tick_params(axis='y', colors='lime')

# Spread bands on the same axis as Close Price
ax2.fill_between(ohlc.index, upper_band, lower_band, color='gray', alpha=0.3, label='Spread Band')

# --- 3) Rolling Volatility (Offset right y-axis) ---
ax3 = ax1.twinx()
ax3.spines["right"].set_position(("outward", 60))  # first offset on right
ax3.plot(rolling_vol.index, rolling_vol, color='orange', linestyle='dashed', linewidth=1, alpha = 0.3, label='Rolling Volatility')
ax3.set_ylabel("Rolling Volatility", color='orange')
ax3.tick_params(axis='y', colors='orange')

# --- 4) Rolling Spread (Further offset right y-axis) ---
ax4 = ax1.twinx()
ax4.spines["right"].set_position(("outward", 120))  # second offset on right, avoids overlap
ax4.plot(rolling_spreads.index, rolling_spreads, color='blue', linewidth=1.5, label='Rolling Spread')
ax4.set_ylabel("Rolling Spread", color='blue')
ax4.tick_params(axis='y', colors='blue')

# Format x-axis for readability
ax1.xaxis.set_major_formatter(mdates.DateFormatter('%Y-%m'))
ax1.xaxis.set_major_locator(mdates.YearLocator())

# Title
fig.suptitle("Rolling Spread, Volume, Close Price, and Rolling Volatility", fontsize=14)

# --- Legend placements (stacked in upper-left corner) ---
ax1.legend(loc="upper left", bbox_to_anchor=(0, 1))       # Volume
ax2.legend(loc="upper left", bbox_to_anchor=(0, 0.85))    # Close Price & Spread Band
ax3.legend(loc="upper left", bbox_to_anchor=(0, 0.7))     # Rolling Vol
ax4.legend(loc="upper left", bbox_to_anchor=(0, 0.55))    # Rolling Spread

plt.show()

import matplotlib.pyplot as plt
import numpy as np

# Compute daily returns and rolling volatility (20-day window)
returns = ohlc['Close'].pct_change()
rolling_vol = returns.rolling(window=20).std()

# --- Set adjustable lag period ---
lag_period = 4  # Change this value to adjust the lag

# Create lagged series.
lagged_spreads = rolling_spreads.shift(lag_period)
lagged_volume = volume.shift(lag_period)
lagged_vol = rolling_vol.shift(lag_period)

# --- Helper function for labeling ---
def label_text(var_name, lag):
    return f"{var_name} (lag={lag})" if lag > 1 else f"Lagged {var_name}"

# Define subplot titles
titles = [
    f"Returns vs {label_text('Spreads', lag_period)}",
    f"Spreads vs {label_text('Volume', lag_period)}",
    f"Returns vs {label_text('Volume', lag_period)}",
    f"Volatility vs {label_text('Spreads', lag_period)}",
    f"{label_text('Volatility', lag_period)} vs Spreads"
]

# --- Prepare data for each plot ---
data_pairs = [
    (lagged_spreads, returns),      # Plot 1: Returns vs Lagged Spreads
    (lagged_volume, rolling_spreads),  # Plot 2: Spreads vs Lagged Volume
    (lagged_volume, returns),       # Plot 3: Returns vs Lagged Volume
    (lagged_spreads, rolling_vol),  # Plot 4: Rolling Volatility vs Lagged Spreads
    (lagged_vol, rolling_spreads)   # Plot 5: Lagged Volatility vs Current Spreads
]

# Colors for each plot
colors = ["blue", "red", "green", "orange", "magenta"]

# Create subplots: 1 row, 5 columns
fig, axes = plt.subplots(1, 5, figsize=(25, 6))
plt.style.use('dark_background')  # Dark theme

for i, ax in enumerate(axes):
    # Mask NaN values
    x, y = data_pairs[i]
    mask = ~np.isnan(x) & ~np.isnan(y)
    ax.scatter(x[mask], y[mask], color=colors[i], alpha=0.7, s=10)  # Scatter plot

    # Titles and labels
    ax.set_title(titles[i], fontsize=12)
    ax.set_xlabel(label_text("Spreads" if i in [0, 3] else "Volume", lag_period), fontsize=10)
    ax.set_ylabel("Returns" if i in [0, 2] else "Spreads" if i in [1, 4] else "Volatility", fontsize=10)

# Adjust layout for better readability
plt.tight_layout()
plt.suptitle(f"Lagged Value Analysis (lag={lag_period})", fontsize=14, y=1.02)

plt.show()

import matplotlib.pyplot as plt
import pandas as pd

# Set parameters
window_corr = 60  # Rolling window for correlation analysis
lag_period = 1    # Adjust this value to set the lag

# Compute lagged series
lagged_volume = volume.shift(lag_period)
lagged_vol = rolling_vol.shift(lag_period)

# Compute rolling correlations using lagged values
rolling_corr_volume_spread = lagged_volume.rolling(window=window_corr).corr(rolling_spreads)
rolling_corr_volatility_spread = lagged_vol.rolling(window=window_corr).corr(rolling_spreads)

# Use dark theme
plt.style.use('dark_background')

# Create figure
fig, ax = plt.subplots(figsize=(25, 6))

# Plot rolling correlation (lagged volume vs spread)
ax.plot(rolling_corr_volume_spread.index, rolling_corr_volume_spread, color='red', linewidth=2,
        label=f"Rolling Corr (Lagged Volume vs Spread, lag={lag_period}, window={window_corr})")

# Plot rolling correlation (lagged volatility vs spread)
ax.plot(rolling_corr_volatility_spread.index, rolling_corr_volatility_spread, color='orange', linewidth=2,
        label=f"Rolling Corr (Lagged Volatility vs Spread, lag={lag_period}, window={window_corr})")

# Labels and title
ax.set_title(f"Rolling Correlations with Lag (Window = {window_corr}, Lag = {lag_period})", fontsize=14)
ax.set_xlabel("Date")
ax.set_ylabel("Correlation")

# Add legend
ax.legend(loc="upper left")

# Improve x-axis readability
plt.xticks(rotation=45)

# Show plot
plt.show()