import gradio as gr
import torch
import torch.nn as nn
import timm
from torchvision import transforms
from PIL import Image

# ------------------------ Model Definition ------------------------
class CycloneIntensityModel(nn.Module):
    def __init__(self):
        super().__init__()
        # use the same attribute name (“model”) that your checkpoint used
        self.model = timm.create_model('resnet18', pretrained=True)
        in_features = self.model.fc.in_features
        self.model.fc = nn.Linear(in_features, 1)

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

# Instantiate & load
model = CycloneIntensityModel()
state = torch.load("cyclone_intensity.pth", map_location='cpu')
model.load_state_dict(state)    # now keys “model.conv1.weight” → self.model.conv1.weight
model.eval()

# ------------------------ Transforms & Classifier ------------------------
transform = transforms.Compose([
    transforms.Resize((224, 224)),
    transforms.ToTensor(),
])

def classify_cyclone(intensity: float) -> str:
    if 74 <= intensity <= 95:
        return "Category 1 – Some damage expected."
    elif 96 <= intensity <= 110:
        return "Category 2 – Extensive damage expected."
    elif 111 <= intensity <= 129:
        return "Category 3 – Devastating damage expected."
    elif 130 <= intensity <= 156:
        return "Category 4 – Catastrophic damage expected."
    else:
        return "Category 5 – Max-severity damage expected."

# ------------------------ Prediction fn ------------------------
def predict_intensity(img: Image.Image):
    img = img.convert("RGB")
    tensor = transform(img).unsqueeze(0)
    with torch.no_grad():
        intensity = model(tensor).item()
    category = classify_cyclone(intensity)
    return f"{intensity:.2f} Knots", category

# ------------------------ Gradio UI ------------------------
with gr.Blocks() as demo:
    gr.Markdown("## Cyclone Intensity Predictor")
    with gr.Row():
        input_img = gr.Image(type="pil", label="Upload Satellite Image")
        with gr.Column():
            btn = gr.Button("Predict")
            output_int = gr.Textbox(label="Predicted Wind Speed")
            output_cat = gr.Textbox(label="Cyclone Category")

    btn.click(
        fn=predict_intensity,
        inputs=input_img,
        outputs=[output_int, output_cat]
    )

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