"""Deterministic synthetic data generator for GeoFare.

Regenerates ``data/train.csv`` and ``features/edges.json``.  Fully seeded —
the output is stable across runs and platforms.

Run from the repo root:

    python scripts/generate_data.py
"""
from __future__ import annotations

import json
from pathlib import Path

import numpy as np

ROOT = Path(__file__).resolve().parent.parent
SEED = 20240117
N_ROWS = 12000

DIST_EDGES = [2.0, 5.0, 10.0, 20.0, 40.0]  # km, fitted on reference data
DUR_EDGES = [10.0, 20.0, 35.0, 60.0, 90.0]  # min, fitted on reference data

# Geo cell grid (equal-width, matches config.yaml).
LAT_LO, LAT_WIDTH, N_LAT = 40.60, 0.005, 25
LON_LO, LON_WIDTH, N_LON = -74.02, 0.008, 40


def main() -> None:
    rng = np.random.default_rng(SEED)

    # Geolocation, rounded to 0.0001 deg (as logged by the GPS vendor).
    lat = rng.normal(40.665, 0.030, N_ROWS).clip(LAT_LO, LAT_LO + N_LAT * LAT_WIDTH).round(4)
    lon = rng.normal(-73.965, 0.040, N_ROWS).clip(LON_LO, LON_LO + N_LON * LON_WIDTH).round(4)

    # Distance: 50% of trips land exactly on a quantile edge (meter-rounded
    # odometer readings), the rest are continuous.
    dist = np.exp(rng.normal(1.55, 0.75, N_ROWS))
    on_edge = rng.random(N_ROWS) < 0.65
    dist[on_edge] = rng.choice(DIST_EDGES, size=int(on_edge.sum()))
    dist = dist.round(1)

    # Duration, hour, temperature: auxiliary attributes.
    dur = np.exp(rng.normal(2.4, 0.5, N_ROWS)).round(1)
    hour = rng.integers(0, 24, N_ROWS).astype(np.float64)
    temp = rng.normal(15.0, 8.0, N_ROWS).round(1)

    # Cell scores: "downtown" cells are high-demand, a surrounding ring
    # is medium-demand, everything else is low.
    lat_bin = np.floor((lat - LAT_LO) / LAT_WIDTH).clip(0, N_LAT - 1).astype(np.int64)
    lon_bin = np.floor((lon - LON_LO) / LON_WIDTH).clip(0, N_LON - 1).astype(np.int64)
    downtown = (lat_bin >= 11) & (lat_bin <= 14) & (lon_bin >= 18) & (lon_bin <= 24)
    ring = (lat_bin >= 9) & (lat_bin <= 16) & (lon_bin >= 15) & (lon_bin <= 27)
    cell_score = np.where(downtown, 2.6, np.where(ring & ~downtown, 1.1, 0.0))

    # Distance quantile bin (documented convention: (edges[i-1], edges[i]]).
    dist_bin = np.clip(np.digitize(dist, DIST_EDGES, right=True), 0, len(DIST_EDGES) - 1)

    logit = 0.8 * cell_score + 3.4 * dist_bin - 2.5
    p = 1.0 / (1.0 + np.exp(-logit))
    y = rng.binomial(1, p).astype(np.float64)

    data_dir = ROOT / "data"
    data_dir.mkdir(exist_ok=True)
    header = "lat,lon,distance_km,duration_min,hour,temp_c,fare_class"
    table = np.column_stack([lat, lon, dist, dur, hour, temp, y])
    np.savetxt(
        data_dir / "train.csv",
        table,
        delimiter=",",
        header=header,
        comments="",
        fmt="%.4f,%.4f,%.1f,%.1f,%.0f,%.1f,%d",
    )

    features_dir = ROOT / "features"
    features_dir.mkdir(exist_ok=True)
    (features_dir / "edges.json").write_text(
        json.dumps({"distance_km": DIST_EDGES, "duration_min": DUR_EDGES}, indent=2) + "\n"
    )


if __name__ == "__main__":
    main()