"""Data Quality Framework — Validation, anomaly detection, and remediation."""

import warnings

warnings.filterwarnings("ignore")

import logging
from datetime import datetime
from pathlib import Path

import numpy as np
import plotly.graph_objects as go
import polars as pl
import structlog
from ml4t.data.anomaly import (
    AnomalyManager,
    PriceStalenessDetector,
    ReturnOutlierDetector,
    VolumeSpikeDetector,
)
from ml4t.data.anomaly.config import (
    AnomalyConfig,
    PriceStalenessConfig,
    ReturnOutlierConfig,
    VolumeSpikeConfig,
)
from ml4t.data.validation import OHLCVValidator

from data import load_us_equities
from utils.paths import get_output_dir
from utils.style import COLORS

# The anomaly/validation libraries log per-symbol detector results through
# structlog, which bypasses stdlib logging; filter at WARNING so the cells show
# call results, not the library's internal debug stream.
structlog.configure(wrapper_class=structlog.make_filtering_bound_logger(logging.WARNING))


def _to_date(value: object) -> object:
    """Normalize a polars timestamp scalar to a datetime.date for printing."""
    return value.date() if hasattr(value, "date") else value

# Five large-cap symbols from the legacy Wiki/Quandl dataset (1962–2018).
# FB (not META) is the canonical Facebook ticker in this vintage. Readers can
# override this universe via Papermill.
SYMBOLS = ["AAPL", "MSFT", "GOOGL", "NVDA", "FB"]

OUTPUT_DIR = get_output_dir(2, "quality")

wiki_df = load_us_equities()
print(
    f"US equities loaded: {len(wiki_df):,} rows; "
    f"{wiki_df['timestamp'].min()} → {wiki_df['timestamp'].max()}"
)

datasets = {
    symbol: (
        wiki_df.lazy()
        .filter(pl.col("symbol") == symbol)
        .select(["timestamp", "symbol", "open", "high", "low", "close", "volume"])
        .collect()
    )
    for symbol in SYMBOLS
}
missing = [s for s, df in datasets.items() if df.is_empty()]
if missing:
    raise RuntimeError(f"Symbols missing from Wiki/Quandl dataset: {missing}")

per_symbol_rows = pl.DataFrame(
    {"symbol": list(datasets), "rows": [len(df) for df in datasets.values()]}
)
per_symbol_rows

validator = OHLCVValidator(
    check_nulls=True,
    check_price_consistency=True,
    negative_price_policy="forbid",
    check_negative_volume=True,
    check_duplicate_timestamps=True,
    check_chronological_order=True,
    check_price_staleness=True,
    check_extreme_returns=True,
    max_return_threshold=0.5,  # flag |return| > 50%
    staleness_threshold=5,  # flag 5+ identical-price days
)

validation_summary = pl.DataFrame(
    [
        {
            "symbol": sym,
            "passed": (r := validator.validate(df)).passed,
            "issues": len(r.issues),
            "critical": r.critical_count,
            "errors": r.error_count,
        }
        for sym, df in datasets.items()
    ]
)
validation_summary

clean_df = datasets["AAPL"]
highs = clean_df["high"].to_numpy().copy()
lows = clean_df["low"].to_numpy().copy()
volumes = clean_df["volume"].to_numpy().copy()

# Fault 1: high < low at rows 10–12
highs[10:13] = lows[10:13] - 1.0
# Fault 2: negative volume at row 20
volumes[20] = -1000
# Fault 3: null close at row 30. Use a genuine null, not NumPy NaN: Polars stores
# np.nan as a float NaN that the null check ignores, so only a real null exercises
# OHLCVValidator's check_nulls.
close_col = pl.Series("close", clean_df["close"].to_numpy()).scatter(30, None)

dirty_df = pl.DataFrame(
    {
        "timestamp": clean_df["timestamp"],
        "symbol": clean_df["symbol"],
        "open": clean_df["open"],
        "high": highs,
        "low": lows,
        "close": close_col,
        "volume": volumes,
    }
)

dirty_result = validator.validate(dirty_df)
print(f"Validation passed: {dirty_result.passed}")
print(f"Critical: {dirty_result.critical_count}, Errors: {dirty_result.error_count}")

dirty_issues = pl.DataFrame(
    [
        {
            "severity": issue.severity.name,
            "check": issue.check,
            "rows": issue.row_count or 0,
            "message": issue.message,
        }
        for issue in dirty_result.issues
    ]
)
dirty_issues

