{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Plot Wflow static maps"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**HydroMT** provides a simple interface to model schematization from which we can make beautiful plots:\n",
    "\n",
    "- Raster layers are saved to the model `staticmaps` component as a  `xarray.Dataset`\n",
    "- Vector layers are saved to the model `staticgeoms` component as a `geopandas.GeoDataFrame`. Note that in case of Wflow these are not used by the model engine, but only for analysis and visualization purposes.\n",
    "\n",
    "We use the [cartopy](https://scitools.org.uk/cartopy/docs/latest/) package to plot maps. This packages provides a simple interface to plot geographic data and add background satellite imagery."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Load dependencies"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import xarray as xr\n",
    "import numpy as np\n",
    "import hydromt\n",
    "from hydromt_wflow import WflowModel"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Read the model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "root = \"wflow_piave_subbasin\"\n",
    "mod = WflowModel(root, mode=\"r\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Plot model schematization base maps"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Here we plot the model basemaps (topography map with rivers, lakes, reservoirs, glaciers and gauges geometries). "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# plot maps dependencies\n",
    "import matplotlib.pyplot as plt\n",
    "from matplotlib import colors\n",
    "import matplotlib.patches as mpatches\n",
    "import cartopy.crs as ccrs\n",
    "import cartopy.io.img_tiles as cimgt\n",
    "\n",
    "# import descartes  # used under the hood to plot polygons"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# read and mask the model elevation\n",
    "da = mod.staticmaps[\"wflow_dem\"].raster.mask_nodata()\n",
    "da.attrs.update(long_name=\"elevation\", units=\"m\")\n",
    "# read/derive river geometries\n",
    "gdf_riv = mod.rivers\n",
    "# read/derive model basin boundary\n",
    "gdf_bas = mod.basins"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.style.use(\"seaborn-whitegrid\")  # set nice style\n",
    "# we assume the model maps are in the geographic CRS EPSG:4326\n",
    "proj = ccrs.PlateCarree()\n",
    "# adjust zoomlevel and figure size to your basis size & aspect\n",
    "zoom_level = 10\n",
    "figsize = (10, 8)\n",
    "shaded = False  # shaded elevation (looks nicer with more pixels (e.g.: larger basins))!\n",
    "\n",
    "# initialize image with geoaxes\n",
    "fig = plt.figure(figsize=figsize)\n",
    "ax = fig.add_subplot(projection=proj)\n",
    "bbox = da.raster.box.to_crs(3857).buffer(5e3).to_crs(da.raster.crs).total_bounds\n",
    "extent = np.array(bbox)[[0, 2, 1, 3]]\n",
    "ax.set_extent(extent, crs=proj)\n",
    "\n",
    "# add sat background image\n",
    "ax.add_image(cimgt.QuadtreeTiles(), zoom_level, alpha=0.5)\n",
    "\n",
    "## plot elevation\\\n",
    "# create nice colormap\n",
    "vmin, vmax = da.quantile([0.0, 0.98]).compute()\n",
    "c_dem = plt.cm.terrain(np.linspace(0.25, 1, 256))\n",
    "cmap = colors.LinearSegmentedColormap.from_list(\"dem\", c_dem)\n",
    "norm = colors.Normalize(vmin=vmin, vmax=vmax)\n",
    "kwargs = dict(cmap=cmap, norm=norm)\n",
    "# plot 'normal' elevation\n",
    "da.plot(\n",
    "    transform=proj, ax=ax, zorder=1, cbar_kwargs=dict(aspect=30, shrink=0.8), **kwargs\n",
    ")\n",
    "# plot elevation with shades\n",
    "if shaded:\n",
    "    ls = colors.LightSource(azdeg=315, altdeg=45)\n",
    "    dx, dy = da.raster.res\n",
    "    _rgb = ls.shade(\n",
    "        da.fillna(0).values,\n",
    "        norm=kwargs[\"norm\"],\n",
    "        cmap=kwargs[\"cmap\"],\n",
    "        blend_mode=\"soft\",\n",
    "        dx=dx,\n",
    "        dy=dy,\n",
    "        vert_exag=200,\n",
    "    )\n",
    "    rgb = xr.DataArray(dims=(\"y\", \"x\", \"rgb\"), data=_rgb, coords=da.raster.coords)\n",
    "    rgb = xr.where(np.isnan(da), np.nan, rgb)\n",
    "    rgb.plot.imshow(transform=proj, ax=ax, zorder=2)\n",
    "\n",
    "# plot rivers with increasing width with stream order\n",
    "gdf_riv.plot(\n",
    "    ax=ax, linewidth=gdf_riv[\"strord\"] / 2, color=\"blue\", zorder=3, label=\"river\"\n",
    ")\n",
    "# plot the basin boundary\n",
    "gdf_bas.boundary.plot(ax=ax, color=\"k\", linewidth=0.3)\n",
    "# plot various vector layers if present\n",
    "if \"gauges\" in mod.staticgeoms:\n",
    "    mod.staticgeoms[\"gauges\"].plot(\n",
    "        ax=ax, marker=\"d\", markersize=25, facecolor=\"k\", zorder=5, label=\"gauges\"\n",
    "    )\n",
    "patches = (\n",
    "    []\n",
    ")  # manual patches for legend, see https://github.com/geopandas/geopandas/issues/660\n",
    "if \"lakes\" in mod.staticgeoms:\n",
    "    kwargs = dict(facecolor=\"lightblue\", edgecolor=\"black\", linewidth=1, label=\"lakes\")\n",
    "    mod.staticgeoms[\"lakes\"].plot(ax=ax, zorder=4, **kwargs)\n",
    "    patches.append(mpatches.Patch(**kwargs))\n",
    "if \"reservoirs\" in mod.staticgeoms:\n",
    "    kwargs = dict(facecolor=\"blue\", edgecolor=\"black\", linewidth=1, label=\"reservoirs\")\n",
    "    mod.staticgeoms[\"reservoirs\"].plot(ax=ax, zorder=4, **kwargs)\n",
    "    patches.append(mpatches.Patch(**kwargs))\n",
    "if \"glaciers\" in mod.staticgeoms:\n",
    "    kwargs = dict(facecolor=\"grey\", edgecolor=\"grey\", linewidth=1, label=\"glaciers\")\n",
    "    mod.staticgeoms[\"glaciers\"].plot(ax=ax, zorder=4, **kwargs)\n",
    "    patches.append(mpatches.Patch(**kwargs))\n",
    "\n",
    "ax.xaxis.set_visible(True)\n",
    "ax.yaxis.set_visible(True)\n",
    "ax.set_ylabel(f\"latitude [degree north]\")\n",
    "ax.set_xlabel(f\"longitude [degree east]\")\n",
    "_ = ax.set_title(f\"wflow base map\")\n",
    "legend = ax.legend(\n",
    "    handles=[*ax.get_legend_handles_labels()[0], *patches],\n",
    "    title=\"Legend\",\n",
    "    loc=\"lower right\",\n",
    "    frameon=True,\n",
    "    framealpha=0.7,\n",
    "    edgecolor=\"k\",\n",
    "    facecolor=\"white\",\n",
    ")\n",
    "\n",
    "# save figure\n",
    "# NOTE create figs folder in model root if it does not exist\n",
    "# fn_out = join(mod.root, \"figs\", \"basemap.png\")\n",
    "# plt.savefig(fn_out, dpi=225, bbox_inches=\"tight\")"
   ]
  }
 ],
 "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.10.2"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}