import gradio as gr
import pandas as pd
from datasets import load_dataset
from sklearn.model_selection import train_test_split
from sklearn.ensemble import RandomForestRegressor

# 1️⃣ Load & prepare data (runs once at startup)
ds = load_dataset("notadib/NASA-Power-Daily-Weather", split="train")
df = pd.DataFrame(ds)[["RH2M", "PRECTOTCORR", "ALLSKY_SFC_SW_DWN", "T2M"]].dropna()

X = df[["RH2M", "PRECTOTCORR", "ALLSKY_SFC_SW_DWN"]]
y = df["T2M"]

# use a small subset so startup stays fast
X_train, X_test, y_train, y_test = train_test_split(
    X, y, test_size=0.2, random_state=42
)

model = RandomForestRegressor(n_estimators=50, random_state=42)
model.fit(X_train, y_train)

# 2️⃣ Define your prediction function
def predict_temperature(rh2m, prectotcorr, solar):
    """Given humidity, precipitation, and solar radiation, predict temperature."""
    val = model.predict([[rh2m, prectotcorr, solar]])[0]
    return round(float(val), 2)

# 3️⃣ Build the Gradio interface
demo = gr.Interface(
    fn=predict_temperature,
    inputs=[
        gr.Number(label="Relative Humidity (%)", value=50, precision=1),
        gr.Number(label="Precipitation (mm)", value=1.0, precision=2),
        gr.Number(label="Solar Radiation (W/m²)", value=200.0, precision=1),
    ],
    outputs=gr.Number(label="Predicted Temp (°C)"),
    title="🌍 ClimatePredict: Daily Temperature Forecast",
    description=(
        "This demo uses a Random Forest model trained on NASA POWER daily weather data. "
        "Adjust the inputs and click **Submit** to see the forecasted temperature."
    ),
    examples=[
        [60, 0.5, 180],
        [30, 2.0, 300],
    ]
)

if __name__ == "__main__":
    demo.launch()