#!/usr/bin/env python
# Wenchang Yang (wenchang@princeton.edu)
# Wed Apr 17 15:34:24 EDT 2024
if __name__ == '__main__':
    import sys,os
    from misc.timer import Timer
    tt = Timer(f'[{os.getcwd()}] start ' + ' '.join(sys.argv))
import sys, os.path, os, glob, datetime
import xarray as xr, numpy as np, pandas as pd, matplotlib.pyplot as plt
#more imports
wython = '/tigress/wenchang/wython'
if wython not in sys.path: sys.path.append(wython); print('added to python path:', wython)
from misc import get_kws_from_argv
#
if __name__ == '__main__':
    tt.check('end import')
#
#start from here
frame = get_kws_from_argv('frame', None)
if frame is not None: frame = int(frame)
das = {}
model = 'FLOR'
#
ifile = 'amoc_FLOR_CTL1860_newdiag_1-4-1xCO2_tigercpu_intelmpi_18_576PE_2001-2280.nc'
expname = 'CTL1860_newdiag_1-4-1xCO2'
da = xr.open_dataarray(ifile).load()
das[expname] = da
#
ifile = 'amoc_FLOR_CTL1860_newdiag_tigercpu_intelmpi_18_576PE_1-3000.nc'
expname = 'CTL1860_newdiag'
da = xr.open_dataarray(ifile).load()
das[expname] = da
#
ifile = 'amoc_FLOR_CTL1990_v201905_tigercpu_intelmpi_18_576PE_1-3000.nc'
expname = 'CTL1990_v201905'
da = xr.open_dataarray(ifile).load()
das[expname] = da
"""
#
ifile = 'amoc_FLOR_co2x4_CTL1860_tigercpu_intelmpi_18_576PE_101-250.nc'
expname = 'co2x4_CTL1860'
da = xr.open_dataarray(ifile).load()
das[expname] = da
#
ifile = 'amoc_FLOR_co2xp25_CTL1860_tigercpu_intelmpi_18_576PE_101-250.nc'
expname = 'co2xp25_CTL1860'
da = xr.open_dataarray(ifile).load()
das[expname] = da
#
ifile = 'amoc_FLOR_p4p0sol_CTL1860_tigercpu_intelmpi_18_576PE_101-400.nc'
expname = 'p4p0sol_CTL1860'
da = xr.open_dataarray(ifile).load()
das[expname] = da
#
ifile = 'amoc_FLOR_m4p0sol_CTL1860_tigercpu_intelmpi_18_576PE_101-800.nc'
expname = 'm4p0sol_CTL1860'
da = xr.open_dataarray(ifile).load()
das[expname] = da
"""
#
ifile = 'amoc_FLOR_p6p0sol_CTL1860_tigercpu_intelmpi_18_576PE_101-600.nc'
expname = 'p6p0sol_CTL1860'
da = xr.open_dataarray(ifile).load()
das[expname] = da
#
ifile = 'amoc_FLOR_m6p0sol_CTL1860_tigercpu_intelmpi_18_576PE_101-597.nc'
expname = 'm6p0sol_CTL1860'
da = xr.open_dataarray(ifile).load()
das[expname] = da
#
ifile = 'amoc_FLOR_co2x2_CTL1860_tigercpu_intelmpi_18_576PE_101-600.nc'
expname = 'co2x2_CTL1860'
da = xr.open_dataarray(ifile).load()
das[expname] = da
#
ifile = 'amoc_FLOR_co2xp5_CTL1860_tigercpu_intelmpi_18_576PE_101-600.nc'
expname = 'co2xp5_CTL1860'
da = xr.open_dataarray(ifile).load()
das[expname] = da
#
ifile = 'amoc_FLOR_p2p0sol_CTL1860_tigercpu_intelmpi_18_576PE_101-400.nc'
expname = 'p2p0sol_CTL1860'
da = xr.open_dataarray(ifile).load()
das[expname] = da
#
ifile = 'amoc_FLOR_m2p0sol_CTL1860_tigercpu_intelmpi_18_576PE_101-400.nc'
expname = 'm2p0sol_CTL1860'
da = xr.open_dataarray(ifile).load()
das[expname] = da




