"""Event Studies — measure abnormal returns around signal triggers and macro announcements."""

from __future__ import annotations

import warnings
from datetime import datetime

import numpy as np
import plotly.graph_objects as go
import polars as pl
from scipy import stats

from utils.reproducibility import set_global_seeds
from utils.style import COLORS  # importing utils.style activates the ml4t Plotly template

warnings.filterwarnings("ignore")

START_DATE = "2018-01-01"
END_DATE = "2024-01-01"
SEED = 42

set_global_seeds(SEED)

from data import load_etfs

etfs = load_etfs()


# Select liquid ETFs for event study
SYMBOLS = ["SPY", "QQQ", "IWM", "TLT", "GLD"]

# Filter
etf_filtered = (
    etfs.filter(pl.col("symbol").is_in(SYMBOLS))
    .filter(
        (pl.col("timestamp") >= datetime.strptime(START_DATE, "%Y-%m-%d"))
        & (pl.col("timestamp") < datetime.strptime(END_DATE, "%Y-%m-%d"))
    )
    .sort(["symbol", "timestamp"])
)

print(f"ETF data: {len(etf_filtered):,} rows")
print(f"Symbols: {etf_filtered['symbol'].n_unique()}")
print(f"Date range: {etf_filtered['timestamp'].min()} to {etf_filtered['timestamp'].max()}")

# Compute daily returns
returns_df = (
    etf_filtered.select(["timestamp", "symbol", "close"])
    .with_columns(pl.col("close").pct_change().over("symbol").alias("return"))
    .drop_nulls()
)

print(f"Returns: {len(returns_df):,} observations")

# Benchmark: SPY as market proxy
benchmark_returns = (
    returns_df.filter(pl.col("symbol") == "SPY")
    .select(["timestamp", "return"])
    .rename({"return": "benchmark_return"})
)

print(f"Benchmark: {len(benchmark_returns):,} days")

def generate_momentum_breakout_events(
    prices: pl.DataFrame,
    lookback: int = 20,
    min_gap_days: int = 21,
) -> pl.DataFrame:
    """
    Generate events when price makes a new N-day high.

    Parameters
    ----------
    prices : DataFrame with timestamp, symbol, close
    lookback : Days to look back for high
    min_gap_days : Minimum gap between events (avoid clustering)

    Returns
    -------
    DataFrame with timestamp, symbol, event_type
    """
    events = []

    for symbol in prices["symbol"].unique().to_list():
        if symbol == "SPY":  # Skip benchmark
            continue

        symbol_data = prices.filter(pl.col("symbol") == symbol).sort("timestamp")
        close_prices = symbol_data["close"].to_numpy()
        timestamps = symbol_data["timestamp"].to_list()

        last_event_idx = -min_gap_days - 1

        for i in range(lookback, len(close_prices) - 30):  # Leave room for event window
            # Check if new high
            if close_prices[i] >= max(close_prices[i - lookback : i]):
                # Check minimum gap
                if i - last_event_idx >= min_gap_days:
                    events.append(
                        {
                            "timestamp": timestamps[i],
                            "symbol": symbol,
                            "event_type": "momentum_breakout",
                        }
                    )
                    last_event_idx = i

    return pl.DataFrame(events)

# Generate events
events_df = generate_momentum_breakout_events(
    etf_filtered.select(["timestamp", "symbol", "close"]),
    lookback=20,
    min_gap_days=30,  # At least 30 days between events per symbol
)

print(f"Generated {len(events_df)} events")
print("\nEvents by symbol:")
events_df.group_by("symbol").len().sort("symbol")

def _estimate_market_model(
    returns_wide: pl.DataFrame,
    event_idx: int,
    symbol: str,
    estimation_window: tuple[int, int],
    min_estimation_obs: int,
) -> tuple[float, float] | None:
    """Estimate market model alpha and beta from estimation window."""
    est_start = event_idx + estimation_window[0]
    est_end = event_idx + estimation_window[1]

    if est_start < 0:
        return None

    est_slice = returns_wide.slice(est_start, est_end - est_start + 1)
    asset_est = est_slice[symbol].to_numpy()
    bench_est = est_slice["benchmark_return"].to_numpy()

    valid = np.isfinite(asset_est) & np.isfinite(bench_est)
    if np.sum(valid) < min_estimation_obs:
        return None

    try:
        slope, intercept, _, _, _ = stats.linregress(bench_est[valid], asset_est[valid])
        return intercept, slope  # alpha, beta
    except Exception:
        return None

