import intake
import uxarray as ux
import cartopy.crs as ccrs
import geoviews.feature as gf
import holoviews as hv

import healpy as hp
import numpy as np

#Calculates the effect of soil moisture anomalies on afternoon precipitation using the ICON dataset for 2020. Inspired by Welty et al 2019 observational study.
#ERIK JANZON
#UNIVERSITY OF NEBRASKA-LINCOLN

#DIGITAL EARTHS HACKATHON
#NCAR NODE
#MAY 2025

#MAKE SURE THE CHANGE THE CONDITIONS BELOW (CELL 42, LINE 36-38) TO REFLECT THE SEASON/HOURS YOU WANT TO LOOK AT 

#MAKE SURE TO ADJUST PLOT TITLES

#Issues:
#1.) Only figured out how to look at one time zone at a time, so "afternoon" is based on CDT in this case. Maybe not a problem since the data is only 
#    3 hourly?
#2.) Cannot calculate moisture convergence. Tried to use uxarray.gradient(), but no luck because one cannot remap to grid cell faces for comparison.
#
#3.) Only one year of data in ICON. Longer model datasets missing soil moisture (unless I missed something...)

#4.) Welty et al 2019 performed more useful statistics. See Yifang Cheng's notebook for example of probabilities.

#5.) HARD CODED STUFF:
    #a.) MAKE SURE THE CHANGE THE CONDITIONS BELOW (CELL 42, LINE 36-38) TO REFLECT THE SEASON/HOURS YOU WANT TO LOOK AT 

    #b.) MAKE SURE TO ADJUST PLOT TITLES

zoom=8

#UTC offset.
offset=6;

#SUBSETS FOR DIFFERENT AREAS.
#CONUS
lon_bounds = (-130, -50)
lat_bounds = (20, 60)

#Central South America
#lon_bounds = (-90, -30)
#lat_bounds = (-40, 0)

#Southern Africa
#lon_bounds = (0, -40)
#lat_bounds = (-30, 0)

from datetime import datetime

#Set up data fetch and list catalogue

url="https://digital-earths-global-hackathon.github.io/catalog/catalog.yaml"

cat = intake.open_catalog(url).NCAR

list(cat)

#Start with mpas_dyamond1

model_run=cat.icon_d3hp003
model_run().describe()

#Play with lower zoom level, fetching data. Increase zoom later.

ds=model_run(zoom=zoom,time='PT3H').to_dask()
uxds=ux.UxDataset.from_healpix(ds)
times=uxds.time

hours=uxds.time.dt.hour
uxds

#Create subsets of soil moisture and precipitation for northern hemisphere south of 60 degrees north across continents

#lon_bounds = (-130, -50)
#lat_bounds = (20, 60)

#lon_bounds = (-90, -30)
#lat_bounds = (-40, 0)

#qv=uxds["ie"].subset.bounding_box(
#    lon_bounds,
#    lat_bounds,
#)

uxsmois = uxds["mrso"].subset.bounding_box(
    lon_bounds,
    lat_bounds,
)

uxrain = uxds["pr"].subset.bounding_box(
    lon_bounds,
    lat_bounds,
)

#conditionafternoon=((uxrain.time.dt.hour >= 12+6) & (uxrain.time.dt.hour <= 17+6))

#conditionmorning=((uxrain.time.dt.hour >= 5+6) & (uxrain.time.dt.hour < 12+6))

#conditionafternoon=((uxrain.time.dt.month >= 12) & (uxrain.time.dt.month <= 2) & (uxrain.time.dt.hour >= 12+offset) & (uxrain.time.dt.hour <= 17+offset))

#conditionmorning=((uxrain.time.dt.month >= 12) & (uxrain.time.dt.month <= 2) & (uxrain.time.dt.hour >= 5+offset) & (uxrain.time.dt.hour < 12+offset))
#wspd=((u**2)+(v**2))**(1/2)
#qvu=qv*wspd


