L10 · Solar Stinger: a Georgia Tech digital twin

"""Campus geometry (OpenStreetMap) and the rooftop solar resource with building-to-building shading.

Pipeline: OSM footprints + heights -> 1 m height raster -> roof sample points -> 36-direction horizon
profile per point (ray-marched, Numba) -> sky-view factor + hourly beam shading over the PVGIS typical
meteorological year -> pvlib Perez transposition + Faiman cell temperature -> kWh per unit efficiency.
"""
import os
os.environ.setdefault("NUMBA_NUM_THREADS", "4")
import json

import numpy as np
import pandas as pd
import pvlib
from matplotlib.path import Path as MPath
from numba import njit, prange

from fetch import DATA, LAT0, LON0, to_xy

RES = 1.0            # height raster resolution, m
SPACING = 3.0        # roof sample spacing, m
NAZ = 36             # horizon directions (10 deg)
TILT, PANEL_Z = 10.0, 0.5    # low-tilt south-facing racking, panel height above roof
USABLE = 0.6         # roof fraction left after setbacks, HVAC, access (NREL rooftop studies use 0.5-0.7)
PR = 0.86            # system losses: soiling, wiring, inverter, mismatch (PVWatts default 14%)
GAMMA = -0.0035      # power temperature coefficient, 1/K (c-Si). ponytail: same for every material
HEIGHT_BY_TYPE = {"house": 7, "detached": 7, "semidetached_house": 7, "garage": 3, "garages": 3, "shed": 3,
                  "roof": 5, "carport": 3, "parking": 12, "apartments": 15, "dormitory": 18, "residential": 10,
                  "university": 16, "college": 16, "school": 12, "commercial": 12, "retail": 7, "office": 20,
                  "church": 12, "stadium": 25, "hotel": 30, "hospital": 25, "industrial": 9}


def join_ways(ways):
    """Chain open member ways of a multipolygon into closed rings (greedy endpoint matching)."""
    ways, out = [list(w) for w in ways if len(w) > 1], []
    while ways:
        ring = ways.pop(0)
        while ring[0] != ring[-1]:
            for i, w in enumerate(ways):
                if w[0] == ring[-1]: ring += w[1:]
                elif w[-1] == ring[-1]: ring += w[::-1][1:]
                else: continue
                ways.pop(i); break
            else: break   # ponytail: unclosable ring (clipped by the bbox) is kept open; Path treats it as closed
        out.append(ring)
    return out


def rings(e):
    if e["type"] == "way":
        return [[(g["lat"], g["lon"]) for g in e["geometry"]]]
    return join_ways([[(g["lat"], g["lon"]) for g in m["geometry"]] for m in e["members"]
                      if m.get("role") == "outer" and "geometry" in m])


def xy_ring(r):
    lat, lon = np.array(r).T
    return to_xy(lat, lon).T


def height(tags, area):
    for key, f in (("height", 1.0), ("building:levels", 3.5)):
        try:
            return float(tags[key].split()[0].replace(",", ".")) * f + (1.5 if f != 1.0 else 0)
        except (KeyError, ValueError):
            pass
    # ponytail: 83% of footprints have no height tag; type/footprint heuristic, LiDAR DSM is the upgrade
    return HEIGHT_BY_TYPE.get(tags.get("building"), 6 if area < 150 else 10 if area < 800 else 16)


def load_geometry():
    els = json.loads((DATA / "osm.json").read_text(encoding="utf-8"))["elements"]
    campus_polys = [MPath(xy_ring(r)) for e in json.loads((DATA / "campus.json").read_text(encoding="utf-8"))["elements"]
                    for r in rings(e)]
    buildings, roads, green = [], [], []
    for e in els:
        t = e.get("tags", {})
        if "building" in t:
            for r in rings(e):
                xy = xy_ring(r)
                if len(xy) < 4: continue
                area = 0.5 * abs(np.dot(xy[:-1, 0], xy[1:, 1]) - np.dot(xy[1:, 0], xy[:-1, 1]))
                c = xy.mean(0)
                buildings.append(dict(xy=xy, h=height(t, area), area=area, name=t.get("name", ""),
                                      kind=t["building"], campus=any(p.contains_point(c) for p in campus_polys)))
        elif "highway" in t and e["type"] == "way":
            roads.append(dict(xy=xy_ring([(g["lat"], g["lon"]) for g in e["geometry"]]), kind=t["highway"]))
        elif e["type"] == "way":
            green.append(dict(xy=xy_ring([(g["lat"], g["lon"]) for g in e["geometry"]]),
                              kind=t.get("leisure") or t.get("landuse") or t.get("natural")))
    return buildings, roads, green, campus_polys


def rasterize(buildings, pad=60):
    allxy = np.vstack([b["xy"] for b in buildings])
    x0, y0 = allxy.min(0) - pad
    nx, ny = (np.ceil((allxy.max(0) + pad - (x0, y0)) / RES)).astype(int)
    H = np.zeros((ny, nx), np.float32)
    pts, owner = [], []
    for i, b in enumerate(buildings):
        lo = np.floor((b["xy"].min(0) - (x0, y0)) / RES).astype(int)
        hi = np.ceil((b["xy"].max(0) - (x0, y0)) / RES).astype(int) + 1
        gx, gy = np.meshgrid(np.arange(lo[0], hi[0]), np.arange(lo[1], hi[1]))
        cx, cy = x0 + (gx + .5) * RES, y0 + (gy + .5) * RES
        inside = MPath(b["xy"]).contains_points(np.c_[cx.ravel(), cy.ravel()]).reshape(gx.shape)
        H[gy[inside], gx[inside]] = np.maximum(H[gy[inside], gx[inside]], b["h"])
        # roof sample points on a SPACING grid (sub-sampled raster cells)
        k = int(SPACING / RES)
        sel = inside & (gx % k == 0) & (gy % k == 0)
        if b["area"] >= 40 and sel.any():
            pts.append(np.c_[cx[sel], cy[sel], np.full(sel.sum(), b["h"] + PANEL_Z)])
            owner.append(np.full(sel.sum(), i))
    return H, (x0, y0), np.vstack(pts), np.concatenate(owner)


@njit(parallel=True, cache=True)
def horizons(H, x0, y0, res, pts, naz, rmax):
    """Max elevation angle of the skyline in naz compass directions (0 = north, clockwise), per point."""
    n = pts.shape[0]
    out = np.zeros((n, naz), np.float32)
    ny, nx = H.shape
    for i in prange(n):
        px, py, pz = pts[i, 0], pts[i, 1], pts[i, 2]
        for k in range(naz):
            az = 2 * np.pi * k / naz
            dx, dy = np.sin(az), np.cos(az)
            best, r = 0.0, 2.0
            while r < rmax:
                ix = int((px + dx * r - x0) / res)
                iy = int((py + dy * r - y0) / res)
                if ix < 0 or iy < 0 or ix >= nx or iy >= ny: break
                a = np.arctan2(H[iy, ix] - pz, r)
                if a > best: best = a
                r += 1.0 + 0.01 * r     # finer steps near the point, where small obstacles matter
            out[i, k] = best
    return out


@njit(parallel=True, cache=True)
def shaded_poa(hor, svf, sun_el, sun_az, poa_dir, poa_sky, poa_gnd):
    """Annual plane-of-array irradiation per roof point (Wh/m^2), beam blocked when the sun is below its skyline."""
    n, naz = hor.shape
    nt = sun_el.shape[0]
    per_pt = np.zeros(n)
    for i in prange(n):
        acc = 0.0
        for t in range(nt):
            if poa_dir[t] + poa_sky[t] <= 0: continue
            f = sun_az[t] / (2 * np.pi) * naz
            k0 = int(f) % naz
            k1 = (k0 + 1) % naz
            hz = hor[i, k0] + (f - int(f)) * (hor[i, k1] - hor[i, k0])
            beam = poa_dir[t] if sun_el[t] > hz else 0.0
            acc += beam + poa_sky[t] * svf[i] + poa_gnd[t]
        per_pt[i] = acc
    return per_pt


@njit(parallel=True, cache=True)
def hourly_total(hor, svf, w, sun_el, sun_az, poa_dir, poa_sky, poa_gnd):
    """Same shading model, summed over points with weights w (panel m^2) per hour -> W. Parallel over hours."""
    n, naz = hor.shape
    nt = sun_el.shape[0]
    out = np.zeros(nt)
    for t in prange(nt):
        if poa_dir[t] + poa_sky[t] <= 0: continue
        f = sun_az[t] / (2 * np.pi) * naz
        k0 = int(f) % naz
        k1 = (k0 + 1) % naz
        fr = f - int(f)
        s = 0.0
        for i in range(n):
            hz = hor[i, k0] + fr * (hor[i, k1] - hor[i, k0])
            s += w[i] * ((poa_dir[t] if sun_el[t] > hz else 0.0) + poa_sky[t] * svf[i] + poa_gnd[t])
        out[t] = s
    return out


def weather():
    tmy = pd.read_csv(DATA / "tmy.csv", index_col=0)
    tmy.index = pd.to_datetime(tmy.index, utc=True)
    tmy = tmy.sort_index()
    # PVGIS TMY stitches months from different years; relabel to one calendar year for day-of-year work
    tmy.index = pd.DatetimeIndex([t.replace(year=2026) for t in tmy.index])
    times = tmy.index + pd.Timedelta(minutes=30)          # hourly averages -> sun at mid-hour
    sp = pvlib.solarposition.get_solarposition(times, LAT0, LON0)
    poa = pvlib.irradiance.get_total_irradiance(
        TILT, 180, sp.apparent_zenith.values, sp.azimuth.values, tmy.dni.values, tmy.ghi.values, tmy.dhi.values,
        dni_extra=pvlib.irradiance.get_extra_radiation(times).values,
        airmass=pvlib.atmosphere.get_relative_airmass(sp.apparent_zenith.values), model="perez")
    return tmy, sp, poa


