#!/usr/bin/env python3

import sys
import os
import grass.script as gs

# set region
gs.run_command("g.region", rows=10, cols=20, res3=1, t=50, b=0, n=50,s=0, w=0,\
 e=200, overwrite=True)

# generate simple raster maps
gs.run_command("r.mapcalc", expression="phead = if(col() == 1, 38.4, 38)", overwrite=True)
gs.run_command("r.mapcalc", expression="status = if(col() == 1 || col() == 20,\
 2, 1)", overwrite=True)
gs.run_command("r.mapcalc", expression="null = 0.0",overwrite=True)
gs.run_command("r.mapcalc", expression="top_unconf = 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.07 * (1.0 + 0.2 * \
(heterogenMap - 46.61)/100 * 6) ",overwrite=True)

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

# Generate the raster maps for transport data
gs.run_command("r.mapcalc", expression="c = if(col() == 2 && row() == 5, 5.0, 0.0)",overwrite=True)
gs.run_command("r.mapcalc", expression="tstatus = if(col() == 1 || col() == 20 , 3, 1)",overwrite=True)
gs.run_command("r.mapcalc", expression="diff = 0.00001 * (1.0 + 0.2 * \
(heterogenMap - 46.61)/100 * 6)",overwrite=True)
gs.run_command("r.mapcalc", expression="R = 1.0",overwrite=True)

gs.run_command("r.mapcalc", expression="PeNr = 10.0 * sqrt(velx^2 + vely^2) / diff",overwrite=True)

# Read coordinates for observation points from file
# Grass GIS base folder is the user home folder. Here, all files were stored
# on the Desktop for simplicity
gs.run_command("v.in.ascii", input = ".\Desktop\week9_points.txt", 
   output="obsPoints", separator="comma")
# Add a attribute table to this vector map
gs.run_command("v.db.addtable", map="obsPoints")
# Add columns to the attribute table of the vector map to store the 
# concentration values of each day
gs.run_command("v.db.addcolumn", map="obsPoints", columns="day1 double precision,\
                                                       day2 double precision,\
                                                       day3 double precision,\
                                                       day4 double precision,\
                                                       day5 double precision,\
                                                       day6 double precision,\
                                                       day7 double precision,\
                                                       day8 double precision,\
                                                       day9 double precision,\
                                                       day10 double precision,\
                                                       day11 double precision,\
                                                       day12 double precision,\
                                                       day13 double precision,\
                                                       day14 double precision,\
                                                       day15 double precision,\
                                                       day16 double precision,\
                                                       day17 double precision,\
                                                       day18 double precision,\
                                                       day19 double precision,\
                                                       day20 double precision,\
                                                       day21 double precision,\
                                                       day22 double precision,\
                                                       day23 double precision,\
                                                       day24 double precision,\
                                                       day25 double precision,\
                                                       day26 double precision,\
                                                       day27 double precision,\
                                                       day28 double precision,\
                                                       day29 double precision,\
                                                       day30 double precision")

# Calculate initial timestep
gs.run_command("r.solute.transport", solver="bicgstab", top="top_unconf",\
    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)

# Calculate subsequent timesteps and define a number of loops as the CFL flag 
# of r.solute.transport can not be called from Python.
# Depending on the chosen resolution, make sure that your timestep (dt/loop)
# is smaller than the maximum timestep allowed by CFL. 10 is for this setting 
# a high number. 2 would be sufficient.
for t in range(0,30,1):
    gs.run_command("r.solute.transport", solver="bicgstab", \
        top="top_unconf", 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.01, at=0.01, overwrite=True, loops=10)

    # save the concentrations at the observation points for this day
    gs.run_command("v.what.rast", map="obsPoints", 
        raster="stresult_"+str(t+1),
        column="day"+str(t+1))

# Export the attribute table of the vector map (=the measurement values at the
# observation points) for further analysis outside of Grass GIS 
gs.run_command("v.db.select",map="obsPoints",separator="tab",
    file=".\Desktop\obsPointsResults.txt",
    overwrite=True)