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

stoday = datetime.datetime.now().strftime('%Y-%m-%d')

figname = 'fig_ts_glbMean.png'
figname = figname.replace('.png', f'__{stoday}.png')

plt.close()
plt.figure(figsize=(8,4))
lw_raw = plt.rcParams['lines.linewidth']
plt.rcParams['lines.linewidth'] = 1.25
n_window = 5 #n_window years lowpass filter
smooth = lambda x: x.where(x.year-x.year[0]<5, other=x.rolling(year=n_window, center=True, min_periods=1).mean())
smooth21 = lambda x: x.where(x.year-x.year[0]<5, other=x.rolling(year=21, center=True, min_periods=1).mean())

# ctl
label = 'CTL'
#ifile = 'data/CTL1860_newdiag_tigercpu_intelmpi_18_576PE.t_surf.glbMean.0001-1000.nc'
ifiles = glob.glob('data/CTL1860_newdiag_tigercpu_intelmpi_18_576PE.t_surf.glbMean.0001-*.nc')
ifiles.sort()
ifile = ifiles[-1]
print(label, ifile)
da = xr.open_dataarray(ifile)
da.groupby('time.year').mean('time').pipe(smooth).plot(color='k', label=label)
ctl = da.groupby('time.year').mean('time').rolling(year=51, center=True).mean()

# m6
label = '$-$6% solar'
color = 'C0'
#ifile = 'data/m6p0sol_CTL1860_tigercpu_intelmpi_18_576PE.t_surf.glbMean.0101-0597.nc'
ifiles = glob.glob('data/m6p0sol_CTL1860_tigercpu_intelmpi_18_576PE.t_surf.glbMean.0101-*.nc')
ifiles.sort()
ifile = ifiles[-1]
print(label, ifile)
da = xr.open_dataarray(ifile)
da.groupby('time.year').mean('time').pipe(smooth).plot(label=label, color=color)

# p6
label = '+6% solar'
color = 'C3'
ls = '-'
#ifile = 'data/p6p0sol_CTL1860_tigercpu_intelmpi_18_576PE.t_surf.glbMean.0101-0400.nc'
ifiles = glob.glob('data/p6p0sol_CTL1860_tigercpu_intelmpi_18_576PE.t_surf.glbMean.0101-*.nc')
ifiles.sort()
ifile = ifiles[-1]
print(label, ifile)
da = xr.open_dataarray(ifile)
da.groupby('time.year').mean('time').pipe(smooth).plot(label=label, color=color, ls=ls)

# p6 scaled
label = '+6% solar: mirrored'
color = 'C3'
ls = '--'
#ifile = 'data/p6p0sol_CTL1860_tigercpu_intelmpi_18_576PE.t_surf.glbMean.0101-0400.nc'
ifiles = glob.glob('data/p6p0sol_CTL1860_tigercpu_intelmpi_18_576PE.t_surf.glbMean.0101-*.nc')
ifiles.sort()
ifile = ifiles[-1]
print(label, ifile)
da = xr.open_dataarray(ifile)
da.groupby('time.year').mean('time').pipe(lambda x: 2*ctl-x) \
    .pipe(smooth).plot(label=label, color=color, ls=ls)

# m6 recover590
label = '$-$6% recover 590'
color = 'C2'
ifiles = glob.glob('data/m6p0sol_CTL1860_recover0591_tigercpu_intelmpi_18_576PE.t_surf.glbMean.0591-????.nc')
ifiles.sort()
ifile = ifiles[-1]
print(label, ifile)
da = xr.open_dataarray(ifile)
da.groupby('time.year').mean('time').pipe(smooth).plot(label=label, color=color)

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

# m6 recover250
label = '$-$6% recover 250'
color = 'C2'
ls = '-'
#ifile = 'data/m6p0sol_CTL1860_recover0251_tigercpu_intelmpi_18_576PE.t_surf.glbMean.0251-0400.nc'
ifiles = glob.glob('data/m6p0sol_CTL1860_recover0251_tigercpu_intelmpi_18_576PE.t_surf.glbMean.0251-????.nc')
ifiles.sort()
ifile = ifiles[-1]
print(label, ifile)
da = xr.open_dataarray(ifile)
da.groupby('time.year').mean('time').pipe(smooth).plot(label=label, color=color, ls=ls)

# m1
label = '$-$1% solar'
color = 'C4'
ls = '-'
#ifile = 'data/m1p0sol_CTL1860_tigercpu_intelmpi_18_576PE.t_surf.glbMean.0101-0600.nc'
ifiles = glob.glob('data/m1p0sol_CTL1860_tigercpu_intelmpi_18_576PE.t_surf.glbMean.0101-*.nc')
ifiles.sort()
ifile = ifiles[-1]
print(label, ifile)
da = xr.open_dataarray(ifile)
da.groupby('time.year').mean('time').pipe(smooth).plot(label=label, color=color, ls=ls)