def _compute_abnormal_returns(
    returns_wide: pl.DataFrame,
    event_idx: int,
    event_date,
    symbol: str,
    alpha: float,
    beta: float,
    event_window: tuple[int, int],
) -> tuple[list[dict], dict | None]:
    """Compute abnormal returns in the event window."""
    evt_start = event_idx + event_window[0]
    evt_end = event_idx + event_window[1]

    if evt_end >= len(returns_wide):
        return [], None

    evt_slice = returns_wide.slice(evt_start, evt_end - evt_start + 1)
    asset_evt = evt_slice[symbol].to_numpy()
    bench_evt = evt_slice["benchmark_return"].to_numpy()
    evt_dates = evt_slice["timestamp"].to_list()

    ars = []
    car = 0.0
    for i, (date, r_actual, r_market) in enumerate(
        zip(evt_dates, asset_evt, bench_evt, strict=False)
    ):
        if not (np.isfinite(r_actual) and np.isfinite(r_market)):
            continue
        r_expected = alpha + beta * r_market
        ar = r_actual - r_expected
        car += ar
        day_relative = event_window[0] + i
        ars.append(
            {
                "event_date": event_date,
                "symbol": symbol,
                "day": day_relative,
                "ar": ar,
                "car_to_day": car,
            }
        )

    event_car = {
        "event_date": event_date,
        "symbol": symbol,
        "car": car,
        "alpha": alpha,
        "beta": beta,
    }
    return ars, event_car

def _aggregate_to_caar(
    all_ars: list[dict],
    event_cars: list[dict],
) -> tuple[pl.DataFrame, pl.DataFrame, pl.DataFrame]:
    """Aggregate abnormal returns to AAR and CAAR."""
    ar_df = pl.DataFrame(all_ars) if all_ars else pl.DataFrame()
    car_df = pl.DataFrame(event_cars) if event_cars else pl.DataFrame()

    if len(ar_df) > 0:
        daily_aar = (
            ar_df.group_by("day")
            .agg(
                [
                    pl.col("ar").mean().alias("aar"),
                    pl.col("ar").std().alias("std"),
                    pl.len().alias("n"),
                ]
            )
            .sort("day")
            .with_columns((pl.col("std") / pl.col("n").sqrt()).alias("se"))
        )
        # CAAR and its standard error
        daily_aar = daily_aar.with_columns(pl.col("aar").cum_sum().alias("caar"))
        daily_aar = daily_aar.with_columns(
            (pl.col("std").pow(2) / pl.col("n")).cum_sum().sqrt().alias("caar_se")
        )
        daily_aar = daily_aar.with_columns(
            [
                (pl.col("aar") / pl.col("se")).alias("t_stat"),
                (pl.col("caar") / pl.col("caar_se")).alias("caar_t_stat"),
            ]
        )
    else:
        daily_aar = pl.DataFrame()

    return ar_df, car_df, daily_aar

def compute_event_study(
    returns_df: pl.DataFrame,
    benchmark_df: pl.DataFrame,
    events_df: pl.DataFrame,
    estimation_window: tuple[int, int] = (-60, -6),
    event_window: tuple[int, int] = (-5, 10),
    min_estimation_obs: int = 30,
) -> dict:
    """
    Compute event study using the three-stage pipeline above.

    Parameters
    ----------
    returns_df : Long-format returns (timestamp, symbol, return)
    benchmark_df : Benchmark returns (timestamp, benchmark_return)
    events_df : Events (timestamp, symbol)
    estimation_window : (start, end) days relative to event
    event_window : (start, end) days relative to event
    min_estimation_obs : Minimum observations in estimation window

    Returns
    -------
    Dict with:
        - abnormal_returns: DataFrame of AR by event-day
        - event_cars: DataFrame of CAR by event
        - daily_aar: DataFrame of AAR by relative day
    """
    # Create wide-format returns for efficient lookup
    returns_wide = returns_df.pivot(on="symbol", index="timestamp", values="return").sort(
        "timestamp"
    )
    returns_wide = returns_wide.join(benchmark_df, on="timestamp", how="inner")

    dates = returns_wide["timestamp"].to_list()
    date_to_idx = {d: i for i, d in enumerate(dates)}
    symbols = [c for c in returns_wide.columns if c not in ["timestamp", "benchmark_return"]]

    all_ars = []
    event_cars = []

    for row in events_df.iter_rows(named=True):
        event_date = row["timestamp"]
        symbol = row["symbol"]
        if symbol not in symbols or event_date not in date_to_idx:
            continue
        event_idx = date_to_idx[event_date]

        # Check event window upper bound
        if event_idx + event_window[1] >= len(dates):
            continue

        # Stage 1: Estimate market model
        model = _estimate_market_model(
            returns_wide, event_idx, symbol, estimation_window, min_estimation_obs
        )
        if model is None:
            continue
        alpha, beta = model

        # Stage 2: Compute abnormal returns
        ars, event_car = _compute_abnormal_returns(
            returns_wide, event_idx, event_date, symbol, alpha, beta, event_window
        )
        all_ars.extend(ars)
        if event_car:
            event_cars.append(event_car)

    # Stage 3: Aggregate
    ar_df, car_df, daily_aar = _aggregate_to_caar(all_ars, event_cars)

    return {
        "abnormal_returns": ar_df,
        "event_cars": car_df,
        "daily_aar": daily_aar,
        "n_events": len(car_df),
    }

