# Author: Simon-Pierre Boucher — simon-pierre.boucher@uqo.ca
"""Unit / integration tests for the SHARP-HPI code base (run: python3 -m pytest -q tests)."""
import sys
from pathlib import Path

sys.path.insert(0, str(Path(__file__).resolve().parents[1] / "src"))
import numpy as np
import polars as pl
import pytest

from sharp_hpi import estimators as E
from sharp_hpi.hedonic import assign_loc_fe, build_features, fit_fe
from sharp_hpi.panel import growth, normalise
from sharp_hpi.sharp import cell_log_index, run_sharp
from sharp_hpi.simulate import Scenario, simulate
from sharp_hpi.statespace import eb_prior, fit_cell, profile_loglik


def test_month_grid():
    g = E.month_grid("2021-11", "2022-02")
    assert g == ["2021-11", "2021-12", "2022-01", "2022-02"]


def test_normalisation_does_not_change_growth():
    months = E.month_grid("2021-01", "2022-12"); li = np.cumsum(np.full(24, 0.01))
    i1, _ = normalise(li, months, "2021"); i2, _ = normalise(li, months, "2021-01")
    g1 = np.diff(np.log(i1)); g2 = np.diff(np.log(i2))
    assert np.allclose(g1, g2) and np.isclose(np.mean(np.log(i1[:12])), np.log(100)) and np.isclose(i2[0], 100)
    d1, p1, p12 = growth(li); assert np.isclose(p12[12], 100 * (np.exp(0.12) - 1))


def test_local_level_recovers_signal_and_tau():
    rng = np.random.default_rng(3); T = 120
    d = np.cumsum(rng.normal(0, 0.01, T)); n = rng.poisson(20, T); y = d + rng.normal(0, 0.2 / np.sqrt(np.maximum(n, 1))); y[n == 0] = np.nan
    f = fit_cell(y, n, 0.04, model="ll", robust=False)
    assert np.sqrt(np.mean((f.smoothed - d) ** 2)) < np.sqrt(np.nanmean((y - d) ** 2))   # smoother beats raw means
    assert 0.004 < np.sqrt(f.tau2) < 0.025
    assert np.all((f.gain >= 0) & (f.gain <= 1)) and np.all(f.gain[n == 0] == 0)


def test_eb_prior_recovers_common_tau():
    rng = np.random.default_rng(5); profs = []
    for _ in range(25):
        d = np.cumsum(rng.normal(0, 0.008, 60)); n = rng.poisson(3, 60); y = d + rng.normal(0, 0.2 / np.sqrt(np.maximum(n, 1))); y[n == 0] = np.nan
        profs.append(profile_loglik(y, n, 0.04))
    m, v = eb_prior(profs)
    assert 0.003 < np.exp(m / 2) < 0.02


def test_simulation_schema_and_truth_shapes():
    s = simulate(Scenario(T=24, n_regions=2, mrc_per_region=2, mun_per_mrc=2, seed=1))
    for c in ["tx_id", "parcel_id", "date", "month", "price", "log_price", "property_class", "region_code", "mrc_code", "mun_code", "floor_area_m2", "age", "grid500", "available_date"]:
        assert c in s.df.columns
    assert s.truth_mun.shape == (8, 24) and s.truth_prov.shape == (24,) and len(s.months) == 24


def test_pooled_td_recovers_known_path():
    s = simulate(Scenario(T=36, n_regions=2, mrc_per_region=2, mun_per_mrc=3, seed=2, size_logmean=2.5, tau_mun=0.0, tau_mrc=0.0, tau_region=0.0, contamination=0.0))
    f = assign_loc_fe(build_features(s.df), min_n=15)
    td = E.pooled_td(f, s.months, robust=True).log_index
    truth = s.truth_prov
    err = np.diff(td) - np.diff(truth)
    assert np.sqrt(np.mean(err ** 2)) < 0.02


def test_sharp_end_to_end_and_hierarchy_consistency():
    s = simulate(Scenario(T=36, n_regions=2, mrc_per_region=2, mun_per_mrc=3, seed=7, size_logmean=1.8))
    f = build_features(s.df)
    r = run_sharp(f, s.months)
    assert len(r.stage1.mu) == 36 and r.stage1.mu[0] == 0.0
    assert any(k[0] == "municipality" for k in r.cells)
    for k in [k for k in r.cells if k[0] == "municipality"][:3]:
        p, v = cell_log_index(r, *k)
        assert len(p) == 36 and np.all(np.isfinite(p)) and np.all(v > 0)
        # child path = parent path + own deviation (nested additivity)
        c = r.cells[k]; parent = cell_log_index(r, "mrc", c.parent_id)[0]
        assert np.allclose(p - parent, c.fit.smoothed)
    # real-time (filtered) path must not depend on later months: truncating the sample reproduces the earlier filtered values up to
    # stage-1 re-estimation; check that the filter is one-sided on fixed inputs
    i = s.months.index("2022-06")
    k = [k for k in r.cells if k[0] == "region"][0]; c = r.cells[k]
    f1 = fit_cell(c.ybar[: i + 1], c.n[: i + 1], c.sigma2_cell, log_tau2=np.log(c.fit.tau2), robust=False)
    f2 = fit_cell(c.ybar, c.n, c.sigma2_cell, log_tau2=np.log(c.fit.tau2), robust=False)
    assert np.allclose(f1.filtered, f2.filtered[: i + 1])


def test_repeat_sales_on_synthetic_pairs():
    s = simulate(Scenario(T=30, n_regions=1, mrc_per_region=1, mun_per_mrc=4, seed=11, size_logmean=3.0, resale_prob=0.3, drift=0.02, tau_mun=0.0, tau_mrc=0.0, tau_region=0.0))
    f = build_features(s.df); rs = E.repeat_sales(f, s.months)
    ok = ~np.isnan(rs.log_index)
    assert rs.meta["pairs"] > 200 and ok.sum() > 20
    assert np.corrcoef(rs.log_index[ok], s.truth_prov[ok])[0, 1] > 0.8


def test_fit_fe_residual_alignment_with_singletons():
    s = simulate(Scenario(T=12, n_regions=1, mrc_per_region=1, mun_per_mrc=2, seed=4, size_logmean=2.0))
    f = build_features(s.df)
    # half of the rows get their own (singleton) fixed effect, the rest share one cell
    f = f.with_columns(pl.when(pl.int_range(pl.len()) % 2 == 0).then(pl.col("tx_id")).otherwise(pl.lit("shared")).alias("loc_fe"))
    res, m = fit_fe(f, "lfa", "loc_fe", robust=False)
    assert len(res.resid) == f.height and res.n == (f.height + 1) // 2 and np.isnan(res.resid).sum() == f.height // 2