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

# In[1]:


import numpy as np


# 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,])
c=np.array([0.31565186, 0.30992895, 0.3301549,  0.423929,
            0.54217166, 0.54365426, 0.41666764, 0.5127122,
            0.5862078 , 0.02054412, 100,   0.])
w=-1*np.ones(np.shape(c))*2.8e-4
w[-1]=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]:


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[21]:


zpos


# In[22]:


zneg


# In[23]:


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


# In[24]:


zbetup


# In[25]:


zbetdo


# In[26]:


#! 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[27]:


pcc


# In[28]:


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


# In[29]:


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


# In[30]:


# 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[31]:


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


# In[32]:


ztra2 # this is before nonosc version


# In[33]:


delc_upwind


# In[34]:


delc_upwind+ztra


# In[35]:


delc_upwind+ztra2


# In[36]:


c


# In[ ]: