from ChemPlugin import *
import sys

def write_line(f, cp, gap, then):
    now = cp[0].Report1("Time", "years")
    if (then - now) <= gap / 1e4:
        f.write("%8.4f" % now)
        for c in cp:
            f.write("\t%8.4f" % c.Report1("concentration Na+", "mmol/kg"))
        f.write("\n")
        f.flush()
        then += gap
    return then

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

print("Model diffusion in one dimension\n")

# Simulation Parameters
nx = 100                     # number of instances along x
length = 100.0               # cm
deltax = length / nx         # cm
deltay = 1.0                 # cm
deltaz = 1.0                 # cm
porosity = 0.25              # volume fraction
diffcoef = 1e-6              # cm2/s
xstable = 0.9                # unitless
time_end = 15.0              # years
delta_years = time_end / 3.  # years
next_output = 0.0            # years

# Open output file and write instance positions on the first line.
f = open("Diffusion.txt", "w")
f.write("   years")
for i in range(nx):
    f.write("\t%8.4f" % ((i + 0.5) * deltax))
f.write("\n")

# Configure and initialize the instances
cmd0 =  "volume = " + str(deltax * deltay * deltaz) + " cm3; "
cmd0 += "porosity = " + str(porosity) + "; "
cmd0 += "time end = " + str(time_end) + " years; "
cmd0 += "Xstable = " + str(xstable)

cmd1 = "Na+ = 1 mmol/kg; Cl- = 1 mmol/kg"
cmd2 = "Na+ = 0.001 mmol/kg; Cl- = 0.001 mmol/kg"

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

for i in range(nx):
    cp[i].Config(cmd0)
    if i < nx / 2:
        cp[i].Config(cmd1)
    else:
        cp[i].Config(cmd2)

    if cp[i].Initialize():
        print("Instance " + str(i) + " failed to initialize")
        end_run(-1)

next_output = write_line(f, cp, delta_years, next_output)

# Link the instances
trans = deltay * deltaz * porosity * diffcoef / deltax
for i in range(1, nx):
    link = cp[i].Link(cp[i - 1])
    link.Transmissivity(trans, "cm3/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.AdvanceChemical(): end_run(-1)

    next_output = write_line(f, cp, delta_years, next_output)