"""Macro Data Alignment — align multi-frequency macro data for daily trading models with PIT correctness."""

import warnings
from datetime import date

import plotly.graph_objects as go
import polars as pl
from plotly.subplots import make_subplots

from data import load_etfs, load_macro
from utils import ML4T_DATA_PATH
from utils.reproducibility import set_global_seeds

# Importing utils.style registers and activates the ML4T Plotly template
# (palette, fonts, backgrounds) as the repo-wide default.
from utils.style import COLORS

warnings.filterwarnings("ignore")

print(f"Data directory: {ML4T_DATA_PATH}")

# Production defaults — Papermill injects overrides for CI
SEED = 42

set_global_seeds(SEED)

# Frequency overview
frequencies = pl.DataFrame(
    {
        "series": ["S&P 500", "Initial Claims", "CPI", "GDP", "10Y Treasury"],
        "frequency": ["Daily", "Weekly", "Monthly", "Quarterly", "Daily"],
        "typical_lag_days": [0, 5, 15, 30, 0],
        "observations_per_year": [252, 52, 12, 4, 252],
    }
)

frequencies

# Load pre-downloaded macro data
macro_df = load_macro()

# Normalize column names to lowercase (test fixture uses uppercase)
macro_df = macro_df.rename({c: c.lower() for c in macro_df.columns})

print(f"Loaded macro data: {macro_df.shape[0]:,} rows × {macro_df.shape[1]} columns")
print(f"Date range: {macro_df['timestamp'].min()} to {macro_df['timestamp'].max()}")
print(f"Columns: {macro_df.columns}")

macro_df.tail(10)

print("Data availability (non-null observations):")
for col in macro_df.columns:
    if col != "timestamp":
        non_null = macro_df[col].drop_nulls().len()
        pct = 100 * non_null / len(macro_df)
        print(f"  {col:12s}: {non_null:>6,} ({pct:>5.1f}%)")

# Release lag configuration (typical values)
RELEASE_LAGS = {
    "dff": 0,  # Real-time (daily)
    "dgs10": 0,  # Real-time (daily)
    "dgs2": 0,  # Real-time (daily)
    "t10y2y": 0,  # Real-time (calculated from daily)
    "vixcls": 0,  # Real-time (market-derived)
    "icsa": 5,  # Weekly claims: ~5 day lag
    "walcl": 7,  # Fed balance sheet: ~7 day lag
    "cpiaucsl": 15,  # CPI: ~15 days after month end
    "cpilfesl": 15,  # Core CPI: ~15 days
    "pcepi": 25,  # PCE: ~25 days
    "unrate": 5,  # Employment report: ~5 days after month
    "payems": 5,  # Nonfarm payrolls: ~5 days
    "civpart": 5,  # Labor force participation: ~5 days
    "indpro": 15,  # Industrial production: ~15 days
    "m2sl": 20,  # Money supply: ~20 days
    "gdp": 30,  # GDP: ~30 days (advance estimate)
    "gdpc1": 30,  # Real GDP: ~30 days
}

print("Release Lags (days after observation period):")
for series, lag in sorted(RELEASE_LAGS.items(), key=lambda x: x[1]):
    print(f"  {series:12s}: {lag:>3} days")

def apply_release_lag(df: pl.DataFrame, lag_days: dict[str, int]) -> pl.DataFrame:
    """
    Apply release lag to create point-in-time correct alignment.

    The forward-filled data already represents what's known at each date,
    but we shift by release lag to account for publication delay.

    For daily series (lag=0), no change needed.
    For monthly series (lag>0), the value shouldn't appear until release date.

    NOTE: This uses row-based shift, which assumes continuous daily data with
    no gaps. For data with gaps (weekends, holidays), a 15-row shift may not
    equal 15 calendar days. For production systems requiring exact calendar-day
    precision, use date arithmetic instead:

        df.with_columns(
            pl.col("timestamp") + pl.duration(days=lag)
        )

    For this educational notebook, the approximation is acceptable since:
    1. Forward-filled macro data has no gaps (every trading day filled)
    2. The lag values are approximate anyway (release schedules vary)
    """
    result = df.clone()

    for col, lag in lag_days.items():
        if col in df.columns and lag > 0:
            # Shift values forward by lag rows (value appears later)
            # Assumes continuous daily data without gaps
            result = result.with_columns(pl.col(col).shift(lag).alias(col))

    return result


