# train_model.py
import numpy as np
import pandas as pd
from sklearn.ensemble import RandomForestRegressor
import joblib, os

def generate_data(n=3000):
    np.random.seed(42)
    dist = np.random.exponential(400, n)
    speed = np.random.uniform(5, 15, n)
    tod = np.random.uniform(0, 86400, n)
    tod_sin = np.sin(2*np.pi * tod/86400)
    tod_cos = np.cos(2*np.pi * tod/86400)
    dow = np.random.randint(0, 7, n)
    eta = dist/(speed*0.277) + np.random.randint(10,120,n)
    df = pd.DataFrame({
        "distance_to_stop_m": dist,
        "speed_mps": speed,
        "tod_sin": tod_sin,
        "tod_cos": tod_cos,
        "day_of_week": dow,
        "eta": eta
    })
    return df

def train():
    os.makedirs("models", exist_ok=True)
    df = generate_data()
    X = df[["distance_to_stop_m","speed_mps","tod_sin","tod_cos","day_of_week"]]
    y = df["eta"]
    model = RandomForestRegressor(n_estimators=40)
    model.fit(X, y)
    joblib.dump(model, "models/model.joblib")
    print("Model saved to models/model.joblib")

if __name__ == "__main__":
    train()