#!/usr/bin/env python
# Wenchang Yang (wenchang@princeton.edu)
# Wed Sep 18 23:01:01 EDT 2019
import os.path, sys, os, glob, datetime
import matplotlib.pyplot as plt
import xarray as xr, numpy as np, pandas as pd

figname = 'fig_netrad_toa_glbMean.png'
stoday = datetime.datetime.now().strftime('%Y-%m-%d')
figname_backup = figname.replace('.png', f'__{stoday}.png')

plt.figure(figsize=(8,4))

# ctl
label = 'CTL'
color = 'k'
ifile = 'data/CTL1860_newdiag_tigercpu_intelmpi_18_576PE.netrad_toa.glbMean.0001-1000.nc'
da = xr.open_dataarray(ifile)
da.groupby('time.year').mean('time').plot(color=color, label=label)

# m6
label = '$-$6% solar'
color = 'C0'
ifile = 'data/m6p0sol_CTL1860_tigercpu_intelmpi_18_576PE.netrad_toa.glbMean.0101-0597.nc'
da = xr.open_dataarray(ifile)
da.groupby('time.year').mean('time').plot(label=label, color=color)

# p6
label = '+6% solar'
color = 'C1'
ifile = 'data/p6p0sol_CTL1860_tigercpu_intelmpi_18_576PE.netrad_toa.glbMean.0101-0400.nc'
da = xr.open_dataarray(ifile)
da.groupby('time.year').mean('time').plot(label=label, color=color)

# m6 recover 0250
label = '$-$6% recover'
color = 'C2'
ifiles = glob.glob('data/m6p0sol_CTL1860_recover0251_tigercpu_intelmpi_18_576PE.netrad_toa.glbMean.0251-????.nc')
ifiles.sort()
ifile = ifiles[-1]
da = xr.open_dataarray(ifile)
da.groupby('time.year').mean('time').plot(label=label, color=color)

# m6 recover 0500
label = '$-$6% recover'
color = 'C2'
#ifile = 'data/m6p0sol_CTL1860_recover0501_tigercpu_intelmpi_18_576PE.netrad_toa.glbMean.0501-0900.nc'
ifiles = glob.glob('data/m6p0sol_CTL1860_recover0501_tigercpu_intelmpi_18_576PE.netrad_toa.glbMean.0501-????.nc')
ifiles.sort()
ifile = ifiles[-1]
da = xr.open_dataarray(ifile)
da.groupby('time.year').mean('time').plot(label=label, color=color)

# m6 recover 0590
label = '$-$6% recover'
color = 'C3'
ifiles = glob.glob('data/m6p0sol_CTL1860_recover0591_tigercpu_intelmpi_18_576PE.netrad_toa.glbMean.0591-????.nc')
ifiles.sort()
ifile = ifiles[-1]
da = xr.open_dataarray(ifile)
da.groupby('time.year').mean('time').plot(label=label, color=color)

plt.xlim(0, 1000)
plt.axvline(590, color='gray', ls='--')
plt.axvline(500, color='gray', ls='--')
plt.axvline(250, color='gray', ls='--')
plt.legend()
plt.ylabel('Global Mean netrad_toa [W m**-2]')
plt.axhline(0, color='gray', ls='--')

plt.tight_layout()
plt.savefig(figname, dpi=128)
if figname is not None:
    if os.path.exists(figname):
        os.rename(figname, figname_backup)
        plt.savefig(figname, dpi=128)
plt.show()