import uxarray as ux
import matplotlib.pyplot as plt
import cartopy.crs as ccrs
import cartopy.feature as cfeature
import os
import OpenVisus as ov
import xarray as xr
from tqdm import tqdm
import time
from sys import getsizeof

# !pip install OpenVisus openvisuspy boto3 xmltodict panel bokeh
# !pip install tqdm

main_dir='/glade/derecho/scratch/digital-earths-hackathon/conus404/native_grid'
home_dir='/glade/u/home/dpanta'
idx_dir='/glade/u/home/dpanta/idx'

st_file='/glade/derecho/scratch/digital-earths-hackathon/conus404/native_grid/2020-10_eastward_wind.nc'


files=[
'2020-12_eastward_wind.nc',
'2020-12_northward_wind.nc',
'2020-12_precipitation_flux.nc',
'2020-12_specific_humidity.nc',
'2020-12_surface_air_pressure.nc',
'2020-12_surface_temperature.nc',
'2020-12_toa_outgoing_longwave_flux.nc',
'2020-12_toa_outgoing_shortwave_flux.nc',
]

ds=xr.open_dataset(st_file)
ds

import datetime

idx_filename=f"/glade/derecho/scratch/dpanta/idx/conus404_all.idx"
fields=['eastward_wind','northward_wind','precipitation_flux','specific_humidity','surface_air_pressure','surface_temperature','toa_outgoing_longwave_flux', 'toa_outgoing_shortwave_flux']
arco="4mb"
fld_idx=[]
for fld in fields:
    print(f'Setting Field::: {fld}')
    st_file=f'/glade/derecho/scratch/digital-earths-hackathon/conus404/native_grid/2020-10_{fld}.nc'

    ds = xr.open_dataset(st_file, group="/", mask_and_scale=False) 
    data=ds[f"{fld}"][0,:,:].values
    if int(len(data.shape))>1:
        
        print(f"{fld}::::{data.shape}")
        try:
            vmin,vmax=np.min(data),np.max(data)
        except:
            vmin,vmax=0,0
        fld=ov.Field.fromString(f"""{fld} {str(data.dtype)} format(row_major) min({vmin}) max({vmax})""")
        fld_idx.append(fld)

T=744
W,H=data.shape

print('Creating IDX Metadata file now...')
db=ov.CreateIdx(
    url=idx_filename, 
	dims=[H,W], 
	fields=fld_idx, 
	compression="raw", 
	arco=arco,
	time=[0,T,"time_%d/"])
print('Created...')

# db=ov.LoadDataset(idx_filename)

print('Writing to idx...')
for i in range(len(fields)):
    print(f'Writing field {fields[i]}...')
    st_file=f'/glade/derecho/scratch/digital-earths-hackathon/conus404/native_grid/2020-10_{fields[i]}.nc'
    ds = xr.open_dataset(st_file, group="/", mask_and_scale=False) 
    data=ds[f"{fields[i]}"][...].values
    for t in range(data.shape[0]):
        data_to_write=data[t,:,:]
        try:
            db.write(data_to_write,time=t,field=fields[i])
        except:
            pass
print('Writing data completed...')
print('Compressing data now...')
db.compressDataset(['zip'])
print('Data compressed!!!')