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

@author: mritchey
"""

#HRRR Zarr
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


with gr.Blocks() 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")
    # input_vars=[date, variable, level]
    output = gr.Numpy(label="Max")
    greet_btn.click(fn=get_hrrr_values,
                    inputs=[date, variable, level],
                    outputs=output,api_name="addition")


# demo = gr.Interface(fn=get_hrrr_values,
#                     inputs=["text","text","text"],
#                     outputs="text",
#                     show_api=True)
demo.launch()