# Author: Simon-Pierre Boucher — simon-pierre.boucher@uqo.ca
"""Baselines on a small synthetic feature matrix built here (independent of the feature pipeline)."""
from __future__ import annotations

import numpy as np
import polars as pl
import pytest

from agentic_valuation.baselines.gbm import AVM, LightGBMModel
from agentic_valuation.baselines.hedonic import HedonicRidge, HierarchicalHedonic
from agentic_valuation.data.synthetic import generate


@pytest.fixture(scope="module")
def fm():
    df = generate(4000, seed=1)
    df = df.with_columns(
        pl.col("price").log().alias("log_price"),
        pl.col("floor_area_m2").log().alias("log_floor_area"),
        (2026 - pl.col("year_built")).alias("age"),
        pl.col("roll_total_value").log().alias("log_roll_total"),
        (pl.col("lat") * 100).floor().cast(pl.Utf8).alias("grid1k"),
        (pl.col("lat") * 20).floor().cast(pl.Utf8).alias("grid5k"),
        pl.col("region_name").alias("region_group"),
        pl.when(pl.col("date") < pl.date(2024, 7, 1)).then(pl.lit("train")).when(pl.col("date") < pl.date(2025, 4, 1)).then(pl.lit("validation")).otherwise(pl.lit("test")).alias("split"),
    ).drop(["exclude_reason", "borough_name", "roll_reference_date", "parcel_id"])
    return df


def _split(fm):
    return (fm.filter(pl.col("split") == s) for s in ("train", "validation", "test"))


def test_hedonic_beats_naive(fm):
    tr, va, te = _split(fm)
    m = HedonicRidge().fit(tr, va)
    p = m.predict(te)
    ae = np.abs(np.log(p.point) - te["log_price"].to_numpy()).mean()
    naive = np.abs(te["log_price"].to_numpy() - tr["log_price"].mean()).mean()
    assert ae < 0.75 * naive
    assert p.quantiles.shape == (te.height, 5)
    assert np.all(np.diff(p.quantiles, axis=1) >= 0)


def test_hierarchical_runs(fm):
    tr, va, te = _split(fm)
    m = HierarchicalHedonic().fit(tr, va)
    p = m.predict(te)
    assert np.isfinite(p.point).all()


def test_lightgbm_and_avm(fm):
    tr, va, te = _split(fm)
    m = LightGBMModel(n_trials=2, max_rounds=300).fit(tr, va)
    p = m.predict(te)
    ae = np.abs(np.log(p.point) - te["log_price"].to_numpy()).mean()
    assert ae < 0.30
    avm = AVM(n_trials=1)
    avm.point.max_rounds = 300
    avm.fit(tr, va)
    # conformal guarantee holds on the calibration partition; under the synthetic trend drift the test-period
    # coverage is lower (documented behaviour, measured in the paper's drift section)
    q = avm.predict(va)
    y = va["log_price"].to_numpy()
    cov = np.mean((np.log(q.quantiles[:, 0]) <= y) & (y <= np.log(q.quantiles[:, 4])))
    assert 0.85 <= cov <= 0.95
    q = avm.predict(te)
    rec = q.to_records("test", "D_avm", "t")
    assert set(rec.columns) == {"tx_id", "split", "arm", "committee", "aggregation", "point", "q05", "q25", "q50", "q75", "q95", "cost_usd", "latency_s", "tokens", "run_id"}