#!/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


# params
data_name = 'PRECT'
data_name_new = 'pr_a'
data_scale = 1000*24*3600 
data_units_new = 'mm/day'
data_long_name = 'Total precipitation anomaly'
'''
data_name = 'TS'
data_name_new = 'ts_a'
data_scale = 1
data_units_new = 'K'
data_long_name = 'Surface temperature anomaly'
'''

months = np.arange(-60, 120)
'''
ens = np.arange(1, 6)
ifiles = [f'data/{data_name}/b.e11.BLMTRC5CN.f19_g16.VOLC_GRA.{i:03d}.cam.h0.{data_name}.085001-200512.nc'
    for i in ens]
'''
ens = np.arange(1, 14)
ifiles = [f'data/{data_name}/b.e11.BLMTRC5CN.f19_g16.{i:03d}.cam.h0.{data_name}.085001-200512.nc'
    for i in ens]
ofile = ifiles[0].split('/')[-1].replace('.001', '.volcs')
ofile = f'data/{ofile}'
ds = xr.open_dataset(ifiles[0])
lats, lons = ds.lat, ds.lon


# In[53]:


# data
ds = xr.open_dataset('data/lme_eruptions_ind.nc')
t_erupt_starts = ds.indexes['volc']
t_starts = ds.indexes['volc'].shift(int(months[0]), 'MS')
t_ends = ds.indexes['volc'].shift(int(months[-1]), 'MS')


# In[54]:


# anomaly from climatology of previous 5 years
das = np.zeros((ens.size, t_starts.size, months.size, lats.size, lons.size))
for i in ens:
    #print(i, 'of', len(ens))
    ifile = ifiles[i-1]
    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


# In[55]:


# ndarray to DataArray
da = xr.DataArray(das, dims=['ien', 'volc', 'month', 'lat', 'lon'],
                 coords={'ien': ens,
                        'volc': t_erupt_starts.values,
                        'month': months,
                        '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(ofile, encoding={'ien': {'dtype': 'int32'},
                       'volc': {'dtype': 'int32'},
                       'month': {'dtype': 'int32'},
                        data_name_new: {'dtype': 'float32'}
                       })