import numpy as np
import sys
from math import *
import matplotlib.pyplot as plt

#global g
#global ICT,AFCA,AFCT
#global TMAX,dt,DTWR
#global PAA,PCI,PCO
#global RWL,TNL,TNA,TNC
#global NST,SEL,SAA,SLB
#global NQT,QTQ,QTI
#global AFCS, QI

def WOUT(dts,zw,vw,qn,fout):
    #print of Time, WL of Surge tank, Velocity of Tunnel, Discharge, k
    global PAA,PCI,PCO
    cf=0.5*(PCI+PCO)*PAA
    k=FK(vw,qn,cf)
    fout.write('{0:15e},{1:15e},{2:15e},{3:15e},{4:15e}\n'.format(dts,zw,vw,qn,k))

def RUNGE(qo,vo,zo,fa,dtt,dts):
    global PAA,PCI,PCO
    wrk=0.0
    dz0=dtt*FZ(vo,qo,fa)
    cf=PCI*PAA              # Water level rising (in-flow from port)
    if dz0>0.0: cf=PCO*PAA  # Water level drop (out-flow from port)
    wrk=FK(vo,qo,cf)
    dv0=dtt*FV(zo,vo,wrk)

    v1=vo+0.5*dv0
    z1=zo+0.5*dz0
    dtw=dts-dtt*0.5
    q1=QNCAL(dtw)
    wrk=FK(v1,q1,cf)
    dv1=dtt*FV(z1,v1,wrk)
    dz1=dtt*FZ(v1,q1,fa)

    v1=vo+0.5*dv1
    z1=zo+0.5*dz1
    wrk=FK(v1,q1,cf)
    dv2=dtt*FV(z1,v1,wrk)
    dz2=dtt*FZ(v1,q1,fa)

    v1=vo+dv2
    z1=zo+dz2
    dtw=dts
    q2=QNCAL(dtw)
    wrk=FK(v1,q2,cf)
    dv3=dtt*FV(z1,v1,wrk)
    dz3=dtt*FZ(v1,q2,fa)

    dv=(dv0+2.0*dv1+2.0*dv2+dv3)/6.0
    dz=(dz0+2.0*dz1+2.0*dz2+dz3)/6.0
    return dv,dz,q2

def FZ(v,q,fa):
    global TNA
    return (q-TNA*v)/fa

def FK(v,q,cf):
    global g,TNA
    return abs((TNA*v-q)/cf)*(TNA*v-q)/cf/2.0/g

def FV(z,v,wrk):
    global g,TNC,TNL
    return (z-TNC*abs(v)*v-wrk)/TNL*g

def QNCAL(tt):
    global AFCA,AFCT
    global NQT,QTQ,QTI
    global AFCS, QI
    qn=QI+AFCA*sin(2.0*pi/AFCT*tt)
    if tt>=AFCS: 
        for i in range(1,NQT):
            if QTI[i-1]<tt and tt<=QTI[i]:
                qn=QTQ[i-1]+(QTQ[i]-QTQ[i-1])/(QTI[i]-QTI[i-1])*(tt-QTI[i-1])
                break
    return qn


# Main routine
g=9.8
param=sys.argv
fnameR=param[1]
fnameW=param[2]
fnameF=param[3]

# Data input
fin=open(fnameR,'r')
text=fin.readline()
text=text.strip()
strcom=text         # Comment

text=fin.readline()
text=text.strip()
text=text.split(',')
ICT =int(text[0])   # Calculation case (1: normal, 2: AFC)
AFCA=float(text[1]) # Discharge of half amplitude of AFC(m3/s)
AFCT=float(text[2]) # Period of AFC(s)

text=fin.readline()
text=text.strip()
text=text.split(',')
TMAX=float(text[0]) # End time of calculation
dt  =float(text[1]) # Time increment for calculation
DTWR=float(text[2]) # Time increment for print out

text=fin.readline()
text=text.strip()
text=text.split(',')
PAA=float(text[0])  # Section area of Port
PCI=float(text[1])  # Discharge coefficient of Port for inflow
PCO=float(text[2])  # Discharge coefficient of Port for outflow

text=fin.readline()
text=text.strip()
text=text.split(',')
RWL=float(text[0])  # Reservoir Water Level
TNL=float(text[1])  # Length of pressure tunnel
TNA=float(text[2])  # Section area of pressure tunnel
TNC=float(text[3])  # Head loss coefficient of pressure tunnel

SAA=[]
SEL=[]
SLB=[]
text=fin.readline()
text=text.strip()
text=text.split(',')
NST=int(text[0])             # Number of sections
for i in range(0,NST):
    text=fin.readline()
    text=text.strip()
    text=text.split(',')
    SAA=SAA+[float(text[0])] # Section area of specified level in the shaft
    SEL=SEL+[float(text[1])] # Level definition of sections of shaft (EL)
    SLB=SLB+[text[2]]        # Label for section (text: for drawing)

