"""Real data: download and preprocess NYC TLC yellow-taxi trips.

Pipeline (all deterministic, seeded):
  1. download one month of trips + zone lookup + zone shapefile
  2. check the zone-ID columns exist, parse timestamps
  3. keep trips inside the month, on weekdays, duration in [1, 120] min,
     positive distance and plausible speed (<= 60 mph)
  4. pick the 20 busiest zones (pickups + drop-offs, unknown zones 264/265 excluded)
     and keep trips with BOTH ends in those zones
  5. subsample ~300k trips with a fixed seed
  6. split by day: Jan 3-19 train, Jan 20-21 validation, Jan 24-31 test
  7. zone-pair mean distance from the TRAIN split only (fallback for unseen pairs:
     linear fit on zone-centroid distance, also fitted on train only)
"""
from __future__ import annotations

import json
import urllib.request
from pathlib import Path

import numpy as np
import pandas as pd

from . import config as C

NEEDED_COLS = ["tpep_pickup_datetime", "tpep_dropoff_datetime",
               "PULocationID", "DOLocationID", "trip_distance"]


# --------------------------------------------------------------------------- download
def download(force: bool = False) -> dict:
    C.DATA_RAW.mkdir(parents=True, exist_ok=True)
    files = {
        "trips": (C.TRIP_URL, C.DATA_RAW / f"yellow_tripdata_{C.MONTH}.parquet"),
        "lookup": (C.ZONE_LOOKUP_URL, C.DATA_RAW / "taxi_zone_lookup.csv"),
        "shapes": (C.ZONE_SHAPES_URL, C.DATA_RAW / "taxi_zones.zip"),
    }
    out = {}
    for key, (url, path) in files.items():
        if force or not path.exists() or path.stat().st_size == 0:
            print(f"  downloading {url}")
            urllib.request.urlretrieve(url, path)
        else:
            print(f"  cached      {path.name} ({path.stat().st_size/1e6:.1f} MB)")
        out[key] = path
    return out


# --------------------------------------------------------------------------- zones
def _ring_centroid(pts: np.ndarray) -> tuple[float, float, float]:
    x, y = pts[:, 0], pts[:, 1]
    x1, y1 = np.roll(x, -1), np.roll(y, -1)
    cross = x * y1 - x1 * y
    a = cross.sum() / 2.0
    if abs(a) < 1e-9:
        return float(x.mean()), float(y.mean()), 0.0
    cx = ((x + x1) * cross).sum() / (6 * a)
    cy = ((y + y1) * cross).sum() / (6 * a)
    return float(cx), float(cy), float(abs(a))


def read_zone_shapes(zip_path: Path) -> dict[int, list[np.ndarray]]:
    """LocationID -> list of polygon rings in km (EPSG:2263 feet * 0.0003048)."""
    import shapefile  # pyshp
    r = shapefile.Reader(str(zip_path))
    fields = [f[0] for f in r.fields[1:]]
    id_col = fields.index("LocationID")
    shapes: dict[int, list[np.ndarray]] = {}
    for sr in r.iterShapeRecords():
        lid = int(sr.record[id_col])
        pts = np.asarray(sr.shape.points, dtype=float) * 0.0003048
        parts = list(sr.shape.parts) + [len(pts)]
        for a, b in zip(parts[:-1], parts[1:]):
            shapes.setdefault(lid, []).append(pts[a:b])
    return shapes


def zone_centroids(zip_path: Path) -> pd.DataFrame:
    rows = []
    for lid, rings in read_zone_shapes(zip_path).items():
        cs = [_ring_centroid(r) for r in rings]
        w = np.array([c[2] for c in cs]) + 1e-12
        cx = float(np.average([c[0] for c in cs], weights=w))
        cy = float(np.average([c[1] for c in cs], weights=w))
        rows.append((lid, cx, cy))
    return pd.DataFrame(rows, columns=["LocationID", "cx_km", "cy_km"])


# --------------------------------------------------------------------------- trips
def load_raw_trips(path: Path) -> pd.DataFrame:
    import pyarrow.parquet as pq
    schema = pq.read_schema(path)
    missing = [c for c in NEEDED_COLS if c not in schema.names]
    if missing:
        raise ValueError(f"zone/time columns missing in {path.name}: {missing}")
    return pq.read_table(path, columns=NEEDED_COLS).to_pandas()


def preprocess(raw: pd.DataFrame, lookup: pd.DataFrame, cents: pd.DataFrame,
               seed: int = C.SEED, n_zones: int = C.N_ZONES,
               target: int = C.TARGET_TRIPS):
    stats = {"raw_rows": int(len(raw))}
    df = raw.rename(columns={"tpep_pickup_datetime": "pickup", "tpep_dropoff_datetime": "dropoff",
                             "PULocationID": "pu_id", "DOLocationID": "do_id",
                             "trip_distance": "distance_mi"})
    month_start = pd.Timestamp(C.MONTH + "-01")
    month_end = month_start + pd.offsets.MonthBegin(1)
    df = df[(df.pickup >= month_start) & (df.pickup < month_end)]
    stats["in_month"] = int(len(df))
    df = df[df.pickup.dt.weekday < 5]
    stats["weekdays"] = int(len(df))
    df = df.assign(duration_min=(df.dropoff - df.pickup).dt.total_seconds() / 60.0)
    df = df[(df.duration_min >= C.MIN_DURATION) & (df.duration_min <= C.MAX_DURATION)]
    stats["duration_1_120"] = int(len(df))
    speed = df.distance_mi / (df.duration_min / 60.0)
    df = df[(df.distance_mi > 0) & (speed <= C.MAX_SPEED_MPH)]
    stats["distance_speed_ok"] = int(len(df))

    known = df[(df.pu_id < 264) & (df.do_id < 264)]
    activity = known.pu_id.value_counts().add(known.do_id.value_counts(), fill_value=0)
    top = [int(z) for z in activity.sort_values(ascending=False).index[:n_zones]]
    top = sorted(top)
    df = df[df.pu_id.isin(top) & df.do_id.isin(top)]
    stats["both_ends_top_zones"] = int(len(df))
    if len(df) > target:
        df = df.sample(n=target, random_state=seed)
    stats["sampled"] = int(len(df))

    idx = {z: i for i, z in enumerate(top)}
    day = df.pickup.dt.normalize()
    split = np.where(day >= pd.Timestamp(C.TEST_START), "test",
                     np.where(day.dt.strftime("%Y-%m-%d").isin(C.VAL_DAYS), "val", "train"))
    trips = pd.DataFrame({
        "pickup": df.pickup.values,
        "date": day.dt.strftime("%Y-%m-%d").values,
        "pu_id": df.pu_id.values.astype(int),
        "do_id": df.do_id.values.astype(int),
        "pu": df.pu_id.map(idx).values.astype(int),
        "do": df.do_id.map(idx).values.astype(int),
        "hour": df.pickup.dt.hour.values.astype(int),
        "weekday": df.pickup.dt.weekday.values.astype(int),
        "distance_mi": df.distance_mi.values.astype(float),
        "duration_min": df.duration_min.values.astype(float),
        "split": split,
    }).sort_values("pickup").reset_index(drop=True)
    for s in ("train", "val", "test"):
        stats[f"n_{s}"] = int((trips.split == s).sum())
        stats[f"days_{s}"] = sorted(trips.loc[trips.split == s, "date"].unique().tolist())

    zones = (lookup[lookup.LocationID.isin(top)]
             .merge(cents, on="LocationID", how="left")
             .sort_values("LocationID").reset_index(drop=True))
    zones.insert(0, "idx", [idx[z] for z in zones.LocationID])
    zones["n_pickups"] = zones.idx.map(trips.pu.value_counts()).fillna(0).astype(int)
    zones["n_dropoffs"] = zones.idx.map(trips["do"].value_counts()).fillna(0).astype(int)
    zones = zones.rename(columns={"Borough": "borough", "Zone": "zone", "service_zone": "service_zone"})
    return trips, zones, stats


def pair_distance_matrix(trips: pd.DataFrame, zones: pd.DataFrame):
    """Mean trip distance per (pu, do) pair from TRAIN rows only.

    Pairs never seen in train get a linear fit on centroid distance (fitted on train pairs).
    Returns (dist[Z,Z] miles, count[Z,Z], fit coefficients).
    """
    Z = len(zones)
    tr = trips[trips.split == "train"]
    g = tr.groupby(["pu", "do"]).distance_mi.agg(["mean", "size"])
    dist = np.full((Z, Z), np.nan)
    cnt = np.zeros((Z, Z), dtype=int)
    for (i, j), r in g.iterrows():
        dist[i, j], cnt[i, j] = r["mean"], r["size"]
    xy = zones.sort_values("idx")[["cx_km", "cy_km"]].to_numpy()
    cd_mi = np.linalg.norm(xy[:, None, :] - xy[None, :, :], axis=2) / 1.609344
    seen = ~np.isnan(dist)
    b, a = np.polyfit(cd_mi[seen], dist[seen], 1)
    dist = np.where(seen, dist, a + b * cd_mi)
    return dist, cnt, {"intercept": float(a), "slope_per_centroid_mile": float(b),
                       "unseen_pairs": int((~seen).sum())}


# --------------------------------------------------------------------------- io
def save_processed(trips, zones, dist, cnt, stats, out: Path = C.DATA_PROC):
    out.mkdir(parents=True, exist_ok=True)
    trips.to_parquet(out / "trips.parquet", index=False)
    zones.to_csv(out / "zones.csv", index=False)
    np.save(out / "pair_distance.npy", dist)
    np.save(out / "pair_count.npy", cnt)
    (out / "preprocess_stats.json").write_text(json.dumps(stats, indent=2))


def load_processed(base: Path = C.DATA_PROC):
    """Returns (trips, zones, pair_distance). Falls back to the committed test fixture."""
    if (base / "trips.parquet").exists():
        trips = pd.read_parquet(base / "trips.parquet")
        zones = pd.read_csv(base / "zones.csv")
        dist = np.load(base / "pair_distance.npy")
        return trips, zones, dist
    return load_fixture()


def write_fixture(trips, zones, dist, n: int = 3000, seed: int = C.SEED, out: Path = C.FIXTURES):
    out.mkdir(parents=True, exist_ok=True)
    parts = [trips[trips.split == s].sample(n=min(n // 3, (trips.split == s).sum()), random_state=seed)
             for s in ("train", "val", "test")]
    sample = pd.concat(parts).sort_values("pickup")
    sample.drop(columns=["pickup"]).to_csv(out / "sample_trips.csv", index=False)
    zones.to_csv(out / "zones.csv", index=False)
    np.savetxt(out / "pair_distance.csv", dist, delimiter=",", fmt="%.4f")


def load_fixture(base: Path = C.FIXTURES):
    trips = pd.read_csv(base / "sample_trips.csv")
    zones = pd.read_csv(base / "zones.csv")
    dist = np.loadtxt(base / "pair_distance.csv", delimiter=",")
    return trips, zones, dist