pit_df = apply_release_lag(macro_df, RELEASE_LAGS)
pit_df.tail(10)

# Demonstrate revision impact (educational example)
# In practice, use FREDProvider.fetch_ohlcv(..., vintage_date=...) for real vintage data

revision_example = pl.DataFrame(
    {
        "quarter": ["Q3 2023", "Q3 2023", "Q3 2023", "Q3 2023"],
        "release_type": ["Advance", "Second", "Third", "Annual Revision"],
        "release_date": ["2023-10-26", "2023-11-29", "2023-12-21", "2024-03-28"],
        "gdp_growth_annualized": [4.9, 5.2, 4.9, 5.2],
    }
)

revision_example

# Create revision history based on documented 2023 GDP revisions
# In production, use: FREDProvider.fetch_ohlcv("GDP", vintage_date=...)
gdp_revisions = pl.DataFrame(
    {
        "quarter": ["Q1", "Q2", "Q3", "Q4"],
        "advance_estimate": [1.3, 2.1, 4.9, 3.3],  # What traders knew at release
        "final_value": [2.0, 2.1, 5.2, 3.4],  # What backtests often use incorrectly
    }
)

# Simple signal: "Exceptional growth" = GDP > 5.0%
gdp_signals = gdp_revisions.with_columns(
    (pl.col("advance_estimate") > 5.0).alias("signal_advance"),
    (pl.col("final_value") > 5.0).alias("signal_final"),
)

signal_match = (
    gdp_signals.select(pl.col("signal_advance") == pl.col("signal_final")).to_series().sum()
)
print(f"Signal agreement: {signal_match}/4 quarters")
gdp_signals

# Visualize the divergence (single cell so the inline backend cannot capture an
# un-laid-out intermediate render of the bars before the layout is applied).
fig_revision = go.Figure()

fig_revision.add_trace(
    go.Bar(
        name="Advance Estimate (PIT-Correct)",
        x=gdp_revisions["quarter"].to_list(),
        y=gdp_revisions["advance_estimate"].to_list(),
        marker_color=COLORS["blue"],
        text=[f"{v:.1f}%" for v in gdp_revisions["advance_estimate"].to_list()],
        textposition="outside",
    )
)

fig_revision.add_trace(
    go.Bar(
        name="Final Value (Look-Ahead Bias)",
        x=gdp_revisions["quarter"].to_list(),
        y=gdp_revisions["final_value"].to_list(),
        marker_color=COLORS["amber"],
        text=[f"{v:.1f}%" for v in gdp_revisions["final_value"].to_list()],
        textposition="outside",
    )
)

fig_revision.add_hline(
    y=5.0,
    line_dash="dash",
    line_color=COLORS["copper"],
    annotation_text="Signal Threshold (5%)",
    annotation_position="top left",
)

fig_revision.update_layout(
    title="Only the revised Q3 GDP crosses the 5% signal line - a trade PIT data never takes",
    barmode="group",
    yaxis_title="GDP Growth (% annualized)",
    xaxis_title="Quarter (2023)",
    height=450,
    legend=dict(orientation="h", yanchor="bottom", y=1.02, xanchor="center", x=0.5),
)

fig_revision.show()