# m1: scaled
label = '$-$1% solar: scaled'
color = 'C4'
ls = '--'
#ifile = 'data/m1p0sol_CTL1860_tigercpu_intelmpi_18_576PE.t_surf.glbMean.0101-0600.nc'
print(label, ifile)
da = xr.open_dataarray(ifile)
da.groupby('time.year').mean('time') \
    .pipe(lambda x: ctl+(x-ctl)*6 ) \
    .pipe(smooth21) \
    .plot(label=label, color=color, ls=ls)

# p1
label = '$+$1% solar'
color = 'C1'
ls = '-'
#ifile = 'data/p1p0sol_CTL1860_tigercpu_intelmpi_18_576PE.t_surf.glbMean.0101-0600.nc'
ifiles = glob.glob('data/p1p0sol_CTL1860_tigercpu_intelmpi_18_576PE.t_surf.glbMean.0101-*.nc')
ifiles.sort()
ifile = ifiles[-1]
print(label, ifile)
da = xr.open_dataarray(ifile)
da.groupby('time.year').mean('time').pipe(smooth).plot(label=label, color=color, ls=ls)

# p1: scaled
label = '$+$1% solar: scaled'
color = 'C1'
ls = '--'
#ifile = 'data/p1p0sol_CTL1860_tigercpu_intelmpi_18_576PE.t_surf.glbMean.0101-0600.nc'
print(label, ifile)
da = xr.open_dataarray(ifile)
da.groupby('time.year').mean('time') \
    .pipe(lambda x: ctl-(x-ctl)*6) \
    .pipe(smooth21) \
    .plot(label=label, color=color, ls=ls)

#co2x2 
label = '2xCO2'
color = 'C6'
ls = '-'
#ifile = 'data/p1p0sol_CTL1860_tigercpu_intelmpi_18_576PE.t_surf.glbMean.0101-0600.nc'
ifiles = glob.glob('data/co2x2_CTL1860_tigercpu_intelmpi_18_576PE.t_surf.glbMean.0101-*.nc')
ifiles.sort()
ifile = ifiles[-1]
print(label, ifile)
da = xr.open_dataarray(ifile)
da.groupby('time.year').mean('time').pipe(smooth).plot(label=label, color=color, ls=ls)


#2xco2: scaled
label = '2xCO2: scaled(4x)'
color = 'C6'
ls = '--'
ifiles = glob.glob('data/co2x2_CTL1860_tigercpu_intelmpi_18_576PE.t_surf.glbMean.0101-*.nc')
ifiles.sort()
ifile = ifiles[-1]
print(label, ifile)
da = xr.open_dataarray(ifile)
da.groupby('time.year').mean('time') \
    .pipe(lambda x: ctl+(x-ctl)*4) \
    .pipe(smooth21).plot(label=label, color=color, ls=ls)

#0.5xco2
label = '0.5xCO2'
color = 'C9'
ls = '-'
#ifile = 'data/p1p0sol_CTL1860_tigercpu_intelmpi_18_576PE.t_surf.glbMean.0101-0600.nc'
ifiles = glob.glob('data/co2xp5_CTL1860_tigercpu_intelmpi_18_576PE.t_surf.glbMean.0101-*.nc')
ifiles.sort()
ifile = ifiles[-1]
print(label, ifile)
da = xr.open_dataarray(ifile)
da.groupby('time.year').mean('time').pipe(smooth).plot(label=label, color=color, ls=ls)

#0.5xco2: scaled
label = '0.5xCO2: scaled(-4x)'
color = 'C9'
ls = '--'
#ifile = 'data/p6p0sol_CTL1860_tigercpu_intelmpi_18_576PE.t_surf.glbMean.0101-0400.nc'
ifiles = glob.glob('data/co2xp5_CTL1860_tigercpu_intelmpi_18_576PE.t_surf.glbMean.0101-*.nc')
ifiles.sort()
ifile = ifiles[-1]
print(label, ifile)
da = xr.open_dataarray(ifile)
da.groupby('time.year').mean('time') \
    .pipe(lambda x: ctl+(ctl-x)*4) \
    .pipe(smooth21).plot(label=label, color=color, ls=ls)

plt.axvline(250, color='gray', ls='--')
plt.axvline(500, color='gray', ls='--')
plt.axvline(590, color='gray', ls='--')
#plt.ylim(245, None)
plt.ylim(270, 295)
#plt.xlim(0,1000)
plt.xlim(100,130)
plt.legend(ncol=3, frameon=False)
plt.ylabel('Global Mean Ts [K]')

plt.tight_layout()
if figname is not None:
    #if os.path.exists(figname):
    #    os.rename(figname, figname_backup)
    plt.savefig(figname, dpi=128)
    print('[saved]:', figname)
plt.show()
plt.rcParams['lines.linewidth'] = lw_raw