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

from datetime import datetime


import xarray as xr
#Calculates the effect of soil moisture anomalies on afternoon precipitation using the MPAS Dyamond 1 dataset for August 2016. Inspired by Welty et al 2019.
#ERIK JANZON
#UNIVERSITY OF NEBRASKA-LINCOLN

#DIGITAL EARTHS HACKATHON
#NCAR NODE
#MAY 2025

#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 MPAS run that has all the needed variables is the Dyamond1 global run, which only has one month of data. Results look neat, but are likely not significant




#Function to remap data onto LatLon Grid (unfinished, based on https://easy.gems.dkrz.de/Processing/healpix/lonlat_remap.html)
#Ultimate goal is to back out structured grid to calculate gradients and, eventually, moisture convergence

def get_nn_lon_lat_index(nside, lons, lats):
    lons2, lats2 = np.meshgrid(lons, lats)
    return xr.DataArray(
        healpy.ang2pix(nside, lons2, lats2, nest=True, lonlat=True),
        coords=[("lat", lats), ("lon", lons)],
    )
zoom=8
#Set up data fetch and list catalogue

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

cat = intake.open_catalog(url).NCAR

#Probably can be commented out. List catalog.
#list(cat)

#Start with mpas_dyamond1

model_run=cat.mpas_dyamond1
#Probably can be commented out. Reporting the available data in each model dataset.
#model_run().describe()

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

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

hours=uxds.time.dt.hour
uxds

#Create subsets of meteorological variables for North America

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

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

u=uxds["uReconstructZonal"].subset.bounding_box(
    lon_bounds,
    lat_bounds,
)

v=uxds["uReconstructMeridional"].subset.bounding_box(
    lon_bounds,
    lat_bounds,
)
qv=qv[:,:,0]
u=u[:,:,0]
v=v[:,:,0]


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

uxrain = uxds["qr"].subset.bounding_box(
    lon_bounds,
    lat_bounds,
)
#Only look at afternoon hours for precipitation and morning hours to filter out low precipitation mornings and for soil moisture
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)
#wspd=((u**2)+(v**2))**(1/2)
#qvu=qv*wspd


#dqu_dx = qv.gradient()  # very rough gradient using wind speed. Probably doesn't work. But placeholder.



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

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

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

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


uxsmois_mean_mor=uxsmois_mor.mean(dim="nSoilLevels") 

#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, August 2016, MPAS Dyamond1",
    **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-3) & (res_rain>1e-3))
agg_rain_smois=agg_rain.where((res_smois>0.4) & (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-3) & (res_rain>1e-3)).mean(dim="time")

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