from ChemPlugin import *
import sys

def write_results(cp, gap, then):
    now = cp[0].Report1("Time", "years")
    if (then - now) < gap / 1e4:
        for i in range(len(cp)):
            cp[i].PrintOutput("RTM.txt", "Instance " + str(i))
        then += gap
    return then

def end_run(code):
    input("Done!")
    sys.exit(code)

print("Model reactive transport in one dimension\n")

# Query user for input file
while True:
    filename = input("Enter RTM input script: ")
    try:
        inputf = open(filename, 'r')
    except:
        print("The input file named " + filename + " doesn't exist")
    else:
        break

# simulation parameters
nx = 100                      # number of instances along x
length = 100.0                # m
deltax = length / nx          # m
deltay = 1.0                  # m
deltaz = 1.0                  # m
time_end = 10.0               # years
delta_years = time_end / 5.0  # years
next_output = 0.0             # years

# Create the inlet and interior instances
cp_inlet = ChemPlugin()
cp_inlet.Console("stdout")

cp = [ChemPlugin() for i in range(nx)]
cp[0].Console("stdout")
cp[0].Config("pluses = banner")

cmd = "volume = " + str(deltax * deltay * deltaz) + " m3; "
cmd += "time end = " + str(time_end) + " years"

for c in cp:
    c.Config(cmd)

# Configure inlet, initial domain
scope = 0
for line in inputf:
    if line.split() == ["go"]:
        break
    elif line.split() == ["scope", "inlet"]:
        scope = 1
    elif line.split() == ["scope", "initial"]:
        scope = 2
    elif scope == 1:
        cp_inlet.Config(line)
    elif scope == 2:
        for c in cp:
            c.Config(line)

# Initialize the inlet and interior instances; write out initial conditions
if cp_inlet.Initialize():
    print("Inlet failed to initialize")
    end_run(-1)

cp_inlet.PrintOutput("RTM.txt", "Inlet fluid")

for c in cp:
    if c.Initialize():
        print("Interior instance failed to initialize\n")
        end_run(-1)

next_output = write_results(cp, delta_years, next_output)

# Query user for velocity; calculate flow rate and transmissivities
veloc_in = float(input("Please enter fluid velocity in m/yr: "))

velocity = veloc_in / 31557600.                         # m/s
porosity = cp[0].Report1("porosity")                    # unitless
flow = deltay * deltaz * porosity * velocity            # m3/s

diffcoef = 1e-10                                        # m2/s
dispersivity = 1.0                                      # m
dispcoef = velocity * dispersivity + diffcoef           # m2/s
trans = deltay * deltaz * porosity * dispcoef / deltax  # m3/s

tcond = 2.0                                             # W/m/K
ttrans = deltay * deltaz * tcond / deltax               # W/K

# Link the instances
link = cp[0].Link(cp_inlet)
link.FlowRate(flow, "m3/s")
link.Transmissivity(trans, "m3/s")
link.HeatTrans(ttrans, "W/K")

for i in range(1, nx):
    link = cp[i].Link(cp[i - 1])
    link.FlowRate(flow, "m3/s")
    link.Transmissivity(trans, "m3/s")
    link.HeatTrans(ttrans, "W/K")

link = cp[nx - 1].Link()
link.FlowRate(-flow, "m3/s")

# Time marching loop.
while True:
    deltat = 1e99
    for c in cp:
        deltat = min(deltat, c.ReportTimeStep())
    for c in cp:
        if c.AdvanceTimeStep(deltat): end_run(0)
    for c in cp:
        if c.AdvanceTransport(): end_run(-1)
    for c in cp:
        if c.AdvanceHeatTransport(): end_run(-1)
    for c in cp:
        if c.AdvanceChemical(): end_run(-1)

    next_output = write_results(cp, delta_years, next_output)