# Author: Simon-Pierre Boucher — simon-pierre.boucher@uqo.ca
"""Leakage-defence suite (synthetic data; ``--real`` re-runs the audit on the processed real data)."""
from __future__ import annotations

from datetime import date

import polars as pl
import pytest

from agentic_valuation.data import leakage as lk
from tests.conftest import SYN_N_BY_SPLIT


def test_split_windows_non_overlapping(cfg: dict) -> None:
    r = lk.check_split_windows(cfg)
    assert r.passed, r.details
    assert r.details["gaps"] == []


def test_split_assignment(syn_analytic: pl.DataFrame, syn_splits: pl.DataFrame, cfg: dict) -> None:
    assert lk.check_split_assignment(syn_analytic, syn_splits, cfg).passed


def test_group_ordering_no_future_sale_in_train(syn_analytic: pl.DataFrame, syn_splits: pl.DataFrame) -> None:
    r = lk.check_group_ordering(syn_analytic, syn_splits)
    assert r.passed and r.details["parcels_in_train_and_later"] > 0


def test_group_ordering_detects_violation(syn_analytic: pl.DataFrame, syn_splits: pl.DataFrame) -> None:
    # later sale B of a repeat parcel into train while its earlier sale A sits in test → must fail
    b = syn_analytic.filter(pl.col("prior_sale_date").is_not_null() & (pl.col("date") >= date(2025, 4, 1))).head(1)
    a = syn_analytic.filter((pl.col("parcel_id") == b["parcel_id"][0]) & (pl.col("date") == b["prior_sale_date"][0])).head(1)
    bad = syn_splits.with_columns(
        pl.when(pl.col("tx_id") == b["tx_id"][0]).then(pl.lit("train"))
        .when(pl.col("tx_id") == a["tx_id"][0]).then(pl.lit("test")).otherwise(pl.col("split")).alias("split")
    )
    r = lk.check_group_ordering(syn_analytic, bad)
    assert not r.passed and r.n_failed == 1


def test_prior_sale_strictly_earlier(syn_analytic: pl.DataFrame) -> None:
    r = lk.check_prior_sale(syn_analytic)
    assert r.passed, r.details


def test_comparables_asof(syn_analytic: pl.DataFrame, syn_subsamples: pl.DataFrame, cfg: dict) -> None:
    subjects = syn_analytic.join(syn_subsamples.select("tx_id"), on="tx_id")
    r = lk.check_comparables(syn_analytic, subjects, cfg)
    assert r.passed, r.details


def test_market_context_asof(syn_analytic: pl.DataFrame, syn_subsamples: pl.DataFrame, cfg: dict) -> None:
    subjects = syn_analytic.join(syn_subsamples.select("tx_id"), on="tx_id").head(15)
    assert lk.check_market_context(syn_analytic, subjects, cfg).passed


def test_feature_matrix_asof(syn_analytic: pl.DataFrame, syn_splits: pl.DataFrame) -> None:
    r = lk.check_feature_matrix_asof(syn_analytic, syn_splits, None, cutoff=date(2025, 4, 1))
    assert r.passed, r.details


def test_feature_matrix_columns(syn_matrix: pl.DataFrame) -> None:
    assert lk.check_feature_matrix_columns(syn_matrix).passed
    bad = syn_matrix.with_columns(pl.col("log_price").alias("price_now"))
    assert not lk.check_feature_matrix_columns(bad).passed


def test_subsample_strata_and_determinism(syn_analytic: pl.DataFrame, syn_splits: pl.DataFrame, syn_subsamples: pl.DataFrame, cfg: dict) -> None:
    r = lk.check_subsample(syn_analytic, syn_splits, syn_subsamples, cfg, n_by_split=SYN_N_BY_SPLIT)
    assert r.passed, r.details
    assert r.details["deterministic_redraw"] is True


def test_full_audit_synthetic(syn_analytic: pl.DataFrame, syn_splits: pl.DataFrame, syn_subsamples: pl.DataFrame, syn_matrix: pl.DataFrame, cfg: dict) -> None:
    rep = lk.audit(syn_analytic, syn_splits, syn_subsamples, syn_matrix, None, cfg, n_subjects=50, n_market_subjects=5, n_by_split=SYN_N_BY_SPLIT, asof_sample=None)
    assert rep["all_passed"], [c for c in rep["checks"] if not c["passed"]]


@pytest.mark.real
def test_real_audit_passes(real_bundle: dict[str, pl.DataFrame], cfg: dict) -> None:
    rep = lk.audit(real_bundle["analytic"], real_bundle["splits"], real_bundle["llm_subsamples"], real_bundle.get("feature_matrix"), real_bundle.get("geo_features"), cfg, n_subjects=200, n_market_subjects=20, asof_sample=20_000)
    assert rep["all_passed"], [c for c in rep["checks"] if not c["passed"]]