#!/usr/bin/env python
"""Annual global total N deposition (Tg N/yr), 1850-2100, from the fndep_wiemip_*_isimip3 files."""
import glob
import os

import matplotlib
matplotlib.use('Agg')
import matplotlib.pyplot as plt
import netCDF4
import numpy as np

HERE = os.path.dirname(os.path.abspath(__file__))
R = 6371.22e3  # m, ELM earth radius
COLORS = {'ssp126': '#2a78d6', 'ssp370': '#eb6834', 'ssp585': '#1baf7a'}

totals = {}
for ssp in COLORS:
    f, = sorted(glob.glob(os.path.join(HERE, f'fndep_wiemip_1850_2100_{ssp}_360x720_isimip3_c*.nc')))
    with netCDF4.Dataset(f) as d:
        lat = d['lat'][:]
        dlat = np.deg2rad(0.5); dlon = np.deg2rad(0.5)
        area = (R ** 2 * dlon * (np.sin(np.deg2rad(lat) + dlat / 2) - np.sin(np.deg2rad(lat) - dlat / 2)))[:, None]
        years = d['YEAR'][:]
        keep = (years >= 1850) & (years <= 2100)
        tot = {}
        for v in ('NDEP_year', 'NDEP_NHx_year', 'NDEP_NOy_year'):
            x = d[v][keep].astype(np.float64)
            tot[v] = (x * area).sum(axis=(1, 2)) / 1e12  # g -> Tg
        totals[ssp] = (years[keep], tot)
    print(ssp, area.sum() / (4 * np.pi * R ** 2))

fig, ax = plt.subplots(figsize=(9, 5), dpi=150)
for s in ('top', 'right'):
    ax.spines[s].set_visible(False)
for s in ('left', 'bottom'):
    ax.spines[s].set_color('#999999')
ax.grid(axis='y', color='#e5e5e5', lw=0.8)
ax.set_axisbelow(True)
for ssp, c in COLORS.items():
    y, tot = totals[ssp]
    ax.plot(y, tot['NDEP_year'], color=c, lw=2, label=ssp.upper())
    ax.annotate(ssp.upper(), (y[-1], tot['NDEP_year'][-1]), xytext=(6, 0), textcoords='offset points',
                va='center', fontsize=9, color='#333333')
ax.axvline(2014.5, color='#999999', lw=0.8, ls='--')
ax.text(2016, ax.get_ylim()[0] + 0.03 * np.ptp(ax.get_ylim()), 'scenario →', fontsize=8, color='#666666')
ax.set_xlim(1850, 2100)
ax.set_xlabel('Year', color='#333333')
ax.set_ylabel('Global total N deposition (Tg N yr$^{-1}$)', color='#333333')
ax.set_title('ISIMIP3 global total N deposition (NHx + NOy), 1850–2100', loc='left', fontsize=11, color='#111111')
ax.legend(frameon=False, loc='upper left')
fig.tight_layout()
out = os.path.join(HERE, 'global_total_ndep_1850-2100_ssp126_ssp370_ssp585.png')
fig.savefig(out)
print('wrote', out)
for ssp in COLORS:
    y, tot = totals[ssp]
    for yr in (1850, 1950, 2000, 2014, 2015, 2050, 2100):
        i = np.where(y == yr)[0][0]
        print(f"{ssp} {yr}: total {tot['NDEP_year'][i]:.1f}  NHx {tot['NDEP_NHx_year'][i]:.1f}  NOy {tot['NDEP_NOy_year'][i]:.1f}")