#!/usr/bin/env python
# coding: utf-8

# In[1]:


import xarray as xr
import cdsapi


# In[2]:


client = cdsapi.Client()


# In[24]:


grid = [0.75, 0.75]
area = [90, -180+0.75/2, -90, 180]
year = ['{}'.format(y) for y in range(1979, 2025)]
month = ['{:02d}'.format(m) for m in range(1, 13)]


# In[12]:


dl_dir = 'tmp/'


# ## Invariant

# In[14]:


filename = 'ERA5_LowRes_Invariant.nc'
dataset = "reanalysis-era5-single-levels-monthly-means"
request = {
    "product_type": ["monthly_averaged_reanalysis"],
    "variable": [
        "geopotential",
        "land_sea_mask"
    ],
    "grid": grid,
    "area": area,
    "year": ["1979"],
    "month": ["01"],
    "time": ["00:00"],
    "data_format": "netcdf",
    "download_format": "unarchived",
}

client.retrieve(dataset, request).download(dl_dir + filename)
print('+++ Done!')


# In[17]:


# Unfortunate necessary cleaning
with xr.open_dataset(dl_dir + filename) as ds:
    # ds = ds.rename({"valid_time": "time"}).drop_vars(["expver", "number"]) #, errors="ignore")
    
    ds = ds.isel(valid_time=0).drop_vars(["expver", "number", "valid_time"])
    ds['orography'] = ds['z']  / 9.80665
    ds['orography'].attrs.update(long_name='Orography', units='m', standard_name='orography')
    ds = ds.drop_vars(["z"])
    ds['lsm'].attrs = {k: v for k, v in ds['lsm'].attrs.items() if not k.startswith("GRIB_")}
    
    ds.attrs = {k: v for k, v in ds.attrs.items() if not k.startswith("GRIB_")}
    ds.to_netcdf(filename)


# ## Temp and precip 

# In[25]:


filename = 'ERA5_LowRes_Monthly_t2m.nc'

dataset = "reanalysis-era5-single-levels-monthly-means"
request = {
    "product_type": ["monthly_averaged_reanalysis"],
    "variable": ["2m_temperature"],
    "grid": grid,
    "area": area,
    "year": year,
    "month": month,
    "time": ["00:00"],
    "data_format": "netcdf",
    "download_format": "unarchived",
}

client.retrieve(dataset, request).download(dl_dir + filename)
print('+++ Done!')


# In[26]:


# Unfortunate necessary cleaning
with xr.open_dataset(dl_dir + filename) as ds:
    ds = ds.rename({"valid_time": "time"}).drop_vars(["expver", "number"]) 
     
    ds.attrs = {k: v for k, v in ds.attrs.items() if not k.startswith("GRIB_")}
    for var in ds.variables.keys():
        ds[var].attrs = {k: v for k, v in ds[var].attrs.items() if not k.startswith("GRIB_")}
    
    ds.to_netcdf(filename)


# In[27]:


filename = 'ERA5_LowRes_Monthly_tp.nc'

dataset = "reanalysis-era5-single-levels-monthly-means"
request = {
    "product_type": ["monthly_averaged_reanalysis"],
    "variable": ["total_precipitation"],
    "grid": grid,
    "area": area,
    "year": year,
    "month": month,
    "time": ["00:00"],
    "data_format": "netcdf",
    "download_format": "unarchived",
}

client.retrieve(dataset, request).download(dl_dir + filename)
print('+++ Done!')


# In[28]:


# Unfortunate necessary cleaning
with xr.open_dataset(dl_dir + filename) as ds:
    ds = ds.rename({"valid_time": "time"}).drop_vars(["expver", "number"]) 
     
    ds.attrs = {k: v for k, v in ds.attrs.items() if not k.startswith("GRIB_")}
    for var in ds.variables.keys():
        ds[var].attrs = {k: v for k, v in ds[var].attrs.items() if not k.startswith("GRIB_")}
    
    ds.to_netcdf(filename)


# In[ ]: