{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Particle Tracking with PRT\n",
    "\n",
    "## Setup\n",
    "\n",
    "### Packages"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from pathlib import Path\n",
    "import flopy as fp\n",
    "import numpy as np\n",
    "import geopandas as gpd\n",
    "import xarray as xr\n",
    "import pandas as pd\n",
    "import matplotlib.pyplot as plt\n",
    "import matplotlib as mpl\n",
    "\n",
    "import nlmod"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Logging"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "nlmod.util.get_color_logger(\"INFO\")\n",
    "nlmod.show_versions()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Load base model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "model_ws = Path(\"../examples/00_model_from_scratch\")\n",
    "model_name = \"from_scratch\"\n",
    "\n",
    "ds = xr.open_dataset(model_ws / f\"{model_name}.nc\")\n",
    "ds.attrs[\"model_ws\"] = model_ws\n",
    "\n",
    "# constants\n",
    "wells = pd.DataFrame(\n",
    "    [[100, -50, -5, -10, -100.0], [200, 150, -20, -30, -300.0]],\n",
    "    columns=[\"x\", \"y\", \"top\", \"botm\", \"Q\"],\n",
    "    index=pd.Index([0, 1], name=\"well no.\"),\n",
    ")\n",
    "xyriv = [\n",
    "    (250.0, -500.0),\n",
    "    (300.0, -300.0),\n",
    "    (275.0, 0.0),\n",
    "    (200.0, 250.0),\n",
    "    (175.0, 500.0),\n",
    "]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## PRT\n",
    "\n",
    "### Forward"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# create new subdir for PRT model\n",
    "prt_ws = Path(model_ws) / \"prt_fw\"\n",
    "\n",
    "# Create a new Simulation object\n",
    "simprt = nlmod.sim.sim(ds, sim_ws=prt_ws)\n",
    "_ = nlmod.sim.tdis(ds, simprt)\n",
    "\n",
    "# Add PRT model\n",
    "prt = nlmod.prt.prt(ds, simprt)\n",
    "\n",
    "# DIS: discretization package\n",
    "_ = nlmod.prt.dis(ds, prt)\n",
    "\n",
    "# MIP: Model Input Package\n",
    "_ = nlmod.prt.mip(ds, prt, porosity=0.3)\n",
    "\n",
    "# PRP: particle release point package\n",
    "# define particle release point  every 11th cell in the first layer\n",
    "pdata = fp.modpath.ParticleData(\n",
    "    [\n",
    "        (k, i, j)\n",
    "        for i in np.arange(0, ds.sizes[\"y\"], 11)\n",
    "        for j in np.arange(0, ds.sizes[\"x\"], 11)\n",
    "        for k in [0]\n",
    "    ],\n",
    "    structured=True,\n",
    ")\n",
    "release_pts = list(pdata.to_prp(prt.modelgrid, global_xy=False))\n",
    "prp = nlmod.prt.prp(ds, prt, packagedata=release_pts, perioddata={0: [\"FIRST\"]})\n",
    "\n",
    "# FMI: flow model interface\n",
    "fmi = nlmod.prt.fmi(ds, prt)\n",
    "\n",
    "# OC: output control\n",
    "trackcsvfile_prt = f\"{prt.name}.trk.csv\"\n",
    "oc = nlmod.prt.oc(\n",
    "    ds,\n",
    "    prt,\n",
    "    trackcsv_filerecord=trackcsvfile_prt,\n",
    ")\n",
    "\n",
    "# EMS: explicit model solution\n",
    "ems = nlmod.sim.ems(simprt, model=prt)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "simprt.write_simulation(silent=False)\n",
    "simprt.run_simulation(silent=False);"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Load pathline data\n",
    "\n",
    "The read_pathlines function is a wrapper around pandas.read_csv but it adds the particle id for you. Additionallly it is usefull as documentation on the different istatus and ireason types."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "nlmod.prt.read_pathlines?"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "pathlines_fw = nlmod.prt.read_pathlines(prt_ws / trackcsvfile_prt)\n",
    "pathlines_fw.head()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Plot results"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "f, ax = plt.subplots(figsize=(8, 8))\n",
    "pmv = fp.plot.PlotMapView(prt, ax=ax)\n",
    "pmv.plot_grid(lw=0.1, color=\"k\")\n",
    "pmv.plot_pathline(\n",
    "    pathlines_fw,\n",
    "    layer=\"all\",\n",
    "    colors=\"0.7\",\n",
    "    lw=1.0,\n",
    ")\n",
    "for ilay in [0, 1, 2]:\n",
    "    mask = pathlines_fw[\"ilay\"] == ilay\n",
    "    pmv.plot_pathline(\n",
    "        pathlines_fw.loc[mask],\n",
    "        layer=ilay - 3,\n",
    "        lw=0.0,\n",
    "        marker=\".\",\n",
    "        ms=4,\n",
    "        markercolor=f\"C{ilay}\",\n",
    "        markerevery=5,\n",
    "    )\n",
    "pmv.plot_endpoint(pathlines_fw, color=\"k\", marker=\".\", direction=\"starting\", zorder=10)\n",
    "\n",
    "\n",
    "# plot river and wells\n",
    "def plot_river_and_wells(ax: plt.Axes):\n",
    "    ax.plot(\n",
    "        [xy[0] for xy in xyriv],\n",
    "        [xy[1] for xy in xyriv],\n",
    "        \"b-\",\n",
    "        lw=2,\n",
    "        label=\"river\",\n",
    "    )\n",
    "    ax.plot(wells[\"x\"].values, wells[\"y\"].values, \"rs\", ms=5)\n",
    "\n",
    "\n",
    "plot_river_and_wells(ax=pmv.ax)\n",
    "handles = [\n",
    "    pmv.ax.plot([], [], \"C0.\", ms=5, label=\"Layer 0\")[0],\n",
    "    pmv.ax.plot([], [], \"C1.\", ms=5, label=\"Layer 1\")[0],\n",
    "    pmv.ax.plot([], [], \"C2.\", ms=5, label=\"Layer 2\")[0],\n",
    "]\n",
    "labels = [h.get_label() for h in handles]\n",
    "pmv.ax.legend(handles, labels, loc=(0, 1), frameon=False, ncol=3)\n",
    "pmv.ax.set_xlabel(\"X [m]\")\n",
    "pmv.ax.set_ylabel(\"Y [m]\");"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Backward and refined\n",
    "\n",
    "#### Refine and run gwf model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "refinement_extent = [0.0, 300.0, -100.0, 200.0]\n",
    "refinement_features = (\n",
    "    [\n",
    "        nlmod.util.extent_to_polygon(refinement_extent),\n",
    "    ],\n",
    "    \"polygon\",  # type of feature\n",
    "    1,  # refinement level\n",
    ")\n",
    "dsr = nlmod.grid.refine(ds, refinement_features=[refinement_features])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "dsr.attrs[\"model_name\"] = f\"{ds.attrs['model_name']}_ref\"\n",
    "dsr.attrs[\"model_ws\"] = f\"./{ds.attrs['model_ws']}_ref\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "sim = nlmod.sim.sim(dsr)\n",
    "tdis = nlmod.sim.tdis(dsr, sim)\n",
    "ims = nlmod.sim.ims(sim, complexity=\"SIMPLE\")\n",
    "gwf = nlmod.gwf.gwf(dsr, sim)\n",
    "disv = nlmod.gwf.disv(dsr, gwf)\n",
    "npf = nlmod.gwf.npf(\n",
    "    dsr, gwf, save_flows=True, save_specific_discharge=True, save_saturation=True\n",
    ")\n",
    "ic = nlmod.gwf.ic(dsr, gwf, starting_head=1.0)\n",
    "oc = nlmod.gwf.oc(dsr, gwf, save_head=True)\n",
    "wel = nlmod.gwf.wells.wel_from_df(wells, gwf)\n",
    "\n",
    "riv_layer = 0  # add to first layer\n",
    "bed_resistance = 0.1  # days\n",
    "riv_stage = 1.0  # m NAP\n",
    "riv_botm = -3.0  # m NAP\n",
    "riv_data = nlmod.gwf.surface_water.rivdata_from_xylist(\n",
    "    gwf, xyriv, riv_layer, riv_stage, bed_resistance, riv_botm\n",
    ")\n",
    "\n",
    "\n",
    "def get_riv_cond_from_cid(icell2d: int, ds: xr.Dataset, bed_resistance: float) -> float:\n",
    "    \"\"\"Get river conductance from cell id.\"\"\"\n",
    "    area = ds.area.sel(icell2d=icell2d).values\n",
    "    # calculate conductance\n",
    "    return area / bed_resistance\n",
    "\n",
    "\n",
    "riv_datac = [\n",
    "    [x[0], x[1], get_riv_cond_from_cid(x[0][1], dsr, x[2]), x[3]] for x in riv_data\n",
    "]\n",
    "riv = fp.mf6.ModflowGwfriv(gwf, stress_period_data={0: riv_datac})\n",
    "\n",
    "nlmod.sim.write_and_run(sim, dsr, silent=False)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Reverse heads and cellbudget files"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "hds = nlmod.gwf.get_headfile(ds=dsr, gwf=gwf, tdis=tdis)\n",
    "hds_bw_path = hds.filename.parent / f\"{hds.filename.stem}_bkwd{hds.filename.suffix}\"\n",
    "hds.reverse(hds_bw_path)\n",
    "\n",
    "cbc = nlmod.gwf.get_cellbudgetfile(ds=dsr, gwf=gwf, tdis=tdis)\n",
    "cbc_bw_path = cbc.filename.parent / f\"{cbc.filename.stem}_bkwd{cbc.filename.suffix}\"\n",
    "cbc.reverse(cbc_bw_path)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Create particle starts list\n",
    "\n",
    "Start a particle in each (active) cell in refined extent"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "dsr[\"z\"] = (dsr[\"top\"] + dsr[\"botm\"]) / 2\n",
    "dsr[\"idomain\"] = (\n",
    "    (\"layer\", \"icell2d\"),\n",
    "    disv.idomain.array,\n",
    ")  # make sure to get idomain otherwhise you might start particles in inactive cells and everything crashes #beentheredonethat\n",
    "\n",
    "# map xy to cellids\n",
    "nodes = dsr[[\"z\", \"idomain\"]].where(\n",
    "    (dsr.x >= refinement_extent[0])\n",
    "    & (dsr.x <= refinement_extent[1])\n",
    "    & (dsr.y >= refinement_extent[2])\n",
    "    & (dsr.y <= refinement_extent[3]),\n",
    "    drop=True,\n",
    ")\n",
    "ix = fp.utils.GridIntersect(gwf.modelgrid)\n",
    "spatial_join = gpd.GeoDataFrame(\n",
    "    geometry=gpd.points_from_xy(nodes.x, nodes.y),\n",
    ").sjoin(\n",
    "    gpd.GeoDataFrame(\n",
    "        {\"cellid\": ix.cellids},\n",
    "        geometry=ix.geoms,\n",
    "    ),\n",
    "    how=\"left\",\n",
    ")\n",
    "\n",
    "# create release points\n",
    "# define particle release point for every cell within refined extent\n",
    "pid = 0\n",
    "release_pts = []\n",
    "for ilay in range(nodes.sizes[\"layer\"]):\n",
    "    nodes_lay = nodes.isel(layer=ilay)\n",
    "    active = (nodes_lay[\"idomain\"] > 0).values\n",
    "    j = spatial_join[\"cellid\"].values.astype(int)[active]\n",
    "    k = np.full_like(j, ilay, dtype=int)\n",
    "    x = nodes_lay[\"x\"].values[active]\n",
    "    y = nodes_lay[\"y\"].values[active]\n",
    "    z = nodes_lay[\"z\"].values[active]\n",
    "    pids = np.arange(pid, pid + len(k), dtype=int)\n",
    "    pid += len(k)\n",
    "    cellid = tuple(zip(k, j))\n",
    "    release_ptsi = list(zip(pids, cellid, x, y, z))\n",
    "    release_pts += release_ptsi\n",
    "\n",
    "print(f\"Number of release points: {len(release_pts)}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# create new subdir for PRT backwards model\n",
    "prt_ws_bw = Path(dsr.attrs[\"model_ws\"]) / \"prt_bw\"\n",
    "\n",
    "# Create a new Simulation object\n",
    "simprt_bw = nlmod.sim.sim(dsr, sim_ws=prt_ws_bw)\n",
    "_ = nlmod.sim.tdis(dsr, simprt_bw)\n",
    "\n",
    "# Add PRT model\n",
    "prt_bw = nlmod.prt.prt(dsr, simprt_bw, modelname=dsr.model_name)\n",
    "\n",
    "# Add DISV model\n",
    "_ = nlmod.prt.disv(dsr, prt_bw)\n",
    "\n",
    "# MIP: Model Input Package\n",
    "_ = nlmod.prt.mip(dsr, prt_bw, porosity=0.3)\n",
    "\n",
    "# PRP: particle release package\n",
    "_ = nlmod.prt.prp(\n",
    "    dsr,\n",
    "    prt_bw,\n",
    "    packagedata=release_pts,\n",
    "    perioddata={0: [\"FIRST\"]},\n",
    ")\n",
    "\n",
    "# OC: output control\n",
    "trackcsvfile_prt_bw = f\"{prt_bw.name}bw.trk.csv\"\n",
    "_ = nlmod.prt.oc(dsr, prt_bw, trackcsv_filerecord=trackcsvfile_prt_bw)\n",
    "\n",
    "# FMI: flow model interface\n",
    "fmi_bw_packagedata = [\n",
    "    (\"GWFHEAD\", hds_bw_path),\n",
    "    (\"GWFBUDGET\", cbc_bw_path),\n",
    "]\n",
    "_ = nlmod.prt.fmi(dsr, prt_bw, packagedata=fmi_bw_packagedata)\n",
    "\n",
    "# EMS: explicit model solution\n",
    "_ = nlmod.sim.ems(simprt_bw, model=prt_bw)\n",
    "\n",
    "# Write and run the simulation\n",
    "simprt_bw.write_simulation()\n",
    "success_bw_prt, _ = simprt_bw.run_simulation()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Load pathline data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "pathlines_bw = nlmod.prt.read_pathlines(\n",
    "    prt_ws_bw / trackcsvfile_prt_bw, icell2d=dsr.icell2d.size\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Get pathlines of particles that started in a well (backward)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "pathlines_bw0 = pathlines_bw.query(\"ireason==0\")\n",
    "well_pids = []\n",
    "for _, row in wel.stress_period_data.dataframe[0].iterrows():\n",
    "    cellid_layer = row[\"cellid_layer\"]\n",
    "    icell2d = row[\"cellid_cell\"]\n",
    "    well_pids.append(\n",
    "        pathlines_bw0.query(f\"ilay=={cellid_layer} and icell2d=={icell2d}\")\n",
    "        .loc[:, \"pid\"]\n",
    "        .to_numpy()[0]\n",
    "    )\n",
    "pathlines_bww = pathlines_bw[pathlines_bw[\"pid\"].isin(well_pids)]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Plot results"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "f, ax = plt.subplots(figsize=(8, 8))\n",
    "for pid, gr in pathlines_bww.groupby(\"pid\"):\n",
    "    for ilay, gr_lay in gr.groupby(\"ilay\"):\n",
    "        ax.plot(gr_lay[\"x\"], gr_lay[\"y\"], lw=1.0, color=f\"C{ilay}\", zorder=1)\n",
    "\n",
    "nlmod.plot.modelgrid(dsr, ax=ax, linewidth=0.1)\n",
    "plot_river_and_wells(ax=ax)\n",
    "ax.set_xlim(ds.extent[0], ds.extent[1])\n",
    "ax.set_ylim(ds.extent[2], ds.extent[3])\n",
    "ax.set_aspect(\"equal\")\n",
    "ax.set_xlabel(\"X [m]\")\n",
    "_ = ax.set_ylabel(\"Y [m]\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "f, axes = plt.subplots(3, 1, figsize=(8, 6), sharex=True, sharey=True)\n",
    "\n",
    "for pid, gr in pathlines_bw.groupby(\"pid\"):\n",
    "    t0 = gr[\"t\"].idxmin()\n",
    "    x0, y0, z0, ilay0 = gr.loc[t0, [\"x\", \"y\", \"z\", \"ilay\"]]\n",
    "    distance = np.sqrt((gr[\"x\"] - x0) ** 2 + (gr[\"y\"] - y0) ** 2 + (gr[\"z\"] - z0) ** 2)\n",
    "    if pid in well_pids:\n",
    "        plot_kwargs = {\"color\": \"k\", \"alpha\": 1.0, \"zorder\": 10, \"label\": f\"well {pid}\"}\n",
    "    else:\n",
    "        plot_kwargs = {\"color\": f\"C{ilay0}\", \"alpha\": 0.01}\n",
    "    axes[ilay0].plot(gr[\"t\"] / 365.25, distance, **plot_kwargs)\n",
    "\n",
    "for i, ax in enumerate(axes):\n",
    "    ax.xaxis.set_major_locator(mpl.ticker.MultipleLocator(250))\n",
    "    ax.yaxis.set_major_locator(mpl.ticker.MultipleLocator(100))\n",
    "    ax.set_ylabel(\"distance [m]\")\n",
    "    ax.set_title(f\"Layer {i}\")\n",
    "    ax.grid(True)\n",
    "    handles, labels = ax.get_legend_handles_labels()\n",
    "    if labels:\n",
    "        ax.legend(handles, labels, loc=\"lower right\", ncol=1, fontsize=8, frameon=True)\n",
    "    ax.set_xlim(0.0, 3000.0)\n",
    "    ax.set_ylim(0.0, 700.0)\n",
    "_ = axes[-1].set_xlabel(\"time [year]\")"
   ]
  }
 ],
 "metadata": {
  "language_info": {
   "name": "python"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}