{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "6ea7a8ea",
   "metadata": {},
   "source": [
    "## Convert Wflow staticmaps netcdf to raster files"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "09ab27a8",
   "metadata": {},
   "source": [
    "In order to inspect or (manually) modify Wflow staticmaps it is convenient to export the maps to a raster format. Here we show how to read the model maps and save to a so-called mapstack (i.e.: a set of raster files with identical grid) using HydroMT.  "
   ]
  },
  {
   "cell_type": "markdown",
   "id": "163edfe3",
   "metadata": {},
   "source": [
    "### Load dependencies"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "d480af50",
   "metadata": {},
   "outputs": [],
   "source": [
    "import xarray as xr\n",
    "from os.path import join\n",
    "import hydromt\n",
    "from hydromt_wflow import WflowModel"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1c078a08",
   "metadata": {},
   "source": [
    "### Read wflow staticmaps"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "fb790d49",
   "metadata": {},
   "source": [
    "HydroMT provides an easy method to read the model schematization trought the Model API."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "15044d4b",
   "metadata": {},
   "outputs": [],
   "source": [
    "root = \"wflow_piave_subbasin\"\n",
    "mod = WflowModel(root, mode=\"r\")\n",
    "ds = mod.staticmaps  # here the staticmaps netcdf is loaded\n",
    "print(ds)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1acdee64",
   "metadata": {},
   "source": [
    "### Write netcdf to mapstack"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "36073d9b",
   "metadata": {},
   "source": [
    "The raster module provides many raster GIS methods throught the **raster** Dataset accessor. To write a Dataset to a mapstack one line with code is sufficient. We only need to provide the output folder in which all raster files are saved. The default output format is *GeoTIFF*, but this can be changed with the `driver` argument. To write to PCRaster map-files it is recommended to have PCRaster python installed."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "07e4d5b7",
   "metadata": {},
   "outputs": [],
   "source": [
    "outdir = join(root, \"staticmaps\")\n",
    "ds.raster.to_mapstack(outdir)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b848f0e9",
   "metadata": {},
   "source": [
    "Now the model files can easily be inspected and modified e.g. QGIS.\n",
    "\n",
    "NOTE: in QGIS, you can also visualize but direct modification is not (yet) possible."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e1f84114",
   "metadata": {},
   "source": [
    "### Create staticmaps netcdf files based on mapstack"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "bbe69373",
   "metadata": {},
   "source": [
    "If you want to update the staticmaps after modification the maps can be read into a Dataset by hydromt. We recommend the following workflow:\n",
    "\n",
    "* read the original model\n",
    "* read the updated mapstack\n",
    "* change the model root to write the updated model to a new directory\n",
    "* update the staticmaps of the model\n",
    "* write the model\n",
    "\n",
    "NOTE: We do not read the forcing as it is probably faster to just copy the file instead of loading it into python and writing it back to netcdf.\n",
    "\n",
    "NOTE: The staticgeoms might be changed because of manual changes in the wflow_river, lakes, reservoir or glacier staticmaps and are therefore not read here. To change these maps we recommend using the hydromt update method to keep the staticgeoms and maps aligned."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "42c8a6f3",
   "metadata": {},
   "outputs": [],
   "source": [
    "# read the original model\n",
    "root = \"wflow_piave_subbasin\"\n",
    "mod = WflowModel(root, mode=\"r\")\n",
    "mod.read_staticmaps()\n",
    "mod.read_config()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "438a1d35",
   "metadata": {},
   "outputs": [],
   "source": [
    "# read the updated mapstack\n",
    "# NOTE: The mapstack does not have to include all staticmaps, only the once that are found will be updated.\n",
    "# The name of the staticmap should however have to be unchanged.\n",
    "ds_updated = hydromt.open_mfraster(join(root, \"staticmaps\", \"*.tif\"))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "98ba4297",
   "metadata": {},
   "outputs": [],
   "source": [
    "# change root to a new directory\n",
    "root_updated = \"wflow_piave_subbasin_updated\"\n",
    "mod.set_root(root_updated, mode=\"w+\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1f9fd886",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Reinitialize and update the new model staticmaps\n",
    "mod._staticmaps = xr.Dataset()\n",
    "mod.set_staticmaps(ds_updated)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2fef2a90",
   "metadata": {},
   "outputs": [],
   "source": [
    "# write the model to the new directory\n",
    "mod.write()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7fdabc58",
   "metadata": {},
   "source": [
    "### Update Wflow staticmaps manually using HydroMT"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e1567262",
   "metadata": {},
   "source": [
    "The previous steps show you how you can easily save a model staticmaps to a GeoTIFF mapstacks and read it again using hydroMT functions **to_mapstack** and **open_mfraster**.\n",
    "\n",
    "However in order to have a fully ready-to-run wflow model, the mapstacks that we create needs to be processed a little more:\n",
    "\n",
    "* The LAI maps have a third *time* dimension (other than x and y coordinates)\n",
    "* The c maps have a third soil *layer* dimension\n",
    "\n",
    "To update manually a wflow model that is then ready to run, we advise to use the following workflow and functions:\n",
    "\n",
    "* read the original model\n",
    "* save the original model staticmaps to formatted PCRaster mapstacks using the **write_staticmaps_pcr** function (based on to_mapstack with c and LAI pre-processing)\n",
    "* manually update the PCRaster maps (eg using QGIS)\n",
    "* read the updated staticmaps using the **read_staticmaps_pcr** function (based on open_mfraster with c and LAI post-processing)\n",
    "* Change the model root to write the updated model to a new directory\n",
    "* Write the model\n",
    "\n",
    "NOTE: We do not read the forcing as it is probably faster to just copy the file instead of loading it into python and writing it back to netcdf.\n",
    "\n",
    "NOTE: The staticgeoms might be changed because of manual changes in the wflow_river, lakes, reservoir or glacier staticmaps and are therefore not read here. To change these maps we recommend using the hydromt update method to keep the staticgeoms and maps aligned."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3464a5c3",
   "metadata": {},
   "source": [
    "#### Save the staticmaps as a PCRaster mapstack"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b06f0f2b",
   "metadata": {},
   "outputs": [],
   "source": [
    "# read the original model\n",
    "root = \"wflow_piave_subbasin\"\n",
    "mod = WflowModel(root, mode=\"r+\")\n",
    "mod.read_staticmaps()\n",
    "mod.read_config()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "330cfb52",
   "metadata": {},
   "outputs": [],
   "source": [
    "# save the staticmaps to PCRaster mapstacks and update them manually where needed\n",
    "from hydromt_wflow.pcrm import write_staticmaps_pcr\n",
    "write_staticmaps_pcr(mod.staticmaps, root=root)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7474fc4e",
   "metadata": {},
   "source": [
    "#### Create a new Wflow model based on the updated PCRaster mapstack"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ca8c57cc",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Import read method for PCRaster files\n",
    "from hydromt_wflow.pcrm import read_staticmaps_pcr\n",
    "# read the updated staticmaps\n",
    "root = \"wflow_piave_subbasin\"\n",
    "mod = WflowModel(root, mode=\"r\")\n",
    "staticmaps = read_staticmaps_pcr(root)\n",
    "mod.set_grid(staticmaps)\n",
    "mod.read_config()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c84202b5",
   "metadata": {},
   "outputs": [],
   "source": [
    "# change root to a new directory\n",
    "root_updated = \"wflow_piave_subbasin_updated\"\n",
    "mod.set_root(root_updated, mode=\"w+\")\n",
    "# re-generate the basins and rivers staticgeoms, for others (gauges, lakes, reservoirs, glaicers) use the hydromt update method\n",
    "mod.basins\n",
    "mod.rivers\n",
    "print(\"Creating a new directory for the updated model\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "04df13a5",
   "metadata": {},
   "outputs": [],
   "source": [
    "# write the model to the new directory\n",
    "mod.write()"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.11.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}