def run():
    buildings, roads, green, campus_polys = load_geometry()
    H, (x0, y0), pts, owner = rasterize(buildings)
    hor = horizons(H, x0, y0, RES, pts.astype(np.float64), NAZ, 400.0)
    svf = (np.cos(hor.astype(np.float64)) ** 2).mean(1)      # isotropic sky view of a horizontal surface
    tmy, sp, poa = weather()
    el = np.radians(sp.apparent_elevation.values)
    az = np.radians(sp.azimuth.values)
    pd_, ps, pg = [np.nan_to_num(np.asarray(poa[k], float)) for k in ("poa_direct", "poa_sky_diffuse", "poa_ground_diffuse")]
    campus = np.array([buildings[o]["campus"] for o in owner])
    w = np.where(campus, SPACING ** 2 * USABLE, 0.0)         # m^2 of panel per sample point, campus roofs only
    per_pt = shaded_poa(hor, svf, el, az, pd_, ps, pg) / 1000          # kWh/m^2/yr POA, shaded
    campus_hourly_poa = hourly_total(hor, svf, w, el, az, pd_, ps, pg)   # W summed over panel m^2
    # temperature derate on the campus-average irradiance (Faiman model), then system losses
    area = w.sum()
    tcell = pvlib.temperature.faiman(campus_hourly_poa / area, tmy.temp_air.values, tmy.wind_speed.values)
    kwh_per_eta = campus_hourly_poa * (1 + GAMMA * (tcell - 25)) * PR / 1000     # kWh per unit efficiency, hourly
    unshaded = (pd_ + ps + pg).sum() / 1000

    cnt = np.bincount(owner, minlength=len(buildings))
    poa_b = np.bincount(owner, per_pt, len(buildings)) / np.maximum(cnt, 1)
    svf_b = np.bincount(owner, svf, len(buildings)) / np.maximum(cnt, 1)
    for i, b in enumerate(buildings):
        b.update(poa=float(poa_b[i]), svf=float(svf_b[i]), n=int(cnt[i]), panel_m2=float(cnt[i] * SPACING ** 2 * USABLE))
    return dict(buildings=buildings, roads=roads, green=green, campus_polys=campus_polys, H=H, origin=(x0, y0),
                pts=pts, owner=owner, hor=hor, svf=svf, per_pt=per_pt, tmy=tmy, sun=sp, kwh_per_eta=kwh_per_eta,
                tcell=tcell, panel_area=float(area), unshaded_poa=float(unshaded), ghi=float(tmy.ghi.sum() / 1000))


def verify(c):
    # 1) analytic horizon: a 20 m wall 40 m north of a point at ground level -> atan(20/40) = 26.6 deg
    H = np.zeros((200, 200), np.float32); H[140:150, :] = 20
    got = np.degrees(horizons(H, 0.0, 0.0, 1.0, np.array([[100.0, 100.0, 0.0]]), NAZ, 300.0)[0, 0])
    print(f"  horizon test: {got:.1f} deg vs analytic {np.degrees(np.arctan2(20, 40.5)):.1f} deg")
    assert abs(got - np.degrees(np.arctan2(20, 40.5))) < 1.5
    # 2) resource vs NREL: Atlanta GHI ~1,650-1,800 kWh/m^2/yr; fixed low-tilt c-Si ~1,300-1,450 kWh/kWp (PVWatts)
    spec = c["unshaded_poa"] * (1 + GAMMA * (np.average(c["tcell"], weights=c["kwh_per_eta"] + 1e-9) - 25)) * PR
    print(f"  TMY GHI {c['ghi']:.0f} kWh/m2/yr, unshaded POA@{TILT:.0f}deg {c['unshaded_poa']:.0f}, specific yield ~{spec:.0f} kWh/kWp")
    assert 1500 < c["ghi"] < 1900 and 1150 < spec < 1600
    loss = 1 - c["per_pt"].mean() / c["unshaded_poa"]
    print(f"  mean roof shading+sky-view loss {100 * loss:.1f}%  over {len(c['pts'])} roof points, "
          f"{sum(b['campus'] for b in c['buildings'])} campus buildings, panel area {c['panel_area'] / 1e3:.0f}k m2")
    assert 0 < loss < 0.3


if __name__ == "__main__":
    import time
    t = time.time(); c = run(); print(f"  campus solar model {time.time() - t:.1f} s")
    verify(c)

"""Export the Materials Project semiconductor set from the Band Gap GNN project to plain JSON.

Runs under the GNN project's own venv (it has pymatgen, which unpickles the Structure objects):
    "../../Band Gap AI Stack/Multimodal GNN/.venv/Scripts/python.exe" extract_materials.py
Read-only on the source project; writes data/materials.json here.
"""
import json
from pathlib import Path

import pandas as pd
import yaml
import pymatgen.core
from pymatgen.analysis.hhi import HHIModel

SRC = Path(r"C:\Users\danma\Documents\Dan\Projects\Band Gap AI Stack\Multimodal GNN\data\materials_dataset_v2_cleaned.pkl")
OUT = Path(__file__).parent / "data" / "materials.json"

df = pd.read_pickle(SRC)
hhi = HHIModel()
colors = yaml.safe_load((Path(pymatgen.core.__file__).parent / "ElementColorSchemes.yaml").read_text())["VESTA"]

rows = []
for r in df.itertuples():
    s = r.structure
    els = sorted({sp.symbol for sp in s.composition.elements})
    hp, hr = hhi.get_hhi(s.composition)
    row = dict(id=r.mp_id, f=r.formula, gap=round(r.band_gap, 4), ef=round(r.formation_energy, 4),
               rho=round(r.density, 3), cs=r.crystal_system, els=els, n=len(s),
               dchi=round(max(e.X for e in s.composition.elements) - min(e.X for e in s.composition.elements), 3),
               hhi_p=round(hp) if hp else None, hhi_r=round(hr) if hr else None)  # None = element not in HHI table
    # ponytail: structures only for PV-relevant gaps and small cells, keeps the file a few MB
    if r.band_gap <= 2.5 and len(s) <= 80:
        row["lat"] = [[round(x, 4) for x in v] for v in s.lattice.matrix]
        row["sp"] = [site.specie.symbol for site in s]
        row["xyz"] = [[round(x, 3) for x in site.coords] for site in s]
    rows.append(row)

