# Author: Simon-Pierre Boucher — simon-pierre.boucher@uqo.ca
"""As-of comparable retrieval (identical for every arm).

Rules (``configs/data.yaml`` → ``comparables``):
* candidate sales have ``date < cutoff`` and ``date >= cutoff − lookback_days`` (strictly before the cutoff);
* the subject's own parcel is excluded (its earlier sale is exposed separately as ``prior_sale_*``);
* same ``property_class`` when ``same_class_required``;
* expanding radius 1 → 3 → 10 km: the first radius with at least ``min_comps`` candidates is used; if even 10 km
  yields fewer than ``min_comps``, the search is extended to 25 km and the result is flagged (``expanded``);
* candidates are ranked by a similarity score (lower = better) combining distance, |Δ log floor area|,
  |Δ year built| and recency; the best ``max_comps`` receive evidence ids C01.. in score order.

Speed: one KD-tree query on EPSG:32198 coordinates per subject (``AnalyticIndex``), then NumPy masks —
about 1–3 ms per subject on the full 632k-row table.
"""
from __future__ import annotations

from dataclasses import dataclass, field
from datetime import date
from typing import Any

import numpy as np
import polars as pl

from ..config import load_config
from ..data.analytic_schema import PROPERTY_CLASSES
from ..schemas import Comparable
from .index import CLASS_CODE, AnalyticIndex, from_days, get_index, subject_to_dict, to_days

EXTENDED_RADIUS_KM = 25.0
#: similarity-score weights (dimensionless; each term is roughly 0–1 in the typical range)
W_DISTANCE, W_AREA, W_YEAR, W_RECENCY = 1.0, 1.5, 0.5, 0.5


@dataclass
class ComparableSet:
    comparables: list[Comparable]
    radius_km: float
    n_candidates: int
    expanded: bool
    notes: list[str] = field(default_factory=list)


def similarity_score(dist_km: np.ndarray, radius_km: float, log_area_diff: np.ndarray, year_diff: np.ndarray, days_before: np.ndarray, lookback_days: int) -> np.ndarray:
    """Lower is more similar. Missing attributes count as a neutral penalty (0.5) rather than being dropped."""
    d = dist_km / max(radius_km, 0.1)
    a = np.where(np.isnan(log_area_diff), 0.5, np.abs(log_area_diff))
    y = np.where(np.isnan(year_diff), 0.5, np.abs(year_diff) / 50.0)
    t = days_before / max(lookback_days, 1)
    return W_DISTANCE * d + W_AREA * a + W_YEAR * y + W_RECENCY * t


def _candidates(idx: AnalyticIndex, x: float, y: float, radius_km: float, cutoff_days: int, lookback_days: int, cls_code: int | None, parcel_id: str | None) -> np.ndarray:
    cand = idx.within_radius(x, y, radius_km * 1000.0)
    if cand.size == 0:
        return cand
    d = idx.date_days[cand]
    mask = (d < cutoff_days) & (d >= cutoff_days - lookback_days)
    if cls_code is not None:
        mask &= idx.cls[cand] == cls_code
    if parcel_id is not None:
        mask &= idx.parcel_id[cand] != parcel_id
    return cand[mask]