def create_macro_features(df: pl.DataFrame) -> pl.DataFrame:
    """Create trading features from aligned macro data."""
    result = df.clone()

    # Yield Curve Features (T10Y2Y)
    if "t10y2y" in df.columns:
        result = result.with_columns(
            [
                # Level
                pl.col("t10y2y").alias("yield_curve_spread"),
                # Regime classification
                pl.when(pl.col("t10y2y") < 0)
                .then(pl.lit("inverted"))
                .when(pl.col("t10y2y") < 0.5)
                .then(pl.lit("flat"))
                .when(pl.col("t10y2y") < 1.5)
                .then(pl.lit("normal"))
                .otherwise(pl.lit("steep"))
                .alias("yield_curve_regime"),
                # Momentum (20-day change)
                (pl.col("t10y2y") - pl.col("t10y2y").shift(20)).alias("yield_curve_momentum_20d"),
                # Rolling z-score (252-day)
                (
                    (pl.col("t10y2y") - pl.col("t10y2y").rolling_mean(252))
                    / pl.col("t10y2y").rolling_std(252)
                ).alias("yield_curve_zscore_252d"),
            ]
        )

    # VIX Volatility Features
    if "vixcls" in df.columns:
        result = result.with_columns(
            [
                pl.col("vixcls").alias("vix"),
                # Regime classification
                pl.when(pl.col("vixcls") < 15)
                .then(pl.lit("low_vol"))
                .when(pl.col("vixcls") < 25)
                .then(pl.lit("normal_vol"))
                .when(pl.col("vixcls") < 35)
                .then(pl.lit("elevated_vol"))
                .otherwise(pl.lit("high_vol"))
                .alias("volatility_regime"),
                # Momentum (5-day change)
                (pl.col("vixcls") - pl.col("vixcls").shift(5)).alias("vix_momentum_5d"),
                # Rolling z-score (63-day / 3-month)
                (
                    (pl.col("vixcls") - pl.col("vixcls").rolling_mean(63))
                    / pl.col("vixcls").rolling_std(63)
                ).alias("vix_zscore_63d"),
            ]
        )

    # Unemployment Features
    if "unrate" in df.columns:
        result = result.with_columns(
            [
                pl.col("unrate").alias("unemployment_rate"),
                # Regime
                pl.when(pl.col("unrate") < 4.0)
                .then(pl.lit("tight"))
                .when(pl.col("unrate") < 6.0)
                .then(pl.lit("normal"))
                .otherwise(pl.lit("recession"))
                .alias("labor_market_regime"),
                # 1-month change (21 trading days)
                (pl.col("unrate") - pl.col("unrate").shift(21)).alias("unemployment_change_1m"),
                # 12-month change
                (pl.col("unrate") - pl.col("unrate").shift(252)).alias("unemployment_change_1y"),
            ]
        )

    # CPI/Inflation Features
    # Important: CPI is a monthly series in levels. If it's been forward-filled to daily,
    # using shift(252) (trading days) is NOT a true year-over-year calculation.
    # A simple approximation on daily grids is 365 calendar days.
    if "cpiaucsl" in df.columns:
        result = result.with_columns(
            [
                # Year-over-year inflation (365 calendar days for monthly forward-filled data)
                (pl.col("cpiaucsl") / pl.col("cpiaucsl").shift(365) - 1).alias("inflation_yoy"),
                # Inflation regime
                pl.when((pl.col("cpiaucsl") / pl.col("cpiaucsl").shift(365) - 1) < 0.02)
                .then(pl.lit("low"))
                .when((pl.col("cpiaucsl") / pl.col("cpiaucsl").shift(365) - 1) < 0.04)
                .then(pl.lit("moderate"))
                .otherwise(pl.lit("high"))
                .alias("inflation_regime"),
            ]
        )

    # Initial Claims Features
    # Note: If weekly claims have been forward-filled to daily, rolling windows over days
    # are smoothing proxies, not true "4-week averages" over weekly observations.
    if "icsa" in df.columns:
        result = result.with_columns(
            [
                pl.col("icsa").alias("initial_claims"),
                # Smoothing proxy on daily grid; for a true 4-week average, compute on weekly data first.
                pl.col("icsa").rolling_mean(20).alias("claims_smooth_20d"),
                # Regime
                pl.when(pl.col("icsa") < 250000)
                .then(pl.lit("strong"))
                .when(pl.col("icsa") < 350000)
                .then(pl.lit("normal"))
                .otherwise(pl.lit("weak"))
                .alias("claims_regime"),
            ]
        )

    return result


# Create features
macro_features = create_macro_features(pit_df)

# List new features
original_cols = set(macro_df.columns)
feature_cols = [c for c in macro_features.columns if c not in original_cols and c != "timestamp"]

print(f"Created {len(feature_cols)} macro features:")
for col in sorted(feature_cols):
    print(f"  - {col}")

