# Test the buoyancy package and the variable density flows between the lake
# and the gwf model.  This model has 4 layers and a lake incised within it.
# The model is transient and has heads in the aquifer higher than the initial
# stage in the lake.  As the model runs, the lake and aquifer equalize and
# should end up at the same level.  The test ensures that the initial and
# final water volumes in the entire system are the same.  There are three
# different cases:
#  1.  No buoyancy package
#  2.  Buoyancy package with lake and aquifer density = 1000.
#  3.  Buoyancy package with lake and aquifer density = 1024.5

import os
import pytest
import sys
import numpy as np

try:
    import flopy
except:
    msg = "Error. FloPy package is not available.\n"
    msg += "Try installing using the following command:\n"
    msg += " pip install flopy"
    raise Exception(msg)

from framework import testing_framework
from simulation import Simulation

ex = ["buy_lak_01a"]  # , 'buy_lak_01b', 'buy_lak_01c']
buy_on_list = [False]  # , True, True]
concbuylist = [0.0]  # , 0., 35.]
exdirs = []
for s in ex:
    exdirs.append(os.path.join("temp", s))


def build_model(idx, dir):
    lx = 7.0
    lz = 4.0
    nlay = 4
    nrow = 1
    ncol = 7
    nper = 1
    delc = 1.0
    delr = lx / ncol
    delz = lz / nlay
    top = 4.0
    botm = [3.0, 2.0, 1.0, 0.0]

    perlen = [10.0]
    nstp = [50]
    tsmult = [1.0]

    Kh = 1.0
    Kv = 1.0

    tdis_rc = []
    for i in range(nper):
        tdis_rc.append((perlen[i], nstp[i], tsmult[i]))

    nouter, ninner = 700, 300
    hclose, rclose, relax = 1e-8, 1e-6, 0.97

    name = ex[idx]

    # build MODFLOW 6 files
    ws = dir
    sim = flopy.mf6.MFSimulation(
        sim_name=name, version="mf6", exe_name="mf6", sim_ws=ws
    )
    # create tdis package
    tdis = flopy.mf6.ModflowTdis(
        sim, time_units="DAYS", nper=nper, perioddata=tdis_rc
    )

    # create gwf model
    gwfname = "gwf_" + name

    gwf = flopy.mf6.ModflowGwf(sim, modelname=gwfname, newtonoptions="NEWTON")

    imsgwf = flopy.mf6.ModflowIms(
        sim,
        print_option="ALL",
        outer_dvclose=hclose,
        outer_maximum=nouter,
        under_relaxation="NONE",
        inner_maximum=ninner,
        inner_dvclose=hclose,
        rcloserecord=rclose,
        linear_acceleration="BICGSTAB",
        scaling_method="NONE",
        reordering_method="NONE",
        relaxation_factor=relax,
        filename="{}.ims".format(gwfname),
    )

    idomain = np.full((nlay, nrow, ncol), 1)
    idomain[0, 0, 1:6] = 0
    idomain[1, 0, 2:5] = 0
    idomain[2, 0, 3:4] = 0
    dis = flopy.mf6.ModflowGwfdis(
        gwf,
        nlay=nlay,
        nrow=nrow,
        ncol=ncol,
        delr=delr,
        delc=delc,
        top=top,
        botm=botm,
        idomain=idomain,
    )

    # initial conditions
    strt = np.zeros((nlay, nrow, ncol), dtype=float)
    strt[0, 0, :] = 3.5
    strt[1, 0, :] = 3.0
    strt[1, 0, 1:6] = 2.5
    strt[2, 0, :] = 2.0
    strt[3, 0, :] = 1.0
    ic = flopy.mf6.ModflowGwfic(gwf, strt=strt)

    # node property flow
    npf = flopy.mf6.ModflowGwfnpf(
        gwf,
        xt3doptions=False,
        save_flows=True,
        save_specific_discharge=True,
        icelltype=1,
        k=Kh,
        k33=Kv,
    )

    sto = flopy.mf6.ModflowGwfsto(gwf, sy=0.3, ss=0.0, iconvert=1)

    c = concbuylist[idx]
    lake_dense = 1000.0 + 0.7 * c
    buy_on = buy_on_list[idx]
    if buy_on:
        pd = [(0, 0.7, 0.0, "none", "none")]
        buy = flopy.mf6.ModflowGwfbuy(
            gwf, packagedata=pd, denseref=1000.0, concentration=c
        )

    nlakeconn = 11  # note: number of connections for this lake
    # pak_data = [lakeno, strt, nlakeconn, dense, boundname]
    pak_data = [(0, 2.25, nlakeconn, lake_dense)]

    connlen = delr / 2.0
    connwidth = delc
    bedleak = "None"
    con_data = [
        # con_data=(lakeno,iconn,(cellid),claktype,bedleak,belev,telev,connlen,connwidth )
        (0, 0, (0, 0, 0), "HORIZONTAL", bedleak, 10, 10, connlen, connwidth),
        (0, 1, (1, 0, 1), "VERTICAL", bedleak, 10, 10, connlen, connwidth),
        (0, 2, (1, 0, 1), "HORIZONTAL", bedleak, 10, 10, connlen, connwidth),
        (0, 3, (2, 0, 2), "VERTICAL", bedleak, 10, 10, connlen, connwidth),
        (0, 4, (2, 0, 2), "HORIZONTAL", bedleak, 10, 10, connlen, connwidth),
        (0, 5, (3, 0, 3), "VERTICAL", bedleak, 10, 10, connlen, connwidth),
        (0, 6, (2, 0, 4), "HORIZONTAL", bedleak, 10, 10, connlen, connwidth),
        (0, 7, (2, 0, 4), "VERTICAL", bedleak, 10, 10, connlen, connwidth),
        (0, 8, (1, 0, 5), "HORIZONTAL", bedleak, 10, 10, connlen, connwidth),
        (0, 9, (1, 0, 5), "VERTICAL", bedleak, 10, 10, connlen, connwidth),
        (0, 10, (0, 0, 6), "HORIZONTAL", bedleak, 10, 10, connlen, connwidth),
    ]

    # period data
    p_data = [
        (0, "STATUS", "ACTIVE"),
    ]

    # note: for specifying lake number, use fortran indexing!
    fname = "{}.lak.obs.csv".format(gwfname)
    lak_obs = {
        fname: [
            ("lakestage", "stage", 1),
            ("lakevolume", "volume", 1),
            ("lak1", "lak", 1, 1),
            ("lak2", "lak", 1, 2),
            ("lak3", "lak", 1, 3),
            ("lak4", "lak", 1, 4),
            ("lak5", "lak", 1, 5),
            ("lak6", "lak", 1, 6),
            ("lak7", "lak", 1, 7),
            ("lak8", "lak", 1, 8),
            ("lak9", "lak", 1, 9),
            ("lak10", "lak", 1, 10),
            ("lak11", "lak", 1, 11),
        ],
        # "digits": 10,
    }

    lak = flopy.mf6.modflow.ModflowGwflak(
        gwf,
        save_flows=True,
        print_input=True,
        print_flows=True,
        print_stage=True,
        stage_filerecord="{}.lak.bin".format(gwfname),
        budget_filerecord="{}.lak.bud".format(gwfname),
        nlakes=len(pak_data),
        ntables=0,
        packagedata=pak_data,
        pname="LAK-1",
        connectiondata=con_data,
        perioddata=p_data,
        observations=lak_obs,
        auxiliary=["DENSITY"],
    )

    # output control
    oc = flopy.mf6.ModflowGwfoc(
        gwf,
        budget_filerecord="{}.cbc".format(gwfname),
        head_filerecord="{}.hds".format(gwfname),
        headprintrecord=[("COLUMNS", 10, "WIDTH", 15, "DIGITS", 6, "GENERAL")],
        saverecord=[("HEAD", "ALL"), ("BUDGET", "ALL")],
        printrecord=[("HEAD", "LAST"), ("BUDGET", "LAST")],
    )

    return sim, None


