#!/usr/bin/env python
"""Build ELM fndep files (1849-2101, 0.5deg) from ISIMIP3 monthly NHx/NOy deposition.

Layout follows ../../WIEMIP/user_inputdata/fndep_wiemip_1850_2000_0.9x1.25_c20211019.nc
(NDEP_year, NDEP_NOy_year, NDEP_NHx_year, YEAR, time; days since 0000-01-01, noleap),
kept on the native ISIMIP3 0.5deg grid with lon 0..360 and lat south->north.
1849 repeats 1850 and 2101 repeats 2100.
"""
import datetime
import os

import netCDF4
import numpy as np

HERE = os.path.dirname(os.path.abspath(__file__))
TEMPLATE = os.path.join(HERE, '../../WIEMIP/user_inputdata/fndep_wiemip_1850_2000_0.9x1.25_c20211019.nc')
HIST = 'ndep-{sp}_histsoc_monthly_1850_2014.nc'
SSP_FILES = {
    'ssp126': 'ndep-{sp}_ssp126soc_monthly_2015_2100.nc',
    'ssp370': 'ndep-{sp}_ssp370_monthly_2015_2100.nc',
    'ssp585': 'ndep-{sp}_ssp585soc_monthly_2015_2100.nc',
}
TODAY = datetime.date.today()
STAMP = TODAY.strftime('c%Y%m%d')


def annual_sum(path, sp):
    """Monthly g N m-2 mon-1 -> annual g N m-2 yr-1, regridded to lon 0..360, lat S->N."""
    with netCDF4.Dataset(path) as d:
        nt = d.dimensions['time'].size
        assert nt % 12 == 0, path
        lat = d['lat'][:]
        lon = d['lon'][:]
        out = np.zeros((nt // 12, lat.size, lon.size), dtype=np.float64)
        for y in range(nt // 12):
            x = d[sp][y * 12:(y + 1) * 12].astype(np.float64)
            assert np.ma.count_masked(x) == 0, f'{path} year index {y} has missing values'
            out[y] = np.ma.getdata(x).sum(axis=0)
    if lat[0] > lat[-1]:
        out = out[:, ::-1, :]
        lat = lat[::-1]
    # -179.75..179.75 -> 0.25..359.75
    lon360 = np.where(lon < 0, lon + 360.0, lon)
    order = np.argsort(lon360)
    return out[:, :, order].astype(np.float32), lat.astype(np.float64), lon360[order].astype(np.float64)


def main():
    tpl = netCDF4.Dataset(TEMPLATE)
    hist = {}
    for sp in ('nhx', 'noy'):
        hist[sp], lat, lon = annual_sum(os.path.join(HERE, HIST.format(sp=sp)), sp)
        print('hist', sp, hist[sp].shape, flush=True)

    years = np.arange(1849, 2102, dtype=np.int32)
    # template time: days since 0000-01-01 noleap, Jan 1 of each YEAR (1849 -> 674885)
    time = years.astype(np.float64) * 365.0

    for ssp, pat in SSP_FILES.items():
        fields = {}
        src = []
        for sp in ('nhx', 'noy'):
            fut, flat, flon = annual_sum(os.path.join(HERE, pat.format(sp=sp)), sp)
            assert np.allclose(flat, lat) and np.allclose(flon, lon)
            a = np.concatenate([hist[sp][:1], hist[sp], fut, fut[-1:]], axis=0)
            assert a.shape[0] == years.size
            fields[sp] = a
            src += [HIST.format(sp=sp), pat.format(sp=sp)]

        fout = os.path.join(HERE, f'fndep_wiemip_1850_2100_{ssp}_360x720_isimip3_{STAMP}.nc')
        with netCDF4.Dataset(fout, 'w', format='NETCDF4_CLASSIC') as o:
            o.title = f'ELM N deposition forcing from ISIMIP3 NHx/NOy, historical 1850-2014 + {ssp} 2015-2100, 0.5deg'
            o.source = ', '.join(os.path.join(HERE, s) for s in src)
            o.template = os.path.abspath(TEMPLATE)
            o.comment = ('1849 repeats 1850 and 2101 repeats 2100. Annual values are sums of ISIMIP3 monthly '
                         'fluxes (g N m-2 mon-1). NDEP_year = NDEP_NHx_year + NDEP_NOy_year.')
            o.history = f'created by {os.path.basename(__file__)} on {TODAY.isoformat()}'

            o.createDimension('time', None)
            o.createDimension('lat', lat.size)
            o.createDimension('lon', lon.size)
            for name, data in (('lon', lon), ('lat', lat)):
                v = o.createVariable(name, 'f8', (name,))
                v.setncatts({k: tpl[name].getncattr(k) for k in tpl[name].ncattrs()})
                v[:] = data
            v = o.createVariable('YEAR', 'i4', ('time',))
            v.setncatts({k: tpl['YEAR'].getncattr(k) for k in tpl['YEAR'].ncattrs()})
            v[:] = years
            v = o.createVariable('time', 'f8', ('time',))
            v.setncatts({k: tpl['time'].getncattr(k) for k in tpl['time'].ncattrs()})
            v[:] = time

            data = {
                'NDEP_year': fields['nhx'] + fields['noy'],
                'NDEP_NOy_year': fields['noy'],
                'NDEP_NHx_year': fields['nhx'],
            }
            for name, arr in data.items():
                attrs = {k: tpl[name].getncattr(k) for k in tpl[name].ncattrs()}
                fill = attrs.pop('_FillValue', None)
                v = o.createVariable(name, 'f4', ('time', 'lat', 'lon'), zlib=True, complevel=4,
                                     chunksizes=(1, lat.size, lon.size), fill_value=fill)
                v.setncatts(attrs)
                v[:] = arr
        print('wrote', fout, flush=True)
    tpl.close()


if __name__ == '__main__':
    main()