import gradio as gr
import requests
from PIL import Image
from io import BytesIO
import numpy as np
import torch
from skillful_nowcasting.dgmr import DGMR
from datetime import datetime, timezone

# Load the pretrained model
model = DGMR.from_pretrained("openclimatefix/dgmr")

def get_latest_radar_image():
    """Fetches the latest radar image from the Iowa State University archive."""
    now = datetime.now(timezone.utc)
    url = f"https://mesonet.agron.iastate.edu/archive/data/{now.year}/{now.month:02d}/{now.day:02d}/GIS/uscomp/n0r_0.png"
    response = requests.get(url)
    img = Image.open(BytesIO(response.content)).convert("RGB")
    return img

def preprocess_image(img):
    """Preprocesses the image for the model."""
    img = img.resize((128, 128))
    img_array = np.array(img) / 255.0
    img_tensor = torch.from_numpy(img_array).permute(2, 0, 1).float()
    # Add batch and sequence dimensions
    return img_tensor.unsqueeze(0).unsqueeze(0)

def postprocess_output(output_tensor):
    """Postprocesses the model output back to an image."""
    # Take the first prediction in the sequence
    output_tensor = output_tensor[0, 0, :, :, :]
    output_array = output_tensor.permute(1, 2, 0).detach().numpy()
    output_array = (output_array * 255).astype(np.uint8)
    return Image.fromarray(output_array)

def predict_nowcast():
    """Fetches the latest radar, runs the nowcasting model, and returns the images."""
    input_image = get_latest_radar_image()
    preprocessed_input = preprocess_image(input_image)
    
    # The DGMR model expects a sequence of 4 input images.
    # For this example, we will feed the same image 4 times.
    model_input = torch.cat([preprocessed_input] * 4, dim=1)
    
    with torch.no_grad():
        output = model(model_input)
    
    predicted_image = postprocess_output(output)
    
    return input_image, predicted_image

# Create the Gradio interface
with gr.Blocks() as demo:
    gr.Markdown("## American Radar Nowcasting with AI")
    gr.Markdown(
        "Click the button to fetch the latest American radar image and generate a nowcast prediction using the `skillful_nowcasting` model."
    )
    
    with gr.Row():
        input_image_display = gr.Image(label="Current Radar Image")
        output_image_display = gr.Image(label="Predicted Nowcast")

    predict_button = gr.Button("Generate Nowcast")
    predict_button.click(
        fn=predict_nowcast,
        inputs=[],
        outputs=[input_image_display, output_image_display]
    )

demo.launch()