from reaktoro import *
import numpy as np
from tqdm.notebook import tqdm
import os

# Import components of bokeh library
from bokeh.io import show, output_notebook
from bokeh.layouts import column
from bokeh.plotting import figure
from bokeh.models import Range1d, ColumnDataSource
from bokeh.layouts import gridplot

second = 1
minute = 60
hour = 60 * minute
day = 24 * hour
week = 7 * day
year = 365 * day

# Discretization parameters
xl = 0.0                # x-coordinate of the left boundary
xr = 100.0              # x-coordinate of the right boundary
ncells = 100            # number of cells in the discretization
nsteps = 1000           # number of steps in the reactive transport simulation
dx = (xr - xl) / ncells # length of the mesh cells (in units of m)
dt = 0.1*day            # time step

# Physical parameters
D = 0               # diffusion coefficient (in units of m2/s)
v = 3e-5            # fluid pore velocity (in units of m/s)
T = 60.0 + 273.15   # temperature (in units of K)
P = 200 * 1.01325 * 1e5   # pressure (in units of Pa)
phi = 0.1           # the porosity

xcells = np.linspace(xl, xr, ncells + 1)  # interval [xl, xr] split into ncells

dirichlet = False  # parameter that determines whether Dirichlet BC must be used

CFL = v * dt / dx
print("CFL = ", CFL)
assert CFL <= 1.0, f"Make sure that CFL = {CFL} is less that 1.0"

output_quantities = """
    pH
    speciesMolality(H+)
    speciesMolality(HS-)
    speciesMolality(S2--)
    speciesMolality(CO3--)
    speciesMolality(HSO4-)
    speciesMolality(H2S(aq))
    speciesMolality(Fe++)
    speciesMolality(Pyrite)
    speciesMolality(Hematite)
""".split()

column_quantities = """
    pH
    Hcation
    HSanion
    S2anion
    CO3anion
    HSO4anion
    H2Saq
    Fe2cation
    pyrite
    hematite
""".split()

columns = ['step', 'x'] + column_quantities
import pandas as pd
df = pd.DataFrame(columns=columns)

folder_results = 'results-rt-scavenging-with-hematite'
def make_results_folders():
    os.system('mkdir -p ' + folder_results)

def define_chemical_system():

    # Construct the chemical system with its phases and species
    db = Database('supcrt07.xml')

    # Set PHREEQC parameters to the model
    dhModel = DebyeHuckelParams()
    dhModel.setPHREEQC()

    # Define the editor
    editor = ChemicalEditor(db)
    # Select species of the of the aqueous phase by providing the list of elements
    editor.addAqueousPhaseWithElements("C Ca Cl Fe H K Mg Na O S").\
        setChemicalModelDebyeHuckel(dhModel)
    # Select mineral phases
    editor.addMineralPhase('Pyrite')
    editor.addMineralPhase('Hematite')
    editor.addMineralPhase('Quartz')

    # Define chemical system
    system = ChemicalSystem(editor)

    return system

def define_initial_condition(system):
    problem_ic = EquilibriumInverseProblem(system)
    problem_ic.setTemperature(T)
    problem_ic.setPressure(P)
    problem_ic.add("H2O", 58.0, "kg")
    problem_ic.add("Cl-", 1122.3e-3, "kg")
    problem_ic.add("Na+", 624.08e-3, "kg")
    problem_ic.add("SO4--", 157.18e-3, "kg")
    problem_ic.add("Mg++", 74.820e-3, "kg")
    problem_ic.add("Ca++", 23.838e-3, "kg")
    problem_ic.add("K+", 23.142e-3, "kg")
    problem_ic.add("HCO3-", 8.236e-3, "kg")
    problem_ic.add("O2(aq)", 58e-12, "kg")
    problem_ic.add("Pyrite", 0.0, "mol")
    problem_ic.add("Hematite", 0.5, "mol") # 10 % of all minerals
    problem_ic.add("Quartz", 0.5 * 9, "mol") # 90 % of all minerals
    problem_ic.pH(8.951, "HCl", "NaOH")

    # Calculate the equilibrium states for the initial conditions
    state_ic = equilibrate(problem_ic)
    state_ic.scalePhaseVolume('Aqueous', phi, 'm3')  # 10% of porosity
    state_ic.scaleVolume(1.0, 'm3')

    # Fetch teh value of the ph in the initial chemical state
    props = state_ic.properties()
    evaluate_pH = ChemicalProperty.pH(system)
    pH = evaluate_pH(props)
    print("ph(IC) = ", pH.val)

    return state_ic

