# Author: Simon-Pierre Boucher — simon-pierre.boucher@uqo.ca
"""Feature matrix for the statistical / ML baselines (arms A–D).

Every feature is observable as of the valuation date (sale date − 1 day):
* subject attributes and the roll in force at the sale (public before the sale);
* calendar position (``month_index``) — the model learns the market level from *training* sales only;
* ``prior_sale_*`` (strictly earlier sale of the same parcel);
* static geo layers (``geo_features.parquet``);
* **as-of comparable summaries** computed on sales strictly before the sale date:
  - ``nb_*``: **all** same-class sales of other parcels within 3 km dated in [sale date − 365 d, sale date − 1]
    (exact; monthly KD-trees of candidate sales): count, median log price, median price/m2, median assessment
    ratio, mean distance;
  - ``mun_*``: municipality × class rolling medians over the 365 days strictly before the sale date
    (polars ``rolling`` with a left-closed window — same-day sales are excluded), exact and cheap.
The target is ``log_price``; ``price`` and ``date`` are carried for evaluation and never used as inputs.
"""
from __future__ import annotations

import warnings
from typing import Any

import numpy as np
import polars as pl
from scipy.spatial import cKDTree

from ..data.analytic_schema import PROPERTY_CLASSES

NB_RADIUS_M = 3000.0
NB_WINDOW_DAYS = 365
MONTH0 = (2021, 1)

ID_COLUMNS = ["tx_id", "split", "log_price", "price", "date", "parcel_id"]
CATEGORICAL = ["region_name", "region_group", "mun_code", "property_class", "subtype", "grid1k", "grid5k", "urban_class"]
GEO_NUMERIC = ["in_flood_zone", "transit_stops_within_1km", "da_population_density", "da_median_income", "da_pct_renters", "distance_to_cbd_km"]


def _month_index(d: pl.Expr) -> pl.Expr:
    return (d.dt.year() - MONTH0[0]) * 12 + (d.dt.month() - MONTH0[1])


def subject_features(analytic: pl.DataFrame) -> pl.DataFrame:
    """Row-wise features from the analytic table alone (no cross-row information)."""
    roll_ref = pl.col("roll_reference_date").str.to_date("%Y-%m-%d", strict=False)
    return analytic.with_columns(
        pl.col("floor_area_m2").log().alias("log_floor_area"),
        pl.col("land_area_m2").log().alias("log_land_area"),
        (pl.col("date").dt.year() - pl.col("year_built")).cast(pl.Float64).alias("age_years"),
        pl.col("storeys").alias("storeys_f"),
        pl.col("n_units").cast(pl.Float64).alias("n_units_f"),
        pl.when(pl.col("roll_total_value") > 0).then(pl.col("roll_total_value").log()).alias("log_roll_total"),
        pl.when(pl.col("roll_land_value") > 0).then(pl.col("roll_land_value").log()).alias("log_roll_land"),
        pl.when(pl.col("roll_building_value") > 0).then(pl.col("roll_building_value").log()).alias("log_roll_building"),
        pl.when(pl.col("roll_total_value") > 0).then(pl.col("roll_land_value") / pl.col("roll_total_value")).alias("roll_land_share"),
        pl.when((pl.col("roll_total_value") > 0) & (pl.col("floor_area_m2") > 0)).then((pl.col("roll_total_value") / pl.col("floor_area_m2")).log()).alias("log_roll_per_m2"),
        (((pl.col("date") - roll_ref).dt.total_days()) / 30.4375).alias("roll_age_months"),
        _month_index(pl.col("date")).cast(pl.Float64).alias("month_index"),
        pl.col("date").dt.month().cast(pl.Float64).alias("sale_month"),
        pl.col("prior_sale_price").log().alias("prior_sale_log_price"),
        (((pl.col("date") - pl.col("prior_sale_date")).dt.total_days()) / 30.4375).alias("months_since_prior_sale"),
        pl.col("new_construction").cast(pl.Float64).alias("new_construction_f"),
    )