# Run event study
result = compute_event_study(
    returns_df.select(["timestamp", "symbol", "return"]),
    benchmark_returns,
    events_df,
    estimation_window=(-60, -6),
    event_window=(-5, 10),
)

print(f"Processed {result['n_events']} events")

if len(result["event_cars"]) > 0:
    print("\nCAR Summary:")
    cars = result["event_cars"]["car"].to_numpy()
    print(f"  Mean CAR: {np.mean(cars) * 100:.2f}%")
    print(f"  Median CAR: {np.median(cars) * 100:.2f}%")
    print(f"  Std CAR: {np.std(cars) * 100:.2f}%")

if len(result["daily_aar"]) > 0:
    daily_aar = result["daily_aar"]

    fig = go.Figure()

    days = daily_aar["day"].to_list()
    caar = daily_aar["caar"].to_numpy() * 100  # Convert to percent
    caar_se = daily_aar["caar_se"].to_numpy() * 100

    # 95% confidence band (Var(CAAR_t) = cumulative sum of daily AAR variances)
    upper = caar + 1.96 * caar_se
    lower = caar - 1.96 * caar_se

    fig.add_trace(
        go.Scatter(
            x=days + days[::-1],
            y=np.concatenate([upper, lower[::-1]]).tolist(),
            fill="toself",
            fillcolor="rgba(10, 22, 40, 0.15)",  # COLORS["blue"] at 15% opacity
            line=dict(width=0),
            name="95% CI",
        )
    )

    # CAAR line drawn on top of the band
    fig.add_trace(
        go.Scatter(
            x=days,
            y=caar,
            mode="lines+markers",
            name="CAAR",
            line=dict(color=COLORS["blue"], width=2),
        )
    )

    # Event day marker and zero reference
    fig.add_vline(
        x=0,
        line_dash="dash",
        line_color=COLORS["amber"],
        annotation_text="Event day",
        annotation_position="top left",
    )
    fig.add_hline(y=0, line_dash="dot", line_color=COLORS["neutral"])

    fig.update_layout(
        title="Momentum breakouts earn ~1.3% abnormal return by the event day, then partly fade",
        xaxis_title="Trading days relative to event",
        yaxis_title="Cumulative average abnormal return (%)",
        height=500,
    )

    fig.show()

if len(result["daily_aar"]) > 0:
    daily_aar = result["daily_aar"]

    days = daily_aar["day"].to_list()
    aar_pct = (daily_aar["aar"] * 100).to_list()
    t_stats = daily_aar["t_stat"].to_list()

    # Color by significance; highlight the event day
    bar_colors = [
        COLORS["amber"] if d == 0 else (COLORS["blue"] if abs(t) > 1.96 else COLORS["silver_muted"])
        for d, t in zip(days, t_stats, strict=True)
    ]

    fig = go.Figure()
    fig.add_trace(go.Bar(x=days, y=aar_pct, marker_color=bar_colors, name="AAR"))
    fig.add_hline(y=0, line_dash="dot", line_color=COLORS["neutral"])
    fig.add_vline(x=0, line_dash="dash", line_color=COLORS["amber"])

    fig.update_layout(
        title="The abnormal return is earned on the breakout day; surrounding days are noise",
        xaxis_title="Trading days relative to event",
        yaxis_title="Average abnormal return (%)",
        height=400,
    )

    fig.show()

