# Author: Simon-Pierre Boucher — contact@spboucher.ai
"""Construction of the estimation sample from the raw DuckDB snapshot.

Steps (documented in the paper's data section):

1. read the raw ``listings`` table (172,019 unique for-sale listings);
2. parse the semi-structured fields (bedrooms, areas, lot, counts);
3. consolidate dwelling type and ownership form; derive the FSA from the
   postal code;
4. keep residential listings with a strictly positive price and non-missing
   core structural fields (living area, bedrooms, bathrooms) and a valid FSA;
5. trim the extreme 1% tails of price and living area;
6. pool FSAs with fewer than 25 listings into a province-level residual
   category.
"""
import duckdb
import numpy as np
import pandas as pd

from . import parsing
from .config import (ANALYSIS_PARQUET, FSA_MIN_LISTINGS, LOT_TAIL_Q, RAW_DB,
                     TRIM_HI, TRIM_LO, ensure_dirs)

_RAW_QUERY = """
select "_category"                        as cat,
       "_province"                        as prov,
       "building.type"                    as btype_raw,
       "ownership.type"                   as own_raw,
       "building.bedrooms"                as bed_raw,
       "building.bathroom_total"          as bath_raw,
       "building.half_bath_total"         as half_raw,
       "building.size_interior"           as size_interior,
       "building.floor_area_measurements" as floor_area,
       "building.stories_total"           as stories_raw,
       "parking.spaces_total"             as parking_raw,
       "land.size_total"                  as lot_raw,
       "location.postal_code"             as postal_code,
       price_cad, lat, lon
from listings
"""

_FSA_PATTERN = r"^[A-Z]\d[A-Z]$"


def load_raw() -> pd.DataFrame:
    """Read the raw listings table from the DuckDB snapshot."""
    with duckdb.connect(str(RAW_DB), read_only=True) as con:
        return con.execute(_RAW_QUERY).df()


def parse_variables(raw: pd.DataFrame) -> pd.DataFrame:
    """Parse raw text fields into numeric/categorical analysis variables."""
    d = raw.copy()
    d["bedrooms"] = d["bed_raw"].map(parsing.parse_bedrooms)
    d["bathrooms"] = d["bath_raw"].map(parsing.parse_count)
    d["half_baths"] = d["half_raw"].map(parsing.parse_count).fillna(0.0)
    d["parking_n"] = d["parking_raw"].map(parsing.parse_count).fillna(0.0)
    d["stories_n"] = d["stories_raw"].map(parsing.parse_count).fillna(0.0)
    d["living_m2"] = (d["floor_area"].map(parsing.parse_floor_area)
                      .fillna(d["size_interior"].map(parsing.parse_size_interior)))

    lot = d["lot_raw"].map(parsing.parse_lot)
    # Very large parsed lots (top 1% of positive values, mostly multi-acre
    # rural acreage strings) are treated as "no usable lot information".
    cap = lot[lot > 0].quantile(LOT_TAIL_Q)
    lot = lot.where(lot <= cap, np.nan)
    d["lot_m2_f"] = lot.fillna(0.0)
    d["has_lot"] = (d["lot_m2_f"] > 0).astype(float)

    d["btype_c"] = d["btype_raw"].map(parsing.consolidate_building_type)
    d["own_c"] = d["own_raw"].map(parsing.consolidate_ownership)

    fsa = d["postal_code"].astype("string").str.upper().str[:3]
    d["fsa"] = fsa.where(fsa.str.match(_FSA_PATTERN, na=False))
    return d


def build_sample(d: pd.DataFrame) -> pd.DataFrame:
    """Apply the sample restrictions and derive the model variables."""
    core = ((d["price_cad"] > 0) & d["living_m2"].notna()
            & d["bedrooms"].notna() & d["bathrooms"].notna() & d["fsa"].notna())
    s = d[core].copy()

    lo_p, hi_p = s["price_cad"].quantile([TRIM_LO, TRIM_HI])
    lo_a, hi_a = s["living_m2"].quantile([TRIM_LO, TRIM_HI])
    s = s[(s["price_cad"].between(lo_p, hi_p)) & (s["living_m2"].between(lo_a, hi_a))].copy()

    s["ln_price"] = np.log(s["price_cad"])
    s["ln_living"] = np.log(s["living_m2"])
    s["ln_lot"] = np.log1p(s["lot_m2_f"])
    s["ppm2"] = s["price_cad"] / s["living_m2"]

    counts = s.groupby("fsa")["fsa"].transform("size")
    s["fsa_c"] = np.where(counts >= FSA_MIN_LISTINGS, s["fsa"], s["prov"] + "_other")

    keep = ["cat", "prov", "btype_c", "own_c", "fsa", "fsa_c",
            "price_cad", "ln_price", "ppm2",
            "living_m2", "ln_living", "bedrooms", "bathrooms", "half_baths",
            "parking_n", "stories_n", "lot_m2_f", "has_lot", "ln_lot",
            "lat", "lon"]
    return s[keep].reset_index(drop=True)


def build_and_save() -> pd.DataFrame:
    """Full raw-to-parquet pipeline; returns the estimation sample."""
    ensure_dirs()
    sample = build_sample(parse_variables(load_raw()))
    sample.to_parquet(ANALYSIS_PARQUET, index=False)
    return sample


def load_sample() -> pd.DataFrame:
    """Load the processed estimation sample (build it first with script 01)."""
    return pd.read_parquet(ANALYSIS_PARQUET)