def municipality_rolling(analytic: pl.DataFrame) -> pl.DataFrame:
    """Municipality × class medians over the 365 days strictly before each sale date
    (tx_id, mun_n_12m, mun_med_log_price_12m, mun_med_ppm2_12m, mun_med_ratio_12m). Row-aligned ``rolling_*_by``
    expressions with a left-closed window: same-day and later sales are never in the window."""
    base = analytic.select(
        "tx_id", "mun_code", "property_class", "date",
        pl.col("log_price"),
        pl.when(pl.col("floor_area_m2") > 0).then(pl.col("price") / pl.col("floor_area_m2")).alias("ppm2"),
        pl.when(pl.col("roll_total_value") > 0).then(pl.col("price") / pl.col("roll_total_value")).alias("ratio"),
    ).sort(["mun_code", "property_class", "date", "tx_id"])
    grp = ["mun_code", "property_class"]
    win = dict(window_size=f"{NB_WINDOW_DAYS}d", closed="left", min_samples=1)
    out = base.with_columns(
        pl.col("log_price").is_not_null().cast(pl.Int32).rolling_sum_by("date", **win).over(grp).fill_null(0).alias("mun_n_12m"),
        pl.col("log_price").rolling_median_by("date", **win).over(grp).alias("mun_med_log_price_12m"),
        pl.col("ppm2").rolling_median_by("date", **win).over(grp).alias("mun_med_ppm2_12m"),
        pl.col("ratio").rolling_median_by("date", **win).over(grp).alias("mun_med_ratio_12m"),
    )
    return out.select("tx_id", "mun_n_12m", "mun_med_log_price_12m", "mun_med_ppm2_12m", "mun_med_ratio_12m")


def _grouped_median(subj: np.ndarray, val: np.ndarray, n_subj: int) -> np.ndarray:
    """Median of ``val`` per subject id (NaN ignored), vectorised through sorting."""
    out = np.full(n_subj, np.nan)
    ok = ~np.isnan(val)
    subj, val = subj[ok], val[ok]
    if subj.size == 0:
        return out
    order = np.lexsort((val, subj))
    subj, val = subj[order], val[order]
    starts = np.flatnonzero(np.r_[True, subj[1:] != subj[:-1]])
    counts = np.diff(np.r_[starts, subj.size])
    lo = starts + (counts - 1) // 2
    hi = starts + counts // 2
    out[subj[starts]] = 0.5 * (val[lo] + val[hi])
    return out


def neighbour_summary(analytic: pl.DataFrame, radius_m: float = NB_RADIUS_M, window_days: int = NB_WINDOW_DAYS) -> pl.DataFrame:
    """Exact as-of neighbourhood summaries: for every sale, all same-class sales within ``radius_m`` whose date lies
    in [date − window_days, date − 1], other parcels only. Implementation: per class, one KD-tree per calendar
    month of *candidate* sales; subjects of month m query the trees of months m−12..m, then exact date masks and
    vectorised grouped medians. Because only earlier sales can enter a window, later sales never influence a
    feature (verified by ``leakage.check_feature_matrix_asof``)."""
    frames = []
    for cls in PROPERTY_CLASSES:
        sub = analytic.filter(pl.col("property_class") == cls).sort(["date", "tx_id"])
        n = sub.height
        if n == 0:
            continue
        x = sub["x_m"].to_numpy(); y = sub["y_m"].to_numpy()
        d = sub["date"].cast(pl.Int32).to_numpy().astype(np.int64)
        month = (sub["date"].dt.year() * 12 + sub["date"].dt.month() - 1).to_numpy().astype(np.int64)
        lp = sub["log_price"].to_numpy()
        fa = sub["floor_area_m2"].fill_null(float("nan")).to_numpy()
        price = sub["price"].to_numpy()
        ppm2 = np.where(fa > 0, price / np.where(fa > 0, fa, 1.0), np.nan)
        rt = sub["roll_total_value"].fill_null(float("nan")).to_numpy()
        ratio = np.where(rt > 0, price / np.where(rt > 0, rt, 1.0), np.nan)
        parcel = sub["parcel_id"].cast(pl.Categorical).to_physical().to_numpy().astype(np.int64)
        xy = np.column_stack([x, y])
        months = np.unique(month)
        trees: dict[int, tuple[cKDTree, np.ndarray]] = {}
        for m in months:
            rows = np.flatnonzero(month == m)
            trees[int(m)] = (cKDTree(xy[rows]), rows)
        n_12m = np.zeros(n, dtype=np.int32)
        med_lp = np.full(n, np.nan); med_ppm2 = np.full(n, np.nan); med_ratio = np.full(n, np.nan); mean_dist = np.full(n, np.nan)
        for m in months:
            subj_rows = np.flatnonzero(month == m)
            pairs_s: list[np.ndarray] = []
            pairs_c: list[np.ndarray] = []
            for mm in range(int(m) - 13, int(m) + 1):
                t = trees.get(mm)
                if t is None:
                    continue
                tree, cand_rows = t
                lists = tree.query_ball_point(xy[subj_rows], r=radius_m, return_sorted=False)
                lens = np.fromiter((len(l) for l in lists), dtype=np.int64, count=len(lists))
                if lens.sum() == 0:
                    continue
                cand = cand_rows[np.concatenate([np.asarray(l, dtype=np.int64) for l in lists if len(l)])]
                subj = np.repeat(subj_rows, lens)
                keep = (d[cand] < d[subj]) & (d[cand] >= d[subj] - window_days) & (parcel[cand] != parcel[subj])
                pairs_s.append(subj[keep]); pairs_c.append(cand[keep])
            if not pairs_s:
                continue
            S = np.concatenate(pairs_s); C = np.concatenate(pairs_c)
            cnt = np.bincount(S, minlength=n)
            n_12m += cnt.astype(np.int32)
            dist = np.hypot(x[C] - x[S], y[C] - y[S]) / 1000.0
            sum_dist = np.bincount(S, weights=dist, minlength=n)
            with np.errstate(all="ignore"):
                mean_dist = np.where(cnt > 0, sum_dist / np.maximum(cnt, 1), mean_dist)
            for arr, target in ((lp, med_lp), (ppm2, med_ppm2), (ratio, med_ratio)):
                med = _grouped_median(S, arr[C], n)
                target[subj_rows] = med[subj_rows]
        frames.append(pl.DataFrame({
            "tx_id": sub["tx_id"], "nb_n_12m": n_12m, "nb_med_log_price": med_lp, "nb_med_ppm2": med_ppm2,
            "nb_med_ratio": med_ratio, "nb_mean_dist_km": mean_dist,
        }))
    return pl.concat(frames)


