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

Residential arm's-length sales (2021–2026) with a high-confidence roll
match and complete structural descriptors from the assessment roll:
floor area, lot area, building age, dwelling class, physical configuration
and coordinates. Writes ``data/processed/analysis.parquet``.
"""
import numpy as np
import pandas as pd

from . import config

LINK_MAP = {"1": "detached", "2": "semi_detached", "3": "row_end",
            "4": "row", "5": "apartment_link"}


def _property_class(df: pd.DataFrame) -> pd.Series:
    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"
    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:
    df = raw if raw is not None else pd.read_parquet(config.RAW_PARQUET)
    log = [("raw snapshot", len(df))]

    df = df[df["role_cubf"].isin(config.RESIDENTIAL_CUBF)]
    log.append(("residential CUBF", len(df)))

    df = df[(df["match_dist_m"] <= config.MATCH_MAX_DIST_M)
            & (df["match_score"] >= config.MATCH_MIN_SCORE)]
    log.append(("high-confidence roll match", len(df)))

    df = df[(df["amount"] >= config.PRICE_MIN)
            & (df["role_aire_etages_m2"] > 20)
            & (df["role_aire_etages_m2"] < 2_000)]
    log.append(("valid price and floor area", len(df)))

    df = df.copy()
    df["sale_date"] = pd.to_datetime(df["date"])
    df["sale_year"] = df["tx_year"].astype(int)
    df["month"] = df["sale_date"].dt.to_period("M").astype(str)
    df["quarter"] = df["sale_date"].dt.to_period("Q").astype(str)
    df["t"] = ((df["sale_date"] - pd.Timestamp("2021-01-01")).dt.days
               / 30.44)  # months since Jan 2021, continuous

    df["price"] = df["amount"].astype(float)
    df["ln_price"] = np.log(df["price"])
    df["area"] = df["role_aire_etages_m2"].astype(float)
    df["ln_area"] = np.log(df["area"])
    df["lot"] = df["role_superficie_terrain_m2"].fillna(0).clip(lower=0)
    df.loc[df["lot"] > df["lot"].quantile(0.995), "lot"] = np.nan
    df["lot"] = df["lot"].fillna(0)
    df["has_lot"] = (df["lot"] > 0).astype(float)
    df["ln_lot"] = np.log(df["lot"].where(df["lot"] > 0, 1.0))

    year_built = pd.to_numeric(df["role_annee_construction"], errors="coerce")
    df["age"] = (df["sale_year"] - year_built).clip(0, config.AGE_MAX)
    df = df[df["age"].notna()]
    log.append(("known building age", len(df)))

    df["floors"] = df["role_nb_etages"].clip(1, 6).fillna(1.0)
    df["units"] = df["role_nb_logements"].clip(1, 12).fillna(1.0)
    df["prop_class"] = _property_class(df)
    df["link"] = df["role_lien_physique"].map(LINK_MAP).fillna("unknown")
    df["muni"] = df["role_code_mun"]
    df["grid"] = ((df["lat"] / config.GRID_DEG).round().astype(int).astype(str)
                  + "_" + (df["lng"] / config.GRID_DEG).round().astype(int).astype(str))
    df["grid5"] = ((df["lat"] / config.GRID5_DEG).round().astype(int).astype(str)
                   + "_" + (df["lng"] / config.GRID5_DEG).round().astype(int).astype(str))

    lo, hi = config.TRIM
    for col in ("price", "area"):
        q = df.groupby("sale_year")[col].quantile([lo, hi]).unstack()
        df = df.join(q.rename(columns={lo: "_qlo", hi: "_qhi"}), on="sale_year")
        df = df[(df[col] >= df["_qlo"]) & (df[col] <= df["_qhi"])]
        df = df.drop(columns=["_qlo", "_qhi"])
    log.append((f"price & area inside [{lo:.0%}, {hi:.0%}] of sale year",
                len(df)))

    df.attrs["selection_log"] = log
    keep = ["id", "sale_date", "sale_year", "month", "quarter", "t",
            "price", "ln_price", "area", "ln_area", "lot", "ln_lot", "has_lot",
            "age", "floors", "units", "prop_class", "link",
            "muni", "grid", "grid5", "lat", "lng", "city",
            "role_municipalite"]
    return df[keep].reset_index(drop=True)


def build_and_save() -> pd.DataFrame:
    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:
    if not config.ANALYSIS_PARQUET.exists():
        return build_and_save()
    return pd.read_parquet(config.ANALYSIS_PARQUET)