def define_boundary_condition(system):

    # Define the boundary condition of the reactive transport modeling problem
    problem_bc = EquilibriumInverseProblem(system)
    problem_bc.setTemperature(T)
    problem_bc.setPressure(P)
    problem_bc.add("H2O", 58.0, "kg")
    problem_bc.add("Cl-", 1122.3e-3, "kg")
    problem_bc.add("Na+", 624.08e-3, "kg")
    problem_bc.add("SO4--", 157.18e-3, "kg")
    problem_bc.add("Mg++", 74.820e-3, "kg")
    problem_bc.add("Ca++", 23.838e-3, "kg")
    problem_bc.add("K+", 23.142e-3, "kg")
    problem_bc.add("HCO3-", 8.236e-3, "kg")
    problem_bc.add("O2(aq)", 58e-12, "kg")
    problem_bc.add("Pyrite", 0.0, "mol")
    problem_bc.add("Hematite", 0.0, "mol")
    problem_bc.add("HS-", 0.0196504, "mol")
    problem_bc.add("H2S(aq)", 0.167794, "mol")
    problem_bc.pH(5.726, "HCl", "NaOH")

    # Calculate the equilibrium states for the boundary conditions
    state_bc = equilibrate(problem_bc)
    # Scale the boundary condition state to 1 m3
    state_bc.scaleVolume(1.0, 'm3')

    # Fetch ph of the evaluated chemical state
    props = state_bc.properties()
    evaluate_pH = ChemicalProperty.pH(system)
    pH = evaluate_pH(props)
    print("ph(BC) = ", pH.val)

    return state_bc

def partition_indices(system):
    nelems = system.numElements()

    ifluid_species = system.indicesFluidSpecies()
    isolid_species = system.indicesSolidSpecies()

    return nelems, ifluid_species, isolid_species

def partition_elements_in_mesh_cell(ncells, nelems, state_ic, state_bc):

    # The concentrations of each element in each mesh cell (in the current time step)
    b = np.zeros((ncells, nelems))

    # Initialize the concentrations (mol/m3) of the elements in each mesh cell
    b[:] = state_ic.elementAmounts()

    # The concentrations (mol/m3) of each element in the fluid partition, in each mesh cell
    bfluid = np.zeros((ncells, nelems))

    # The concentrations (mol/m3) of each element in the solid partition, in each mesh cell
    bsolid = np.zeros((ncells, nelems))

    # Initialize the concentrations (mol/m3) of each element on the boundary
    b_bc = state_bc.elementAmounts()

    return b, bfluid, bsolid, b_bc

def transport(states, bfluid, bsolid, b, b_bc, nelems, ifluid_species, isolid_species):
    # Collect the amounts of elements from fluid and solid partitions
    for icell in range(ncells):
        bfluid[icell] = states[icell].elementAmountsInSpecies(ifluid_species)
        bsolid[icell] = states[icell].elementAmountsInSpecies(isolid_species)

    # Get the porosity of the boundary cell
    bc_cell = 0
    phi_bc = states[bc_cell].properties().fluidVolume().val / states[bc_cell].properties().volume().val

    # Transport each element in the fluid phase
    for j in range(nelems):
        transport_fullimplicit(bfluid[:, j], dt, dx, v, D, phi_bc * b_bc[j])

    # Update the amounts of elements in both fluid and solid partitions
    b[:] = bsolid + bfluid

    return bfluid, bsolid, b

def transport_fullimplicit(u, dt, dx, v, D, ul):

    # Number of DOFs
    n = len(u)
    alpha = D * dt / dx ** 2
    beta = v * dt / dx

    # Upwind finite volume scheme
    a = np.full(n, -beta - alpha)
    b = np.full(n, 1 + beta + 2 * alpha)
    c = np.full(n, -alpha)

    # Set the boundary condition on the left cell
    if dirichlet:
        # Use Dirichlet BC boundary conditions
        b[0] = 1.0
        c[0] = 0.0
        u[0] = ul

    else:
        # Flux boundary conditions (implicit scheme for the advection)
        # Left boundary
        b[0] = 1 + alpha + beta
        c[0] = -alpha  # stays the same as it is defined -alpha
        u[0] += beta * ul  # = dt/dx * v * g, flux that we prescribe is equal v * ul

    # Right boundary is free
    a[-1] = - beta
    b[-1] = 1 + beta

    # Solve a tridiagonal matrix equation
    thomas(a, b, c, u)

