"""Point-in-Time Data Validation."""

import os
from datetime import datetime

import plotly.graph_objects as go
import polars as pl
from ml4t.data.providers import FREDProvider

from data import load_etfs, load_macro
from utils.style import COLORS

# Production defaults — Papermill injects overrides for CI

spy = load_etfs().filter(pl.col("symbol") == "SPY").sort("timestamp")
print(f"SPY data: {len(spy)} rows; {spy['timestamp'].min()} → {spy['timestamp'].max()}")
spy.head()

spy_ma = (
    spy.with_columns(
        pl.col("close").rolling_mean(window_size=20).alias("ma20_trailing"),
        pl.col("close").rolling_mean(window_size=20, center=True).alias("ma20_centered"),
    )
    .drop_nulls()
    .tail(120)
)

fig = go.Figure()
fig.add_trace(
    go.Scatter(
        x=spy_ma["timestamp"].to_list(),
        y=spy_ma["close"].to_list(),
        name="SPY close",
        line=dict(color=COLORS["blue"], width=1),
    )
)
fig.add_trace(
    go.Scatter(
        x=spy_ma["timestamp"].to_list(),
        y=spy_ma["ma20_trailing"].to_list(),
        name="20d trailing MA (PIT-correct)",
        line=dict(color=COLORS["slate"], width=2),
    )
)
fig.add_trace(
    go.Scatter(
        x=spy_ma["timestamp"].to_list(),
        y=spy_ma["ma20_centered"].to_list(),
        name="20d centered MA (uses future)",
        line=dict(color=COLORS["copper"], width=2, dash="dash"),
    )
)
fig.update_layout(
    title="Trailing vs centered moving average — centered knows the future",
    xaxis_title="Date",
    yaxis_title="SPY price",
    height=420,
)
fig.show()

def naive_leakage_corr(df: pl.DataFrame, feature: str, price: str = "close") -> float:
    """Correlation between a level feature and the next-period price return.

    Reproduces the textbook-warning heuristic that *seems* like it should
    catch lookahead but rarely does because it compares incompatible scales.
    """
    enriched = df.with_columns(
        (pl.col(price).shift(-1) / pl.col(price) - 1).alias("next_ret")
    ).drop_nulls([feature, "next_ret"])
    return float(enriched.select(pl.corr(feature, "next_ret")).item())

naive_features = spy.with_columns(
    pl.col("close").rolling_mean(window_size=5).alias("ma5_trailing"),
    pl.col("close").rolling_mean(window_size=5, center=True).alias("ma5_centered"),
    pl.col("close").shift(-1).alias("tomorrow_close"),
).drop_nulls()

naive_results = pl.DataFrame(
    [
        {"feature": f, "naive_corr_with_next_return": naive_leakage_corr(naive_features, f)}
        for f in ["ma5_trailing", "ma5_centered", "tomorrow_close"]
    ]
)
naive_results

def leakage_corr_scale_invariant(df: pl.DataFrame, feature: str, price: str = "close") -> float:
    """Correlation between a feature's own return and the next price return.

    Most useful when `feature` is a price-like (level) series — for return-
    shaped features the pct_change is unstable near zero, so the caller
    should pass them through directly instead.
    """
    enriched = (
        df.with_columns(
            pl.col(feature).pct_change().alias("feat_ret"),
            (pl.col(price).shift(-1) / pl.col(price) - 1).alias("next_ret"),
        )
        .drop_nulls(["feat_ret", "next_ret"])
        .filter(pl.col("feat_ret").is_finite() & pl.col("next_ret").is_finite())
    )
    if enriched.is_empty():
        return float("nan")
    return float(enriched.select(pl.corr("feat_ret", "next_ret")).item())

fixed_results = pl.DataFrame(
    [
        {
            "feature": f,
            "scale_invariant_corr": leakage_corr_scale_invariant(naive_features, f),
        }
        for f in ["ma5_trailing", "ma5_centered", "tomorrow_close"]
    ]
)
fixed_results

def validate_signal_trade_lag(signals: pl.DataFrame, execution_lag: int = 1) -> pl.DataFrame:
    """Tag each signal with the earliest tradable date `execution_lag` rows ahead."""
    return signals.sort("timestamp").with_columns(
        pl.col("timestamp").alias("signal_date"),
        pl.col("timestamp").shift(-execution_lag).alias("earliest_execution"),
    )

spy_with_signal = spy.with_columns(
    (pl.col("close") / pl.col("close").shift(5) - 1).alias("momentum_signal")
)
validate_signal_trade_lag(spy_with_signal, execution_lag=1).select(
    ["timestamp", "close", "momentum_signal", "signal_date", "earliest_execution"]
).head(8)

macro = load_macro()
macro = macro.rename({col: col.lower() for col in macro.columns})

