# Author: Simon-Pierre Boucher — simon-pierre.boucher@uqo.ca
"""Analytic contract, derived columns, splits and subsamples (synthetic data)."""
from __future__ import annotations

from datetime import date

import numpy as np
import polars as pl
import pytest

from agentic_valuation.data.analytic_schema import (
    ANALYTIC_COLUMNS, DERIVED_COLUMNS, REQUIRED_NON_NULL, SYNTHETIC_COLUMNS, add_prior_sale, coerce_analytic, haversine_km,
    project_xy, region_group_expr, validate_analytic,
)
from agentic_valuation.data.build import roll_reference_at_sale_expr, urban_class_expr
from agentic_valuation.data.splits import SPLIT_ORDER, assign_splits, proportional_allocation, split_windows, stratified_subsample
from tests.conftest import SYN_N_BY_SPLIT


def test_contract_shape() -> None:
    assert len(SYNTHETIC_COLUMNS) == 24
    assert len(DERIVED_COLUMNS) == 9
    assert list(ANALYTIC_COLUMNS)[:24] == list(SYNTHETIC_COLUMNS)
    assert set(REQUIRED_NON_NULL) <= set(ANALYTIC_COLUMNS)


def test_synthetic_analytic_validates(syn_analytic: pl.DataFrame) -> None:
    validate_analytic(syn_analytic)
    assert syn_analytic.columns == list(ANALYTIC_COLUMNS)
    assert syn_analytic.height == 20_000
    # no NaN survives coercion (nulls only)
    for c, t in ANALYTIC_COLUMNS.items():
        if t == pl.Float64:
            assert syn_analytic[c].is_nan().fill_null(False).sum() == 0, c


def test_validate_raises_on_violations(syn_analytic: pl.DataFrame) -> None:
    with pytest.raises(ValueError, match="missing columns"):
        validate_analytic(syn_analytic.drop("grid1k"))
    with pytest.raises(ValueError, match="dtype"):
        validate_analytic(syn_analytic.with_columns(pl.col("price").cast(pl.Int64)))
    with pytest.raises(ValueError, match="nulls"):
        validate_analytic(syn_analytic.with_columns(pl.when(pl.arange(0, pl.len()) == 0).then(None).otherwise(pl.col("parcel_id")).alias("parcel_id")))
    with pytest.raises(ValueError, match="unexpected columns"):
        validate_analytic(syn_analytic.with_columns(pl.lit(1).alias("extra")))
    validate_analytic(syn_analytic.with_columns(pl.lit(1).alias("extra")), allow_extra=True)
    with pytest.raises(ValueError, match="prior_sale_date"):
        bad = syn_analytic.with_columns(pl.col("date").alias("prior_sale_date"))
        validate_analytic(bad)


def test_coerce_adds_missing_and_orders(syn_analytic: pl.DataFrame) -> None:
    shuffled = syn_analytic.select(list(reversed(syn_analytic.columns))).drop("borough_name")
    out = coerce_analytic(shuffled)
    assert out.columns == list(ANALYTIC_COLUMNS)
    assert out["borough_name"].null_count() == out.height
    assert out.schema["date"] == pl.Date


def test_projection_matches_epsg32198() -> None:
    x, y = project_xy(np.array([46.973765046]), np.array([-69.799217111]))
    assert abs(x[0] - (-98672.84)) < 1.0 and abs(y[0] - 331793.78) < 1.0  # value from the WP 22 table
    assert abs(haversine_km(45.5017, -73.5673, 46.8139, -71.2080) - 233) < 3  # Montréal–Québec ≈ 233 km


def test_region_group_rule() -> None:
    df = pl.DataFrame({
        "region_name": ["Montréal", "Montérégie", "Montérégie", "Capitale-Nationale", "Capitale-Nationale", "Outaouais", "Estrie"],
        "lat": [45.50, 45.40, 45.00, 46.80, 47.40, 45.45, 45.40],
        "lng": [-73.57, -73.30, -72.50, -71.22, -71.00, -75.72, -71.90],
    }).with_columns(region_group_expr().alias("rg"))
    assert df["rg"].to_list() == ["Montréal CMA", "Montréal CMA", "other", "Québec CMA", "other", "Outaouais", "other"]


def test_prior_sale_strictly_earlier_and_same_day_invisible() -> None:
    df = pl.DataFrame({
        "tx_id": ["a", "b", "c", "d", "e"],
        "parcel_id": ["P1", "P1", "P1", "P2", "P1"],
        "date": [date(2021, 1, 1), date(2021, 1, 1), date(2022, 6, 1), date(2021, 3, 3), date(2023, 1, 1)],
        "price": [100.0, 120.0, 200.0, 50.0, 300.0],
    })
    out = add_prior_sale(df).sort("tx_id")
    assert out["prior_sale_price"].to_list() == [None, None, 120.0, None, 200.0]
    assert out["prior_sale_date"].to_list() == [None, None, date(2021, 1, 1), None, date(2022, 6, 1)]


