#!/usr/bin/env python
# Wenchang Yang (wenchang@princeton.edu)
# Thu Dec 18 01:48:25 PM EST 2025
if __name__ == '__main__':
    import sys,os
    try:
        from misc.timer import Timer
        tt = Timer(f'[{os.getcwd()}] start ' + ' '.join(sys.argv))
    except:
        pass
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
from modelout import update_modelout_data
from misc import roll_lon
#
if __name__ == '__main__':
    try:
        tt.check('end import')
    except:
        pass
#
#start from here
cinMask = True if 'cinMask' in sys.argv else False
cinModel = True if 'cinModel' in sys.argv else False
def stp(ds):
    """assume ds is from atmos_daily"""
    ifile = ds.attrs['ifile']
    ifile_month = ifile.replace('atmos_daily', 'atmos_month')
    #ds = xr.open_dataset(ifile)
    #cape
    cape = ds.cape.load()
    fcape = cape/1500
    #for lcl
    ps = ds.ps.load()
    plcl = ds.pcldb.load()
    #srh
    srh = ds.srh01.load()
    #for wind shear
    u500 = ds.u500.load()
    v500 = ds.v500.load()
    u200 = ds.u200.load()
    v200 = ds.u200.load()
    h500 = ds.h500.load()
    h200 = ds.h200.load()
    u0 = ds.u_ref.load()
    v0 = ds.v_ref.load()
    #
    #ifile = f'/projects/GEOCLIM/wenchang/MODEL_OUT/AM2.5/tiger3/CTL1990_tiger3_intelmpi_24_1080PE/POSTP/{year:04d}0101.atmos_month.nc'
    h0 = xr.open_dataset(ifile_month)['zsurf'].load()
    
    #lcl
    p500 = 500*100
    sblcl = (h500-h0) * np.log(ps/plcl) / np.log(ps/p500)
    sblcl = sblcl.clip(0)
    flcl = (2000 - sblcl)/1500
    flcl = flcl.clip(0)
    
    #srh
    srh = srh.where(srh.grid_yt>=0, other=-srh)
    fsrh = srh.clip(0)/150
    
    #shear
    u = u500 - u0 + (u200 - u500) * (h0 + 6000 - h500) / (h200 - h500)
    v = v500 - v0 + (v200 - v500) * (h0 + 6000 - h500) / (h200 - h500)
    bwd = np.sqrt(u**2 + v**2)
    fbwd = bwd/20
    
    #stp
    stp = fcape * flcl * fsrh * fbwd
    #
    if cinModel:
        cin = ds.cin.load()
        fcin = (200 - cin)/150
        fcin = fcin.clip(0)
        stp = stp*fcin #https://scholarworks.sfasu.edu/cgi/viewcontent.cgi?article=1292&context=urc
    elif cinMask:
        cin = ds.cin.load()
        stp = stp.where(cin<50, other=0)

    #annnual
    stp = stp.resample(time='AS').mean('time', keep_attrs=True)
    stp = stp.assign_attrs(long_name='annual mean significant tonado parameter')
    #monthly
    #stp = stp.resample(time='MS').mean('time', keep_attrs=True)
    #stp = stp.assign_attrs(long_name='monthly mean significant tonado parameter')
    return stp

model = 'AM2.5C360'
expname = 'CTL2010s_tiger3_intel24ifort_openmpi_2430PE'
func = stp
funcname = 'stp'
if cinModel:
    funcname += '_cinModel'
elif cinMask:
    funcname += '_cminMask'
#funcname = 'stp_monthly'
dsname = 'atmos_daily'
da = update_modelout_data(daname=None, model=model, expname=expname, func=func, funcname=funcname, dsname=dsname) 
da = da.resample(time='AS').mean('time', keep_attrs=True) #monthly to annual
da2010 = da.pipe(roll_lon)
nyears_2010 = da.time.size
expname = 'CTL1980s_tiger3_intel24ifort_openmpi_2430PE'
da = update_modelout_data(daname=None, model=model, expname=expname, func=func, funcname=funcname, dsname=dsname) 
da = da.resample(time='AS').mean('time', keep_attrs=True) #monthly to annual
da1980 = da.pipe(roll_lon)
nyears_1980 = da.time.size
expname = 'CTL1980s_plus2K_tiger3_intel24ifort_openmpi_2430PE'
da = update_modelout_data(daname=None, model=model, expname=expname, func=func, funcname=funcname, dsname=dsname) 
da = da.resample(time='AS').mean('time', keep_attrs=True) #monthly to annual
da1980plus2K = da.pipe(roll_lon)
nyears_1980plus2K = da.time.size
expname = 'CTL1980s_plusPattern_tiger3_intel24ifort_openmpi_2430PE'
da = update_modelout_data(daname=None, model=model, expname=expname, func=func, funcname=funcname, dsname=dsname) 
da = da.resample(time='AS').mean('time', keep_attrs=True) #monthly to annual
da1980plusPattern = da.pipe(roll_lon)
nyears_1980plusPattern = da.time.size

