#!/usr/bin/env python3

import grass.script as gs

# set region
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)

# generate simple raster maps
gs.run_command("r.mapcalc", expression="phead = if(col() == 1, 39, 38)", overwrite=True)
gs.run_command("r.mapcalc", expression="status = if(col() == 1 || col() == 100,\
 2, 1)", overwrite=True)
gs.run_command("r.mapcalc", expression="null = 0.0",overwrite=True)
gs.run_command("r.mapcalc", expression="top = 45",overwrite=True)

# generate artificial heterogeneity
gs.run_command("r.random.surface", output="heterogenMap", distance=30.0, 
flat=10.0, seed=500, high=100, overwrite=True)

# apply heterogeneity on maps
# use the statistics tool to calculate the mean of heterogenMap (46.61), 
# subtract the mean and divide by 100 to obtain a distribution of cell values
# between -0.5 and 0.5. Then apply a factor (6) to achieve the desired 
# standard deviation of 0.000075 m/s. Use the statistics tool to verify
gs.run_command("r.mapcalc", expression="hydcond = 0.0015 * (1.0 + 0.2 *  \
(heterogenMap - 46.61)/100 * 6) ",overwrite=True)

gs.run_command("r.mapcalc", expression="poros = 0.25",overwrite=True)

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

gs.run_command("r.mapcalc", expression="c = if(col() == 5 && row() == 50, 50.0, 0.0)",overwrite=True)
gs.run_command("r.mapcalc", expression="tstatus = if(col() == 1 || col() == 100 , 3, 1)",overwrite=True)
gs.run_command("r.mapcalc", expression="diff = 0.00001",overwrite=True)
gs.run_command("r.mapcalc", expression="R = 1.0",overwrite=True)

# Calculate initial timestep
gs.run_command("r.solute.transport", top="top",\
    bottom="null", phead="gwresult", status="tstatus", hc_x="hydcond",\
    hc_y="hydcond", q="null", rd="R", cs="null", nf="poros", output="stresult_0", \
    dt=3600, diff_x="diff", diff_y="diff", c="c", al=0.01, at=0.01,\
    overwrite=True)

for t in range(0,30,1):
    gs.run_command("r.solute.transport", \
        top="top", bottom="null", phead="gwresult", \
        status="tstatus", q="null",hc_x="hydcond", hc_y="hydcond", rd="R", \
        cs="null", nf="poros", output="stresult_" + str(t + 1), \
        dt=86400, diff_x="diff", diff_y="diff", c="stresult_" + str(t), \
        al=0.1, at=0.1, overwrite=True)