#!/usr/bin/env python
# coding: utf-8

# In[1]:


import numpy as np
from matplotlib import pyplot as plt

get_ipython().run_line_magic('matplotlib', 'inline')


# In[2]:


#w=np.array([-.1,-.1,-.1,-.1,-.1, 0])
#w=np.array([-1,-1,-1,-1,-1, 0])
#c=np.array([  1,  1,  1,.01,100, 1e10])*1e-6
e1e2t=214298.32535430687
e3t=np.array([25.70331479, 26.28684983, 26.59728865, 26.75965336, 
              26.84381704, 26.88724213, 26.90959407, 26.92108493,
              26.9269885,  26.93002054, 26.93157752, 26.93237697,])
# step 24:
c=np.array([2.65121251e-01, 1.90705761e-01, 1.62778839e-01, 1.63109794e-01, 
                   1.41612262e-01, 1.20875522e-01, 1.02854975e-01, 9.03575793e-02,
                   6.69111013e-02, 3.21567655e-02, 2.10060075e-01, 0.0])
# step 0:
#c=np.array([0.20331727, 0.16865452, 0.15772118, 0.14340848,
#            0.11826674, 0.10116881, 0.08674008, 0.07820327,
#            0.06499595, 0.04768646, 0.1204031 , 0.])
#w=-1*np.ones(np.shape(c))*2.8e-4
#w[-1]=0
w=np.array([ 3.871025,  -96.635605, -40.600544,  20.52446,  
            -6.9390144, -10.459614, -70.14894,  -30.304188,  
            -17.269821, -46.92218,  -61.723824,  0.      ]) 
tmask=np.ones(np.shape(c))
tmask[-1]=0


# In[3]:


#1. bottom value zero
zwz=np.zeros(np.shape(c))
#2. upstream advection
zwz[1:]=w[1:]*c[:-1]
zwz


# In[4]:


dt=40*2


# In[5]:


ztra=np.zeros(np.shape(c))
ztra[:-1]=-1*(zwz[:-1]-zwz[1:])/e1e2t
delc_upwind=ztra.copy() # ** added to tra at this step
zwi=(e3t*c+dt*ztra)/e3t*tmask # set fse3t's to 1
zwi # upwind increases low point from .01 to .0298 in one time step


# In[6]:


# 3. antidiffusive flux
zwz_sav=zwz.copy()


# In[7]:


jnzts=20 # same as SSC


# In[8]:


cb=c.copy()


# In[9]:


ca=np.zeros(np.shape(c))
zwzts=np.zeros(np.shape(c))


# In[10]:


for jl in range(0,jnzts):
    if jl==0:
        jtaken=np.mod(jnzts+1,2)
        print(jtaken)
        zts=dt/jnzts # z_rzts=1/jnzts
        cn=c.copy() # this makes first step euler
    elif jl==1:
        zts=2*dt/jnzts
    zwz[1:]=.5*w[1:]*(cn[1:]+cn[:-1])
    if jtaken==0:
        zwzts[1:]=zwzts[1:]+zwz[1:]*zts
    jtaken=np.mod(jtaken+1,2) # switch on/off
    
    zbtr=1/(e1e2t*e3t)
    ztra[1:-1]=-1*zbtr[1:-1]*(zwz[1:-1]-zwz[2:])
    ca[1:-1]=cb[1:-1]+zts*ztra[1:-1] # step forward
    # swap for next step:
    cb=cn.copy()
    cn=ca.copy()
    ca=np.zeros(np.shape(c))


# In[11]:


cn


# In[12]:


zwzts


# In[13]:


# anti-diffusive vertical flux using average flux from sub-timestepping
zwz=zwzts/dt-zwz_sav


# In[14]:


zwz


# In[15]:


zbig=1e40
zrtrn=1e-15


# In[16]:


# 4. nonosc(pbef=ptb,zwx,zwy,pcc=zwz,paft=zwi,p2dt)
# zwi is after value based on upwind
pbef=c.copy()
paft=zwi.copy()
# search local extrema
zbup=np.maximum(pbef*tmask-zbig*(1-tmask),paft*tmask-zbig*(1-tmask))
zbup


# In[17]:


zbdo=np.minimum(pbef*tmask+zbig*(1-tmask),paft*tmask+zbig*(1-tmask))
zbdo


# In[18]:


