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.log_alpha = nn.Parameter(torch.tensor(0.0))

    def forward(self, x):
        eta = self.linear(x)
        mu = torch.exp(eta).squeeze(-1)
        alpha = torch.exp(self.log_alpha)
        return mu, alpha

def negbinom_loss(y, mu, alpha):
    log_prob = (
        torch.lgamma(y + 1/alpha)
        - torch.lgamma(1/alpha)
        - torch.lgamma(y + 1)
        + (1/alpha) * torch.log(1 / (1 + alpha * mu))
        + y * torch.log((alpha * mu) / (1 + alpha * mu))
    )
    return -torch.mean(log_prob)
    
model = NegBinomialModel(17)
model.load_state_dict(torch.load("model_weights.pt"))
model.eval()

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

    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),
        "Num Banks": round(float(mu_pred), 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()