"""The ml4t Library Ecosystem - unified data, feature, and diagnostic libraries for Chapters 7-12."""

from __future__ import annotations

import warnings
from datetime import datetime

import numpy as np
import polars as pl
from IPython.display import display
from ml4t.diagnostic.signal import analyze_signal
from ml4t.engineer import compute_features
from ml4t.engineer.core.registry import get_registry

from data import load_etfs

warnings.filterwarnings("ignore")

# Production defaults - Papermill injects overrides for CI
SPY_START_DATE = "2015-01-01"

etfs = load_etfs()
print(f"ETF universe: {len(etfs):,} rows, {etfs['symbol'].n_unique()} symbols")
print(f"Columns: {etfs.columns}")
print(f"Date range: {etfs['timestamp'].min()} to {etfs['timestamp'].max()}")

# For demos: SPY only
spy = etfs.filter(
    (pl.col("symbol") == "SPY")
    & (pl.col("timestamp") >= datetime.fromisoformat(SPY_START_DATE).date())
).sort("timestamp")

print(f"SPY: {len(spy):,} rows from {spy['timestamp'].min()} to {spy['timestamp'].max()}")

registry = get_registry()
all_features = registry.list_all()

# Features by category
categories = {}
for name in all_features:
    metadata = registry.get(name)
    categories.setdefault(metadata.category, []).append(name)

print(f"Total features available: {len(all_features)}\n")
print("Features by category:")
for cat, feats in sorted(categories.items()):
    examples = ", ".join(feats[:4])
    suffix = ", ..." if len(feats) > 4 else ""
    print(f"  {cat:20s}: {len(feats):3d}  ({examples}{suffix})")

for feature_name in ["rsi", "atr", "garman_klass_volatility"]:
    meta = registry.get(feature_name)
    print(f"\n{'=' * 50}")
    print(f"Feature:     {meta.name}")
    print(f"Category:    {meta.category}")
    print(f"Description: {meta.description}")
    print(f"Formula:     {meta.formula or '(not recorded)'}")
    print(f"Parameters:  {meta.parameters}")
    print(f"Input type:  {meta.input_type}")

result = compute_features(spy, ["rsi", "sma", "ema", "atr"])

new_cols = [c for c in result.columns if c not in spy.columns]
print(f"Computed {len(new_cols)} features: {new_cols}")
display(result.select(["timestamp", "close"] + new_cols).tail(5))

# Distinct features, each with an explicit parameter set. Bollinger bands take the
# two deviations separately, following TA-Lib, so the upper and lower band need not
# sit the same distance from the moving average:
parameterized = [
    {"name": "rsi", "params": {"period": 10}},
    {"name": "atr", "params": {"period": 20}},
    {"name": "bollinger_bands", "params": {"period": 20, "nbdevup": 2.0, "nbdevdn": 2.0}},
]

result = compute_features(spy, parameterized)

new_cols = [c for c in result.columns if c not in spy.columns]
print(f"Parameterized features ({len(new_cols)} columns): {', '.join(new_cols)}")

# Multiple horizons of one indicator: one call per period, suffix, then join.
rsi_multi = spy.select(["timestamp"])
for period in (10, 21, 63):
    horizon = compute_features(spy, [{"name": "rsi", "params": {"period": period}}])
    rsi_multi = rsi_multi.join(
        horizon.select(["timestamp", pl.col("rsi").alias(f"rsi_{period}")]),
        on="timestamp",
    )

rsi_cols = [c for c in rsi_multi.columns if c != "timestamp"]
print(f"Multi-horizon RSI columns: {rsi_cols}")
display(rsi_multi.tail(5))

def manual_rsi(df: pl.DataFrame, period: int = 14) -> pl.DataFrame:
    """Manual RSI using EWM span (not Wilder's smoothing)."""
    return (
        df.with_columns(pl.col("close").diff().alias("delta"))
        .with_columns(
            pl.when(pl.col("delta") > 0).then(pl.col("delta")).otherwise(0).alias("gain"),
            pl.when(pl.col("delta") < 0).then(-pl.col("delta")).otherwise(0).alias("loss"),
        )
        .with_columns(
            pl.col("gain").ewm_mean(span=period, adjust=False).alias("avg_gain"),
            pl.col("loss").ewm_mean(span=period, adjust=False).alias("avg_loss"),
        )
        .with_columns(
            (100 - 100 / (1 + pl.col("avg_gain") / pl.col("avg_loss"))).alias("rsi_manual"),
        )
    )