zup=np.zeros(np.shape(c))
zdo=zup.copy()
zup[0]=np.maximum(zbup[0],zbup[1])
zup[1:-1]=np.maximum(np.maximum(zbup[:-2],zbup[1:-1]),zbup[2:])
#zup[-1]=np.maximum(zbup[-2],zbup[-1]) this is never calculated
zup


# In[19]:


zdo[0]=np.minimum(zbdo[0],zbdo[1])
zdo[1:-1]=np.minimum(np.minimum(zbdo[:-2],zbdo[1:-1]),zbdo[2:])
#zdo[-1]=np.minimum(zbdo[-2],zbdo[-1])
zdo


# In[20]:


fig,ax=plt.subplots(1,2,figsize=(12,6))
ax[0].plot(zup,np.arange(0,len(zbup)),'k*')
ax[0].set_title('zup')
ax[0].set_ylim(12,0)
ax[1].plot(zdo,np.arange(0,len(zdo)),'k*')
ax[1].set_title('zdo')
ax[1].set_ylim(12,0)


# In[21]:


zeros=np.zeros(np.shape(c))
pcc=zwz.copy()
zpos=np.zeros(np.shape(c))
zneg=zpos.copy()
zpos[:-1]=np.maximum(zeros[1:],pcc[1:])-np.minimum(zeros[:-1],pcc[:-1])
zneg[:-1]=np.maximum(zeros[:-1],pcc[:-1])-np.minimum(zeros[1:],pcc[1:])


# In[22]:


zpos


# In[23]:


zneg


# In[24]:


fig,ax=plt.subplots(1,2,figsize=(12,6))
ax[0].plot(zpos,np.arange(0,len(zpos)),'k*')
ax[0].set_title('zpos')
ax[0].set_ylim(12,0)
ax[1].plot(zneg,np.arange(0,len(zneg)),'k*')
ax[1].set_title('zneg')
ax[1].set_ylim(12,0)


# In[25]:


#! up & down beta terms
zbt = e1e2t*e3t / dt
zbetup = ( zup - paft) / ( zpos + zrtrn ) * zbt
zbetdo = ( paft - zdo ) / ( zneg + zrtrn ) * zbt


# In[26]:


zbetup


# In[27]:


zbetdo


# In[28]:


fig,ax=plt.subplots(1,2,figsize=(12,6))
ax[0].plot(zbetup,np.arange(0,len(zbetup)),'k*')
ax[0].set_title('zbetup')
ax[0].set_ylim(12,0)
ax[1].plot(zbetdo,np.arange(0,len(zbetdo)),'k*')
ax[1].set_title('zbetdo')
ax[1].set_ylim(12,0)


# In[29]:


#! monotonic flux in the k direction, i.e. pcc
ones=np.ones(np.shape(c))
za=zeros.copy()
zb=zeros.copy()
zc=zeros.copy()
za[:-1] = np.minimum(np.minimum( ones[1:], zbetdo[1:]), zbetup[:-1] )
zb[:-1] = np.minimum(np.minimum( ones[1:], zbetup[1:]), zbetdo[:-1] )
zc[:-1]=np.where(pcc[1:]<0,0,1)
#zc[:-1]=.5+np.abs(pcc[1:])/pcc[1:]*.5
#zc =       ( 0.5  + SIGN( 0.5 , pcc(ji,jj,jk+1) ) )
# SIGN(A,B) returns value of A with sign of B
pcc[1:]=pcc[1:]*(zc[:-1]*za[:-1]+(1-zc[:-1])*zb[:-1])
#pcc(ji,jj,jk+1) = pcc(ji,jj,jk+1) * ( zc * za + ( 1._wp - zc) * zb )


# In[30]:


pcc


# In[31]:


# before nonosc, zwz contains difference between centered substepped leapfrog and upwind:
zwz


# In[32]:


# after nonosc, zwz (pcc) is adjusted so that only one element changes from upwind 


# In[33]:


# 5. final trend with corrected fluxes
ztra=0*ztra
ztra2=ztra.copy()
ztra[:-1]=-1*zbtr[:-1]*(pcc[:-1]-pcc[1:])
ztra2[:-1]=-1*zbtr[:-1]*(zwz[:-1]-zwz[1:])


# In[34]:


# total change is ztra+delc, ztra is after nonosc adjustment
ztra


# In[35]:


ztra2 # this is before nonosc version


# In[36]:


delc_upwind


# In[37]:


delc_upwind+ztra


# In[38]:


delc_upwind+ztra2


# In[39]:


c


# In[ ]: