"""Listing-day risk set with time-varying local competition (Phase 3, analysis_plan §2).

Competitor set of listing i on day t: listings of the same make & model, model year
within ±1, at a *different* dealer, within radius R km, active on day t, drawn from
the reliable-source universe (S2L: direct dealers, full portals, banners; not
oscillating). Radii 10 / 25 / 50 km. Day index: UTC calendar days since the first
crawl day (2026-08-12 = day 0).
"""
from __future__ import annotations

import json
import math
import time
from pathlib import Path

import numpy as np
import polars as pl

from . import config as C

RADII = (10.0, 25.0, 50.0)
DAY0 = math.floor(C.WINDOW_START_TS / 86400)


def day_of(ts: np.ndarray) -> np.ndarray:
    return (np.floor(ts / 86400) - DAY0).astype(np.int32)


def haversine_km(lat1, lng1, lat2, lng2):
    R = 6371.0
    p1, p2 = np.radians(lat1), np.radians(lat2)
    dphi = p2[None, :] - p1[:, None]
    dl = np.radians(lng2)[None, :] - np.radians(lng1)[:, None]
    a = np.sin(dphi / 2) ** 2 + np.cos(p1)[:, None] * np.cos(p2)[None, :] * np.sin(dl / 2) ** 2
    return 2 * R * np.arcsin(np.sqrt(np.clip(a, 0, 1)))


def daily_price_matrix(ev: pl.DataFrame, ids: list[str], e_day: np.ndarray, T: int) -> tuple[np.ndarray, np.ndarray]:
    """Step-function log price per listing-day (nan when no price yet) and daily cut indicator."""
    n = len(ids)
    P = np.full((n, T), np.nan, dtype=np.float32)
    K = np.zeros((n, T), dtype=np.float32)
    idx = {l: i for i, l in enumerate(ids)}
    g = ev.filter(pl.col("listing_id").is_in(ids)).sort(["listing_id", "ts"]).group_by("listing_id", maintain_order=True).agg([pl.col("ts"), pl.col("price"), pl.col("kind")])
    days = np.arange(T)
    for lid, ts, pr, kd in g.iter_rows():
        i = idx[lid]
        d = day_of(np.asarray(ts, dtype=np.float64))
        d = np.minimum(d, T - 1)
        # price at end of each day = last event with day <= t
        pos = np.searchsorted(d, days, side="right") - 1
        valid = pos >= 0
        P[i, valid] = np.log(np.asarray(pr, dtype=np.float64))[pos[valid]]
        for dd, k in zip(d, kd):
            if k == "cut" and 0 <= dd < T:
                K[i, dd] = 1.0
    return P, K


def build(listings_full: Path = C.INTERIM / "listings_full.parquet", events: Path = C.PROCESSED / "price_events.parquet",
          out: Path = C.PROCESSED / "risk_set.parquet") -> dict:
    t0 = time.time()
    L = pl.read_parquet(listings_full)
    E = pl.read_parquet(events)
    T = int(day_of(np.array([C.SNAPSHOT_TS]))[0]) + 1
    U = L.filter(pl.col("s2l_dynamic_with_stock") & pl.col("lat").is_not_null()).select(
        ["listing_id", "dealer_id", "make", "model", "year", "lat", "lng", "entry_ts", "end_ts", "exit_confirmed", "s2_dynamic", "left_censored", "first_cut_ts"])
    U = U.with_columns([
        pl.Series("e_day", day_of(U["entry_ts"].to_numpy())),
        pl.Series("x_day", np.minimum(day_of(U["end_ts"].to_numpy()), T - 1)),
    ])
    parts = []
    n_groups = 0
    for (mk, md), G in U.group_by(["make", "model"]):
        n_groups += 1
        ids = G["listing_id"].to_list()
        n = len(ids)
        e = G["e_day"].to_numpy(); x = G["x_day"].to_numpy(); y = G["year"].to_numpy()
        lat = G["lat"].to_numpy(); lng = G["lng"].to_numpy()
        dealer = pl.Series(G["dealer_id"]).cast(pl.Categorical).to_physical().to_numpy()
        A = ((np.arange(T)[None, :] >= e[:, None]) & (np.arange(T)[None, :] <= x[:, None])).astype(np.float32)
        P, K = daily_price_matrix(E, ids, e, T)
        Pf = np.where(np.isnan(P), 0.0, P).astype(np.float32)
        has_p = (~np.isnan(P)).astype(np.float32) * A
        D = haversine_km(lat, lng, lat, lng)
        same_year = np.abs(y[:, None] - y[None, :]) <= 1
        other_dealer = dealer[:, None] != dealer[None, :]
        base = same_year & other_dealer
        cols = {}
        for R in RADII:
            M = (base & (D <= R)).astype(np.float32)
            dens = M @ A                       # n x T competitor counts
            ccut = M @ K                       # competitor cuts on day t
            sp = M @ (Pf * has_p); cnt = M @ has_p
            mean_lp = np.where(cnt > 0, sp / np.maximum(cnt, 1), np.nan)
            tag = f"{int(R)}"
            cols[f"dens_{tag}"] = dens; cols[f"ccut_{tag}"] = ccut; cols[f"cmlp_{tag}"] = mean_lp
        # within-dealer same-model stock (own-dealer substitutes)
        Mo = (same_year & ~other_dealer).astype(np.float32) - np.eye(n, dtype=np.float32)
        own = Mo @ A
        # explode to listing-days
        rows_i, rows_t = [], []
        for i in range(n):
            tt = np.arange(e[i], x[i] + 1)
            rows_i.append(np.full(len(tt), i)); rows_t.append(tt)
        ii = np.concatenate(rows_i); tt = np.concatenate(rows_t)
        part = {"listing_id": np.array(ids, dtype=object)[ii], "t": tt.astype(np.int16), "listing_age": (tt - e[ii]).astype(np.int16),
                "logp": P[ii, tt], "cut": K[ii, tt].astype(np.int8), "own_stock": own[ii, tt].astype(np.int16)}
        for k, v in cols.items():
            part[k] = v[ii, tt]
        parts.append(pl.DataFrame(part))
    RS = pl.concat(parts)
    # exit indicator: last day of the spell for confirmed exits
    last = U.select(["listing_id", "x_day", "exit_confirmed", "e_day", "s2_dynamic", "left_censored", "first_cut_ts"])
    RS = RS.join(last, on="listing_id", how="left").with_columns([
        ((pl.col("t") == pl.col("x_day")) & pl.col("exit_confirmed")).cast(pl.Int8).alias("exit"),
        (pl.col("t") == pl.col("x_day")).cast(pl.Int8).alias("last_day"),
    ])
    # risk set for the first markdown: days up to and including the first cut day
    fc_day = pl.when(pl.col("first_cut_ts").is_not_null()).then(((pl.col("first_cut_ts") / 86400).floor() - DAY0).cast(pl.Int32)).otherwise(None)
    RS = RS.with_columns(fc_day.alias("first_cut_day"))
    RS = RS.with_columns([
        (pl.col("first_cut_day").is_null() | (pl.col("t") <= pl.col("first_cut_day"))).alias("in_first_cut_risk"),
        ((pl.col("t") == pl.col("first_cut_day"))).fill_null(False).cast(pl.Int8).alias("first_cut"),   # NULL first_cut_day -> 0 (never cut), see E-016
    ])
    for tag in ("10", "25", "50"):
        RS = RS.with_columns((pl.col("logp") - pl.col(f"cmlp_{tag}")).alias(f"pricepos_{tag}"))
    RS = RS.drop(["x_day", "first_cut_ts"]).sort(["listing_id", "t"])
    RS.write_parquet(out, compression="zstd")
    stats = {"listing_days": RS.height, "listings": RS["listing_id"].n_unique(), "groups": n_groups, "T": T,
             "exits": int(RS["exit"].sum()), "first_cuts": int(RS["first_cut"].sum()), "seconds": round(time.time() - t0, 1)}
    (C.PROCESSED / "risk_set_stats.json").write_text(json.dumps(stats, indent=2))
    return stats


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