def thomas(a, b, c, d):
    n = len(d)
    c[0] /= b[0]
    for i in range(1, n - 1):
        c[i] /= b[i] - a[i] * c[i - 1]
    d[0] /= b[0]
    for i in range(1, n):
        d[i] = (d[i] - a[i] * d[i - 1]) / (b[i] - a[i] * c[i - 1])
    x = d
    for i in reversed(range(0, n - 1)):
        x[i] -= c[i] * x[i + 1]
    return x

def reactive_chemistry(solver, states, b):
    # Equilibrating all cells with the updated element amounts
    for icell in range(ncells):
        solver.solve(states[icell], T, P, b[icell])
    return states

def outputstate_df(step, system, states):
    # Define the instance of ChemicalQuantity class
    quantity = ChemicalQuantity(system)

    # Create the list with empty values to populate with chemical properties
    values = [None] * len(columns)
    for state, x in zip(states, xcells):

        # Populate values with number of reactive transport step and spacial coordinates
        values[0] = step
        values[1] = x

        # Update the
        quantity.update(state)
        for quantity_name, i in zip(output_quantities, range(2, len(states))):
            values[i] = quantity.value(quantity_name)
        df.loc[len(df)] = values

make_results_folders()

system = define_chemical_system()

state_ic = define_initial_condition(system)

state_bc = define_boundary_condition(system)

nelems, ifluid_species, isolid_species = partition_indices(system)

b, bfluid, bsolid, b_bc = partition_elements_in_mesh_cell(ncells, nelems, state_ic, state_bc)

states = [state_ic.clone() for _ in range(ncells + 1)]

solver = EquilibriumSolver(system)

step = 0  # the current step number
t = 0.0  # the current time (in seconds)

# Output the initial state of the reactive transport calculation
outputstate_df(step, system, states)

with tqdm(total=nsteps, desc="Reactive transport simulations") as pbar:
    while step < nsteps:
        # Perform transport calculations
        bfluid, bsolid, b = transport(states, bfluid, bsolid, b, b_bc, nelems, ifluid_species, isolid_species)

        # Perform reactive chemical calculations
        states = reactive_chemistry(solver, states, b)


        # Increment time step and number of time steps
        t += dt
        step += 1

        # Output the current state of the reactive transport calculation
        outputstate_df(step, system, states)

        # Update a progress bar
        pbar.update(1)

df

df.to_csv(folder_results + '/rt.scavenging-with-hematite.csv', index=False)

def titlestr(t):
    t = t / minute  # Convert from seconds to minutes
    h = int(t) / 60  # The number of hours
    m = int(t) % 60  # The number of remaining minutes
    return 'Time: %2dh %2dm' % (h, m)

def plot_figures_ph(steps):

    plots = []
    for i in steps:
        print("On pH figure at time step: {}".format(i))
        t = i * dt
        source = ColumnDataSource(df[df['step'] == i])

        p = figure(plot_width=600, plot_height=250)
        p.line(x='x', y='pH', color='teal', line_width=2, legend_label='pH', source=source)
        p.x_range = Range1d(xl - 1, xr - 1)
        p.y_range = Range1d(5.0, 10.0)
        p.xaxis.axis_label = 'Distance [m]'
        p.yaxis.axis_label = 'pH'
        p.legend.location = 'center_right'
        p.title.text = titlestr(t)

        plots.append([p])

    grid = gridplot(plots)
    show(grid)

