#!/usr/bin/env python
# Wenchang Yang (wenchang@princeton.edu)
# Wed Jun  9 20:52:58 EDT 2021
if __name__ == '__main__':
    from misc.timer import Timer
    tt = Timer(f'start {__file__}')
import sys, os.path, os, glob, datetime
import xarray as xr, numpy as np, pandas as pd, matplotlib.pyplot as plt
#more imports
from mystats import p2t
#
if __name__ == '__main__':
    tt.check('end import')
#
#start from here
dfile = __file__.replace('.py', '.nc')
if os.path.exists(dfile):
    ds = xr.open_dataset(dfile)
    print('[loaded]:', dfile)
else:
    ds = xr.open_dataset('wyfig01_lines.nc')
    #Ndt
    daname = 'N'
    da = ds[daname].sel(sphere='GL').sum('time')/12
    units, long_name = 'W m$^{-2}$ yr', r'$\int$Ndt'
    da = da.assign_attrs(units=units, long_name=long_name)
    Ndt = da
    #Fdt
    daname = 'F'
    da = ds[daname].sel(sphere='GL').sum('time')/12
    units, long_name = 'W m$^{-2}$ yr', r'$\int$Ndt'
    da = da.assign_attrs(units=units, long_name=long_name)
    Fdt = da
    #Tsdt
    daname = 'tsa'
    da = ds[daname].sel(sphere='GL').sum('time')/12
    units, long_name = 'K yr', r'$\int$T$_s$dt'
    da = da.assign_attrs(units=units, long_name=long_name)
    Tsdt = da
    
    #save
    ds = xr.Dataset(dict(Ndt=Ndt, Fdt=Fdt, Tsdt=Tsdt))
    ds.to_netcdf(dfile)
    print('[saved]:', dfile)


 
def wyplot(ax=None, volc='Pinatubo', color='k'):
    if ax is None:
        fig, ax = plt.subplots(figsize=(6, 5.5))
    alpha = .5
    xmin, xmax = -2, 0
    ymin, ymax = -6.5, 0
    capsize = 3
    marker = 's'
    markersize = 8
    Fmarker = 'D'
    Fmarkersize = 8
    
    # 
    #volc = 'Pinatubo'
    #color = 'k'
    
    daname = 'Tsdt'
    dax = ds[daname].sel(volc=volc)
    xmean_err = p2t(0.05, dax.size-1) * dax.std() * (dax.size)**(-1/2)
    
    daname = 'Ndt'
    day = ds[daname].sel(volc=volc)
    ymean_err = p2t(0.05, day.size-1) * day.std() * (day.size)**(-1/2)
    daname = 'Fdt'
    dayF = ds[daname].sel(volc=volc, en=range(1,11))
    Fmean_err = p2t(0.05, dayF.size-1) * dayF.std() * (dayF.size)**(-1/2)
    dayF = dayF.mean('en')
    
    ax.scatter(dax, day, label=volc, color=color, alpha=alpha)
    ax.errorbar(dax.mean(), day.mean(), 
                 xerr=xmean_err, yerr=ymean_err, 
                 color=color, marker=marker, ms=markersize, alpha=alpha, capsize=capsize)
    ax.axline([dax.mean(), day.mean()], [0, 0], ls=':', color=color, alpha=alpha)
    ax.errorbar(dax.mean(), dayF.mean(), 
                 xerr=xmean_err, yerr=Fmean_err, 
                 color=color, marker=Fmarker, ms=Fmarkersize, alpha=alpha, capsize=capsize)
    ax.axline([dax.mean(), dayF], [0, 0], ls='--', color=color, alpha=alpha)

    ax.set_xlim(xmin, xmax)
    ax.set_ylim(ymin, ymax)
    ax.set_xlabel(f'$\int$ T$_s$dt [K yr]')
    ax.set_ylabel(f'$\int$ Ndt [W m$^{-2}$ yr]')
    plt.legend(frameon=True, loc='lower right')

    print(volc)
    print('sensitivity', dax.mean().item()/dayF.item(), 'K/(W $^{-2}$)')
    print('beta', (dayF - day.mean()).item()/dax.mean().item())
    print('lambda', day.mean().item()/dax.mean().item())
    print()
 
if __name__ == '__main__':
    from wyconfig import * #my plot settings
    fig, ax = plt.subplots(figsize=(6, 5))
    wyplot(ax=ax, volc='Pinatubo', color='k')
    wyplot(ax=ax, volc='Agung', color='C0')
    wyplot(ax=ax, volc='StMaria', color='C1')

    #savefig
    if 'savefig' in sys.argv:
        figname = __file__.replace('.py', f'.png')
        wysavefig(figname)
    tt.check(f'**Done**')
    plt.show()