sample_df = datasets["AAPL"]
methods = ["mad", "zscore", "iqr"]

method_summary = pl.DataFrame(
    [
        {
            "method": m,
            "anomalies": len(
                ReturnOutlierDetector(
                    config=ReturnOutlierConfig(method=m, threshold=3.0, min_samples=20)
                ).detect(sample_df, symbol="AAPL")
            ),
        }
        for m in methods
    ]
)
method_summary

fig = go.Figure(
    go.Bar(
        x=[m.upper() for m in method_summary["method"].to_list()],
        y=method_summary["anomalies"].to_list(),
        marker_color=COLORS["slate"],
        text=method_summary["anomalies"].to_list(),
        textposition="outside",
    )
)
fig.update_layout(
    title="AAPL return outliers by method (same threshold = 3.0)",
    xaxis_title="Method",
    yaxis_title="Events flagged",
    height=400,
)
fig.show()

mad_anomalies = ReturnOutlierDetector(
    config=ReturnOutlierConfig(method="mad", threshold=3.0, min_samples=20)
).detect(sample_df, symbol="AAPL")

top_mad = pl.DataFrame(
    [
        {
            "date": _to_date(a.timestamp),
            "return_pct": float(a.value),  # value is already a percentage
        }
        for a in sorted(mad_anomalies, key=lambda x: abs(x.value), reverse=True)[:5]
    ]
)
top_mad

returns_df = sample_df.with_columns((pl.col("close").pct_change() * 100).alias("return_pct"))
anomaly_dates = {_to_date(a.timestamp) for a in mad_anomalies}
flagged = returns_df.with_columns(
    pl.col("timestamp")
    .map_elements(lambda t: _to_date(t) in anomaly_dates, return_dtype=pl.Boolean)
    .alias("flagged")
)
flagged_pts = flagged.filter(pl.col("flagged"))

fig = go.Figure()
fig.add_trace(
    go.Scatter(
        x=returns_df["timestamp"].to_list(),
        y=returns_df["return_pct"].to_list(),
        mode="lines",
        line=dict(color=COLORS["slate"], width=0.6),
        name="Daily return (%)",
    )
)
fig.add_trace(
    go.Scatter(
        x=flagged_pts["timestamp"].to_list(),
        y=flagged_pts["return_pct"].to_list(),
        mode="markers",
        marker=dict(color=COLORS["copper"], size=5),
        name="MAD-flagged",
    )
)
fig.update_layout(
    title="AAPL daily returns with MAD-flagged outliers (splits + the 2000 crash)",
    xaxis_title="Date",
    yaxis_title="Return (%)",
    height=420,
)
fig.show()

volume_anomalies = VolumeSpikeDetector(
    config=VolumeSpikeConfig(window=20, threshold=3.0, min_volume=0, min_samples=20)
).detect(sample_df, symbol="AAPL")

top_volume = pl.DataFrame(
    [
        {
            "date": _to_date(a.timestamp),
            "volume": int(a.value),
            "ratio_vs_avg": (
                a.value / a.metadata["average_volume"] if a.metadata.get("average_volume") else None
            ),
        }
        for a in sorted(volume_anomalies, key=lambda x: x.value, reverse=True)[:5]
    ]
)
print(f"Found {len(volume_anomalies)} AAPL volume spikes")
top_volume

stale_anomalies = PriceStalenessDetector(
    config=PriceStalenessConfig(max_unchanged_days=3, check_close_only=False)
).detect(sample_df, symbol="AAPL")
print(f"AAPL stale-price periods (≥4 consecutive days): {len(stale_anomalies)}")

anomaly_config = AnomalyConfig(
    enabled=True,
    report_severity_threshold="warning",
    return_outliers=ReturnOutlierConfig(method="mad", threshold=3.0),
    volume_spikes=VolumeSpikeConfig(window=20, threshold=3.0),
    price_staleness=PriceStalenessConfig(max_unchanged_days=5),
)
manager = AnomalyManager(config=anomaly_config)
reports = manager.analyze_batch(datasets)

batch_summary = pl.DataFrame(
    [
        {
            "symbol": sym,
            "total_anomalies": len(rep.anomalies),
            "critical": len(rep.get_critical_anomalies()),
        }
        for sym, rep in reports.items()
    ]
)
batch_summary