def plot_figures_minerals(steps):
    plots = []
    for i in steps:
        print("On pyrite-hematite-siderite figure at time step: {}".format(i))
        t = i * dt
        source = ColumnDataSource(df[df['step'] == i])

        p = figure(plot_width=600, plot_height=250)
        p.line(x='x', y='pyrite', color='blue', line_width=2, legend_label='Pyrite',
               muted_color='blue', muted_alpha=0.2, source=source)
        p.line(x='x', y='hematite', color='orange', line_width=2, legend_label='Hematite',
               muted_color='orange', muted_alpha=0.2, source=source)
        p.x_range = Range1d(xl - 1, xr - 1)
        p.y_range = Range1d(-1e-3, 0.012)
        p.xaxis.axis_label = 'Distance [m]'
        p.yaxis.axis_label = 'Concentration [mol/m3]'
        p.legend.location = 'top_right'
        p.title.text = titlestr(t)
        p.legend.click_policy = 'mute'
        plots.append([p])

    grid = gridplot(plots)
    show(grid)

def plot_figures_pyrite(steps):
    plots = []
    for i in steps:
        print("On pyrite figure at time step: {}".format(i))
        t = i * dt
        source = ColumnDataSource(df[df['step'] == i])

        p = figure(plot_width=600, plot_height=250)
        p.line(x='x', y='pyrite', color='blue', line_width=2, legend_label='Pyrite',
               muted_color='blue', muted_alpha=0.2, source=source)
        p.x_range = Range1d(xl - 1, xr + 1)
        p.xaxis.axis_label = 'Distance [m]'
        p.yaxis.axis_label = 'Concentration [mol/m3]'
        p.legend.location = 'top_right'
        p.title.text = titlestr(t)
        p.legend.click_policy = 'mute'
        plots.append([p])

    grid = gridplot(plots)
    show(grid)

def plot_figures_hematite(steps):
    plots = []
    for i in steps:
        print("On hematite figure at time step: {}".format(i))
        t = i * dt
        source = ColumnDataSource(df[df['step'] == i])

        p = figure(plot_width=600, plot_height=250)
        p.line(x='x', y='hematite', color='orange', line_width=2, legend_label='Hematite',
               muted_color='orange', muted_alpha=0.2, source=source)


        p.x_range = Range1d(xl - 1, xr + 1)
        p.xaxis.axis_label = 'Distance [m]'
        p.yaxis.axis_label = 'Concentration [mol/m3]'
        p.legend.location = 'top_right'
        p.title.text = titlestr(t)
        p.legend.click_policy = 'mute'
        plots.append([p])

    grid = gridplot(plots)
    show(grid)

def plot_figures_aqueous_species(steps):
    plots = []
    for i in steps:
        print("On aqueous species figure at time step: {}".format(i))
        source = ColumnDataSource(df[df['step'] == i])
        t = dt * i

        p = figure(plot_width=600, plot_height=300, y_axis_type = 'log',)
        p.line(x='x', y='HSanion', color='darkcyan', line_width=2, legend_label='HS-', source=source)
        p.line(x='x', y='S2anion', color='darkorange', line_width=2, legend_label='S2--', source=source)
        p.line(x='x', y='CO3anion', color='seagreen', line_width=2, legend_label='CO3--', source=source)
        p.line(x='x', y='HSO4anion', color='indianred', line_width=2, legend_label='HSO4-', source=source)
        p.line(x='x', y='H2Saq', color='gray', line_width=2, legend_label='H2S(aq)', source=source)
        p.line(x='x', y='Hcation', color='darkviolet', line_width=2, legend_label='H+', source=source)
        p.line(x='x', y='Fe2cation', color='darkblue', line_width=2, legend_label='Fe++', source=source)
        p.x_range = Range1d(xl - 1, xr - 1)
        p.y_range = Range1d(1e-16, 1e-2)
        p.xaxis.axis_label = 'Distance [m]'
        p.yaxis.axis_label = 'Concentration [molal]'
        p.legend.location = 'top_right'
        p.title.text = titlestr(t)
        p.legend.click_policy = 'mute'
        plots.append([p])

    grid = gridplot(plots)
    show(grid)

selected_steps_to_plot = [120, 480, 960]
assert all(step <= nsteps for step in selected_steps_to_plot), f"Make sure that selected steps are less than " \
                                                               f"total amount of steps {nsteps}"

output_notebook()

plot_figures_ph(selected_steps_to_plot)

plot_figures_minerals(selected_steps_to_plot)

plot_figures_aqueous_species(selected_steps_to_plot)

step = 25