def eval_results(sim):
    print("evaluating results...")

    # calculate volume of water and make sure it is conserved
    name = ex[sim.idxsim]
    gwfname = "gwf_" + name
    fname = gwfname + ".lak.bin"
    fname = os.path.join(sim.simpath, fname)
    assert os.path.isfile(fname)
    bobj = flopy.utils.HeadFile(fname, text="STAGE")
    stage = bobj.get_alldata().flatten()
    # print(stage)

    fname = gwfname + ".hds"
    fname = os.path.join(sim.simpath, fname)
    assert os.path.isfile(fname)
    hobj = flopy.utils.HeadFile(fname)
    head = hobj.get_data()
    # print(head)

    # calculate initial water volume
    v0 = 3.5 * 2  # outermost columns
    v0 += 2.5 * 2  # next innermost columns
    v0 += 2.0 * 2  # next innermost columns
    v0 += 1.0 * 1  # middle column
    v0 = v0 * 0.3  # specific yield
    v0 = v0 + (2.25 - 2.0) * 2 + (2.25 - 1.0)
    print("initial volume of water in model = {}".format(v0))

    # calculate ending water volume in model
    h = head[0, 0, 0]
    s = stage[-1]
    v = h * 4 + 2.0 * 2 + 1.0 * 1
    v = v * 0.3  # specific yield
    v = v + (s - 2.0) * 2 + (s - 1.0)
    print("final volume of water in model = {}".format(v))

    # check to make sure starting water volume same as equalized final volume
    errmsg = "initial and final water volume not equal: {} {}".format(v0, v)
    assert np.allclose(v0, v)

    # todo: add a better check of the lake concentrations
    # assert False