if __name__ == '__main__':
    from wyconfig import * #my plot settings
    import xfilter
    lowpass = lambda x: x.filter.lowpass(1/50, dim='year', padtype='even')
    fig,ax = plt.subplots()
    # depth/lat selection according to Curtis and Fedorov 2024: https://doi.org/10.1007/s00382-024-07185-3
    z0 = None# 750 # m, below this depth
    lat0 = None# 30 #N, north of this lat
    alpha = 0.3
    lw = 1
    ls = '-'
    iframe = 0
    iframe += 1
    if frame is None or frame>=iframe:
        expname = 'CTL1860_newdiag'
        da = das[expname] 
        da = da.groupby('time.year').mean('time', keep_attrs=True) \
            .sel(depth=slice(z0, None), lat=slice(lat0, None)) \
            .max(['depth', 'lat'], keep_attrs=True)
        da.plot(lw=lw, alpha=alpha, color=f'C{iframe-1}')
        da.pipe(lowpass).assign_attrs(das[expname].attrs).plot(label=expname, color=f'C{iframe-1}', ls=ls)

    iframe += 1
    if frame is None or frame>=iframe:
        expname = 'CTL1990_v201905'
        da = das[expname] 
        da = da.groupby('time.year').mean('time', keep_attrs=True) \
            .sel(depth=slice(z0, None), lat=slice(lat0, None)) \
            .max(['depth', 'lat'], keep_attrs=True)
        da.plot(lw=lw, alpha=alpha, color=f'C{iframe-1}')
        da.pipe(lowpass).assign_attrs(das[expname].attrs).plot(label=expname, color=f'C{iframe-1}', ls=ls)

    iframe += 1
    if frame is None or frame>=iframe:
        expname = 'co2x2_CTL1860'
        da = das[expname] 
        da = da.groupby('time.year').mean('time', keep_attrs=True) \
            .sel(depth=slice(z0, None), lat=slice(lat0, None)) \
            .max(['depth', 'lat'], keep_attrs=True)
        da.plot(lw=lw, alpha=alpha, color=f'C{iframe-1}')
        da.pipe(lowpass).assign_attrs(das[expname].attrs).plot(label=expname, color=f'C{iframe-1}', ls=ls)

    iframe += 1
    if frame is None or frame>=iframe:
        expname = 'co2xp5_CTL1860'
        da = das[expname] 
        da = da.groupby('time.year').mean('time', keep_attrs=True) \
            .sel(depth=slice(z0, None), lat=slice(lat0, None)) \
            .max(['depth', 'lat'], keep_attrs=True)
        da.plot(lw=lw, alpha=alpha, color=f'C{iframe-1}')
        da.pipe(lowpass).assign_attrs(das[expname].attrs).plot(label=expname, color=f'C{iframe-1}', ls=ls)
    """
    iframe += 1
    if frame is None or frame>=iframe:
        expname = 'co2x4_CTL1860'
        da = das[expname] 
        da = da.groupby('time.year').mean('time', keep_attrs=True) \
            .sel(depth=slice(z0, None), lat=slice(lat0, None)) \
            .max(['depth', 'lat'], keep_attrs=True)
        da.plot(lw=lw, alpha=alpha, color=f'C{iframe-1}')
        da.pipe(lowpass).assign_attrs(das[expname].attrs).plot(label=expname, color=f'C{iframe-1}', ls=ls)

    iframe += 1
    if frame is None or frame>=iframe:
        expname = 'co2xp25_CTL1860'
        da = das[expname] 
        da = da.groupby('time.year').mean('time', keep_attrs=True) \
            .sel(depth=slice(z0, None), lat=slice(lat0, None)) \
            .max(['depth', 'lat'], keep_attrs=True)
        da.plot(lw=lw, alpha=alpha, color=f'C{iframe-1}')
        da.pipe(lowpass).assign_attrs(das[expname].attrs).plot(label=expname, color=f'C{iframe-1}', ls=ls)

    iframe += 1
    if frame is None or frame>=iframe:
        expname = 'p4p0sol_CTL1860'
        da = das[expname] 
        da = da.groupby('time.year').mean('time', keep_attrs=True) \
            .sel(depth=slice(z0, None), lat=slice(lat0, None)) \
            .max(['depth', 'lat'], keep_attrs=True)
        da.plot(lw=lw, alpha=alpha, color=f'C{iframe-1}')
        da.pipe(lowpass).assign_attrs(das[expname].attrs).plot(label=expname, color=f'C{iframe-1}', ls=ls)

    iframe += 1
    if frame is None or frame>=iframe:
        expname = 'm4p0sol_CTL1860'
        da = das[expname] 
        da = da.groupby('time.year').mean('time', keep_attrs=True) \
            .sel(depth=slice(z0, None), lat=slice(lat0, None)) \
            .max(['depth', 'lat'], keep_attrs=True)
        da.plot(lw=lw, alpha=alpha, color=f'C{iframe-1}')
        da.pipe(lowpass).assign_attrs(das[expname].attrs).plot(label=expname, color=f'C{iframe-1}', ls=ls)
    """

    iframe += 1
    if frame is None or frame>=iframe:
        expname = 'p2p0sol_CTL1860'
        da = das[expname] 
        da = da.groupby('time.year').mean('time', keep_attrs=True) \
            .sel(depth=slice(z0, None), lat=slice(lat0, None)) \
            .max(['depth', 'lat'], keep_attrs=True)
        da.plot(lw=lw, alpha=alpha, color=f'C{iframe-1}')
        da.pipe(lowpass).assign_attrs(das[expname].attrs).plot(label=expname, color=f'C{iframe-1}', ls=ls)

    iframe += 1
    if frame is None or frame>=iframe:
        expname = 'm2p0sol_CTL1860'
        da = das[expname] 
        da = da.groupby('time.year').mean('time', keep_attrs=True) \
            .sel(depth=slice(z0, None), lat=slice(lat0, None)) \
            .max(['depth', 'lat'], keep_attrs=True)
        da.plot(lw=lw, alpha=alpha, color=f'C{iframe-1}')
        da.pipe(lowpass).assign_attrs(das[expname].attrs).plot(label=expname, color=f'C{iframe-1}', ls=ls)

    iframe += 1
    if frame is None or frame>=iframe:
        expname = 'p6p0sol_CTL1860'
        da = das[expname] 
        da = da.groupby('time.year').mean('time', keep_attrs=True) \
            .sel(depth=slice(z0, None), lat=slice(lat0, None)) \
            .max(['depth', 'lat'], keep_attrs=True)
        da.plot(lw=lw, alpha=alpha, color=f'C{iframe-1}')
        da.pipe(lowpass).assign_attrs(das[expname].attrs).plot(label=expname, color=f'C{iframe-1}', ls=ls)

    iframe += 1
    if frame is None or frame>=iframe:
        expname = 'm6p0sol_CTL1860'
        da = das[expname] 
        da = da.groupby('time.year').mean('time', keep_attrs=True) \
            .sel(depth=slice(z0, None), lat=slice(lat0, None)) \
            .max(['depth', 'lat'], keep_attrs=True)
        da.plot(lw=lw, alpha=alpha, color=f'C{iframe-1}')
        da.pipe(lowpass).assign_attrs(das[expname].attrs).plot(label=expname, color=f'C{iframe-1}', ls=ls)

    iframe += 1
    if frame is None or frame>=iframe:
        expname = 'CTL1860_newdiag_1-4-1xCO2'
        da = das[expname] 
        da = da.groupby('time.year').mean('time', keep_attrs=True) \
            .sel(depth=slice(z0, None), lat=slice(lat0, None)) \
            .max(['depth', 'lat'], keep_attrs=True)
        da.plot(lw=lw, alpha=alpha, color=f'C{iframe-1}')
        da.pipe(lowpass).assign_attrs(das[expname].attrs).plot(label=expname, color=f'C{iframe-1}', ls=ls)

    ax.set_title(f'AMOC from {model}')
    ax.legend(ncol=2)
    ax.set_ylim(10,50)
    
    #savefig
    if 'savefig' in sys.argv or 's' in sys.argv:
        figname = __file__.replace('.py', f'.png')
        if z0 is not None: figname = figname.replace('.png', f'__{z0}m.png')
        if lat0 is not None: figname = figname.replace('.png', f'__{lat0}N.png')
        if frame is not None: 
            dirname = os.path.dirname(figname)
            basename = os.path.basename(figname)
            figname = os.path.join(dirname, 'frames',  basename.replace('.png', f'__frame{frame:02d}.png'))
        if 'overwritefig' in sys.argv or 'o' in sys.argv:
            wysavefig(figname, overwritefig=True)
        else:
            wysavefig(figname)
 
    tt.check(f'**Done**')
    print()
    if 'notshowfig' in sys.argv or 'n' in sys.argv:
        pass
    else:
        if 'plt' in globals(): plt.show()