# Author: Simon-Pierre Boucher — simon-pierre.boucher@uqo.ca
"""Repeat sales: a transparent monthly deflator and the resale-implied value of a sale.

Definitions (paper §sec:fundamental, notation of theory.tex §subsec:values)
---------------------------------------------------------------------------
* Monthly index. For municipality m, property class c and calendar month τ, let q_{mcτ} be the median of
  log(price / floor area) over the sales of that cell. Cells with fewer than `min_n` sales (default 5) are replaced by
  the region-group × class cell. The series is smoothed with a centred 3-month rolling median and *chained*: the
  log index is the cumulative sum of month-to-month changes of the smoothed series, each change taken from the
  municipality series when both months are available there and from the region series otherwise. Chaining means that
  a switch between the municipality and the region source never introduces a level jump. The index is an ex-post
  deflator (it uses all sales of the snapshot) and is never shown to any arm.
* Resale-implied value. For a sale (i, t1) at price P1 whose parcel resells at (t2 > t1) for P2,
      log V^R_{t1} = log P2 − [ I(m, c, τ2) − I(m, c, τ1) ],
  i.e. the later price brought back to the first sale date by the index. Under the repeat-sales assumptions of
  Bailey–Muth–Nourse (1963) and Case–Shiller (1989) — unchanged quality, market movement common to the cell — the
  difference log V^R − log P1 is the difference of the two transaction noises, ε_{t2} − ε_{t1}, so an estimate at t1
  that is closer to V^R than P1 was captured a component of value that the first transaction's noise obscured.
* Annualised holding return: (log P2 − log P1) × 12 / months_held.

Selection (documented, never corrected here): parcels that resell within the snapshot window (2021–2026) are not a
random sample — flips, forced sales, renovations and loss-averse listing behaviour (Genesove & Mayer 2001) correlate
the resale decision with ε_{t1}. Holding periods shorter than 12 months are flagged (`short_hold`); renovations are
unobserved because the roll attributes are a single snapshot for both sales.
"""
from __future__ import annotations

from typing import Any

import numpy as np
import pandas as pd

from ..evaluation._common import to_pandas

MIN_CELL_N = 5
SMOOTH_WINDOW = 3
SHORT_HOLD_MONTHS = 12
DAYS_PER_MONTH = 365.25 / 12
# rows of the raw snapshot that cannot serve as either leg of a repeat-sales pair (data-quality exclusions of WP 22)
STRUCTURAL_EXCLUSIONS = {"exact_duplicate", "same_parcel_same_date_diff_price", "shared_parcel", "multi_parcel_candidate",
                         "land", "non_residential", "built_after_sale"}
INDEX_KEY = ["mun_code", "property_class"]


# ------------------------------------------------------------------------------------------------ monthly index

def _month_id(date: pd.Series) -> pd.Series:
    d = pd.to_datetime(date)
    return (d.dt.year * 12 + d.dt.month - 1).astype(int)