date_col = "timestamp" if "timestamp" in macro.columns else "date"
macro_summary = pl.DataFrame(
    [
        {
            "series": s,
            "non_null": macro.filter(pl.col(s).is_not_null()).height,
            "first": macro.filter(pl.col(s).is_not_null())[date_col].min(),
            "last": macro.filter(pl.col(s).is_not_null())[date_col].max(),
        }
        for s in ["dgs10", "unrate", "vixcls"]
        if s in macro.columns
    ]
)
macro_summary

unrate = (
    macro.select([date_col, "unrate"])
    .filter(pl.col(date_col).dt.year() >= 2023)
    .with_columns(pl.col("unrate").forward_fill().alias("unrate_filled"))
    .filter(
        (pl.col(date_col) >= datetime(2023, 1, 1)) & (pl.col(date_col) <= datetime(2023, 3, 31))
    )
)
unrate.head(10)

FRED_API_KEY = os.getenv("FRED_API_KEY")
if not FRED_API_KEY:
    raise RuntimeError(
        "FRED_API_KEY not set. Get a free key at "
        "https://fred.stlouisfed.org/docs/api/api_key.html and export it."
    )

provider = FREDProvider()

gdp_early = provider.fetch_ohlcv(
    "GDP",
    "2023-01-01",
    "2023-09-30",
    frequency="quarterly",
    vintage_date="2023-11-01",
)
gdp_late = provider.fetch_ohlcv(
    "GDP",
    "2023-01-01",
    "2023-09-30",
    frequency="quarterly",
    vintage_date="2024-06-01",
)
provider.close()

revisions = (
    gdp_early.select(["timestamp", pl.col("close").alias("vintage_2023_11_01")])
    .join(
        gdp_late.select(["timestamp", pl.col("close").alias("vintage_2024_06_01")]),
        on="timestamp",
    )
    .with_columns(
        (
            (pl.col("vintage_2024_06_01") - pl.col("vintage_2023_11_01"))
            / pl.col("vintage_2023_11_01")
            * 100
        )
        .round(3)
        .alias("revision_pct")
    )
)
revisions

class PITValidator:
    """Point-in-time validation for daily price-and-feature panels."""

    LEAK_THRESHOLD = 0.30  # |corr| above which we flag

    def __init__(self, df: pl.DataFrame, date_col: str = "timestamp"):
        self.df = df
        self.date_col = date_col
        self.violations: list[dict] = []

    def check_future_leakage(self, feature_col: str, price_col: str = "close") -> dict:
        corr = leakage_corr_scale_invariant(self.df, feature_col, price_col)
        severity = "HIGH" if abs(corr) > self.LEAK_THRESHOLD else "LOW"
        result = {
            "feature": feature_col,
            "scale_invariant_corr": round(corr, 4),
            "severity": severity,
            "violation": severity == "HIGH",
        }
        if result["violation"]:
            self.violations.append(result)
        return result

    def check_date_gaps(self, max_gap_days: int = 5) -> dict:
        gaps = (
            self.df.sort(self.date_col)
            .with_columns(pl.col(self.date_col).diff().dt.total_days().alias("gap_days"))
            .filter(pl.col("gap_days") > max_gap_days)
        )
        return {
            "rows": len(self.df),
            "gaps_above_threshold": len(gaps),
            "max_gap_days": int(gaps["gap_days"].max()) if len(gaps) > 0 else 0,
        }

    def check_monotonic_dates(self) -> dict:
        diffs = self.df[self.date_col].diff().drop_nulls().dt.total_days()
        is_sorted = bool((diffs >= 0).all())
        return {"is_monotonic": is_sorted, "violation": not is_sorted}

spy_features = spy.with_columns(
    pl.col("close").rolling_mean(window_size=5).alias("ma5_trailing"),
    pl.col("close").rolling_mean(window_size=20).alias("ma20_trailing"),
    pl.col("close").rolling_mean(window_size=5, center=True).alias("ma5_centered_BAD"),
    pl.col("close").shift(-1).alias("tomorrow_close_BAD"),
).drop_nulls()

validator = PITValidator(spy_features)
validation_table = pl.DataFrame(
    [
        validator.check_future_leakage(f)
        for f in [
            "ma5_trailing",
            "ma20_trailing",
            "ma5_centered_BAD",
            "tomorrow_close_BAD",
        ]
    ]
)
validation_table

gap_check = validator.check_date_gaps(max_gap_days=5)
mono_check = validator.check_monotonic_dates()
print(
    f"Date gaps over 5 calendar days: {gap_check['gaps_above_threshold']} "
    f"(max {gap_check['max_gap_days']} d); monotonic={mono_check['is_monotonic']}"
)
print(f"Violations flagged: {len(validator.violations)}")