# app.py — Single-file Transit Tracker API for Hugging Face (No extra files required)
from fastapi import FastAPI, HTTPException
from pydantic import BaseModel
from typing import Dict, Any
import numpy as np
import joblib
import os
import math
import time

app = FastAPI(title="Transit Tracker API - Single File Version")

# -----------------------------
# In-memory Vehicle Locations
# -----------------------------
VEHICLES: Dict[str, Dict[str, Any]] = {}

# -----------------------------
# Train a lightweight ETA model
# -----------------------------
def train_and_save_model():
    from sklearn.ensemble import RandomForestRegressor

    np.random.seed(42)
    N = 3000

    dist = np.random.exponential(400, N)
    speed = np.random.uniform(5, 15, N)
    tod = np.random.uniform(0, 86400, N)
    dow = np.random.randint(0, 7, N)

    tod_sin = np.sin(2*np.pi * tod/86400)
    tod_cos = np.cos(2*np.pi * tod/86400)

    eta = dist/(speed*0.277) + np.random.randint(10, 120, N)

    X = np.column_stack([dist, speed, tod_sin, tod_cos, dow])
    y = eta

    model = RandomForestRegressor(n_estimators=30)
    model.fit(X, y)
    joblib.dump(model, "model.joblib")
    print("MODEL TRAINED AND SAVED ✔")

# -----------------------------
# Load or Train Model
# -----------------------------
def load_model():
    if not os.path.exists("model.joblib"):
        train_and_save_model()
    return joblib.load("model.joblib")

MODEL = load_model()

# -----------------------------
# JSON Schemas
# -----------------------------
class VehicleUpdate(BaseModel):
    vehicle_id: str
    lat: float
    lon: float
    speed: float = 0.0

class PredictRequest(BaseModel):
    distance_to_stop_m: float
    speed_mps: float = 4.0


# -----------------------------
# Utils
# -----------------------------
def haversine(lat1, lon1, lat2, lon2):
    R = 6371000
    p1, p2 = math.radians(lat1), math.radians(lat2)
    dphi = math.radians(lat2 - lat1)
    dl = math.radians(lon2 - lon1)
    a = math.sin(dphi / 2) ** 2 + math.cos(p1) * math.cos(p2) * math.sin(dl / 2) ** 2
    return R * 2 * math.atan2(math.sqrt(a), math.sqrt(1 - a))


# -----------------------------
# ROOT
# -----------------------------
@app.get("/")
def root():
    return {"message": "Transit Tracker API Running", "vehicles": len(VEHICLES)}


# -----------------------------
# LIVE VEHICLE LOCATION APIs
# -----------------------------
@app.post("/vehicle/update")
def update_vehicle(v: VehicleUpdate):
    VEHICLES[v.vehicle_id] = {
        "lat": v.lat,
        "lon": v.lon,
        "speed": v.speed,
        "timestamp": time.time()
    }
    return {"success": True, "vehicle": VEHICLES[v.vehicle_id]}


@app.get("/vehicles")
def get_all_vehicles():
    return VEHICLES


# -----------------------------
# ETA PREDICTION
# -----------------------------
@app.post("/predict")
def predict_eta(req: PredictRequest):
    dist = req.distance_to_stop_m
    if dist < 0:
        raise HTTPException(400, "distance cannot be negative")

    t = time.gmtime()
    seconds = t.tm_hour * 3600 + t.tm_min * 60 + t.tm_sec
    tod_sin = np.sin(2*np.pi * seconds/86400)
    tod_cos = np.cos(2*np.pi * seconds/86400)
    dow = t.tm_wday

    X = np.array([[dist, req.speed_mps, tod_sin, tod_cos, dow]])
    eta_sec = float(MODEL.predict(X)[0])

    return {"eta_seconds": eta_sec, "eta_minutes": eta_sec / 60}