# Display sample of macro features
display_cols = ["timestamp", "yield_curve_spread", "yield_curve_regime", "vix", "volatility_regime"]
if "unemployment_rate" in macro_features.columns:
    display_cols.extend(["unemployment_rate", "labor_market_regime"])

macro_features.select([c for c in display_cols if c in macro_features.columns]).tail(10)

# Prepare data for visualization (2020 onwards for readability)
VIZ_START_DATE = date(2020, 1, 1)

viz_df = (
    macro_features.filter(pl.col("timestamp") >= VIZ_START_DATE)
    .filter(pl.col("yield_curve_spread").is_not_null())
    .to_pandas()
)

# Build the three-panel macro-indicator figure in a SINGLE cell so the
# inline backend cannot capture a yield-curve-only intermediate render
# (feedback_split_cell_figure_bug).
fig = make_subplots(
    rows=3,
    cols=1,
    shared_xaxes=True,
    subplot_titles=("Yield Curve Spread (10Y-2Y)", "Unemployment Rate", "VIX Volatility Index"),
    vertical_spacing=0.08,
)

# Yield curve
fig.add_trace(
    go.Scatter(
        x=viz_df["timestamp"],
        y=viz_df["yield_curve_spread"],
        mode="lines",
        name="10Y-2Y Spread",
        line={"color": COLORS["blue"]},
    ),
    row=1,
    col=1,
)
fig.add_hline(y=0, line_dash="dash", line_color=COLORS["negative"], row=1, col=1)

# Unemployment
if "unemployment_rate" in viz_df.columns:
    fig.add_trace(
        go.Scatter(
            x=viz_df["timestamp"],
            y=viz_df["unemployment_rate"],
            mode="lines",
            name="Unemployment",
            line={"color": COLORS["slate"]},
        ),
        row=2,
        col=1,
    )

# VIX
if "vix" in viz_df.columns:
    fig.add_trace(
        go.Scatter(
            x=viz_df["timestamp"],
            y=viz_df["vix"],
            mode="lines",
            name="VIX",
            line={"color": COLORS["copper"]},
        ),
        row=3,
        col=1,
    )
    fig.add_hline(y=20, line_dash="dash", line_color=COLORS["neutral"], row=3, col=1)

fig.update_layout(
    height=700,
    title_text="Yield-curve inversion, unemployment, and VIX trace the 2020-2025 macro regimes",
    showlegend=False,
)
fig.update_yaxes(title_text="Spread (%)", row=1, col=1)
fig.update_yaxes(title_text="Rate (%)", row=2, col=1)
fig.update_yaxes(title_text="Index", row=3, col=1)

fig.show()

# Analyze regime distributions
print("Regime Distributions (2020+):\n")

regime_cols = [c for c in macro_features.columns if c.endswith("_regime")]
analysis_df = macro_features.filter(pl.col("timestamp") >= VIZ_START_DATE)

for col in regime_cols:
    if col in analysis_df.columns:
        counts = (
            analysis_df.filter(pl.col(col).is_not_null())
            .group_by(col)
            .len()
            .sort("len", descending=True)
        )
        total = counts["len"].sum()
        print(f"{col}:")
        for row in counts.iter_rows():
            pct = 100 * row[1] / total
            print(f"  {row[0]:15s}: {row[1]:>5,} ({pct:>5.1f}%)")
        print()

def combine_with_prices(
    macro_df: pl.DataFrame, prices_df: pl.DataFrame, date_col: str = "timestamp"
) -> pl.DataFrame:
    """Combine macro features with price data via left join on date."""
    return prices_df.join(macro_df, on=date_col, how="left")


# Load real SPY closes and align to the macro window
price_data = (
    load_etfs(symbols=["SPY"], start_date=str(VIZ_START_DATE))
    .sort("timestamp")
    .with_columns(spy_return=pl.col("close").pct_change())
    .select(timestamp="timestamp", spy_close="close", spy_return="spy_return")
)

# Combine: SPY trading days on the left, daily-aligned macro features attached
combined = combine_with_prices(macro_features, price_data)

cols = ["timestamp", "spy_close", "spy_return", "yield_curve_regime", "volatility_regime"]
if "labor_market_regime" in combined.columns:
    cols.append("labor_market_regime")
combined.select([c for c in cols if c in combined.columns]).tail(10)