"""Cleaning, normalisation, linkage keys and sample construction (Phase 3).

Produces in data/processed:
  listings.parquet        one row per listing (anonymised ids, normalised fields, sample flags)
  price_events.parquet    sustained price events (entry / cut / rise / minor)
  exclusions.csv          auditable exclusion log (uid, rule, value)
  sample_construction.csv attrition table for the paper
  sources_registry.csv    public registry of sources with class and quality flags
  id_maps/ (interim)      uid -> listing_id, dealer key -> dealer_id (restricted)
"""
from __future__ import annotations

import csv
import hashlib
import json
import math
import re
from pathlib import Path

import numpy as np
import polars as pl

from . import config as C
from . import pricepaths

# ----------------------------------------------------------------------------- helpers

def _norm(s: str | None) -> str:
    return (s or "").strip().lower()


def powertrain(fuel: str | None, engine: str | None, title: str | None, make: str | None) -> str:
    f = _norm(fuel)
    if "rechargeable" in f or "plug" in f:
        return "PHEV"
    if "hybride" in f or "hybrid" in f:
        return "HEV"
    if "électrique" in f or "electrique" in f or f in {"ev", "electric"}:
        return "BEV"
    if f.startswith("diesel"):
        return "Diesel"
    if f.startswith("essence") or "flex" in f or "polycarb" in f or "éthanol" in f or "propane" in f:
        return "ICE"
    if _norm(make) in {"tesla", "polestar"}:
        return "BEV"
    e = _norm(engine) + " " + _norm(title)
    if re.search(r"\b(électrique|electric|\bev\b|kwh)", e):
        return "BEV"
    return "unknown"


_BODY = [
    (r"vus|sport utility|suv|crossover|utilitaire", "SUV"),
    (r"camionnette|pickup|pick-up|truck|camion", "Pickup"),
    (r"fourgon|minivan|van", "Van"),
    (r"berline|sedan|4dr car|4 portes|intermédiaire|compact", "Sedan"),
    (r"hayon|hatch|5 portes", "Hatchback"),
    (r"coupé|coupe|2 portes|2dr|voiture 2 portes", "Coupe"),
    (r"cabriolet|convertible|décapotable", "Convertible"),
    (r"familiale|wagon", "Wagon"),
]


def body_segment(body: str | None) -> str | None:
    b = _norm(body)
    if not b or b in {"n/d", "autre", "other", "custom", "n.d."}:
        return None
    for pat, lab in _BODY:
        if re.search(pat, b):
            return lab
    return None


def transmission_norm(t: str | None) -> str:
    s = _norm(t)
    if not s or s in {"n/a", "0", "null", "na", "unknown", "other", "transmission"}:
        return "unknown"
    if "manuel" in s or "manual" in s:
        return "Manual"
    return "Automatic"


def drivetrain_norm(d: str | None) -> str:
    s = _norm(d)
    if not s or s in {"n/a", "null", "other", "not defined"}:
        return "unknown"
    if "intégrale" in s or "4x4" in s or "four" in s or "4 roues" in s or "quatre" in s or "awd" in s:
        return "AWD/4x4"
    if "traction" in s or "fwd" in s or "4x2" in s:
        return "FWD"
    if "propulsion" in s or "rwd" in s:
        return "RWD"
    return "unknown"


def source_class(src: str) -> str:
    if src in C.SAMPLED_PORTALS:
        return "sampled_portal"
    if src in C.PRIVATE_MARKETPLACE:
        return "private_marketplace"
    if src in C.FULL_PORTALS:
        return "full_portal"
    if src in C.BANNERS:
        return "banner"
    if src in C.SMALL_PORTALS:
        return "small_portal"
    return "direct_dealer"


def _hid(prefix: str, key: str, n: int = 10) -> str:
    return prefix + hashlib.sha1(key.encode()).hexdigest()[:n]