if __name__ == '__main__':
    from wyconfig import * #my plot settings
    from geoplots import mapplot
    from misc.landmask import whereland
    from misc.comp2samp import comp2samp
    import salem
    vmax_anom = 0.2
    df = salem.read_shapefile('/tigress/wenchang/analysis/climate_and_disease/data/US_states/cb_2016_us_state_5m/cb_2016_us_state_5m.shp')

    #import cmocean
    levels = None
    NA = True if 'NA' in sys.argv else False
    func = lambda x: (x*365).pipe(whereland).assign_attrs(units='days')
    func_region = lambda x: x.sel(lat=slice(20, 60), lon=slice(-120, -60)) if NA else x
    da = da2010
    nyears = da.time.size
    #da.mean('time').pipe(func).pipe(func_region).plot.contourf(vmax=2, levels=levels)
    da.mean('time').pipe(func).pipe(func_region).plot(vmax=2, levels=levels)
    mapplot()
    ax = plt.gca()
    ax.set_title(f'annual STP from {model} CTL2010s, {nyears} years')
    if NA: df.boundary.plot(color='gray', lw=0.5, ax=ax)
    #savefig
    if 'savefig' in sys.argv or 's' in sys.argv:
        figname = __file__.replace('.py', f'__world_2010.png')
        if NA: figname = figname.replace('__world', '__NA')
        if 'overwritefig' in sys.argv or 'o' in sys.argv:
            wysavefig(figname, overwritefig=True)
        else:
            wysavefig(figname)

    #CTL2010 - CTL1980
    plt.figure()
    da = da2010.mean('time') - da1980.mean('time')
    daa_total = da
    #da.pipe(func).pipe(func_region).plot.contourf(vmax=vmax_anom, extend='both', levels=levels)#, cmap=cmocean.cm.balance)
    da.pipe(func).pipe(func_region).plot(vmax=vmax_anom, extend='both', levels=levels)#, cmap=cmocean.cm.balance)
    ds = comp2samp(da2010, da1980, dim='time')
    #significance info
    p = ds.pvalue
    p.pipe(whereland).pipe(func_region).where(p<0.05).plot.contourf(colors='none', hatches=['......'], add_colorbar=False)
    da = da2010.mean('time')
    if NA: da.pipe(func).pipe(func_region).plot.contour(levels=[1.5], colors='g', linewidths=0.5)
    mapplot()
    ax = plt.gca()
    ax.set_title(f'annual STP from {model}: CTL2010s({nyears_2010}) $-$ CTL1980s ({nyears_1980})')
    if NA: df.boundary.plot(color='gray', lw=0.5, ax=ax)
    #savefig
    if 'savefig' in sys.argv or 's' in sys.argv:
        figname = __file__.replace('.py', f'__world_2010-1980.png')
        if NA: figname = figname.replace('__world', '__NA')
        if 'overwritefig' in sys.argv or 'o' in sys.argv:
            wysavefig(figname, overwritefig=True)
        else:
            wysavefig(figname)
 
    #plus2K -CTL
    plt.figure()
    da = da1980plus2K.mean('time') - da1980.mean('time')
    da = da/2*0.27 #scaled to the mean warming from 1980s to 2010s
    daa_uniform = da #save for later use
    #da.pipe(func).pipe(func_region).plot.contourf(vmax=vmax_anom, extend='both', levels=21)#, cmap=cmocean.cm.balance)
    da.pipe(func).pipe(func_region).plot(vmax=vmax_anom, extend='both', levels=levels)#, cmap=cmocean.cm.balance)
    ds = comp2samp(da1980plus2K, da1980, dim='time')
    #significance info
    p = ds.pvalue
    p.pipe(whereland).pipe(func_region).where(p<0.05).plot.contourf(colors='none', hatches=['......'], add_colorbar=False)
    da = da2010.mean('time')
    if NA: da.pipe(func).pipe(func_region).plot.contour(levels=[1.5], colors='g', linewidths=0.5)
    mapplot()
    ax = plt.gca()
    ax.set_title(f'annual STP from {model}: +2K({nyears_1980plus2K}) $-$ CTL1980s({nyears_1980}), scaled')
    if NA: df.boundary.plot(color='gray', lw=0.5, ax=ax)
    #savefig
    if 'savefig' in sys.argv or 's' in sys.argv:
        figname = __file__.replace('.py', f'__world_1980plus2K-1980.png')
        if NA: figname = figname.replace('__world', '__NA')
        if 'overwritefig' in sys.argv or 'o' in sys.argv:
            wysavefig(figname, overwritefig=True)
        else:
            wysavefig(figname)
 
    #plusPattern -CTL
    plt.figure()
    da = da1980plusPattern.mean('time') - da1980.mean('time')
    daa_pattern = da #save for later use
    #da.pipe(func).pipe(func_region).plot.contourf(vmax=vmax_anom, extend='both', levels=21)#, cmap=cmocean.cm.balance)
    da.pipe(func).pipe(func_region).plot(vmax=vmax_anom, extend='both', levels=levels)#, cmap=cmocean.cm.balance)
    ds = comp2samp(da1980plusPattern, da1980, dim='time')
    #significance info
    p = ds.pvalue
    p.pipe(whereland).pipe(func_region).where(p<0.05).plot.contourf(colors='none', hatches=['......'], add_colorbar=False)
    da = da2010.mean('time')
    if NA: da.pipe(func).pipe(func_region).plot.contour(levels=[1.5], colors='g', linewidths=0.5)
    mapplot()
    ax = plt.gca()
    ax.set_title(f'annual STP from {model}: +Pattern({nyears_1980plusPattern}) $-$ CTL1980s({nyears_1980})')
    if NA: df.boundary.plot(color='gray', lw=0.5, ax=ax)
    #savefig
    if 'savefig' in sys.argv or 's' in sys.argv:
        figname = __file__.replace('.py', f'__world_1980plusPattern-1980.png')
        if NA: figname = figname.replace('__world', '__NA')
        if 'overwritefig' in sys.argv or 'o' in sys.argv:
            wysavefig(figname, overwritefig=True)
        else:
            wysavefig(figname)

    #uniform + pattern
    plt.figure()
    da = daa_uniform + daa_pattern
    #da.pipe(func).pipe(func_region).plot.contourf(vmax=vmax_anom, extend='both', levels=21)#, cmap=cmocean.cm.balance)
    da.pipe(func).pipe(func_region).plot(vmax=vmax_anom, extend='both', levels=levels)#, cmap=cmocean.cm.balance)
    da = da2010.mean('time')
    if NA: da.pipe(func).pipe(func_region).plot.contour(levels=[1.5], colors='g', linewidths=0.5)
    mapplot()
    ax = plt.gca()
    ax.set_title(f'annual STP from {model}: uniform warming + pattern')
    if NA: df.boundary.plot(color='gray', lw=0.5, ax=ax)
    #savefig
    if 'savefig' in sys.argv or 's' in sys.argv:
        figname = __file__.replace('.py', f'__world_uniform_and_pattern.png')
        if NA: figname = figname.replace('__world', '__NA')
        if 'overwritefig' in sys.argv or 'o' in sys.argv:
            wysavefig(figname, overwritefig=True)
        else:
            wysavefig(figname)

    #residual
    plt.figure()
    da = daa_total - (daa_uniform + daa_pattern)
    #da.pipe(func).pipe(func_region).plot.contourf(vmax=vmax_anom, extend='both', levels=21)#, cmap=cmocean.cm.balance)
    da.pipe(func).pipe(func_region).plot(vmax=vmax_anom, extend='both', levels=levels)#, cmap=cmocean.cm.balance)
    da = da2010.mean('time')
    if NA: da.pipe(func).pipe(func_region).plot.contour(levels=[1.5], colors='g', linewidths=0.5)
    mapplot()
    ax = plt.gca()
    ax.set_title(f'annual STP from {model}: residual')
    if NA: df.boundary.plot(color='gray', lw=0.5, ax=ax)
    #savefig
    if 'savefig' in sys.argv or 's' in sys.argv:
        figname = __file__.replace('.py', f'__world_residual.png')
        if NA: figname = figname.replace('__world', '__NA')
        if 'overwritefig' in sys.argv or 'o' in sys.argv:
            wysavefig(figname, overwritefig=True)
        else:
            wysavefig(figname)

    try:
        tt.check(f'**Done**')
    except:
        pass
    print()
    if 'notshowfig' in sys.argv or 'n' in sys.argv:
        pass
    else:
        if 'plt' in globals(): plt.show()