def modify_doc(doc):
    # Initialize the data by the initial chemical state
    source = ColumnDataSource(df[df['step'] == 0])

    # Auxiliary function that returns a string for the title of a figure in the format Time: #h##m
    def titlestr(t):
        t = t / minute  # Convert from seconds to minutes
        h = int(t) / 60  # The number of hours
        m = int(t) % 60  # The number of remaining minutes
        return 'Time: %2dh %2dm' % (h, m)

    # Plot for ph
    p1 = figure(plot_width=500, plot_height=250)
    p1.line(x='x', y='pH', color='teal', line_width=2, legend_label='pH', source=source)
    p1.x_range = Range1d(xl - 1, xr - 1)
    p1.y_range = Range1d(5.0, 10.0)
    p1.xaxis.axis_label = 'Distance [m]'
    p1.yaxis.axis_label = 'pH'
    p1.legend.location = 'center_right'
    p1.title.text = titlestr(0 * dt)

    # Plot for calcite and dolomite

    p2 = figure(plot_width=500, plot_height=250)
    p2.line(x='x', y='pyrite', color='blue', line_width=2, legend_label='Pyrite', muted_color='blue',
            muted_alpha=0.2, source=source)
    p2.line(x='x', y='hematite', color='orange', line_width=2, legend_label='Hematite', muted_color='orange',
            muted_alpha=0.2, source=source)
    p2.x_range = Range1d(xl - 1, xr - 1)
    p2.y_range = Range1d(-1e-3, 0.012)
    p2.xaxis.axis_label = 'Distance [m]'
    p2.yaxis.axis_label = 'Concentrations [mol/m3]'
    p2.legend.location = 'top_right'
    p2.title.text = titlestr(0 * dt)
    p2.legend.click_policy = 'mute'

    p3 = figure(plot_width=500, plot_height=250, y_axis_type='log')
    p3.line(x='x', y='HSanion', color='darkcyan', line_width=2, legend_label='HS-', source=source)
    p3.line(x='x', y='S2anion', color='darkorange', line_width=2, legend_label='S2--', source=source)
    p3.line(x='x', y='CO3anion', color='seagreen', line_width=2, legend_label='CO3--', source=source)
    p3.line(x='x', y='HSO4anion', color='indianred', line_width=2, legend_label='HSO4-', source=source)
    p3.line(x='x', y='H2Saq', color='gray', line_width=2, legend_label='H2S(aq)', source=source)
    p3.line(x='x', y='Hcation', color='darkviolet', line_width=2, legend_label='H+', source=source)
    p3.line(x='x', y='Fe2cation', color='darkblue', line_width=2, legend_label='Fe++', source=source)
    p3.x_range = Range1d(xl - 1, xr - 1)
    p3.y_range = Range1d(1e-16, 1e-2)
    p3.xaxis.axis_label = 'Distance [m]'
    p3.yaxis.axis_label = 'Concentration [molal]'
    p3.legend.location = 'top_right'
    p3.title.text = titlestr(0 * dt)
    p3.legend.click_policy = 'mute'

    layout = column(p1, p2, p3)

    # Function that return the data dictionary with provided index of the file
    def update():

        if source.data['step'][0] + 1 <= nsteps:
            step_number = source.data['step'][0] + step
        else:
            step_number = 0

        new_source = ColumnDataSource(df[df['step'] == step_number])
        new_data = dict(index=np.linspace(0, ncells, ncells + 1, dtype=int),
                        step=new_source.data['step'],
                        x=new_source.data['x'],
                        pH=new_source.data['pH'],
                        pyrite=new_source.data['pyrite'],
                        hematite=new_source.data['hematite'],
                        HSanion=new_source.data['HSanion'],
                        S2anion=new_source.data['S2anion'],
                        CO3anion=new_source.data['CO3anion'],
                        HSO4anion=new_source.data['HSO4anion'],
                        H2Saq=new_source.data['H2Saq'],
                        Hcation=new_source.data['Hcation'],
                        Fe2cation=new_source.data['Fe2cation'])

        p1.title.text = titlestr(step_number * dt)
        p2.title.text = titlestr(step_number * dt)
        p3.title.text = titlestr(step_number * dt)

        source.stream(new_data, rollover=ncells+1)

    doc.add_periodic_callback(update, 500)
    doc.add_root(layout)

output_notebook()
show(modify_doc, notebook_url="http://localhost:8888")