QTQ=[]
QTI=[]
text=fin.readline()
text=text.strip()
text=text.split(',')
NQT=int(text[0])             # #Number of discharge input points
for i in range(0,NQT):
    text=fin.readline()
    text=text.strip()
    text=text.split(',')
    QTQ=QTQ+[float(text[0])] # Discharge at specified time point
    QTI=QTI+[float(text[1])] # Time at specified time point
fin.close()

# Preprocessing
QI=QTQ[0]
if ICT==1: AFCS=-1.0
if ICT==2: AFCS=QTI[1]
_list = [[-1]*3 for i in range(0,NST)]
for i in range(0,NST):
    _list[i][0]=SEL[i]
    _list[i][1]=SAA[i]
    _list[i][2]=SLB[i]
_list.sort()         # Sort section numbers in small order of elevation
for i in range(0,NST):
    SEL[i]=_list[i][0]
    SAA[i]=_list[i][1]
    SLB[i]=_list[i][2]
del _list
SAA=SAA+[SAA[NST-1]] # Dummy for judgement of top level (section area)
SEL=SEL+[SEL[NST-1]] # Dummy for judgement of top level (EL)

# Main calculation
fout=open(fnameW,'w')
fout.write('Time,WL of Surge tank,Velocity of Tunnel,Discharge,k\n')
nnp=int(DTWR/dt+0.00001) # Step number of print out
dto=dt                   # Time increment
qn=QI                    # Initial discharge
vi=qn/TNA                # Initial velocity of tunnel discharge
zi=TNC*vi*vi             # Initial head loss due to initial velocity
if vi<0.0: zi=-TNC*vi*vi # Definition of sign of initial head loss (direction of flow)
zo=RWL-zi                # Initial water level in Surge Tank
# Print out initial values (time=0)
dts=0.0
zw=zo
vw=vi
zi=zi
qn=qn
WOUT(dts,zw,vw,qn,fout)
#To search the section number with initial water level
for j in range(NST-1,-1,-1):
    if SEL[j]<zo: break
npe=j      #Number of section with initial water level
npw=0      #Initialization of counter for print out
dts=0.0    #Initialization of elapsed time
iret=0
while dts<=TMAX:
    dts=dts+dto  #The time found solution
    npw=npw+1
    fa=SAA[npe]
    dv,dz,q1=RUNGE(qn,vi,zi,fa,dto,dts)
    zw=RWL-zi-dz
    iw=npe+1
    if iw>NST-1: iw=NST-1
    if zw<SEL[npe] or SEL[iw]<zw:
        dti=0.1*dto
        for ic in range(1,11):
            dtw=dts-dto+dti*float(ic)
            zw=RWL-zi
            if SEL[npe+1]<zw: npe=npe+1
            if zw<SEL[npe]: npe=npe-1
            if NST-1<npe: iret=1 # water level greater than USWL
            if npe<0:     iret=2 # water level less than DSWL
            fa=SAA[npe]
            dv,dz,q2=RUNGE(q1,vi,zi,fa,dti,dtw)
            vi=vi+dv
            zi=zi+dz
            q1=q2
    else:
        vi=vi+dv
        zi=zi+dz
        qn=q1
    # Print out results
    if npw>=nnp:
        npw=0
        zw=RWL-zi
        vw=vi
        qn=q1
        WOUT(dts,zw,vw,qn,fout)
fout.close()
print('iret=%1d'%iret)

data=np.loadtxt(fnameW,delimiter=',', skiprows=1, usecols=(0,1))
TI=data[:,0]
WL=data[:,1]
xmin=0.0
xmax=TMAX
fig=plt.figure(figsize=(10, 5))
plt.plot(TI,WL,color='black',lw=2.0,label='WL in the shaft')
plt.hlines(RWL,xmin,xmax,alpha=0.3,color='black',lw=1.5,linestyles='-')
plt.text(xmax-10,RWL,'RWL',fontsize=10,color='black',ha='right',va='bottom')
for i in range(0,NST):
    plt.hlines(SEL[i],xmin,xmax,alpha=0.3,color='black',lw=1.5,linestyles='--')
    plt.text(xmax-10,SEL[i],SLB[i],fontsize=10,color='black',ha='right',va='bottom')
n1=np.argmax(WL)
n2=np.argmin(WL)
plt.text(TI[n1],WL[n1],'max:%.3f(%.1fsec)'%(WL[n1],TI[n1]),fontsize=10,color='black',ha='left',va='bottom')
plt.text(TI[n2],WL[n2],'min:%.3f(%.1fsec)'%(WL[n2],TI[n2]),fontsize=10,color='black',ha='left',va='top')
plt.xlabel('Time (second)')
plt.ylabel('Elevation (EL.m)')
plt.xlim(xmin,xmax)
plt.grid()
plt.savefig(fnameF,dpi=200)
#plt.show()
plt.close()