import xarray as xr
import numpy as np
import matplotlib.pyplot as plt
import gradio as gr

# Charger les données
sst = xr.open_dataset("data/sst.nc")
pr = xr.open_dataset("data/pr_senegal.nc")

def compute_corr(start_year, end_year):

    sst_sel = sst.sel(time=slice(f"{start_year}", f"{end_year}"))
    pr_sel = pr.sel(time=slice(f"{start_year}", f"{end_year}"))

    # moyenne SST région
    sst_mean = sst_sel.mean(dim=["lat","lon"])

    # corrélation
    corr = xr.corr(sst_mean, pr_sel, dim="time")

    fig, ax = plt.subplots()

    ax.plot(sst_mean.time, sst_mean, label="SST")
    ax.plot(pr_sel.time, pr_sel, label="Precip Senegal")

    ax.legend()
    ax.set_title(f"Correlation = {float(corr.values):.2f}")

    return fig

interface = gr.Interface(
    fn=compute_corr,
    inputs=[
        gr.Slider(1980,2025,value=2000,label="Start year"),
        gr.Slider(1980,2025,value=2020,label="End year")
    ],
    outputs="plot",
    title="SST Atlantic NE vs Senegal Rainfall"
)

interface.launch()