used = {e for r in rows for e in r["els"]}
OUT.parent.mkdir(exist_ok=True)
OUT.write_text(json.dumps({"materials": rows, "colors": {e: colors[e] for e in used if e in colors}}))
print(len(rows), "materials,", sum("xyz" in r for r in rows), "with structures ->", OUT, f"{OUT.stat().st_size/1e6:.1f} MB")

"""Download and cache every external input of the twin into data/. Safe to re-run: cached files are skipped.

Sources: bus.gatech.edu (routes/stops), OpenStreetMap Overpass (buildings/roads/green),
PVGIS TMY (NREL NSRDB-based typical year), Open-Meteo elevation (Copernicus GLO-90 DEM),
local HTEM thin-film library data.  Materials come from extract_materials.py (separate venv).
"""
import json
import os
import time
from pathlib import Path

import numpy as np
import pandas as pd
import pvlib
import requests

DATA = Path(__file__).parent / "data"
DATA.mkdir(exist_ok=True)
LAT0, LON0 = 33.7765, -84.3990            # local origin: roughly Tech Tower / campus centre
ROUTES = {20: "Red", 21: "Blue", 17: "Green", 29: "Gold", 28: "Clough"}  # campus routes (Emory/weekend routes leave campus)
UA = {"User-Agent": "SimulationLab-GT-student-project/1.0"}
HTEM_CSV = Path(r"C:\Users\danma\Documents\Dan\Projects\HTEM DB\Data\HTEM Libraries\Datasets\htem_220library_dataset_cleaned.csv")


def decode_polyline(s):
    """Google encoded-polyline algorithm -> [(lat, lon)]."""
    pts, i, lat, lon = [], 0, 0, 0
    while i < len(s):
        for which in (0, 1):
            shift = result = 0
            while True:
                b = ord(s[i]) - 63; i += 1
                result |= (b & 0x1F) << shift; shift += 5
                if b < 0x20: break
            d = ~(result >> 1) if result & 1 else result >> 1
            if which == 0: lat += d
            else: lon += d
        pts.append((lat / 1e5, lon / 1e5))
    return pts


def cached(name, fn):
    p = DATA / name
    if not p.exists():
        print("fetching", name)
        fn(p)
    return p


def routes(p):
    # needs GT_BUS_API_KEY: the public map key bus.gatech.edu's own web page sends (see its network requests)
    r = requests.get("https://bus.gatech.edu/Services/JSONPRelay.svc/GetRoutesForMapWithScheduleWithEncodedLine",
                     params={"apiKey": os.environ["GT_BUS_API_KEY"]}, timeout=30, headers=UA)
    out = []
    for rt in r.json():
        if rt["RouteID"] not in ROUTES: continue
        out.append(dict(id=rt["RouteID"], name=rt["Description"], color=rt["MapLineColor"],
                        line=decode_polyline(rt["EncodedPolyline"]),
                        stops=[dict(id=s["RouteStopID"], name=s["Name"], lat=s["Latitude"], lon=s["Longitude"],
                                    dwell=s["SecondsAtStop"], to_next=s["SecondsToNextStop"]) for s in rt["Stops"]]))
    p.write_text(json.dumps(out), encoding="utf-8")


def bbox(margin_deg=0.0025):
    lat, lon = np.array([pt for rt in json.loads(cached("routes.json", routes).read_text(encoding="utf-8")) for pt in rt["line"]]).T
    return lat.min() - margin_deg, lon.min() - margin_deg, lat.max() + margin_deg, lon.max() + margin_deg


def osm(p):
    s, w, n, e = bbox()
    q = f"""[out:json][timeout:120];(
      way["building"]({s},{w},{n},{e}); relation["building"]["type"="multipolygon"]({s},{w},{n},{e});
      way["highway"]({s},{w},{n},{e});
      way["leisure"~"park|pitch|garden|stadium|track"]({s},{w},{n},{e});
      way["landuse"~"grass|recreation_ground|forest"]({s},{w},{n},{e}); way["natural"~"wood|scrub"]({s},{w},{n},{e});
    );out geom;"""
    r = requests.post("https://overpass-api.de/api/interpreter", data={"data": q}, timeout=180, headers=UA)
    r.raise_for_status()
    p.write_text(r.text, encoding="utf-8")


def campus(p):
    b = ",".join(map(str, bbox()))
    q = f"""[out:json][timeout:60];(way["amenity"="university"]({b}); relation["amenity"="university"]({b}););out geom;"""
    r = requests.post("https://overpass-api.de/api/interpreter", data={"data": q}, timeout=120, headers=UA)
    r.raise_for_status()
    p.write_text(r.text, encoding="utf-8")


def tmy(p):
    res = pvlib.iotools.get_pvgis_tmy(LAT0, LON0, map_variables=True, timeout=60)
    res[0].to_csv(p)  # first element is the hourly frame in every pvlib version


def elevation(p):
    # route polylines resampled every ~25 m; Open-Meteo takes 100 points per call
    pts = []
    for rt in json.loads(cached("routes.json", routes).read_text(encoding="utf-8")):
        ll = np.array(rt["line"])
        xy = to_xy(ll[:, 0], ll[:, 1])
        s = np.r_[0, np.cumsum(np.hypot(*np.diff(xy, axis=1)))]
        si = np.arange(0, s[-1], 25.0)
        pts += [(rt["id"], a, b, c) for a, b, c in zip(si, np.interp(si, s, ll[:, 0]), np.interp(si, s, ll[:, 1]))]
    elev = []
    for k in range(0, len(pts), 100):
        chunk = pts[k:k + 100]
        while "elevation" not in (j := requests.get("https://api.open-meteo.com/v1/elevation", timeout=30, headers=UA, params={
                "latitude": ",".join(f"{c[2]:.6f}" for c in chunk), "longitude": ",".join(f"{c[3]:.6f}" for c in chunk)}).json()):
            print("  elevation API:", j.get("reason"), "- waiting 65 s"); time.sleep(65)  # free tier is rate limited per minute
        elev += j["elevation"]
    pd.DataFrame(pts, columns=["route", "s", "lat", "lon"]).assign(elev=elev).to_csv(p, index=False)