def test_price_tercile_balanced_by_class(syn_analytic: pl.DataFrame) -> None:
    t = syn_analytic.group_by("property_class", "price_tercile").len().sort(["property_class", "price_tercile"])
    for cls, g in t.group_by("property_class"):
        counts = g["len"].to_list()
        assert max(counts) - min(counts) <= 1, cls
    # monotone in price
    q = syn_analytic.group_by("property_class", "price_tercile").agg(pl.col("price").min().alias("lo"), pl.col("price").max().alias("hi")).sort(["property_class", "price_tercile"])
    for cls, g in q.group_by("property_class"):
        his = g.sort("price_tercile")["hi"].to_list(); los = g.sort("price_tercile")["lo"].to_list()
        assert his[0] <= los[1] and his[1] <= los[2]


def test_urban_class_and_roll_reference_rules() -> None:
    df = pl.DataFrame({"mamh_population": [250000, 15000, 900, None], "roll_reference_date": ["2024-07-01", "2023-07-01", "2022-07-01", None],
                       "date": [date(2026, 2, 1), date(2024, 12, 31), date(2021, 3, 1), date(2022, 1, 1)]})
    out = df.with_columns(urban_class_expr().alias("u"), roll_reference_at_sale_expr().alias("r"))
    assert out["u"].to_list() == ["metro", "suburban", "rural", "rural"]
    # 2024-07-01 roll in force from 2026-01-01 → k=0 ; 2023-07-01 roll in force 2025-01-01 > 2024-12-31 → step back 1 cycle ;
    # 2022-07-01 roll in force 2024-01-01, sale 2021-03 → back 1 cycle (2019-07-01 in force 2021-01-01) ; null stays null
    assert out["r"].to_list() == ["2024-07-01", "2020-07-01", "2019-07-01", None]


def test_splits_windows_and_assignment(syn_analytic: pl.DataFrame, syn_splits: pl.DataFrame, cfg: dict) -> None:
    w = split_windows(cfg)
    for a, b in zip(SPLIT_ORDER, SPLIT_ORDER[1:]):
        assert w[a][1] < w[b][0]
    joined = syn_analytic.select("tx_id", "date").join(syn_splits, on="tx_id")
    for name, (a, b) in w.items():
        part = joined.filter(pl.col("split") == name)
        assert part.height > 0
        assert part["date"].min() >= a and part["date"].max() <= b
    assert set(syn_splits["split"].unique()) <= set(SPLIT_ORDER) | {"unused"}
    assert syn_splits["spatial_holdout_region"].dtype == pl.Boolean
    assert syn_splits.filter(pl.col("spatial_holdout_region")).height == syn_analytic.filter(pl.col("region_name").is_in(cfg["spatial_holdout"]["regions"])).height


def test_overlapping_windows_rejected(syn_analytic: pl.DataFrame, cfg: dict) -> None:
    bad = {**cfg, "splits": {**cfg["splits"], "validation": {"start": "2024-06-01", "end": "2025-03-31"}}}
    with pytest.raises(ValueError, match="overlaps"):
        assign_splits(syn_analytic, bad)


def test_proportional_allocation() -> None:
    alloc = proportional_allocation({"a": 100, "b": 50, "c": 2}, 30)
    assert sum(alloc.values()) == 30 and alloc["c"] <= 2 and alloc["a"] > alloc["b"]
    alloc = proportional_allocation({"a": 3, "b": 3}, 10)
    assert alloc == {"a": 3, "b": 3}


def test_subsample_counts_strata_and_determinism(syn_analytic: pl.DataFrame, syn_splits: pl.DataFrame, syn_subsamples: pl.DataFrame, cfg: dict) -> None:
    counts = {k: int(v) for k, v in syn_subsamples.group_by("split").len().iter_rows()}
    assert counts == SYN_N_BY_SPLIT
    assert syn_subsamples["tx_id"].n_unique() == syn_subsamples.height
    assert syn_subsamples.columns == ["tx_id", "split", "stratum", "draw_seed"]
    again = stratified_subsample(syn_analytic, syn_splits, cfg, n_by_split=SYN_N_BY_SPLIT)
    assert again.equals(syn_subsamples)
    other = stratified_subsample(syn_analytic, syn_splits, cfg, n_by_split=SYN_N_BY_SPLIT, seed=1)
    assert not other.equals(syn_subsamples)
    # each stratum label = region_group|property_class|price_tercile of the row
    chk = syn_subsamples.join(syn_analytic.select("tx_id", "region_group", "property_class", "price_tercile"), on="tx_id")
    lab = chk.select(pl.concat_str([pl.col("region_group"), pl.col("property_class"), pl.col("price_tercile").cast(pl.String)], separator="|"))
    assert lab.to_series().to_list() == chk["stratum"].to_list()


@pytest.mark.real
def test_real_analytic_contract(real_bundle: dict[str, pl.DataFrame]) -> None:
    validate_analytic(real_bundle["analytic"])
    assert real_bundle["analytic"].height > 600_000