def calculate_psi(
    baseline: pl.Series, current: pl.Series, n_bins: int = 10, epsilon: float = 1e-6
) -> tuple[float, pl.DataFrame]:
    """Population Stability Index between baseline and current distributions."""
    baseline_clean = baseline.drop_nulls()
    current_clean = current.drop_nulls()

    percentiles = [i * 100 / n_bins for i in range(n_bins + 1)]
    bin_edges = [baseline_clean.quantile(p / 100) for p in percentiles]

    # Ensure unique edges
    unique_edges = [bin_edges[0]]
    for edge in bin_edges[1:]:
        if edge <= unique_edges[-1]:
            edge = unique_edges[-1] + epsilon
        unique_edges.append(edge)

    baseline_counts = np.histogram(baseline_clean.to_numpy(), bins=unique_edges)[0]
    current_counts = np.histogram(current_clean.to_numpy(), bins=unique_edges)[0]
    baseline_pct = np.maximum(baseline_counts / len(baseline_clean), epsilon)
    current_pct = np.maximum(current_counts / len(current_clean), epsilon)

    psi_values = (current_pct - baseline_pct) * np.log(current_pct / baseline_pct)
    breakdown = pl.DataFrame(
        {
            "bin": list(range(1, n_bins + 1)),
            "baseline_pct": baseline_pct.round(4),
            "current_pct": current_pct.round(4),
            "psi_contribution": psi_values.round(4),
        }
    )
    return float(np.sum(psi_values)), breakdown

df = datasets["AAPL"].with_columns((pl.col("close").pct_change() * 100).alias("return_pct"))
midpoint = len(df) // 2
baseline_returns = df["return_pct"][:midpoint]
current_returns = df["return_pct"][midpoint:]

psi_value, psi_breakdown = calculate_psi(baseline_returns, current_returns)
psi_severity = (
    "no significant change"
    if psi_value < 0.1
    else "moderate shift"
    if psi_value < 0.25
    else "significant shift"
)
print(
    f"AAPL daily-return PSI: {psi_value:.4f} ({psi_severity})\n"
    f"Baseline: {_to_date(df['timestamp'][0])} → {_to_date(df['timestamp'][midpoint])}\n"
    f"Current:  {_to_date(df['timestamp'][midpoint])} → {_to_date(df['timestamp'][-1])}"
)
psi_breakdown

fig = go.Figure()
fig.add_trace(
    go.Bar(
        x=psi_breakdown["bin"].to_list(),
        y=(psi_breakdown["baseline_pct"] * 100).to_list(),
        name="Baseline %",
        marker_color=COLORS["slate"],
    )
)
fig.add_trace(
    go.Bar(
        x=psi_breakdown["bin"].to_list(),
        y=(psi_breakdown["current_pct"] * 100).to_list(),
        name="Current %",
        marker_color=COLORS["copper"],
    )
)
fig.add_hline(y=10, line_dash="dash", line_color=COLORS["amber"], line_width=1)
fig.update_layout(
    title="AAPL return-decile membership, baseline vs current",
    xaxis_title="Baseline decile bin",
    yaxis_title="Share of observations (%)",
    barmode="group",
    height=420,
)
fig.show()

def detect_gaps(df: pl.DataFrame, max_gap_days: int = 5) -> pl.DataFrame:
    """Return rows whose gap from the previous timestamp exceeds `max_gap_days`."""
    return (
        df.sort("timestamp")
        .with_columns(
            pl.col("timestamp").diff().dt.total_days().alias("days_since_prev"),
            pl.col("timestamp").shift(1).alias("prev_timestamp"),
        )
        .filter(pl.col("days_since_prev") > max_gap_days)
        .select(["prev_timestamp", "timestamp", "days_since_prev"])
    )

gap_rows = []
for symbol, df in datasets.items():
    gaps = detect_gaps(df, max_gap_days=5)
    for row in gaps.iter_rows(named=True):
        gap_rows.append(
            {
                "symbol": symbol,
                "prev_date": _to_date(row["prev_timestamp"]),
                "next_date": _to_date(row["timestamp"]),
                "days": int(row["days_since_prev"]),
            }
        )
gap_table = (
    pl.DataFrame(gap_rows)
    if gap_rows
    else pl.DataFrame({"symbol": [], "prev_date": [], "next_date": [], "days": []})
)
gap_table

sample = datasets["AAPL"].head(100)
df_with_dups = pl.concat([sample, sample.head(10)]).sort("timestamp")
print(
    f"Original: {len(df_with_dups)} rows; keep first: "
    f"{df_with_dups.unique(subset=['timestamp'], keep='first').shape[0]} rows; "
    f"keep last: {df_with_dups.unique(subset=['timestamp'], keep='last').shape[0]} rows"
)