from ml4t.diagnostic.config import EventConfig
from ml4t.diagnostic.config.event_config import WindowSettings
from ml4t.diagnostic.evaluation import EventStudyAnalysis

# Prepare data in library format
# Returns: date, asset, return
lib_returns = returns_df.select(
    pl.col("timestamp").alias("date"),
    pl.col("symbol").alias("asset"),
    pl.col("return"),
)

# Benchmark: date, return
lib_benchmark = benchmark_returns.rename({"timestamp": "date", "benchmark_return": "return"})

# Events: date, asset
lib_events = events_df.select(
    pl.col("timestamp").alias("date"),
    pl.col("symbol").alias("asset"),
)

# Configure event study
config = EventConfig(
    window=WindowSettings(
        estimation_start=-60,
        estimation_end=-6,
        event_start=-5,
        event_end=10,
    ),
    model="market_model",
    min_estimation_obs=30,
)

# Run library event study
lib_analysis = EventStudyAnalysis(
    returns=lib_returns,
    events=lib_events,
    benchmark=lib_benchmark,
    config=config,
)
lib_result = lib_analysis.run()

# Compare results
print("=== Library EventStudyAnalysis Results ===\n")
print(lib_result.summary())

# The library defaults to the Boehmer et al. (1991) BMP test, which is robust to
# event-induced variance - the manual t-test above assumes constant variance.
print("\n=== Robust significance test (library) ===")
print(f"Test: {lib_result.test_name}")
print(f"Test statistic: {lib_result.test_statistic:.2f}")
print(f"P-value: {lib_result.p_value:.4f}")
print(f"Significant at 5%: {'Yes' if lib_result.p_value < 0.05 else 'No'}")

if len(result["event_cars"]) > 0:
    cars = result["event_cars"]["car"].to_numpy() * 100

    fig = go.Figure()

    fig.add_trace(
        go.Histogram(
            x=cars,
            nbinsx=25,
            marker_color=COLORS["blue"],
            name="CAR distribution",
        )
    )

    # Zero reference and mean (mean highlighted in amber)
    fig.add_vline(
        x=0,
        line_dash="dot",
        line_color=COLORS["neutral"],
        annotation_text="Zero",
        annotation_position="top left",
    )
    fig.add_vline(
        x=float(np.mean(cars)),
        line_dash="dash",
        line_color=COLORS["amber"],
        annotation_text=f"Mean: {np.mean(cars):.2f}%",
        annotation_position="top right",
    )

    fig.update_layout(
        title="Breakout CARs skew positive, with a mean of +1.1% over the 16-day window",
        xaxis_title="Cumulative abnormal return over event window (%)",
        yaxis_title="Number of events",
        height=400,
    )

    fig.show()

    # Statistical test: Mean CAR = 0
    t_stat, p_value = stats.ttest_1samp(cars, 0)

    print("\nStatistical Test (H0: Mean CAR = 0):")
    print(f"  Mean CAR: {np.mean(cars):.2f}%")
    print(f"  Median CAR: {np.median(cars):.2f}%")
    print(f"  T-statistic: {t_stat:.2f}")
    print(f"  P-value: {p_value:.4f}")
    print(f"  Significant at 5%: {'Yes' if p_value < 0.05 else 'No'}")

if len(result["abnormal_returns"]) > 0:
    ar_df = result["abnormal_returns"]

    # Create unique event identifier (date + symbol can have multiple events)
    ar_df = ar_df.with_columns(
        (pl.col("event_date").dt.strftime("%Y-%m-%d") + "_" + pl.col("symbol")).alias("event_id")
    )

    # Pivot to wide format (event x day)
    ar_pivot = ar_df.pivot(on="day", index="event_id", values="ar").sort("event_id")

    # Convert to numpy for heatmap
    day_cols = sorted([c for c in ar_pivot.columns if c != "event_id"], key=lambda x: int(x))
    ar_matrix = ar_pivot.select(day_cols).to_numpy() * 100

    # Limit to 30 events for readability
    if len(ar_matrix) > 30:
        ar_matrix = ar_matrix[:30]
        event_ids = ar_pivot["event_id"].to_list()[:30]
    else:
        event_ids = ar_pivot["event_id"].to_list()

    # ML4T diverging scale centered at zero (red = negative, green = positive AR)
    diverging_scale = [
        [0.0, COLORS["negative"]],
        [0.5, COLORS["silver_muted"]],
        [1.0, COLORS["positive"]],
    ]

    fig = go.Figure(
        data=go.Heatmap(
            z=ar_matrix,
            x=day_cols,
            y=event_ids,
            colorscale=diverging_scale,
            zmid=0,
            colorbar=dict(title="AR (%)"),
        )
    )

    fig.add_vline(x=0, line_dash="dash", line_color=COLORS["amber"], line_width=2)

    fig.update_layout(
        title="Event-day abnormal returns stand out across individual breakout events",
        xaxis_title="Trading days relative to event",
        yaxis_title="Event (date + symbol)",
        height=600,
    )

    fig.show()