def month_id_to_timestamp(m: np.ndarray | pd.Series) -> pd.Series:
    m = pd.Series(np.asarray(m, dtype=int))
    return pd.to_datetime({"year": m // 12, "month": m % 12 + 1, "day": 1})


def _centred_rolling_median(wide: np.ndarray, window: int) -> np.ndarray:
    """Centred rolling median along axis 1 ignoring NaN (min 1 valid value); NaN where no neighbour is valid."""
    half = window // 2
    n_rows, n_cols = wide.shape
    stack = np.full((window, n_rows, n_cols), np.nan)
    for k, shift in enumerate(range(-half, half + 1)):
        if shift < 0:
            stack[k, :, -shift:] = wide[:, :shift]
        elif shift > 0:
            stack[k, :, :-shift] = wide[:, shift:]
        else:
            stack[k] = wide
    with np.errstate(all="ignore"):
        return np.nanmedian(stack, axis=0)


def monthly_index(analytic: Any, *, min_n: int = MIN_CELL_N, window: int = SMOOTH_WINDOW) -> pd.DataFrame:
    """Municipality × class monthly chained log-price index (long format).

    Columns: mun_code, property_class, region_group, month_id, month (Timestamp), n_mun, n_region, raw_mun, raw_region,
    smooth_mun, smooth_region, delta, source ("mun" | "region" | "carry"), log_index (0 at the first month).
    Every (municipality, class) that has at least one sale gets a complete monthly series from the first to the last
    month of the table.
    """
    df = to_pandas(analytic)
    need = INDEX_KEY + ["region_group", "date", "price", "floor_area_m2"]
    missing = [c for c in need if c not in df.columns]
    if missing:
        raise ValueError(f"analytic table is missing {missing}")
    df = df[need].dropna(subset=["floor_area_m2", "price"])
    df = df[(df["floor_area_m2"] > 0) & (df["price"] > 0)].copy()
    if df.empty:
        raise ValueError("no sale with a positive floor area and price")
    df["lp"] = np.log(df["price"].astype(float) / df["floor_area_m2"].astype(float))
    df["month_id"] = _month_id(df["date"])
    months = np.arange(int(df["month_id"].min()), int(df["month_id"].max()) + 1)
    n_months = len(months)

    # municipality → region group (modal, in case the snapshot has a few inconsistent rows)
    mun_region = df.groupby("mun_code")["region_group"].agg(lambda s: s.value_counts().index[0])

    def _wide(frame: pd.DataFrame, keys: list[str]) -> tuple[pd.DataFrame, np.ndarray, np.ndarray]:
        g = frame.groupby(keys + ["month_id"])["lp"].agg(["median", "size"])
        med = g["median"].unstack("month_id").reindex(columns=months)
        n = g["size"].unstack("month_id").reindex(columns=months).fillna(0)
        return med, med.to_numpy(dtype=float), n.to_numpy(dtype=float)

    mun_med, mun_raw, mun_n = _wide(df, INDEX_KEY)
    reg_med, reg_raw, reg_n = _wide(df, ["region_group", "property_class"])
    mun_raw = np.where(mun_n >= min_n, mun_raw, np.nan)
    reg_raw = np.where(reg_n >= min_n, reg_raw, np.nan)
    # the smoother uses the valid neighbours of a month but never *creates* a value for a month below `min_n`
    mun_s = np.where(np.isfinite(mun_raw), _centred_rolling_median(mun_raw, window), np.nan)
    reg_s = np.where(np.isfinite(reg_raw), _centred_rolling_median(reg_raw, window), np.nan)

    # align each municipality row with its region row
    reg_pos = {k: i for i, k in enumerate(reg_med.index)}
    mun_keys = list(mun_med.index)  # (mun_code, property_class)
    rows_reg = np.array([reg_pos.get((mun_region[m], c), -1) for m, c in mun_keys])
    reg_s_al = np.where(rows_reg[:, None] >= 0, reg_s[np.clip(rows_reg, 0, None)], np.nan)
    reg_raw_al = np.where(rows_reg[:, None] >= 0, reg_raw[np.clip(rows_reg, 0, None)], np.nan)
    reg_n_al = np.where(rows_reg[:, None] >= 0, reg_n[np.clip(rows_reg, 0, None)], 0.0)

    d_mun = np.full_like(mun_s, np.nan)
    d_mun[:, 1:] = mun_s[:, 1:] - mun_s[:, :-1]
    d_reg = np.full_like(reg_s_al, np.nan)
    d_reg[:, 1:] = reg_s_al[:, 1:] - reg_s_al[:, :-1]
    delta = np.where(np.isfinite(d_mun), d_mun, np.where(np.isfinite(d_reg), d_reg, 0.0))
    delta[:, 0] = 0.0
    source = np.where(np.isfinite(d_mun), "mun", np.where(np.isfinite(d_reg), "region", "carry"))
    source[:, 0] = "base"
    log_index = np.cumsum(delta, axis=1)

    n_rows = len(mun_keys)
    out = pd.DataFrame({
        "level": "mun",
        "mun_code": np.repeat([str(k[0]) for k in mun_keys], n_months),
        "property_class": np.repeat([k[1] for k in mun_keys], n_months),
        "region_group": np.repeat([mun_region[k[0]] for k in mun_keys], n_months),
        "month_id": np.tile(months, n_rows),
        "n_mun": mun_n.ravel(), "n_region": reg_n_al.ravel(),
        "raw_mun": mun_raw.ravel(), "raw_region": reg_raw_al.ravel(),
        "smooth_mun": mun_s.ravel(), "smooth_region": reg_s_al.ravel(),
        "delta": delta.ravel(), "source": source.ravel(), "log_index": log_index.ravel(),
    })
    # region × class series (level="region", keyed by the region group in `mun_code`): the fallback when a
    # municipality × class has no series at all (e.g. every sale of the cell lacks a floor area)
    d_reg_only = np.full_like(reg_s, np.nan)
    d_reg_only[:, 1:] = reg_s[:, 1:] - reg_s[:, :-1]
    d_reg_only = np.where(np.isfinite(d_reg_only), d_reg_only, 0.0)
    d_reg_only[:, 0] = 0.0
    reg_keys = list(reg_med.index)
    reg_out = pd.DataFrame({
        "level": "region",
        "mun_code": np.repeat([str(k[0]) for k in reg_keys], n_months),
        "property_class": np.repeat([k[1] for k in reg_keys], n_months),
        "region_group": np.repeat([k[0] for k in reg_keys], n_months),
        "month_id": np.tile(months, len(reg_keys)),
        "n_mun": np.nan, "n_region": reg_n.ravel(), "raw_mun": np.nan, "raw_region": reg_raw.ravel(),
        "smooth_mun": np.nan, "smooth_region": reg_s.ravel(), "delta": d_reg_only.ravel(),
        "source": np.concatenate([np.full((len(reg_keys), 1), "base"),
                                  np.where(np.isfinite(reg_s[:, 1:] - reg_s[:, :-1]), "region", "carry")], axis=1).ravel(),
        "log_index": np.cumsum(d_reg_only, axis=1).ravel(),
    })
    out = pd.concat([out, reg_out], ignore_index=True)
    out["month"] = month_id_to_timestamp(out["month_id"]).to_numpy()
    out.attrs["min_n"] = min_n
    out.attrs["window"] = window
    return out


def index_change(index: pd.DataFrame, mun_code: Any, property_class: Any, date_from: Any, date_to: Any,
                 region_group: Any | None = None) -> np.ndarray:
    """Vectorised I(m, c, τ_to) − I(m, c, τ_from). When the (municipality, class) series does not exist and
    `region_group` is given, the region × class series is used; NaN otherwise (never a silent zero)."""
    key = pd.MultiIndex.from_arrays([index["level"].astype(str), index["mun_code"].astype(str), index["property_class"].astype(str),
                                     index["month_id"].astype(int)])
    lookup = pd.Series(index["log_index"].to_numpy(dtype=float), index=key)
    mun = pd.Series(mun_code).astype(str).to_numpy()
    cls = pd.Series(property_class).astype(str).to_numpy()
    m_from = _month_id(pd.Series(date_from)).to_numpy()
    m_to = _month_id(pd.Series(date_to)).to_numpy()
    lvl = np.full(len(mun), "mun")
    v_from = lookup.reindex(pd.MultiIndex.from_arrays([lvl, mun, cls, m_from])).to_numpy(dtype=float)
    v_to = lookup.reindex(pd.MultiIndex.from_arrays([lvl, mun, cls, m_to])).to_numpy(dtype=float)
    if region_group is not None:
        reg = pd.Series(region_group).astype(str).to_numpy()
        lvl_r = np.full(len(mun), "region")
        r_from = lookup.reindex(pd.MultiIndex.from_arrays([lvl_r, reg, cls, m_from])).to_numpy(dtype=float)
        r_to = lookup.reindex(pd.MultiIndex.from_arrays([lvl_r, reg, cls, m_to])).to_numpy(dtype=float)
        missing = ~(np.isfinite(v_from) & np.isfinite(v_to))
        v_from = np.where(missing, r_from, v_from)
        v_to = np.where(missing, r_to, v_to)
    return v_to - v_from


# ------------------------------------------------------------------------------------------------ pairs

def resale_pairs(full: Any) -> pd.DataFrame:
    """Consecutive sales of the same parcel in the full snapshot (one row per (first sale, next sale) pair).

    Rows with a structural data-quality exclusion (duplicates, shared or multi-parcel deeds, land, non-residential,
    built after the sale) and rows without parcel id or positive price cannot be either leg. Pairs on the same date
    are dropped. The later leg keeps its `exclude_reason` (e.g. assessment-ratio outlier) in `resale_exclude_reason`
    and its headline-sample membership in `resale_in_headline` so that users can tighten the sample.
    """
    df = to_pandas(full)
    need = ["tx_id", "parcel_id", "date", "price"]
    missing = [c for c in need if c not in df.columns]
    if missing:
        raise ValueError(f"full table is missing {missing}")
    df = df.dropna(subset=["parcel_id", "price"])
    df = df[df["price"] > 0]
    if "exclude_reason" in df.columns:
        df = df[~df["exclude_reason"].isin(STRUCTURAL_EXCLUSIONS)]
    df = df.assign(date=pd.to_datetime(df["date"])).sort_values(["parcel_id", "date", "tx_id"]).reset_index(drop=True)
    opt = {c: df[c] if c in df.columns else pd.Series(pd.NA, index=df.index) for c in ("property_class", "exclude_reason", "in_headline_sample")}
    nxt = df.groupby("parcel_id", sort=False).shift(-1)
    pairs = pd.DataFrame({
        "tx_id": df["tx_id"].astype(str),
        "parcel_id": df["parcel_id"].astype(str),
        "date": df["date"],
        "price": df["price"].astype(float),
        "property_class": opt["property_class"],
        "resale_tx_id": nxt["tx_id"],
        "resale_date": nxt["date"],
        "resale_price": nxt["price"],
        "resale_property_class": nxt["property_class"] if "property_class" in df.columns else pd.NA,
        "resale_exclude_reason": nxt["exclude_reason"] if "exclude_reason" in df.columns else pd.NA,
        "resale_in_headline": nxt["in_headline_sample"] if "in_headline_sample" in df.columns else pd.NA,
    })
    pairs = pairs.dropna(subset=["resale_tx_id"])
    pairs = pairs[pairs["resale_date"] > pairs["date"]].copy()
    pairs["resale_tx_id"] = pairs["resale_tx_id"].astype(str)
    pairs["resale_price"] = pairs["resale_price"].astype(float)
    return pairs.reset_index(drop=True)


def _months_between(d1: pd.Series, d2: pd.Series) -> np.ndarray:
    return (pd.to_datetime(d2) - pd.to_datetime(d1)).dt.days.to_numpy(dtype=float) / DAYS_PER_MONTH


def resale_targets(analytic: Any, full: Any | None = None, index: pd.DataFrame | None = None, *, min_n: int = MIN_CELL_N) -> pd.DataFrame:
    """Resale-implied value at the first sale date for every analytic sale with a later resale of the same parcel.

    Returns one row per analytic `tx_id` that has a later sale: tx_id, resale_tx_id, resale_price, resale_date,
    months_held, index_change, resale_implied_log_value, resale_implied_value, log_price, gap (= implied − log price),
    annualised_log_return, short_hold, class_changed, resale_flagged, mun_code, property_class, region_group, date, price.
    `full` defaults to the analytic table itself (pairs then only use headline sales); `index` defaults to
    `monthly_index(analytic)`.
    """
    a = to_pandas(analytic)
    a["tx_id"] = a["tx_id"].astype(str)
    if index is None:
        index = monthly_index(a, min_n=min_n)
    pairs = resale_pairs(full if full is not None else a)
    keep = ["tx_id", "parcel_id", "date", "price", "mun_code", "property_class", "region_group"]
    keep = [c for c in keep if c in a.columns]
    t = pairs.drop(columns=["date", "price", "property_class", "parcel_id"]).merge(a[keep], on="tx_id", how="inner", validate="one_to_one")
    t["log_price"] = np.log(t["price"].astype(float))
    t["months_held"] = _months_between(t["date"], t["resale_date"])
    t["index_change"] = index_change(index, t["mun_code"], t["property_class"], t["date"], t["resale_date"], t.get("region_group"))
    t["resale_implied_log_value"] = np.log(t["resale_price"]) - t["index_change"]
    t["resale_implied_value"] = np.exp(t["resale_implied_log_value"])
    t["gap"] = t["resale_implied_log_value"] - t["log_price"]
    t["annualised_log_return"] = (np.log(t["resale_price"]) - t["log_price"]) * 12.0 / t["months_held"]
    t["short_hold"] = t["months_held"] < SHORT_HOLD_MONTHS
    t["class_changed"] = t["resale_property_class"].notna() & (t["resale_property_class"].astype("string") != t["property_class"].astype("string"))
    t["resale_flagged"] = t["resale_exclude_reason"].notna() | (t["resale_in_headline"].fillna(True).astype(bool) == False)  # noqa: E712
    cols = ["tx_id", "parcel_id", "date", "price", "log_price", "mun_code", "property_class", "region_group",
            "resale_tx_id", "resale_date", "resale_price", "resale_property_class", "resale_exclude_reason", "resale_in_headline",
            "months_held", "index_change", "resale_implied_log_value", "resale_implied_value", "gap", "annualised_log_return",
            "short_hold", "class_changed", "resale_flagged"]
    return t[[c for c in cols if c in t.columns]].sort_values(["date", "tx_id"]).reset_index(drop=True)


def prior_targets(analytic: Any, full: Any | None = None, index: pd.DataFrame | None = None, *, min_n: int = MIN_CELL_N) -> pd.DataFrame:
    """Backward direction: for every analytic sale with an *earlier* sale of the same parcel, the earlier price
    brought forward to the sale date by the index (`prior_implied_log_value`). Same selection caveats; unlike
    `prior_sale_price` in the analytic table this uses the full snapshot and the deflator."""
    a = to_pandas(analytic)
    a["tx_id"] = a["tx_id"].astype(str)
    if index is None:
        index = monthly_index(a, min_n=min_n)
    pairs = resale_pairs(full if full is not None else a).rename(columns={
        "tx_id": "prior_tx_id", "date": "prior_date", "price": "prior_price", "property_class": "prior_property_class",
        "resale_tx_id": "tx_id", "exclude_reason": "prior_exclude_reason"})
    pairs = pairs.drop(columns=["resale_date", "resale_price", "resale_property_class", "resale_exclude_reason", "resale_in_headline", "parcel_id"])
    keep = [c for c in ["tx_id", "date", "price", "mun_code", "property_class", "region_group"] if c in a.columns]
    t = pairs.merge(a[keep], on="tx_id", how="inner", validate="one_to_one")
    t["log_price"] = np.log(t["price"].astype(float))
    t["months_held"] = _months_between(t["prior_date"], t["date"])
    t["index_change"] = index_change(index, t["mun_code"], t["property_class"], t["prior_date"], t["date"], t.get("region_group"))
    t["prior_implied_log_value"] = np.log(t["prior_price"]) + t["index_change"]
    t["gap"] = t["prior_implied_log_value"] - t["log_price"]
    t["short_hold"] = t["months_held"] < SHORT_HOLD_MONTHS
    return t.sort_values(["date", "tx_id"]).reset_index(drop=True)


# ------------------------------------------------------------------------------------------------ description

def describe_resale_sample(targets: pd.DataFrame, analytic: Any | None = None) -> dict[str, Any]:
    """Counts, holding periods, returns and the gap between the resale-implied value and the price, overall and by
    class / region group / holding-period bucket. Pure description; no selection correction."""
    t = targets.copy()
    t["hold_bucket"] = pd.cut(t["months_held"], [0, 12, 24, 36, 72], labels=["<12", "12-24", "24-36", "36+"], right=False)

    def block(g: pd.DataFrame) -> dict[str, float]:
        gap = g["gap"].to_numpy(dtype=float)
        ok = np.isfinite(gap)
        return {
            "n": int(len(g)),
            "share_short_hold": float(g["short_hold"].mean()) if len(g) else np.nan,
            "share_flagged": float(g["resale_flagged"].mean()) if len(g) else np.nan,
            "months_held_p25": float(g["months_held"].quantile(0.25)) if len(g) else np.nan,
            "months_held_median": float(g["months_held"].median()) if len(g) else np.nan,
            "months_held_p75": float(g["months_held"].quantile(0.75)) if len(g) else np.nan,
            "annualised_log_return_median": float(g["annualised_log_return"].median()) if len(g) else np.nan,
            "index_change_median": float(g["index_change"].median()) if len(g) else np.nan,
            "gap_mean": float(gap[ok].mean()) if ok.any() else np.nan,
            "gap_median": float(np.median(gap[ok])) if ok.any() else np.nan,
            "gap_sd": float(gap[ok].std(ddof=1)) if ok.sum() > 1 else np.nan,
            "gap_abs_median": float(np.median(np.abs(gap[ok]))) if ok.any() else np.nan,
        }

    out: dict[str, Any] = {"overall": block(t), "by_class": {}, "by_region": {}, "by_hold_bucket": {}, "by_split": {}}
    for k, g in t.groupby("property_class"):
        out["by_class"][str(k)] = block(g)
    for k, g in t.groupby("region_group"):
        out["by_region"][str(k)] = block(g)
    for k, g in t.groupby("hold_bucket", observed=True):
        out["by_hold_bucket"][str(k)] = block(g)
    if "split" in t.columns:
        for k, g in t.groupby("split"):
            out["by_split"][str(k)] = block(g)
    if analytic is not None:
        a = to_pandas(analytic)
        out["n_analytic"] = int(len(a))
        out["share_with_resale"] = float(len(t) / len(a)) if len(a) else np.nan
    return out