def retrieve_comparables_detailed(analytic: pl.DataFrame, subject: dict[str, Any] | pl.DataFrame, cutoff: date, cfg: dict[str, Any] | None = None) -> ComparableSet:
    """Full retrieval result (comparables + radius used + flags)."""
    cfg = cfg or load_config("data")
    c = cfg["comparables"] if "comparables" in cfg else cfg
    radii = [float(r) for r in c["radius_km"]]
    lookback = int(c["lookback_days"])
    max_comps, min_comps = int(c["max_comps"]), int(c["min_comps"])
    same_class = bool(c.get("same_class_required", True))

    s = subject_to_dict(subject)
    idx = get_index(analytic)
    if "x_m" in s and s.get("x_m") is not None:
        x, y = float(s["x_m"]), float(s["y_m"])
    else:
        from ..data.analytic_schema import project_xy

        xx, yy = project_xy(np.array([s["lat"]]), np.array([s["lng"]]))
        x, y = float(xx[0]), float(yy[0])
    cls_code = CLASS_CODE.get(s.get("property_class")) if same_class else None
    parcel = s.get("parcel_id")
    cutoff_days = to_days(cutoff)
    notes: list[str] = []
    expanded = False
    used_radius = radii[-1]
    cand = np.empty(0, dtype=np.int64)
    for r in radii:
        cand = _candidates(idx, x, y, r, cutoff_days, lookback, cls_code, parcel)
        used_radius = r
        if cand.size >= min_comps:
            break
    if cand.size < min_comps:
        cand = _candidates(idx, x, y, EXTENDED_RADIUS_KM, cutoff_days, lookback, cls_code, parcel)
        used_radius = EXTENDED_RADIUS_KM
        expanded = True
        notes.append(f"fewer than {min_comps} same-class sales within {radii[-1]:g} km in the last {lookback} days; search extended to {EXTENDED_RADIUS_KM:g} km")
        if cand.size < min_comps:
            notes.append(f"only {cand.size} comparable(s) found within {EXTENDED_RADIUS_KM:g} km (thin market)")
    if cand.size == 0:
        return ComparableSet([], used_radius, 0, expanded, notes)

    dist_km = np.hypot(idx.x[cand] - x, idx.y[cand] - y) / 1000.0
    s_area = s.get("floor_area_m2")
    with np.errstate(all="ignore"):
        log_area_diff = np.log(idx.floor_area[cand]) - np.log(float(s_area)) if s_area else np.full(cand.size, np.nan)
    s_year = s.get("year_built")
    year_diff = (idx.year_built[cand] - float(s_year)) if s_year else np.full(cand.size, np.nan)
    days_before = (cutoff_days - idx.date_days[cand]).astype(float)
    score = similarity_score(dist_km, used_radius, log_area_diff, year_diff, days_before, lookback)
    order = np.lexsort((idx.tx_id[cand].astype(str), score))[:max_comps]

    comps: list[Comparable] = []
    for rank, j in enumerate(order, start=1):
        i = int(cand[j])
        fa = idx.floor_area[i]
        rt = idx.roll_total[i]
        comps.append(
            Comparable(
                evidence_id=f"C{rank:02d}",
                sale_date=from_days(idx.date_days[i]),
                price=float(idx.price[i]),
                distance_km=round(float(dist_km[j]), 3),
                property_class=PROPERTY_CLASSES[idx.cls[i]],
                subtype=str(idx.subtype[i]) if idx.subtype[i] in ("detached", "semi", "row", "integrated") else "unknown",
                municipality=str(idx.mun_name[i]),
                floor_area_m2=None if np.isnan(fa) else float(fa),
                land_area_m2=None if np.isnan(idx.land_area[i]) else float(idx.land_area[i]),
                year_built=None if np.isnan(idx.year_built[i]) else int(idx.year_built[i]),
                storeys=None if np.isnan(idx.storeys[i]) else float(idx.storeys[i]),
                n_units=None if np.isnan(idx.n_units[i]) else int(idx.n_units[i]),
                roll_total_value=None if np.isnan(rt) else float(rt),
                price_per_m2=None if np.isnan(fa) or fa <= 0 else round(float(idx.price[i] / fa), 1),
                assessment_ratio=None if np.isnan(rt) or rt <= 0 else round(float(idx.price[i] / rt), 4),
                months_before_valuation=round(float(days_before[j]) / 30.4375, 2),
            )
        )
    return ComparableSet(comps, used_radius, int(cand.size), expanded, notes)


def retrieve_comparables(analytic: pl.DataFrame, subject: dict[str, Any] | pl.DataFrame, cutoff: date, cfg: dict[str, Any] | None = None) -> list[Comparable]:
    """Comparables C01..C12 for ``subject`` using only sales strictly before ``cutoff`` (see module docstring)."""
    return retrieve_comparables_detailed(analytic, subject, cutoff, cfg).comparables