def vehicle_age_years(ts: float, year: int) -> float:
    """Age at observation: years since 1 October of (model year - 1), floored at 0."""
    import datetime as dt
    ref = dt.datetime(year - 1, 10, 1, tzinfo=dt.timezone.utc).timestamp()
    return max(0.0, (ts - ref) / (365.25 * 86400))


# ----------------------------------------------------------------------------- region recovery

def build_city_region_map(v: pl.DataFrame) -> dict[str, str]:
    d = (v.filter((pl.col("region") != "") & pl.col("region").is_not_null() & (pl.col("city") != "") & pl.col("city").is_not_null())
         .with_columns(pl.col("city").str.to_lowercase().str.strip_chars())
         .group_by(["city", "region"]).len().sort(["city", "len"], descending=[False, True])
         .unique(subset=["city"], keep="first"))
    return dict(zip(d["city"].to_list(), d["region"].to_list()))


def nearest_city_region(lat: float, lng: float, cities: list[tuple[str, float, float]], cmap: dict[str, str]) -> str | None:
    best, bd = None, 1e9
    for name, clat, clng in cities:
        d = (lat - clat) ** 2 + ((lng - clng) * math.cos(math.radians(lat))) ** 2
        if d < bd:
            bd, best = d, name
    if best is None or bd > (0.5 ** 2):  # > ~50 km from any known city -> leave missing
        return None
    return cmap.get(best)


REGION_FIX = {"Saguenay-Lac-Saint-Jean": "Saguenay–Lac-Saint-Jean"}


# ----------------------------------------------------------------------------- main

