#!/usr/bin/env python3
# Copyleft Rostislav Kouznetsov 10.2022- 

import matplotlib.pyplot as plt
import numpy as np
import scipy.integrate as SI
import os
import sys
#plt.style.use('seaborn-poster')
import pandas as pd

np.seterr(invalid='raise')

#
# X -- vector of rho and v along the pipe
#

km = int(sys.argv[1]) ##Pipe length in kilometers
L = km*1000.
N = km // 2     ## Points along the pipe

#L = 1.23e6 ##m pipe length full
#L = 1.03e6 ##m pipe length long end
#L = 0.2e6 ##m pipe length short end
#L = 0.3e6 ##m pipe length short end
# L comes from the command line
D = 1.153 ## m pipe diameter
RToverMu = 8.31 * 278 / 0.016  ## +5C assumed

S = np.pi*(D/2)**2

VdW = True ## VanDerWaals equation of state
if VdW:
    ## van der Waals gas
    rho_out = 4.9  ##
    rho0 = 90.
else:
    # Ideal-gas approximation
    ## Methane density at standard conditions 0.66 kg/m3
    rho_out = 7*0.66; ## density at the open end
    rho0 = 105*0.66; ## density at full pressure

rho_atm = 0.66;

seconds=0

#https://www.nuclear-power.com/nuclear-engineering/fluid-dynamics/major-head-loss-friction-loss/friction-factor-turbulent-flow-colebrook/
#                                      10m/s          100 m/s
# @1  bar  nu = 15e-6  m2/s            7e6             7e7
# @10 bar  nu = 1.5e-6 m2/s            7e5             7e6
# @100 bar nu = 0.2e-6 m2/s            5e4             5e5 
## nu = 15e-6 * rho_atm/rho

#Blasius formula:   f = (100 Re)^-¼
#f =  0.02    (100 bar, 10m/s)
#f =  0.006    (10 bar, 100m/s)
#f = 0.01  ## geometric mean

delta = L/(N-1)

def F(t, Y): ## Derivative over time
    global seconds
    
    rho = Y[0:N]
    v = Y[N:]

#    ## Boundary conditions
    rho[-1]  = rho_out
    v[0] = 0
    
    out = np.empty(Y.shape)

    totv = np.sum(v*v)
    totrhodelta = np.sum((rho - rho_out)**2)
    if (t > seconds):
        #print(rho[-5:])
        print("t=%10.4f, vnorm = %g, rhonorm= %g"%(t/3600., totv/N, totrhodelta/N))
        seconds += 3600

    ## Stop iterating when velocity is small 
    if totv < N*1e-6 and totrhodelta < N*1e-6: ## <1m/s
        out[:] = 0.
        return out


    drhodx = np.gradient(rho, delta)
    dvdx = np.gradient(v, delta)


    ## 1e5 to keep f finite at v=0
    try:
        f =  np.power((100 * np.abs(v) /(15e-6 * rho_atm/np.maximum(rho, rho_atm) ) + 1e5), -0.25)
        ### densities below rho_out should be allowed
    except:
        print(rho[-5:])
        __import__('pdb').set_trace()
    # f = 0.01


    out[0:N] = - v * drhodx - rho * dvdx   ## drho/dt

    if VdW:
        rho_l = 374. ## kg/m3 "liquid density"
        a_over_mu2 = 898 ## J m3 / kg2  
        #out[N:] = 1./rho * ( - v * out[0:N] - (RToverMu / ((1.- rho/rho_l)**2) - 2* rho * a_over_mu2 )* drhodx - rho*v*np.abs(v) * f/(2*D))
        out[N:] = 1./rho * (  - (RToverMu / ((1.- rho/rho_l)**2) - 2* rho * a_over_mu2 )* drhodx - rho*v*np.abs(v) * f/(2*D))
    else:
        out[N:] = 1./rho * ( - v * out[0:N] - RToverMu * drhodx - rho*v*np.abs(v) * f/(2*D))

    return out

#Initial condition
X0 = np.zeros((2*N,))
X0[0:N] = rho0


ndays=8
step= 15
maxt = 24*3600*ndays
t_eval = np.linspace(0,maxt, num = maxt//step+1)
sol = SI.solve_ivp(F, (0,maxt+1), X0, t_eval=t_eval) ## Solve the initial-value problem
#sol = SI.solve_ivp(F, (0,maxt+1), X0, t_eval=t_eval, method="LSODA", atol=1e-3) ## Solve the initial-value problem



x = np.arange(N) * L /(N-1) / 1000  ## C
seconds = np.around(sol.t).astype(np.int32)
kilometers = ["%d"%(m) for m in x ] ## Labels in km
iL = int(np.round(L/1000))  ##Length in km (to mark the case)


fig = plt.figure(figsize = (12, 12))
axrho = fig.add_subplot(211)
axv = fig.add_subplot(212)
#Plot every hour
axrho.plot(x, sol.y[0:N,::int(3600/step)])
axrho.grid()
axv.plot(x, sol.y[N:,::int(3600/step)]  )
axv.set_xlabel('X, km')
axv.set_ylabel('V, m/s')
axrho.set_ylabel('Rho kg/m3')
axv.grid()
fig.tight_layout()
fig.savefig('case%04dRhoV.png'%(iL))
#os.system("xli 1.png")


fig = plt.figure(figsize = (12, 4))
axrate = fig.add_subplot(111)
flow = sol.y[N-2,:]* sol.y[2*N-2,:]*S #Rho*v #last point is unstable...
axrate.plot(sol.t/3600, flow)
axrate.set_xlabel('t, hours')
axrate.set_ylabel('Discharge, kg/s')
fig.tight_layout()
fig.savefig('case%04dDischarge.png'%(iL))
gasTot = rho0*L*S*1e-6 #kt
gasLeaked = np.sum(flow)*step*1e-6 # in kt
gasLeft = (np.sum(sol.y[1:N-1,-1]) + 0.5*(sol.y[0,-1]+sol.y[N-1,-1]))*S*delta*1e-6 ##kt  # Half-step at ends
 #  
print("Len= %4d km, gasTot= %.2f kt, gasLeft = %.2f kt, gasLeaked= %.2f kt, delta = %.2f kt"%(L/1000,gasTot, gasLeft, gasLeaked, gasTot- gasLeaked-gasLeft))


tslice, = np.where((seconds < 600)*(seconds%60 == 0) + (seconds < 3600)*(seconds%600 == 0) + (seconds%3600 == 0)) ##Every minute, then every 10 minutes, then every hour
columnheads = ["%dm"%(ss//60) if ss < 3600 else "%dh"%(ss//3600)  for ss in seconds[tslice] ]


##
##
pd.DataFrame(sol.y[:N,tslice],index=kilometers, columns=columnheads).to_csv("case%04dRho.csv"%(iL))
pd.DataFrame(sol.y[N:,tslice],index=kilometers, columns=columnheads).to_csv("case%04dV.csv"%(iL))
pd.DataFrame(sol.y[N:,tslice]*sol.y[0:N,tslice]*S,index=kilometers, columns=columnheads).to_csv("case%04dRhoVS.csv"%(iL))
pd.DataFrame(flow,index=seconds, columns=["Discharge(kg/s)"]).to_csv("case%04dDischarge.csv"%(iL))
pd.DataFrame(sol.y[1,:],index=seconds, columns=["Rho0(kg/m3)"]).to_csv("case%04dRho0.csv"%(iL))