def htem(p):
    df = pd.read_csv(HTEM_CSV, usecols=["sample_id", "library_id", "deposition_sample_time_min", "deposition_initial_temp_c", "thickness"])
    df.to_csv(p, index=False)


def to_xy(lat, lon):
    """Local tangent-plane metres (east, north). ponytail: equirectangular, <0.1% error over a 3 km campus."""
    return np.array([(np.asarray(lon) - LON0) * 111320 * np.cos(np.radians(LAT0)), (np.asarray(lat) - LAT0) * 110574])


if __name__ == "__main__":
    for name, fn in [("routes.json", routes), ("osm.json", osm), ("campus.json", campus), ("tmy.csv", tmy), ("elevation.csv", elevation), ("htem_films.csv", htem)]:
        print(cached(name, fn), f"{(DATA / name).stat().st_size / 1e3:.0f} kB")

"""The Stinger fleet, electrified on paper.

Real GPS polls (Stinger Delay project, March 2026) -> map-matched to the official route polylines ->
distance-along-route time series -> PCHIP drive cycles (1 Hz speed/accel) -> longitudinal-dynamics battery
power (inertia, rolling, aero, grade from a DEM, regen, HVAC) -> kWh/mi, loop energies, service profile ->
8760-h fleet demand -> SimPy charging-infrastructure design study.
"""
import json
from pathlib import Path

import numpy as np
import pandas as pd
import simpy
from scipy.interpolate import PchipInterpolator

from fetch import DATA, to_xy

RAW = Path(r"C:\Users\danma\Documents\Dan\Projects\Stinger Proj\stinger-delay\Analysis\model_comparison\raw_data")
SOURCES = {20: ("red/red_line_data.csv", "snapshot_time_utc"),           # UTC with offset
           29: ("gold/gps_only_scraper/route_29_data.csv", "snapshot_time"),  # UTC with offset
           17: ("green/green_bus_*.csv", "snapshot_time"),                 # naive local
           28: ("clough/clough_bus_data.csv", "snapshot_time")}            # "... EST" local
TZ = "America/New_York"
MAX_GAP, OFF_ROUTE = 120.0, 60.0     # s between polls to still interpolate; m from polyline to count as in service
V_MAX = 16.0                          # m/s (36 mph): campus streets are 25-35 mph; faster "moves" are GPS jumps

# 35-ft battery-electric transit bus (Proterra/Gillig/New Flyer class). All assumptions, all tunable.
MASS = 13_000 + 15 * 75     # curb kg + average load (15 riders)
CRR, CDA, RHO, G = 0.009, 5.0, 1.2, 9.81
ROT = 1.05                  # rotating-mass inertia factor
ETA_DRIVE, ETA_REGEN = 0.85, 0.60     # battery->wheel, wheel->battery (braking energy recovered)
# Opportunity-charger sites, in priority order. A route uses the first one it passes. The Kendeda Building is GT's
# net-positive Living Building and a shared Red/Gold/Green stop, so three routes queue for the same chargers there.
SITES = ["Kendeda Building", "5th Street Bridge EB", "North Avenue Apts - To West Campus"]
LAYOVER = 300.0             # s scheduled recovery time at the hub each loop
P_CHARGER = 300.0           # kW pantograph opportunity charger


def aux_kw(temp_c):
    """Base auxiliaries + HVAC. ponytail: linear in |T - 20 C|, capped; a heat-pump COP(T) model is the upgrade."""
    return 4.0 + np.minimum(0.9 * np.abs(np.asarray(temp_c) - 20.0), 25.0)


def battery_kw(v, a, grade):
    f = MASS * ROT * a + MASS * G * CRR * np.cos(np.arctan(grade)) + MASS * G * np.sin(np.arctan(grade)) + 0.5 * RHO * CDA * v ** 2
    pw = f * v / 1000
    return np.where(pw > 0, pw / ETA_DRIVE, pw * ETA_REGEN)


# ---------------------------------------------------------------- geometry
def route_geoms():
    routes = json.loads((DATA / "routes.json").read_text(encoding="utf-8"))
    elev = pd.read_csv(DATA / "elevation.csv")
    out = {}
    for rt in routes:
        xy = to_xy(*np.array(rt["line"]).T).T
        if np.hypot(*(xy[0] - xy[-1])) > 1: xy = np.vstack([xy, xy[:1]])        # routes are loops
        seg = np.diff(xy, axis=0)
        slen = np.hypot(*seg.T)
        keep = slen > 0.01
        a, seg, slen = xy[:-1][keep], seg[keep], slen[keep]
        s0 = np.r_[0, np.cumsum(slen)[:-1]]
        e = elev[elev.route == rt["id"]]
        z = np.convolve(np.pad(e.elev.values, 3, mode="wrap"), np.ones(7) / 7, "valid")   # 175 m smoothing of a 90 m DEM
        grade = np.clip(np.gradient(z, e.s.values), -0.12, 0.12)
        stops_xy = to_xy(np.array([s["lat"] for s in rt["stops"]]), np.array([s["lon"] for s in rt["stops"]])).T
        out[rt["id"]] = dict(id=rt["id"], name=rt["name"], color=rt["color"], xy=xy, a=a, d=seg, len=slen, s0=s0,
                             L=float(slen.sum()), grade_s=e.s.values, grade=grade, elev=z, stops=rt["stops"],
                             stop_s=project(stops_xy, a, seg, slen, s0)[0])
    return out