def detect_corporate_actions(df: pl.DataFrame, threshold: float = 0.25) -> pl.DataFrame:
    """Flag overnight returns whose magnitude exceeds `threshold`."""
    return (
        df.with_columns(pl.col("close").shift(1).alias("prev_close"))
        .with_columns(((pl.col("open") / pl.col("prev_close")) - 1).alias("overnight_return"))
        .filter(pl.col("overnight_return").abs() > threshold)
        .select(["timestamp", "prev_close", "open", "overnight_return"])
    )

events_rows = []
for symbol, df in datasets.items():
    for row in detect_corporate_actions(df, threshold=0.25).iter_rows(named=True):
        events_rows.append(
            {
                "symbol": symbol,
                "date": _to_date(row["timestamp"]),
                "prev_close": round(row["prev_close"], 2),
                "open": round(row["open"], 2),
                "overnight_pct": round(row["overnight_return"] * 100, 1),
            }
        )
events_table = pl.DataFrame(events_rows)
print(f"Potential corporate-action events across {len(datasets)} symbols: {len(events_table)}")
events_table.head(10)

fig = go.Figure(
    go.Scatter(
        x=events_table["date"].to_list(),
        y=events_table["overnight_pct"].to_list(),
        mode="markers",
        marker=dict(color=COLORS["slate"], size=8),
        text=events_table["symbol"].to_list(),
        hovertemplate="%{text}<br>%{x|%Y-%m-%d}: %{y:.1f}%<extra></extra>",
        name="Flagged event",
    )
)
fig.add_hline(y=25, line_dash="dash", line_color=COLORS["amber"], line_width=1)
fig.add_hline(y=-25, line_dash="dash", line_color=COLORS["amber"], line_width=1)
fig.update_layout(
    title="Flagged corporate actions are large discrete jumps, not data noise",
    xaxis_title="Date",
    yaxis_title="Overnight return (%)",
    height=420,
)
fig.show()

def quality_check_pipeline(
    df: pl.DataFrame,
    symbol: str,
    quarantine_dir: Path,
    anomaly_cfg: AnomalyConfig | None = None,
) -> tuple[pl.DataFrame, dict]:
    """Run validation → anomaly detection → dedup; quarantine on critical issues."""
    results: dict = {
        "symbol": symbol,
        "timestamp": datetime.now().isoformat(),
        "input_rows": len(df),
        "actions": [],
    }

    validation = OHLCVValidator().validate(df)
    results["validation_issues"] = len(validation.issues)
    results["actions"].append(
        f"{'PASS' if validation.passed else 'FAIL'} Validation: {len(validation.issues)} issues"
    )

    if not validation.passed:
        critical = [i for i in validation.issues if i.severity.value == "critical"]
        if critical:
            quarantine_dir.mkdir(parents=True, exist_ok=True)
            path = quarantine_dir / f"{symbol}_{datetime.now():%Y%m%d_%H%M%S}.parquet"
            df.write_parquet(path)
            results["actions"].append(f"QUARANTINED: {path.name}")

    mgr = AnomalyManager(config=anomaly_cfg) if anomaly_cfg else AnomalyManager()
    report = mgr.analyze(df, symbol)
    results["anomaly_count"] = len(report.anomalies)
    results["actions"].append(f"Anomalies: {len(report.anomalies)}")

    duplicates = len(df) - df["timestamp"].n_unique()
    if duplicates > 0:
        df = df.unique(subset=["timestamp"], keep="last")
        results["actions"].append(f"Removed {duplicates} duplicates")

    results["output_rows"] = len(df)
    results["status"] = (
        "PASS" if validation.passed and not report.get_critical_anomalies() else "REVIEW"
    )
    return df, results

quarantine_dir = OUTPUT_DIR / "quarantine"
pipeline_rows = []
for symbol, df in datasets.items():
    _, result = quality_check_pipeline(df, symbol, quarantine_dir, anomaly_config)
    pipeline_rows.append(
        {
            "symbol": symbol,
            "status": result["status"],
            "input_rows": result["input_rows"],
            "output_rows": result["output_rows"],
            "validation_issues": result["validation_issues"],
            "anomalies": result["anomaly_count"],
        }
    )
pipeline_summary = pl.DataFrame(pipeline_rows)
pipeline_summary