manual_result = manual_rsi(spy.select(["timestamp", "close"]))
library_result = compute_features(spy, [{"name": "rsi", "params": {"period": 14}}])

comparison = manual_result.join(
    library_result.select(["timestamp", "rsi"]),
    on="timestamp",
    how="inner",
).select(["timestamp", "rsi_manual", "rsi"])

print("RSI comparison (manual EWM vs library Wilder's):")
display(comparison.filter(pl.col("rsi").is_not_null()).tail(10))

# Both columns can carry a leading NaN/null from the warm-up window; in Polars
# NaN is distinct from null, so guard against both before averaging.
diff = comparison.filter(
    pl.col("rsi").is_not_null()
    & pl.col("rsi").is_not_nan()
    & pl.col("rsi_manual").is_not_null()
    & pl.col("rsi_manual").is_not_nan()
).with_columns((pl.col("rsi") - pl.col("rsi_manual")).abs().alias("abs_diff"))
print(f"Mean absolute difference: {diff['abs_diff'].mean():.4f}")
print("Differences are due to smoothing method (EWM span vs Wilder's).")

# Compute one momentum factor across the whole ETF cross-section, then evaluate.
panel = etfs.sort(["symbol", "timestamp"])
factor = (
    panel.group_by("symbol", maintain_order=True)
    .map_groups(lambda g: compute_features(g, [{"name": "mom", "params": {"period": 21}}]))
    .select(["timestamp", "symbol", pl.col("mom").alias("factor")])
    .drop_nulls()
)
prices = panel.select(["timestamp", "symbol", pl.col("close").alias("price")])

signal = analyze_signal(
    factor,
    prices,
    periods=(1, 5, 21),
    quantiles=5,
    date_col="timestamp",
    asset_col="symbol",
    price_col="price",
    factor_col="factor",
)

print(f"Evaluated on {signal.n_assets} assets over {signal.n_dates:,} dates")
for horizon in ("1D", "5D", "21D"):
    print(
        f"  {horizon:>3s}  IC={signal.ic[horizon]:+.4f}  "
        f"t-stat={signal.ic_t_stat[horizon]:+.2f}  "
        f"quantile spread={signal.spread[horizon]:+.4f}"
    )

multi_symbol = etfs.filter(pl.col("symbol").is_in(["SPY", "QQQ", "IWM"])).sort(
    ["symbol", "timestamp"]
)

features = multi_symbol.with_columns(
    pl.col("close").pct_change().over("symbol").alias("ret_1d"),
    (pl.col("close").pct_change().rolling_std(21).over("symbol") * np.sqrt(252)).alias("vol_21d"),
    pl.col("close").rolling_mean(21).over("symbol").alias("sma_21"),
)

print("Window function features:")
display(features.select(["symbol", "timestamp", "close", "ret_1d", "vol_21d"]).tail(10))

trades = spy.select(["timestamp", "close"]).rename({"close": "trade_price"}).head(1000)

quotes = (
    spy.select(["timestamp", "close"])
    .rename({"close": "quote_price"})
    .with_columns(pl.col("timestamp").cast(pl.Date))
    .head(1000)
)

matched = trades.sort("timestamp").join_asof(
    quotes.sort("timestamp"),
    on="timestamp",
    strategy="backward",
)

print("ASOF join result (most recent quote for each trade):")
display(matched.head(5))

# Demonstrate lazy API pattern with in-memory data
# (In production, use pl.scan_parquet on the actual file)
spy_lazy = (
    load_etfs(symbols=["SPY"], start_date="2020-01-01")
    .lazy()
    .select(["timestamp", "close", "volume"])
    .with_columns(pl.col("close").pct_change().alias("returns"))
)

result = spy_lazy.collect()
print(f"Lazy query: {len(result):,} rows")
print(f"\nQuery plan:\n{spy_lazy.explain()}")