# -*- coding: utf-8 -*-
"""
Created on Tue Feb 21 11:44:34 2023

@author: mritchey
"""

import gradio as gr
from joblib import Parallel, delayed
import s3fs 
import xarray as xr 
import numpy as np

s3 = s3fs.S3FileSystem(anon=True)

def lookup(path):
	return s3fs.S3Map(path, s3=s3)

def get_values(date_time,variable,level):
    date,time=date_time
    path = f"hrrrzarr/sfc/{date}/{date}_{time[-2:]}z_anl.zarr/{level}/{variable}" 
    ds = xr.open_mfdataset([lookup(path), lookup(f"{path}/{level}")], 
       engine="zarr") 
    return ds[variable].values

def get_hrrr_values(date,variable,level):
    times=[f'0{t}'[-2:] for t in range(0,24)]
    date_times=[(date,t)  for t in times]
    results=Parallel(n_jobs=24, prefer="threads")(delayed(get_values)(i,variable,level) for i in date_times)
    ds_all=np.stack(results)
    ds_max=ds_all.max(axis=0)
    return ds_max

def multiple_dates(dates,variable,level):
    dates=dates.split(";")
    results=Parallel(n_jobs=6)(delayed(get_hrrr_values)(i,variable,level) for i in dates)
    ds_all=np.stack(results)
    return ds_all
    

with gr.Blocks(show_api=True) as demo:
    date = gr.Textbox(label="Date")
    variable = gr.Textbox(label="Variable",placeholder='GUST')
    level = gr.Textbox(label="Level",placeholder='surface')
    greet_btn = gr.Button("Get Max")
    output = gr.Numpy(label="Max")
    greet_btn.click(fn=multiple_dates, 
                    inputs=[date, variable, level], 
                    outputs=output,api_name="addition")



demo.launch()