#!/usr/bin/env python

# coding: utf-8

# # Anomalies for each eruption
# * Wenchang Yang (wenchang@prnceton.edu)
# * Department of Geosciences, Princeton University

# In[46]:


# init
import numpy as np, xarray as xr


# In[48]:


# misc
from dask.distributed import Client
client = Client('127.0.0.1:8982')
client


# In[49]:


# params
data_name = 'TS'
data_name_new = 'ts_a'
data_scale = 1
data_units_new = 'K'
data_long_name = 'Surface temperature anomaly'
ens = np.arange(1, 6)
nmonths = 120
ifile = 'data/PRECT/b.e11.BLMTRC5CN.f19_g16.VOLC_GRA.001.cam.h0.PRECT.085001-200512.nc'
ds = xr.open_dataset(ifile)
lats, lons = ds.lat, ds.lon


# In[53]:


# data
ds = xr.open_dataset('eruptions_855_2000_large.nc')     .rename(ivolc='time')
ds = ds.assign_coords(time=ds.t_erupt_start.values)
volcs = ds

t_starts = volcs.indexes['time'].shift(-60, 'MS')
t_ends = volcs.indexes['time'].shift(59, 'MS')


# In[54]:


# anomaly from climatology of previous 5 years
das = np.zeros((ens.size, t_starts.size, nmonths, lats.size, lons.size))
for i in ens:
    ifile = f'data/{data_name}/b.e11.BLMTRC5CN.f19_g16.VOLC_GRA.{i:03d}.cam.h0.{data_name}.085001-200512.nc'
    # print(ifile)
    ds = xr.open_dataset(ifile)
    da = ds[data_name]
    for itime in range(t_starts.size):# t0,t1 in zip(t_starts, t_ends):
        t0, t1 = t_starts[itime], t_ends[itime]
        da_ = da.sel(time=slice(t0, t1))             .pipe(lambda x: x.groupby('time.month') 
                  - x.isel(time=slice(0,60)).groupby('time.month').mean('time') )
        das[i-1, itime, :, :, :] = da_.data
das


# In[55]:


# ndarray to DataArray
da = xr.DataArray(das, dims=['ien', 'volc', 'month', 'lat', 'lon'],
                 coords={'ien': ens,
                        'volc': volcs['t_erupt_start'].values,
                        'month': range(-60, 60),
                        'lat': lats,
                        'lon': lons})
da


# In[56]:


# to netcdf
da.pipe(lambda x: x*data_scale)     .assign_attrs(long_name=data_long_name, units=data_units_new)     .to_dataset(name=data_name_new)     .to_netcdf(f'data/{data_name_new}.nc',
              encoding={'ien': {'dtype': 'int32'},
                       'volc': {'dtype': 'int32'},
                       'month': {'dtype': 'int32'},
                        data_name_new: {'dtype': 'float32'}
                       })