def project(p, a, d, slen, s0):
    """Nearest-point projection of points p (N,2) onto every segment -> (s along route, distance) of the best one."""
    t = np.clip(((p[:, None, :] - a[None]) * d[None]).sum(-1) / slen ** 2, 0, 1)
    dist = np.hypot(*(p[:, None, :] - (a[None] + t[..., None] * d[None])).transpose(2, 0, 1))
    k = dist.argmin(1)
    return s0[k] + t[np.arange(len(p)), k] * slen[k], dist[np.arange(len(p)), k]


def grade_at(g, s):
    return np.interp(np.mod(s, g["L"]), g["grade_s"], g["grade"], period=g["L"])


# ---------------------------------------------------------------- GPS -> distance along route
def load_gps(rid):
    pat, col = SOURCES[rid]
    df = pd.concat(pd.read_csv(f) for f in RAW.glob(pat))
    ts = df[col].astype(str).str.replace(r" E[SD]T$", "", regex=True)
    t = pd.to_datetime(ts, format="mixed")
    t = t.dt.tz_convert(TZ) if t.dt.tz is not None else t.dt.tz_localize(TZ, ambiguous="NaT", nonexistent="NaT")
    df = df.assign(t=t).dropna(subset=["t", "vehicle_lat"]).drop_duplicates(["t", "vehicle_id"]).sort_values(["vehicle_id", "t"])
    return df[["t", "vehicle_id", "vehicle_lat", "vehicle_lon", "vehicle_speed"]].reset_index(drop=True)


def track(g, df):
    """Tracking map-matcher: each poll may only move forward along the loop by what the elapsed time allows.
    Handles the out-and-back streets a pure nearest-point match gets wrong. Returns df + s, dist, D (cumulative m)."""
    p = to_xy(df.vehicle_lat.values, df.vehicle_lon.values).T
    t = (df.t - df.t.iloc[0]).dt.total_seconds().values
    veh = df.vehicle_id.values
    L = g["L"]
    tt = np.clip(((p[:, None, :] - g["a"][None]) * g["d"][None]).sum(-1) / g["len"] ** 2, 0, 1)
    dist = np.hypot(*(p[:, None, :] - (g["a"][None] + tt[..., None] * g["d"][None])).transpose(2, 0, 1))
    sall = g["s0"][None] + tt * g["len"][None]
    s, dmin, D, seg_id = np.zeros(len(p)), np.zeros(len(p)), np.zeros(len(p)), np.zeros(len(p), int)
    prev, sid = None, -1
    for i in range(len(p)):
        cont = prev is not None and veh[i] == veh[prev] and t[i] - t[prev] < MAX_GAP and dmin[prev] < OFF_ROUTE
        di = dist[i]
        if cont:
            ahead = np.mod(sall[i] - s[prev], L)
            di = np.where((ahead < V_MAX * (t[i] - t[prev]) + 30) | (ahead > L - 40), di, np.inf)
        k = di.argmin()
        if not np.isfinite(di[k]):          # impossible jump: restart tracking from a global match
            cont, k = False, dist[i].argmin()
        s[i], dmin[i] = sall[i, k], dist[i, k]
        if cont and dmin[i] < OFF_ROUTE:
            step = np.mod(s[i] - s[prev], L)
            D[i] = D[prev] + (step if step < L - 40 else 0.0)      # small backward jitter -> standing still
            seg_id[i] = sid
        else:
            sid += 1
            D[i], seg_id[i] = s[i], sid
        prev = i
    return df.assign(s=s, dist=dmin, D=D, seg=seg_id, tsec=t)


def drive_cycles(g, m):
    """1 Hz speed/accel/grade per continuous in-service segment (PCHIP keeps distance monotone -> v >= 0)."""
    cycles = []
    for _, sg in m[m.dist < OFF_ROUTE].groupby("seg"):
        if len(sg) < 8 or sg.tsec.iloc[-1] - sg.tsec.iloc[0] < 300: continue
        T, D = sg.tsec.values, np.maximum.accumulate(sg.D.values)
        T, idx = np.unique(T, return_index=True)
        f = PchipInterpolator(T, D[idx])
        tt = np.arange(T[0], T[-1], 1.0)
        v = np.clip(f(tt, 1), 0, V_MAX)
        a = np.gradient(v)
        cycles.append(dict(vehicle=sg.vehicle_id.iloc[0], t=sg.t.iloc[0] + pd.to_timedelta(tt - T[0], "s"),
                           tsec=tt, D=f(tt), v=v, a=a, grade=grade_at(g, f(tt)), api_speed=sg.vehicle_speed.values,
                           poll_v=np.gradient(D[idx], T), poll_t=T))
    return cycles


# ---------------------------------------------------------------- per-route energy statistics
def analyse_route(g, m, cycles):
    km = sum(c["D"][-1] - c["D"][0] for c in cycles) / 1000
    hours = sum(len(c["tsec"]) for c in cycles) / 3600
    p = np.concatenate([battery_kw(c["v"], c["a"], c["grade"]) for c in cycles])
    v = np.concatenate([c["v"] for c in cycles])
    trac_kwh = p.sum() / 3600
    parts = {  # energy at the wheel by physical mechanism, kWh (positive = spent)
        "inertia (accelerating)": sum((MASS * ROT * np.clip(c["a"], 0, None) * c["v"]).sum() for c in cycles) / 3.6e6,
        "rolling resistance": sum((MASS * G * CRR * c["v"]).sum() for c in cycles) / 3.6e6,
        "aerodynamic drag": sum((0.5 * RHO * CDA * c["v"] ** 3).sum() for c in cycles) / 3.6e6,
        "climbing": sum((MASS * G * np.clip(c["grade"], 0, None) * c["v"]).sum() for c in cycles) / 3.6e6,
        "regen recovered": -sum(np.clip(-battery_kw(c["v"], c["a"], c["grade"]), 0, None).sum() for c in cycles) / 3600,
    }
    # loops: time and traction energy to cover one route length L
    loops = []
    for c in cycles:
        pc = np.cumsum(battery_kw(c["v"], c["a"], c["grade"])) / 3600
        marks = np.arange(c["D"][0] + g["L"], c["D"][-1], g["L"])
        idx = np.searchsorted(c["D"], np.r_[c["D"][0], marks])
        loops += [(c["tsec"][j] - c["tsec"][i], pc[j] - pc[i]) for i, j in zip(idx[:-1], idx[1:])]
    loops = np.array(loops) if loops else np.zeros((0, 2))
    # service profile: mean buses on route per local hour, over hours the scraper was actually running
    mm = m[m.dist < OFF_ROUTE].assign(date=lambda d: d.t.dt.date, hour=lambda d: d.t.dt.hour)
    active = m.assign(date=m.t.dt.date, hour=m.t.dt.hour).groupby(["date", "hour"]).size().index
    per = mm.groupby(["date", "hour"]).vehicle_id.nunique().reindex(active, fill_value=0)
    profile = per.groupby(level="hour").mean().reindex(range(24), fill_value=0.0)
    # poll-level speed check: map-matched speed vs the speed field the API reports (mph)
    pv = np.concatenate([c["poll_v"] for c in cycles]) * 2.23694
    av = np.concatenate([c["api_speed"][:len(c["poll_v"])] for c in cycles])
    ok = np.isfinite(av) & np.isfinite(pv)
    return dict(km=km, hours=hours, mean_mph=float(v.mean() * 2.23694), trac_kwh=trac_kwh,
                trac_kwh_per_km=trac_kwh / km, trac_kw_per_bus=trac_kwh / hours, stopped=float((v < 0.3).mean()),
                parts=parts, loops=loops, profile=profile.values, n_polls=len(m),
                speed_corr=float(np.corrcoef(pv[ok], av[ok])[0, 1]) if ok.sum() > 10 else float("nan"))


def replay_day(g, m, cycles):
    """Busiest local weekday: per vehicle, distance along route on a 10 s grid (NaN = not in service) + traction kWh."""
    days = m[m.dist < OFF_ROUTE].assign(date=m.t.dt.date)
    days = days[pd.to_datetime(days.date).dt.dayofweek < 5].groupby("date").size()
    day = days.idxmax()
    grid = np.arange(0, 86400, 10.0)
    veh = {}
    for c in cycles:
        if c["t"][0].date() != day: continue
        sec = (c["t"] - c["t"][0].normalize()).total_seconds().values
        inside = (grid >= sec[0]) & (grid <= sec[-1])
        rec = veh.setdefault(c["vehicle"], dict(D=np.full(len(grid), np.nan), E=np.full(len(grid), np.nan), off=0.0))
        e = rec["off"] + np.cumsum(battery_kw(c["v"], c["a"], c["grade"])) / 3600   # traction kWh since pull-out
        rec["off"] = e[-1]
        rec["D"][inside] = np.interp(grid[inside], sec, c["D"])
        rec["E"][inside] = np.interp(grid[inside], sec, e)
    return str(day), grid, veh


# ---------------------------------------------------------------- 8760-h demand
def annual_demand(stats, tmy):
    """Hourly fleet battery energy (kWh) for a TMY year on the campus-route schedule.
    ponytail: semester weekdays full service, summer weekdays half, weekends/holidays none (weekend routes excluded)."""
    t = tmy.index.tz_convert(TZ)
    doy, dow, hr = t.dayofyear.values, t.dayofweek.values, t.hour.values
    semester = ((doy >= 12) & (doy <= 121)) | ((doy >= 229) & (doy <= 344))       # ~Jan 12-May 1, Aug 17-Dec 10
    factor = np.where(dow >= 5, 0.0, np.where(semester, 1.0, 0.5))
    buses = sum(s["profile"][hr] for s in stats.values())
    trac = sum(s["profile"][hr] * s["trac_kw_per_bus"] for s in stats.values())
    drive = factor * (trac + buses * aux_kw(tmy.temp_air.values))
    # overnight depot charging: each day's energy delivered 00:00-06:00 local the same morning
    day_id = t.normalize()
    daily = pd.Series(drive, index=t).groupby(day_id).transform("sum").values
    night = np.where(hr < 6, daily / 6, 0.0)
    return dict(drive=drive, night=night, factor=factor, buses=buses, local=t)