def validate_signal_with_event_study(
    returns_df: pl.DataFrame,
    benchmark_df: pl.DataFrame,
    signal_df: pl.DataFrame,
    signal_column: str = "signal",
    threshold: float = 2.0,
    event_window: tuple[int, int] = (-5, 10),
) -> dict:
    """
    Validate a trading signal using event study methodology.

    Parameters
    ----------
    returns_df : Long-format returns
    benchmark_df : Benchmark returns
    signal_df : Signal values (timestamp, symbol, signal)
    signal_column : Column name for signal
    threshold : Z-score threshold for event trigger
    event_window : Days around event to analyze

    Returns
    -------
    Dict with long and short event study results
    """
    # Z-score signals cross-sectionally
    signal_zscored = signal_df.with_columns(
        (
            (pl.col(signal_column) - pl.col(signal_column).mean().over("timestamp"))
            / pl.col(signal_column).std().over("timestamp")
        ).alias("zscore")
    )

    # Generate events from extreme signals
    long_events = (
        signal_zscored.filter(pl.col("zscore") > threshold)
        .select(["timestamp", "symbol"])
        .with_columns(pl.lit("long_signal").alias("event_type"))
    )

    short_events = (
        signal_zscored.filter(pl.col("zscore") < -threshold)
        .select(["timestamp", "symbol"])
        .with_columns(pl.lit("short_signal").alias("event_type"))
    )

    print(f"Long signal events: {len(long_events)}")
    print(f"Short signal events: {len(short_events)}")

    results = {}

    if len(long_events) > 10:
        results["long"] = compute_event_study(
            returns_df, benchmark_df, long_events, event_window=event_window
        )

    if len(short_events) > 10:
        results["short"] = compute_event_study(
            returns_df, benchmark_df, short_events, event_window=event_window
        )

    return results

# Example: Validate momentum signal
# Create momentum signal
prices_wide = (
    etf_filtered.select(["timestamp", "symbol", "close"])
    .pivot(on="symbol", index="timestamp", values="close")
    .sort("timestamp")
)

symbols = [c for c in prices_wide.columns if c != "timestamp"]

# 21-day momentum
momentum = prices_wide.select(
    pl.col("timestamp"), *[(pl.col(s) / pl.col(s).shift(21) - 1).alias(s) for s in symbols]
)

# Melt to long format
momentum_long = (
    momentum.unpivot(index="timestamp", variable_name="symbol", value_name="momentum")
    .drop_nulls()
    .filter(pl.col("momentum").is_finite())
)

print(f"Momentum signal: {len(momentum_long):,} observations")

# Validate momentum signal
validation = validate_signal_with_event_study(
    returns_df.select(["timestamp", "symbol", "return"]),
    benchmark_returns,
    momentum_long,
    signal_column="momentum",
    threshold=1.5,
)

if "long" in validation and len(validation["long"]["daily_aar"]) > 0:
    print("\nLong Signal Validation:")
    aar_long = validation["long"]["daily_aar"]
    final_caar = aar_long["caar"].to_numpy()[-1] * 100
    final_t = aar_long["caar_t_stat"].to_numpy()[-1]
    print(f"  Final CAAR: {final_caar:.2f}%")
    print(f"  CAAR t-stat: {final_t:.2f}")
    print(f"  Significant: {'Yes' if abs(final_t) > 1.96 else 'No'}")

if "short" in validation and len(validation["short"]["daily_aar"]) > 0:
    print("\nShort Signal Validation:")
    aar_short = validation["short"]["daily_aar"]
    final_caar = aar_short["caar"].to_numpy()[-1] * 100
    final_t = aar_short["caar_t_stat"].to_numpy()[-1]
    print(f"  Final CAAR: {final_caar:.2f}%")
    print(f"  CAAR t-stat: {final_t:.2f}")
    print(f"  Significant: {'Yes' if abs(final_t) > 1.96 else 'No'}")