def run(interim: Path = C.INTERIM, out: Path = C.PROCESSED, registry: Path = C.RAW / "registry") -> dict:
    out.mkdir(parents=True, exist_ok=True)
    (interim / "id_maps").mkdir(exist_ok=True)
    v = pl.read_parquet(interim / "vehicles.parquet")
    pl_log = pl.read_parquet(interim / "price_log.parquet")
    sync = pl.read_parquet(interim / "sync_log.parquet")
    snap_ts = C.SNAPSHOT_TS
    excl: list[tuple[str, str, str]] = []

    # --- sources registry ---------------------------------------------------
    reg = json.loads((registry / "sources.json").read_text())["sources"]
    reg_rows = []
    for s in reg:
        reg_rows.append({"source": s["id"], "name": s.get("name"), "type": s.get("type", "direct"),
                         "connector": s.get("connector"), "platform": s.get("platform"), "region": s.get("region"),
                         "city": s.get("city"), "status": s.get("status"), "class": source_class(s["id"]),
                         "in_dynamic_sample": source_class(s["id"]) in {"direct_dealer", "full_portal", "banner"} and s["id"] not in C.FLICKER_OR_FAILED})
    # sync quality per source
    sq = (sync.group_by("source").agg([pl.len().alias("n_syncs"), pl.col("ok").sum().alias("n_ok"),
                                       pl.col("found").mean().alias("mean_found"), pl.col("added").sum().alias("added"),
                                       pl.col("removed").sum().alias("removed")]))
    inact = v.group_by("source").agg([(pl.col("active") == 0).sum().alias("n_inactive"), pl.len().alias("n_listings")])
    sq = sq.join(inact, on="source", how="left").with_columns(
        ((pl.col("removed") - pl.col("n_inactive")) / pl.col("removed").clip(lower_bound=1)).alias("reactivation_share"))
    regdf = pl.DataFrame(reg_rows).join(sq, on="source", how="left")
    regdf.write_csv(out / "sources_registry.csv")

    # first non-empty crawl per source: entries within 24 h of it are the source's start-up stock (left-censored, D-014)
    ff = sync.filter(pl.col("found") > 0).group_by("source").agg(pl.col("ts").min().alias("first_found"))
    first_found = dict(zip(ff["source"].to_list(), ff["first_found"].to_list()))

    # crawl interval per source (median seconds) for interval-censoring midpoints
    iv = (sync.sort(["source", "ts"]).with_columns((pl.col("ts") - pl.col("ts").shift(1).over("source")).alias("dt"))
          .filter(pl.col("dt").is_not_null()).group_by("source").agg(pl.col("dt").median().alias("crawl_median_s")))
    med_iv = dict(zip(iv["source"].to_list(), iv["crawl_median_s"].to_list()))

    # --- price paths ---------------------------------------------------------
    events, psum = pricepaths.build(pl_log, v)

    # --- region recovery -----------------------------------------------------
    cmap = build_city_region_map(v)
    cmap = {k: REGION_FIX.get(r, r) for k, r in cmap.items()}
    gps = json.loads((registry / "villes_gps.json").read_text())
    cities = [(name, c[0], c[1]) for name, c in gps.items() if name in cmap]

    # --- per-listing table ---------------------------------------------------
    v = v.join(psum, on="uid", how="left")
    rows = []
    dealer_keys: dict[str, str] = {}
    vin_re = re.compile(r"^[A-HJ-NPR-Z0-9]{17}$")
    for r in v.iter_rows(named=True):
        uid, src = r["uid"], r["source"]
        cls = source_class(src)
        region = REGION_FIX.get(r["region"] or "", r["region"] or "") or None
        city = (r["city"] or "").strip() or None
        if region is None and city:
            region = cmap.get(city.lower())
        region_recovered = False
        if region is None and r["lat"] is not None and r["lng"] is not None:
            region = nearest_city_region(r["lat"], r["lng"], cities, cmap)
            region_recovered = region is not None
        dealer_key = src if cls in {"direct_dealer", "banner", "small_portal"} else f"{src}|{(r['dealer_name'] or '').strip().lower()}"
        if cls == "private_marketplace":
            dealer_key = "kijiji|private"
        dealer_id = dealer_keys.setdefault(dealer_key, _hid("d", dealer_key, 8))
        vin = (r["vin"] or "").strip().upper()
        vin_ok = bool(vin_re.match(vin))
        year = r["year"]
        entry_ts = r["first_seen"]
        entry_price = r["entry_price"] if r["entry_price"] is not None else r["price"]
        # exit variables (D-004)
        active = int(r["active"] or 0)
        last_seen = r["last_seen"]
        exit_confirmed = (active == 0) and (last_seen <= snap_ts - C.EXIT_CONFIRM_DAYS * 86400)
        half_iv = 0.5 * med_iv.get(src, 3.3 * 3600)
        if exit_confirmed:
            end_ts = min(last_seen + half_iv, snap_ts)
        else:
            end_ts = last_seen if active == 0 else snap_ts
        rows.append({
            "listing_id": _hid("l", uid), "uid": uid, "title": r["title"], "source": src, "source_class": cls,
            "dealer_id": dealer_id, "kind": r["kind"], "make": (r["make"] or "").strip() or None,
            "model": (r["model"] or "").strip() or None, "trim": (r["trim"] or "").strip() or None, "year": year,
            "price_raw": r["price"], "entry_price": entry_price, "last_price": r["last_price"] if r["last_price"] is not None else r["price"],
            "mileage_km": r["mileage_km"], "transmission": transmission_norm(r["transmission"]),
            "drivetrain": drivetrain_norm(r["drivetrain"]), "powertrain": powertrain(r["fuel"], r["engine"], r["title"], r["make"]),
            "body": body_segment(r["body_type"]), "exterior_color": (r["exterior_color"] or "").strip().lower() or None,
            "region": region, "region_recovered": region_recovered, "city": city,
            "lat_r": round(r["lat"], 2) if r["lat"] is not None else None, "lng_r": round(r["lng"], 2) if r["lng"] is not None else None,
            "lat": r["lat"], "lng": r["lng"],  # exact coords stay in interim-only columns (dropped from release)
            "vin_ok": vin_ok, "vin_hash": _hid("v", vin, 12) if vin_ok else None,
            "has_stock_number": bool((r["stock_number"] or "").strip()), "is_canonical": r["dup_of"] is None,
            "n_dup_sources": len(json.loads(r["dup_sources"])) if r["dup_sources"] else 0,
            "entry_ts": entry_ts, "last_seen_ts": last_seen, "end_ts": end_ts, "active": active,
            "exit_confirmed": exit_confirmed,
            "left_censored": (entry_ts < C.LEFT_CENSOR_CUTOFF_TS) or (entry_ts < first_found.get(src, C.WINDOW_START_TS) + C.SOURCE_STARTUP_SECONDS),
            "duration_days": (end_ts - entry_ts) / 86400,
            "n_cuts": r["n_cuts"] or 0, "n_rises": r["n_rises"] or 0, "n_transients": r["n_transients"] or 0,
            "oscillating": bool(r["oscillating"]) if r["oscillating"] is not None else False,
            "first_cut_ts": r["first_cut_ts"], "first_cut_pct": r["first_cut_pct"], "total_change_pct": r["total_change_pct"],
            "age_years_entry": vehicle_age_years(entry_ts, int(year)) if year and 1980 <= year <= 2027 else None,
        })
    L = pl.DataFrame(rows)

    # --- coordinates outside the Quebec bounding box are errors -> null (flagged) ---
    inbox = pl.col("lat").is_between(44.0, 63.0) & pl.col("lng").is_between(-80.0, -56.0)
    L = L.with_columns([
        (pl.col("lat").is_not_null() & ~inbox).alias("coords_out_of_bbox"),
        pl.when(inbox).then(pl.col("lat")).otherwise(None).alias("lat"),
        pl.when(inbox).then(pl.col("lng")).otherwise(None).alias("lng"),
    ])

    # --- coordinates: fill missing from the dealer's median location (flagged) ---
    dm = (L.filter(pl.col("lat").is_not_null()).group_by("dealer_id")
          .agg(pl.col("lat").median().alias("_dlat"), pl.col("lng").median().alias("_dlng")))
    L = L.join(dm, on="dealer_id", how="left").with_columns([
        (pl.col("lat").is_null() & pl.col("_dlat").is_not_null()).alias("coords_from_dealer"),
        pl.when(pl.col("lat").is_null()).then(pl.col("_dlat")).otherwise(pl.col("lat")).alias("lat"),
        pl.when(pl.col("lng").is_null()).then(pl.col("_dlng")).otherwise(pl.col("lng")).alias("lng"),
    ]).drop(["_dlat", "_dlng"])

    # --- coordinates: fill missing from city centroid (flagged) --------------
    gps_l = {k: v for k, v in gps.items()}
    L = L.with_columns([
        pl.col("city").str.to_lowercase().str.strip_chars().alias("_city_l")])
    fill_lat = pl.col("_city_l").map_elements(lambda c: gps_l.get(c, [None, None])[0] if c else None, return_dtype=pl.Float64)
    fill_lng = pl.col("_city_l").map_elements(lambda c: gps_l.get(c, [None, None])[1] if c else None, return_dtype=pl.Float64)
    L = L.with_columns([
        (pl.col("lat").is_null() & fill_lat.is_not_null()).alias("coords_from_city"),
        pl.when(pl.col("lat").is_null()).then(fill_lat).otherwise(pl.col("lat")).alias("lat"),
        pl.when(pl.col("lng").is_null()).then(fill_lng).otherwise(pl.col("lng")).alias("lng"),
    ]).drop("_city_l")
    L = L.with_columns([pl.col("lat").round(2).alias("lat_r"), pl.col("lng").round(2).alias("lng_r")])

    # --- powertrain recovery: title keywords, then make-model(-trim) modal (>= 90 % agreement) ---
    kw = pl.col("title").fill_null("").str.to_lowercase()
    L = L.with_columns(
        pl.when(pl.col("powertrain") != "unknown").then(pl.col("powertrain"))
        .when(kw.str.contains("rechargeable|plug-in|phev")).then(pl.lit("PHEV"))
        .when(kw.str.contains("hybrid")).then(pl.lit("HEV"))
        .when(kw.str.contains(r"électrique|electric|\bev\b")).then(pl.lit("BEV"))
        .when(kw.str.contains("diesel|tdi")).then(pl.lit("Diesel"))
        .otherwise(pl.lit("unknown")).alias("powertrain"))
    known = L.filter(pl.col("powertrain") != "unknown")
    def modal(keys):
        g = known.group_by(keys + ["powertrain"]).len()
        tot = g.group_by(keys).agg(pl.col("len").sum().alias("tot"))
        g = g.join(tot, on=keys).with_columns((pl.col("len") / pl.col("tot")).alias("sh")).filter(pl.col("sh") >= 0.9)
        return g.select(keys + [pl.col("powertrain").alias("pt_fill")])
    L = L.with_columns(pl.col("powertrain").alias("_pt0"))
    for keys in (["make", "model", "trim", "year"], ["make", "model", "year"], ["make", "model"]):
        m = modal(keys)
        L = L.join(m, on=keys, how="left").with_columns(
            pl.when((pl.col("powertrain") == "unknown") & pl.col("pt_fill").is_not_null()).then(pl.col("pt_fill")).otherwise(pl.col("powertrain")).alias("powertrain")).drop("pt_fill")
    L = L.with_columns(((pl.col("_pt0") == "unknown") & (pl.col("powertrain") != "unknown")).alias("powertrain_recovered")).drop("_pt0")

    # --- body recovery from make-model modal body ----------------------------
    mm = (L.filter(pl.col("body").is_not_null()).group_by(["make", "model", "body"]).len()
          .sort(["make", "model", "len"], descending=[False, False, True]).unique(subset=["make", "model"], keep="first")
          .rename({"body": "body_mm"}).select(["make", "model", "body_mm"]))
    L = L.join(mm, on=["make", "model"], how="left").with_columns(
        pl.when(pl.col("body").is_null()).then(pl.col("body_mm")).otherwise(pl.col("body")).alias("body_filled"),
        pl.col("body").is_null().alias("body_recovered")).drop("body_mm")

    # --- sample flags + exclusion log ---------------------------------------
    def flag(cond: pl.Expr, name: str):
        return pl.when(cond).then(pl.lit(name)).otherwise(pl.lit(None))

    L = L.with_columns([
        (pl.col("kind") == "auto").alias("f_auto"),
        pl.col("is_canonical").alias("f_canonical"),
        pl.col("year").is_between(C.YEAR_MIN, C.YEAR_MAX).fill_null(False).alias("f_year"),
        pl.col("mileage_km").is_between(C.KM_MIN, C.KM_MAX).fill_null(False).alias("f_km"),
        pl.col("entry_price").is_between(C.PRICE_MIN, C.PRICE_MAX).fill_null(False).alias("f_price"),
        (pl.col("make").is_not_null() & pl.col("model").is_not_null()).alias("f_mm"),
        (~pl.col("uid").is_in(list(C.UID_COLLISIONS))).alias("f_nocollision"),
        (~pl.col("source").is_in(list(C.SAMPLED_PORTALS | C.PRIVATE_MARKETPLACE))).alias("f_not_sampled_or_private"),
        (~pl.col("source").is_in(list(C.FLICKER_OR_FAILED | C.SMALL_PORTALS))).alias("f_source_reliable"),
    ])
    L = L.with_columns([
        (pl.col("f_auto") & pl.col("f_canonical") & pl.col("f_year") & pl.col("f_km") & pl.col("f_price") & pl.col("f_mm") & pl.col("f_nocollision")).alias("s1_all_sources"),
    ])
    L = L.with_columns([
        (pl.col("s1_all_sources") & pl.col("f_not_sampled_or_private")).alias("s1_hedonic"),
    ])
    L = L.with_columns([
        (pl.col("s1_hedonic") & pl.col("f_source_reliable") & ~pl.col("oscillating") & ~pl.col("left_censored")).alias("s2_dynamic"),
        (pl.col("s1_hedonic") & pl.col("f_source_reliable") & ~pl.col("oscillating")).alias("s2l_dynamic_with_stock"),
    ])

    # exclusion log (sequential attrition, S0 -> S1 -> S2)
    steps = [
        ("kind != auto", ~pl.col("f_auto")), ("VIN duplicate (non-canonical)", ~pl.col("f_canonical")),
        (f"model year outside [{C.YEAR_MIN},{C.YEAR_MAX}] or missing", ~pl.col("f_year")),
        (f"mileage outside [{C.KM_MIN:.0f},{C.KM_MAX:.0f}] km or missing", ~pl.col("f_km")),
        (f"entry price outside [{C.PRICE_MIN:.0f},{C.PRICE_MAX:.0f}] $ or missing/placeholder", ~pl.col("f_price")),
        ("make or model missing", ~pl.col("f_mm")), ("uid identity collision", ~pl.col("f_nocollision")),
        ("sampled portal (autotrader/cargurus) or private marketplace (kijiji)", ~pl.col("f_not_sampled_or_private")),
        ("flickering/failed connector or small portal", ~pl.col("f_source_reliable")),
        ("oscillating price path", pl.col("oscillating")), ("left-censored stock (day 1 or source start-up day)", pl.col("left_censored")),
    ]
    remaining = L
    attrition = [("S0 all listings in snapshot", len(L))]
    for name, cond in steps:
        dropped = remaining.filter(cond)
        for uid in dropped["uid"].to_list():
            excl.append((uid, name, ""))
        remaining = remaining.filter(~cond)
        attrition.append((f"− {name}", len(remaining)))
    with open(out / "exclusions.csv", "w", newline="") as f:
        w = csv.writer(f); w.writerow(["uid", "rule", "value"]); w.writerows(excl)
    with open(out / "sample_construction.csv", "w", newline="") as f:
        w = csv.writer(f); w.writerow(["step", "listings_remaining"]); w.writerows(attrition)
    stats = {"S0": len(L), "S1_all_sources": int(L["s1_all_sources"].sum()), "S1_hedonic": int(L["s1_hedonic"].sum()),
             "S2_dynamic": int(L["s2_dynamic"].sum()), "S2L": int(L["s2l_dynamic_with_stock"].sum()),
             "exit_confirmed_in_S2": int(L.filter(pl.col("s2_dynamic"))["exit_confirmed"].sum()),
             "region_recovered": int(L["region_recovered"].sum()), "region_missing": int(L["region"].is_null().sum()),
             "body_recovered": int(L["body_recovered"].sum()), "body_missing_after": int(L["body_filled"].is_null().sum()),
             "oscillating": int(L["oscillating"].sum()), "n_dealer_ids": L["dealer_id"].n_unique(),
             "powertrain_recovered": int(L["powertrain_recovered"].sum()), "powertrain_unknown_after": int((L["powertrain"] == "unknown").sum()),
             "coords_from_city": int(L["coords_from_city"].sum()), "coords_from_dealer": int(L["coords_from_dealer"].sum()), "coords_out_of_bbox": int(L["coords_out_of_bbox"].sum()), "coords_missing_after": int(L["lat"].is_null().sum())}
    # restricted id maps stay in interim
    pl.DataFrame({"dealer_key": list(dealer_keys.keys()), "dealer_id": list(dealer_keys.values())}).write_csv(interim / "id_maps" / "dealers.csv")
    L.select(["uid", "listing_id"]).write_csv(interim / "id_maps" / "listings.csv")
    L.write_parquet(interim / "listings_full.parquet")           # includes exact lat/lng, uid
    L.drop(["lat", "lng", "uid", "city", "title"]).write_parquet(out / "listings.parquet")  # release-grade
    events = events.join(L.select(["uid", "listing_id"]), on="uid", how="left").drop("uid")
    events.write_parquet(out / "price_events.parquet")
    (out / "cleaning_stats.json").write_text(json.dumps(stats, indent=2))
    return stats


if __name__ == "__main__":
    print(json.dumps(run(), indent=2))