import gradio as gr
import torch
from torch import nn

import numpy as np
import pandas as pd

from utils import compute_features


class NegBinomialModel(nn.Module):
    def __init__(self, in_features):
        super().__init__()
        self.linear = nn.Linear(in_features, 1)
        self.alpha = nn.Parameter(torch.tensor(0.5))

    def forward(self, x):
        # safer activation than exp()
        mu = torch.exp(torch.clamp(self.linear(x), min=-5, max=5))
        alpha = torch.clamp(self.alpha, min=1e-3, max=10)
        return mu.squeeze(), alpha

    
model = NegBinomialModel(16)
model.load_state_dict(torch.load("model_weights.pt", map_location='cpu'))
model.eval()

def predict_score(lat, lon):
    # Convert input to tensor
    # inputs = torch.tensor([[lat, lon]], dtype=torch.float32)
    inputs = compute_features((lat,lon))
    num_banks = inputs.get("num_banks_in_radius", 0)

    inputs = torch.tensor([lat,lon] + list(inputs.values()))

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

    # Unpack into respective values
    mu_pred, alpha = outputs

    score = (1 * np.abs(mu_pred + 0.1)) * 100

    # You can apply any post-processing here
    return {
        "Score": round(float(score), 3),
        "Number of current ATMs": round(float(mu_pred), 3),
        "Number of ideal ATMs" : num_banks_in_radius
        # "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 Current Banks"),
        gr.Number(label="Num Ideal 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()