#!/usr/bin/env python
# Wenchang Yang (wenchang@princeton.edu)
# Thu Jul  2 14:17:46 EDT 2020
if __name__ == '__main__':
    from misc.timer import Timer
    tt = Timer(f'start {__file__}')
import sys, os.path, os, glob
import xarray as xr, numpy as np, pandas as pd
#import matplotlib.pyplot as plt
#more imports
#
if __name__ == '__main__':
    tt.check('end import')
#
#start from here
ifile = 'ts.historical_ssp585.1850-2100.ens26.ocean.nc'
years_base = slice('1981','2000')
years_warm = slice('2081','2100')
 
 
if __name__ == '__main__':
    from wyconfig import * #my plot settings
    with xr.open_dataset(ifile) as ds:
        da = ds['ts']
    da_base = da.sel(time=years_base).mean('time')
    da_warm = da.sel(time=years_warm).mean('time')
    vmin = 0
    vmax = 6#(da_warm - da_base).sel(lat=slice(-30,30)).max().item()
    """ 
    fig, ax = plt.subplots()
    da = (da_warm - da_base).mean('model')
    da.pipe(np.log2).assign_attrs(long_name='$\log_2(\Delta{}SST)$').plot(levels=21)

    #CIESM
    plt.figure()
    da = (da_warm - da_base).sel(model='CIESM').sel(lat=slice(-30,30))
    da.plot.contourf(levels=21, vmin=vmin, vmax=vmax, cmap='turbo', extend='neither')

    #EC-Earth3
    plt.figure()
    da = (da_warm - da_base).sel(model='EC-Earth3').sel(lat=slice(-30,30))
    da.plot.contourf(levels=21, vmin=vmin, vmax=vmax, cmap='turbo', extend='neither')
    """
    da = (da_warm - da_base).sel(lat=slice(-30,30))
    for model in da.model:
        plt.figure(figsize=(6,2.4))
        long_name = f'$\Delta$SST: {years_warm.start}-{years_warm.stop} minus {years_base.start}-{years_base.stop}'
        da.sel(model=model).assign_attrs(long_name=long_name).plot.contourf(levels=21, vmin=vmin, vmax=vmax, cmap='turbo', extend='both')
        plt.xlabel('')
        plt.ylabel('')
        
        figname = __file__.replace('.py', f'.map.dsst.{model.item()}.{tt.today()}.png')
        figname = os.path.join('figs', figname)
        plt.savefig(figname)
        print('**saved**:', figname)

    tt.check(f'**Done**')
    plt.show()