# Author: Simon-Pierre Boucher — contact@spboucher.ai
"""Estimation-sample construction for WP10.

Reads the raw matched transaction–roll snapshot, restricts to residential
arm's-length sales with a high-confidence roll match, builds the assessment
ratio and every derived regressor, applies the trims, and writes
``data/processed/analysis.parquet``.

Key variable definitions
------------------------
ratio        AV / SP where AV = ``role_valeur_immeuble`` (total assessed value
             on the roll in force at the sale date) and SP = ``amount``.
ln_ratio     log(ratio).
cell         municipality × roll vintage × sale year — the market-timing cell
             inside which assessed values share a common reference date, so
             the roll lag is constant and ratio comparisons are clean.
lag_months   months elapsed between the roll's market-condition reference
             date (July 1, by statute) and the sale date.
land_share   assessed land value / total assessed value.
age          sale year minus year built (from the roll where available).
"""
import numpy as np
import pandas as pd

from . import config


def _property_class(df: pd.DataFrame) -> pd.Series:
    """Harmonised dwelling class from the listing type, CUBF code and unit count."""
    cls = pd.Series("other", index=df.index, dtype="object")
    cls[df["role_cubf"] == "1100"] = "cottage"
    cls[df["role_cubf"] == "1211"] = "mobile"
    is_dwelling = df["role_cubf"] == "1000"
    cls[is_dwelling & (df["propertyType"] == "condo")] = "condo"
    cls[is_dwelling & (df["propertyType"] == "plex")] = "plex"
    cls[is_dwelling & (df["propertyType"] == "unifamilial")] = "single_family"
    # dwellings the feed left untyped: use the roll's unit count
    untyped = is_dwelling & (cls == "other")
    cls[untyped & (df["role_nb_logements"] == 1)] = "single_family"
    cls[untyped & (df["role_nb_logements"].between(2, 5))] = "plex"
    return cls


def build(raw: pd.DataFrame | None = None) -> pd.DataFrame:
    """Apply every sample restriction and return the analysis DataFrame."""
    df = raw if raw is not None else pd.read_parquet(config.RAW_PARQUET)
    n0 = len(df)
    log = [("raw snapshot", n0)]

    # -------------------------------------------------- residential use codes
    df = df[df["role_cubf"].isin(config.RESIDENTIAL_CUBF)]
    log.append(("residential CUBF (1000/1100/1211/1990)", len(df)))

    # -------------------------------------------------- match confidence
    df = df[(df["match_dist_m"] <= config.MATCH_MAX_DIST_M)
            & (df["match_score"] >= config.MATCH_MIN_SCORE)]
    log.append((f"match dist ≤ {config.MATCH_MAX_DIST_M:.0f} m & score ≥ "
                f"{config.MATCH_MIN_SCORE:.0f}", len(df)))

    # -------------------------------------------------- valid AV and SP
    df = df[(df["amount"] >= config.PRICE_MIN)
            & (df["role_valeur_immeuble"] > 5_000)
            & df["role_date_cond_marche"].notna()]
    log.append(("positive AV, SP ≥ 50k, dated roll", len(df)))

    # -------------------------------------------------- derived variables
    df = df.copy()
    df["sale_date"] = pd.to_datetime(df["date"])
    df["sale_year"] = df["tx_year"].astype(int)
    df["ratio"] = df["role_valeur_immeuble"] / df["amount"]
    df["ln_ratio"] = np.log(df["ratio"])
    df["ln_price"] = np.log(df["amount"].astype(float))
    df["ln_av"] = np.log(df["role_valeur_immeuble"])
    df["lag_months"] = ((df["sale_date"] - df["role_date_cond_marche"]).dt.days
                        / 30.44)
    df["land_share"] = (df["role_valeur_terrain"]
                        / df["role_valeur_immeuble"]).clip(0, 1)
    df["age"] = (df["sale_year"]
                 - pd.to_numeric(df["role_annee_construction"], errors="coerce"))
    df.loc[(df["age"] < 0) | (df["age"] > 300), "age"] = np.nan
    df["prop_class"] = _property_class(df)
    df["muni"] = df["role_code_mun"]
    df["roll"] = df["role_anrole"]
    df["cell"] = (df["muni"] + "_" + df["roll"] + "_"
                  + df["sale_year"].astype(str))

    # -------------------------------------------------- ratio trim (by roll vintage,
    # so the mechanical drift of ratios across reference dates is not trimmed away)
    lo, hi = config.RATIO_TRIM
    q = df.groupby("roll")["ratio"].quantile([lo, hi]).unstack()
    df = df.join(q.rename(columns={lo: "_qlo", hi: "_qhi"}), on="roll")
    df = df[(df["ratio"] >= df["_qlo"]) & (df["ratio"] <= df["_qhi"])]
    df = df.drop(columns=["_qlo", "_qhi"])
    log.append((f"ratio inside [{lo:.0%}, {hi:.0%}] of its roll vintage", len(df)))

    # -------------------------------------------------- market-timing cells
    counts = df.groupby("cell")["cell"].transform("size")
    df = df[counts >= config.CELL_MIN_OBS]
    log.append((f"cells (muni × roll × year) with ≥ {config.CELL_MIN_OBS} sales",
                len(df)))

    df.attrs["selection_log"] = log
    keep = ["id", "sale_date", "sale_year", "amount", "ln_price",
            "role_valeur_immeuble", "role_valeur_terrain", "role_valeur_batiment",
            "totalArValue", "previousValue", "ln_av", "ratio", "ln_ratio",
            "lag_months", "land_share", "age", "prop_class",
            "muni", "roll", "cell", "city", "role_municipalite",
            "lat", "lng", "role_superficie_terrain_m2", "role_aire_etages_m2",
            "role_nb_logements", "match_dist_m", "match_score",
            "match_valeur_exacte", "ownerType", "role_cubf"]
    return df[keep].reset_index(drop=True)


def build_and_save() -> pd.DataFrame:
    """Build the sample, print the selection log, persist to parquet."""
    config.ensure_dirs()
    s = build()
    for step, n in s.attrs["selection_log"]:
        print(f"  {n:>9,}  after: {step}")
    s.to_parquet(config.ANALYSIS_PARQUET, index=False)
    return s


def load() -> pd.DataFrame:
    """Load the processed analysis sample (build it first if missing)."""
    if not config.ANALYSIS_PARQUET.exists():
        return build_and_save()
    return pd.read_parquet(config.ANALYSIS_PARQUET)