def build_feature_matrix(analytic: pl.DataFrame, splits: pl.DataFrame, geo: pl.DataFrame | None = None, comps_summary: bool = True) -> pl.DataFrame:
    """Assemble the baseline design matrix (one row per sale in ``splits``)."""
    df = subject_features(analytic).join(splits.select("tx_id", "split"), on="tx_id", how="inner")
    numeric = [
        "log_floor_area", "log_land_area", "age_years", "storeys_f", "n_units_f", "log_roll_total", "log_roll_land", "log_roll_building",
        "roll_land_share", "log_roll_per_m2", "roll_age_months", "lat", "lng", "x_m", "y_m", "month_index", "sale_month",
        "prior_sale_log_price", "months_since_prior_sale", "new_construction_f",
    ]
    extra: list[str] = []
    if comps_summary:
        df = df.join(municipality_rolling(analytic), on="tx_id", how="left").join(neighbour_summary(analytic), on="tx_id", how="left")
        extra += ["mun_n_12m", "mun_med_log_price_12m", "mun_med_ppm2_12m", "mun_med_ratio_12m",
                  "nb_n_12m", "nb_med_log_price", "nb_med_ppm2", "nb_med_ratio", "nb_mean_dist_km"]
    if geo is not None:
        g = geo.select(["tx_id", *[c for c in GEO_NUMERIC if c in geo.columns]])
        if "in_flood_zone" in g.columns:
            g = g.with_columns(pl.col("in_flood_zone").cast(pl.Float64))
        df = df.join(g, on="tx_id", how="left")
        extra += [c for c in GEO_NUMERIC if c in g.columns]
    cols = ID_COLUMNS + numeric + extra + CATEGORICAL
    return df.select(cols).with_columns([pl.col(c).cast(pl.Float64) for c in numeric + extra]).sort(["date", "tx_id"])


def feature_columns(matrix: pl.DataFrame) -> tuple[list[str], list[str]]:
    """(numeric feature names, categorical feature names) of a built matrix."""
    num = [c for c in matrix.columns if c not in ID_COLUMNS and c not in CATEGORICAL]
    cat = [c for c in CATEGORICAL if c in matrix.columns]
    return num, cat