import gradio as gr
import torch
import numpy as np

# ======== Example Dummy Model ========
# Replace this with your own model loading
class DummyModel(torch.nn.Module):
    def __init__(self):
        super().__init__()
        self.linear = torch.nn.Linear(2, 3)  # 2 inputs (lat, lon) → 3 outputs

    def forward(self, x):
        return self.linear(x)

# load your trained model (replace with torch.load if you have .pt file)
model = DummyModel()
model.eval()

# ======== Prediction Function ========
def predict_score(lat, lon):
    # Convert input to tensor
    inputs = torch.tensor([[lat, lon]], dtype=torch.float32)

    # Get model output
    with torch.no_grad():
        outputs = model(inputs).numpy().flatten()

    # Unpack into respective values
    score, num_banks, normal_score = outputs

    # You can apply any post-processing here
    return {
        "Score": round(float(score), 3),
        "Num Banks": round(float(num_banks), 3),
        "Normal Score": round(float(normal_score), 3),
    }

# ======== Gradio Interface ========
interface = gr.Interface(
    fn=predict_score,
    inputs=[
        gr.Number(label="Latitude"),
        gr.Number(label="Longitude"),
    ],
    outputs=[
        gr.Number(label="Score"),
        gr.Number(label="Num Banks"),
        gr.Number(label="Normal Score"),
    ],
    title="Bank Location Scoring Model",
    description="Enter latitude and longitude to get the predicted score, number of banks, and normalized score.",
)


interface.launch()