# ---------------------------------------------------------------- SimPy charging design study
def fleet_day(stats, n_chargers, batt_kwh, temps, seed, start_h=7, end_h=23):
    """One service day. Each bus loops its route; at its charging site it may take an opportunity charge during
    layover, but only if a charger frees up before the layover ends (reneging: the schedule wins over the battery)."""
    rng = np.random.default_rng(seed)
    env = simpy.Environment()
    hubs = {site: simpy.Resource(env, n_chargers) for site in SITES} if n_chargers else {}
    out = dict(min_soc=[], stranded=0, charges=0, skipped=0, busy=0.0, arrivals=0)

    def bus(st, k, n):
        soc = batt_kwh
        yield env.timeout(start_h * 3600 + k * np.median(st["loops"][:, 0]) / n)      # staggered departures
        lo = soc
        while env.now < end_h * 3600:
            dur, trac = st["loops"][rng.integers(len(st["loops"]))]
            yield env.timeout(dur)
            soc -= trac + aux_kw(temps[int(env.now // 3600) % 24]) * dur / 3600
            lo = min(lo, soc)
            if soc < 0:
                out["stranded"] += 1; break
            if st["site"] in hubs:
                out["arrivals"] += 1
                arrive = env.now
                with hubs[st["site"]].request() as req:
                    got = yield req | env.timeout(LAYOVER)
                    if req in got:
                        t_charge = LAYOVER - (env.now - arrive)
                        soc = min(batt_kwh, soc + P_CHARGER * t_charge / 3600)
                        out["charges"] += 1; out["busy"] += t_charge
                        yield env.timeout(t_charge)
                    else:
                        out["skipped"] += 1
            else:
                yield env.timeout(LAYOVER)
        out["min_soc"].append(lo / batt_kwh)

    for st in stats.values():
        n = max(1, int(round(st["profile"][start_h:end_h].max())))
        for k in range(n):
            env.process(bus(st, k, n))
    env.run(until=(end_h + 2) * 3600)
    used = len({st["site"] for st in stats.values()} & set(hubs))
    out["util"] = out["busy"] / (used * n_chargers * (end_h - start_h + 2) * 3600) if n_chargers else 0.0
    return out


def design_study(stats, days, reps=30):
    rows = []
    for dname, temps in days.items():
        for batt in (100, 150, 200, 300, 450):
            for k in (0, 1, 2):
                res = [fleet_day(stats, k, batt, temps, seed) for seed in range(reps)]     # same seeds = CRN across designs
                ms = np.array([min(r["min_soc"]) for r in res])
                rows.append(dict(day=dname, batt=batt, chargers=k, p_fail=float(np.mean(ms < 0.10)),
                                 min_soc=float(ms.mean()), min_soc_ci=float(1.96 * ms.std(ddof=1) / np.sqrt(reps)),
                                 util=float(np.mean([r["util"] for r in res])),
                                 skipped=float(np.mean([r["skipped"] / max(r["arrivals"], 1) for r in res]))))
    return pd.DataFrame(rows)


def run(tmy):
    geoms = route_geoms()
    site_xy = {name: next(to_xy(s["lat"], s["lon"]) for g in geoms.values() for s in g["stops"] if s["name"] == name) for name in SITES}

    def site_of(g):
        return next((n for n in SITES if project(np.array([site_xy[n]]), g["a"], g["d"], g["len"], g["s0"])[1][0] < 60), None)
    stats, replays, matched = {}, {}, {}
    for rid in SOURCES:
        g = geoms[rid]
        m = track(g, load_gps(rid))
        cyc = drive_cycles(g, m)
        st = analyse_route(g, m, cyc)
        st["site"] = site_of(g)
        stats[rid], replays[rid], matched[rid] = st, replay_day(g, m, cyc), (m, cyc)
    # Blue has no GPS logs: it is Red's counter-rotating twin, so borrow Red's per-km energy and fleet size
    red = stats[20]
    b = geoms[21]
    stats[21] = {**red, "site": site_of(b),
                 "loops": red["loops"] * [b["L"] / geoms[20]["L"], b["L"] / geoms[20]["L"]], "synthetic": True}
    return dict(geoms=geoms, stats=stats, replays=replays, matched=matched, demand=annual_demand(stats, tmy), sites=site_xy)


def verify(f):
    for rid, st in f["stats"].items():
        if st.get("synthetic"): continue
        kwh_mi = (st["trac_kwh"] + aux_kw(18) * st["hours"]) / (st["km"] / 1.609)
        print(f"  {f['geoms'][rid]['name']:7s} {st['km']:6.0f} km in service, mean {st['mean_mph']:4.1f} mph, "
              f"stopped {100 * st['stopped']:3.0f}%, {kwh_mi:.2f} kWh/mi at 18C, speed corr vs API {st['speed_corr']:.2f}")
        assert 1.2 < kwh_mi < 3.5, "outside the NREL in-service BEB range (~1.7-2.5 kWh/mi, wider for slow campus duty)"
    # physics check: around a closed loop the DEM elevations must close (net climb ~ 0)
    for g in f["geoms"].values():
        net = np.trapezoid(g["grade"], g["grade_s"])
        assert abs(net) < 3.0, f"{g['name']} loop does not close: {net:.1f} m"
    # DES check (Little's-law style): charger busy time = charges * mean charge time
    r = fleet_day(f["stats"], 1, 200, np.full(24, 20.0), 0)
    print(f"  DES check: {r['charges']} charges, utilisation {100 * r['util']:.0f}%, skipped {r['skipped']}")
    assert r["charges"] * LAYOVER >= r["busy"] > 0