#!/usr/bin/env python3

import grass.script as gs

# Set the region for now arbitrary to 100 x 100 m. We might need to specify
# the region later depending on how far the contaminant traveled over the years
gs.run_command("g.region", rows=100, cols=100, res3=1, t=50, b=0, n=100,s=0, w=0, e=100, overwrite=True)

# Define the boundary conditions following to the hydraulic gradient of 1% in
# west-east and 2% in south-north direction with a mark of 35m in the SW corner
gs.run_command("r.mapcalc",expression="bottomBC = if(row() == nrows(), 35 + 0.01*x(), 35)")
gs.run_command("r.mapcalc",expression="leftBC = if(col() == 1 && row()>1 && row()<nrows(), 0.02*y(), 0)")
gs.run_command("r.mapcalc",expression="topBC = if(row() == 1, 0.02*y() + 0.01*x(), 0)")
gs.run_command("r.mapcalc",expression="rightBC = if(col() == ncols() && row()>1 && row()<nrows(), 0.01*x() + 0.02*y(), 0)")
gs.run_command("r.mapcalc",expression="phead = topBC + bottomBC + leftBC + rightBC")

# Define parameter raster maps
gs.run_command("r.mapcalc", expression="status = if(col() == 1 || col() == ncols() || row() == 1 || row() == nrows(), 2, 1)", overwrite=True)
gs.run_command("r.mapcalc", expression="null = 0.0",overwrite=True)
gs.run_command("r.mapcalc", expression="top = 43",overwrite=True)
gs.run_command("r.mapcalc", expression="hydcond = 0.00005",overwrite=True)
gs.run_command("r.mapcalc", expression="poros = 0.12",overwrite=True)

# Compute a steady state groundwater flow as initial flow pattern
gs.run_command("r.gwflow", top="top", bottom="null",phead="phead",
    status="status", hc_x="hydcond", hc_y="hydcond", s="poros",
    output="gwresult_0", dt=8640000, type="unconfined",
    error="0.00000001",
    maxit="100000000000",
    vx="velx", vy="vely", overwrite=True)

# Define parameters for solute transport
gs.run_command("r.mapcalc", expression="tstatus = if(col() == 1 || col() == ncols() || row() == 1 || row() == nrows(), 3, 1)", overwrite=True)
gs.run_command("r.mapcalc", expression="diff = 0.0000001",overwrite=True)
gs.run_command("r.mapcalc", expression="cin = if(col()==ncols()-15 && row()==15, -0.001, 0)",overwrite=True)
gs.run_command("r.mapcalc", expression="rd = 1.6",overwrite=True)
gs.run_command("r.mapcalc", expression="stresult_0 = 0",overwrite=True)
gs.run_command("r.mapcalc", expression="q = if(col()==ncols()-15 && row()==15, 0.01, 0)",overwrite=True)

# Calculate flow and concentration for the first 30days of solute injection
# with pumping 
for t in range(0,30,1):
    gs.run_command("r.gwflow", top="top", bottom="null",phead="gwresult_"+str(t),
        status="status", hc_x="hydcond", hc_y="hydcond", s="poros",
        output="gwresult_"+str(t+1), dt=86400, type="unconfined",
        error="0.0000001",
        maxit="100000000000",
        q="q",
        vx="velx", vy="vely", overwrite=True)

    gs.run_command("r.solute.transport", \
        top="top", bottom="null", phead="gwresult", \
        status="tstatus", q="q",hc_x="hydcond", hc_y="hydcond", rd="rd", \
        cs="null", nf="poros", output="stresult_" + str(t + 1), \
        cin="cin",\
        dt=86400, diff_x="diff", diff_y="diff", c="stresult_" + str(t), \
        al=0.01, at=0.01, overwrite=True, loops=2)

# subsequent 10 days of pumping without solute injection
for t in range(30,40,1):
    gs.run_command("r.gwflow", top="top", bottom="null",phead="gwresult_"+str(t),
        status="status", hc_x="hydcond", hc_y="hydcond", s="poros",
        output="gwresult_"+str(t+1), dt=86400, type="unconfined",
        error="0.0000001",
        maxit="100000000000",
        q="q",
        vx="velx", vy="vely", overwrite=True)

    gs.run_command("r.solute.transport", \
        top="top", bottom="null", phead="gwresult", \
        status="tstatus", q="q",hc_x="hydcond", hc_y="hydcond", rd="rd", \
        cs="null", nf="poros", output="stresult_" + str(t + 1), \
        cin="null",\
        dt=86400, diff_x="diff", diff_y="diff", c="stresult_" + str(t), \
        al=0.01, at=0.01, overwrite=True, loops=2)

# subsequent 30 days regular transport in groundwater flow
for t in range(40,70,1):
    gs.run_command("r.gwflow", top="top", bottom="null",phead="gwresult_"+str(t),
        status="status", hc_x="hydcond", hc_y="hydcond", s="poros",
        output="gwresult_"+str(t+1), dt=86400, type="unconfined",
        error="0.0000001",
        maxit="100000000000",
        q="null",
        vx="velx", vy="vely", overwrite=True)

    gs.run_command("r.solute.transport", \
        top="top", bottom="null", phead="gwresult", \
        status="tstatus", q="null",hc_x="hydcond", hc_y="hydcond", rd="rd", \
        cs="null", nf="poros", output="stresult_" + str(t + 1), \
        cin="null",\
        dt=86400, diff_x="diff", diff_y="diff", c="stresult_" + str(t), \
        al=0.01, at=0.01, overwrite=True, loops=2)

# Estimating locations of pumping and injection borehole for a pump-&-treat setup assuming
# steady state flow conditions.
gs.run_command("r.mapcalc",expression="pump=if(col()==70 && row()==nrows()-57, -0.015,0)")
gs.run_command("r.mapcalc",expression="inje=if(col()==50 && row()==nrows()-30, 0.015,0)")
gs.run_command("r.mapcalc",expression="wells = pump + inje")

gs.run_command("r.gwflow", top="top", bottom="null",phead="gwresult_70",
    status="status", hc_x="hydcond", hc_y="hydcond", s="poros",
    output="gwresult_ph", dt=86400000, type="unconfined",
    q="wells", vx="velx", vy="vely", overwrite=True)