#MAKE SURE TO CHANGE THIS WHEN LOOKING AT NORTH/SOUTH HEMISPHERE!!!!!!
conditionafternoon=(((uxrain.time.dt.month >= 6) | (uxrain.time.dt.month <= 8)) & (uxrain.time.dt.hour >= 12+offset) & (uxrain.time.dt.hour <= 17+offset))

conditionmorning=(((uxrain.time.dt.month >= 6) | (uxrain.time.dt.month <= 8)) & (uxrain.time.dt.hour >= 5+offset) & (uxrain.time.dt.hour < 12+offset))




uxrain_aft=uxrain[conditionafternoon]
uxrain_mor=uxrain[conditionmorning]
uxsmois_mor=uxsmois[conditionmorning]
#qv_mor=qv[conditionmorning]
uxrain_aft.time

plot_opts = {"width": 700, "height": 350}

#uxrain_mean_aft=uxrain_aft.sum(dim="nVertLevels") 

#uxrain_mean_mor=uxrain_mor.sum(dim="nVertLevels") 

uxrain_mean_mor=uxrain_mor

uxrain_mean_aft=uxrain_aft

uxsmois_mean_mor=uxsmois_mor.mean(dim="soil_level") 
#uxsmois_mean_mor=uxsmois_mean_mor

#uxsmois_mean_mor=uxsmois_mor.mean(dim="soil_level") 
uxrain_plot=uxrain_mean_aft.mean(dim="time")



features = gf.coastline(
    projection=ccrs.PlateCarree(), line_width=1, scale="50m"
) * gf.states(projection=ccrs.PlateCarree(), line_width=1, scale="50m")

uxrain_plot.plot(
    rasterize=True,
    periodic_elements="exclude",
    title="Mean Afternoon Total Column Rain Mixing Ratio, JJA, ERA5",
    **plot_opts,
) * features

# Condition: values greater than 0.5
#uxrain_max=uxrain_mean_aft.max(dim="n_face")




# Apply the condition to get the values
#sum_hours_rain=uxrain_mean_aft.groupby_bins("n_face",[.01e-6,100]).groups
#sum_hours_rain

daily_rain=uxrain_mean_aft.resample(time='D')

daily_mor_rain=uxrain_mean_mor.resample(time='D')
daily_smois=uxsmois_mean_mor.resample(time='D')
#dailyqv=qv_mor.resample(time='D')
#condition = daily_rain > 0.01e-6

#uxrain1=daily_rain.where(condition).sum(dim="time")
res_rain=daily_rain.sum()
res_mor_rain=daily_mor_rain.sum()
res_smois=daily_smois.sum()
#res_qv=dailyqv.mean()

res_smois.mean(dim="time").plot(
    rasterize=True,
    periodic_elements="exclude",
    title="Mean Total Soil Moisture (sum of layers), Morning",
    **plot_opts,
) * features

agg_rain=res_rain.where((res_mor_rain<0.02e-5) & (res_rain>1e-5))
agg_rain_smois=agg_rain.where((res_smois>0.1) & (res_smois<=1)).count(dim="time")
agg_rain_smois.plot(
    rasterize=True,
    periodic_elements="exclude",
    clim=(1,15),
    title="Afternoon Rainfall, with no morning rainfall",
    **plot_opts,
) * features

#Soil moisture anomalies
#Remove water values

smois_valid=res_smois.where((res_smois<=1))

smois_anom=smois_valid-smois_valid.mean(dim="time")

mean_smois_anom_rain=smois_anom.where((res_mor_rain<0.02e-5) & (res_rain>1e-5)).mean(dim="time")

mean_smois_anom_rain.plot(
    rasterize=True,
    periodic_elements="exclude",
    title="JJA ICON Morning Soil Moisture Anomaly (with Afternoon Rainfall)",
    **plot_opts,
) * features