# - No need to change any code below
@pytest.mark.parametrize(
    "idx, dir",
    list(enumerate(exdirs)),
)
def test_mf6model(idx, dir):
    # initialize testing framework
    test = testing_framework()

    # build the model
    test.build_mf6_models(build_model, idx, dir)

    # run the test model
    test.run_mf6(Simulation(dir, exfunc=eval_results, idxsim=idx))


def main():
    # initialize testing framework
    test = testing_framework()

    # run the test model
    for idx, dir in enumerate(exdirs):
        test.build_mf6_models(build_model, idx, dir)
        sim = Simulation(dir, exfunc=eval_results, idxsim=idx)
        test.run_mf6(sim)


if __name__ == "__main__":
    # print message
    print("standalone run of {}".format(os.path.basename(__file__)))

    # run main routine
    main()

[((1693, 1772), 'flopy.mf6.MFSimulation', 'flopy.mf6.MFSimulation', ([], {'sim_name': 'name', 'version': '"""mf6"""', 'exe_name': '"""mf6"""', 'sim_ws': 'ws'}), "(sim_name=name, version='mf6', exe_name='mf6', sim_ws=ws)\n", (1715, 1772), False, 'import flopy\n'), ((1824, 1900), 'flopy.mf6.ModflowTdis', 'flopy.mf6.ModflowTdis', (['sim'], {'time_units': '"""DAYS"""', 'nper': 'nper', 'perioddata': 'tdis_rc'}), "(sim, time_units='DAYS', nper=nper, perioddata=tdis_rc)\n", (1845, 1900), False, 'import flopy\n'), ((1978, 2046), 'flopy.mf6.ModflowGwf', 'flopy.mf6.ModflowGwf', (['sim'], {'modelname': 'gwfname', 'newtonoptions': '"""NEWTON"""'}), "(sim, modelname=gwfname, newtonoptions='NEWTON')\n", (1998, 2046), False, 'import flopy\n'), ((2508, 2538), 'numpy.full', 'np.full', (['(nlay, nrow, ncol)', '(1)'], {}), '((nlay, nrow, ncol), 1)\n', (2515, 2538), True, 'import numpy as np\n'), ((2630, 2754), 'flopy.mf6.ModflowGwfdis', 'flopy.mf6.ModflowGwfdis', (['gwf'], {'nlay': 'nlay', 'nrow': 'nrow', 'ncol': 'ncol', 'delr': 'delr', 'delc': 'delc', 'top': 'top', 'botm': 'botm', 'idomain': 'idomain'}), '(gwf, nlay=nlay, nrow=nrow, ncol=ncol, delr=delr,\n    delc=delc, top=top, botm=botm, idomain=idomain)\n', (2653, 2754), False, 'import flopy\n'), ((2867, 2908), 'numpy.zeros', 'np.zeros', (['(nlay, nrow, ncol)'], {'dtype': 'float'}), '((nlay, nrow, ncol), dtype=float)\n', (2875, 2908), True, 'import numpy as np\n'), ((3040, 3078), 'flopy.mf6.ModflowGwfic', 'flopy.mf6.ModflowGwfic', (['gwf'], {'strt': 'strt'}), '(gwf, strt=strt)\n', (3062, 3078), False, 'import flopy\n'), ((3115, 3240), 'flopy.mf6.ModflowGwfnpf', 'flopy.mf6.ModflowGwfnpf', (['gwf'], {'xt3doptions': '(False)', 'save_flows': '(True)', 'save_specific_discharge': '(True)', 'icelltype': '(1)', 'k': 'Kh', 'k33': 'Kv'}), '(gwf, xt3doptions=False, save_flows=True,\n    save_specific_discharge=True, icelltype=1, k=Kh, k33=Kv)\n', (3138, 3240), False, 'import flopy\n'), ((3311, 3367), 'flopy.mf6.ModflowGwfsto', 'flopy.mf6.ModflowGwfsto', (['gwf'], {'sy': '(0.3)', 'ss': '(0.0)', 'iconvert': '(1)'}), '(gwf, sy=0.3, ss=0.0, iconvert=1)\n', (3334, 3367), False, 'import flopy\n'), ((6681, 6713), 'os.path.join', 'os.path.join', (['sim.simpath', 'fname'], {}), '(sim.simpath, fname)\n', (6693, 6713), False, 'import os\n'), ((6725, 6746), 'os.path.isfile', 'os.path.isfile', (['fname'], {}), '(fname)\n', (6739, 6746), False, 'import os\n'), ((6758, 6799), 'flopy.utils.HeadFile', 'flopy.utils.HeadFile', (['fname'], {'text': '"""STAGE"""'}), "(fname, text='STAGE')\n", (6778, 6799), False, 'import flopy\n'), ((6902, 6934), 'os.path.join', 'os.path.join', (['sim.simpath', 'fname'], {}), '(sim.simpath, fname)\n', (6914, 6934), False, 'import os\n'), ((6946, 6967), 'os.path.isfile', 'os.path.isfile', (['fname'], {}), '(fname)\n', (6960, 6967), False, 'import os\n'), ((6979, 7006), 'flopy.utils.HeadFile', 'flopy.utils.HeadFile', (['fname'], {}), '(fname)\n', (6999, 7006), False, 'import flopy\n'), ((7813, 7831), 'numpy.allclose', 'np.allclose', (['v0', 'v'], {}), '(v0, v)\n', (7824, 7831), True, 'import numpy as np\n'), ((8097, 8116), 'framework.testing_framework', 'testing_framework', ([], {}), '()\n', (8114, 8116), False, 'from framework import testing_framework\n'), ((8342, 8361), 'framework.testing_framework', 'testing_framework', ([], {}), '()\n', (8359, 8361), False, 'from framework import testing_framework\n'), ((1127, 1150), 'os.path.join', 'os.path.join', (['"""temp"""', 's'], {}), "('temp', s)\n", (1139, 1150), False, 'import os\n'), ((3532, 3610), 'flopy.mf6.ModflowGwfbuy', 'flopy.mf6.ModflowGwfbuy', (['gwf'], {'packagedata': 'pd', 'denseref': '(1000.0)', 'concentration': 'c'}), '(gwf, packagedata=pd, denseref=1000.0, concentration=c)\n', (3555, 3610), False, 'import flopy\n'), ((8232, 8280), 'simulation.Simulation', 'Simulation', (['dir'], {'exfunc': 'eval_results', 'idxsim': 'idx'}), '(dir, exfunc=eval_results, idxsim=idx)\n', (8242, 8280), False, 'from simulation import Simulation\n'), ((8494, 8542), 'simulation.Simulation', 'Simulation', (['dir'], {'exfunc': 'eval_results', 'idxsim': 'idx'}), '(dir, exfunc=eval_results, idxsim=idx)\n', (8504, 8542), False, 'from simulation import Simulation\n'), ((8658, 8684), 'os.path.basename', 'os.path.basename', (['__file__'], {}), '(__file__)\n', (8674, 8684), False, 'import os\n')]