{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1f7daec4",
   "metadata": {},
   "outputs": [],
   "source": [
    "import progressbar as pgb\n",
    "import geopandas as gpd\n",
    "import xarray as xr\n",
    "import numpy as np\n",
    "import functools\n",
    "import atlite\n",
    "import logging\n",
    "import time\n",
    "import matplotlib.pyplot as plt\n",
    "from _helpers import mock_snakemake\n",
    "\n",
    "from _helpers import configure_logging\n",
    "\n",
    "logger = logging.getLogger(__name__)\n",
    "\n",
    "if __name__ == \"__main__\":\n",
    "    if \"snakemake\" not in globals():\n",
    "        from _helpers import mock_snakemake\n",
    "\n",
    "        snakemake = mock_snakemake(\n",
    "            \"determine_availability_matrix_MD_UA\", technology=\"solar\"\n",
    "        )\n",
    "    configure_logging(snakemake)\n",
    "    pgb.streams.wrap_stderr()\n",
    "\n",
    "    nprocesses = snakemake.config[\"atlite\"].get(\"nprocesses\")\n",
    "    noprogress = not snakemake.config[\"atlite\"].get(\"show_progress\", True)\n",
    "    config = snakemake.config[\"renewable\"][snakemake.wildcards.technology]\n",
    "\n",
    "    cutout = atlite.Cutout(snakemake.input[\"cutout\"])\n",
    "    regions = (\n",
    "        gpd.read_file(snakemake.input.regions).set_index(\"name\").rename_axis(\"bus\")\n",
    "    )\n",
    "    buses = regions.index\n",
    "\n",
    "    excluder = atlite.ExclusionContainer(crs=3035, res=100)\n",
    "\n",
    "    corine = config.get(\"corine\", {})\n",
    "    if \"grid_codes\" in corine:\n",
    "\n",
    "        # Land cover codes to emulate CORINE results\n",
    "        if snakemake.wildcards.technology == \"solar\":\n",
    "            codes = [20, 30, 40, 50, 60, 90, 100]\n",
    "        elif snakemake.wildcards.technology == \"onwind\":\n",
    "            codes = [\n",
    "                20,\n",
    "                30,\n",
    "                40,\n",
    "                60,\n",
    "                100,\n",
    "                111,\n",
    "                112,\n",
    "                113,\n",
    "                114,\n",
    "                115,\n",
    "                116,\n",
    "                121,\n",
    "                122,\n",
    "                123,\n",
    "                124,\n",
    "                125,\n",
    "                126,\n",
    "            ]\n",
    "        elif snakemake.wildcards.technology == \"offshore-ac\":\n",
    "            codes = [80, 200]\n",
    "        elif snakemake.wildcards.technology == \"offshore-dc\":\n",
    "            codes = [80, 200]\n",
    "        else:\n",
    "            assert False, \"technology not supported\"\n",
    "\n",
    "        excluder.add_raster(\n",
    "            snakemake.input.copernicus, codes=codes, invert=True, crs=\"EPSG:4326\"\n",
    "        )\n",
    "    if \"distance\" in corine and corine.get(\"distance\", 0.0) > 0.0:\n",
    "        # Land cover codes to emulate CORINE results\n",
    "        if snakemake.wildcards.technology == \"onwind\":\n",
    "            codes = [50]\n",
    "        else:\n",
    "            assert False, \"technology not supported\"\n",
    "\n",
    "        buffer = corine[\"distance\"]\n",
    "        excluder.add_raster(\n",
    "            snakemake.input.copernicus, codes=codes, buffer=buffer, crs=\"EPSG:4326\"\n",
    "        )\n",
    "\n",
    "    if \"max_depth\" in config:\n",
    "        # lambda not supported for atlite + multiprocessing\n",
    "        # use named function np.greater with partially frozen argument instead\n",
    "        # and exclude areas where: -max_depth > grid cell depth\n",
    "        func = functools.partial(np.greater, -config[\"max_depth\"])\n",
    "        excluder.add_raster(snakemake.input.gebco, codes=func, crs=4236, nodata=-1000)\n",
    "\n",
    "    if \"min_shore_distance\" in config:\n",
    "        buffer = config[\"min_shore_distance\"]\n",
    "        excluder.add_geometry(snakemake.input.country_shapes, buffer=buffer)\n",
    "\n",
    "    if \"max_shore_distance\" in config:\n",
    "        buffer = config[\"max_shore_distance\"]\n",
    "        excluder.add_geometry(\n",
    "            snakemake.input.country_shapes, buffer=buffer, invert=True\n",
    "        )\n",
    "\n",
    "    kwargs = dict(nprocesses=nprocesses, disable_progressbar=noprogress)\n",
    "    if noprogress:\n",
    "        logger.info(\"Calculate landuse availabilities...\")\n",
    "        start = time.time()\n",
    "        availability = cutout.availabilitymatrix(regions, excluder, **kwargs)\n",
    "        duration = time.time() - start\n",
    "        logger.info(f\"Completed availability calculation ({duration:2.2f}s)\")\n",
    "    else:\n",
    "        availability = cutout.availabilitymatrix(regions, excluder, **kwargs)\n",
    "\n",
    "    # Limit results only to buses for UA and MD\n",
    "    buses = regions.loc[regions[\"country\"].isin([\"UA\", \"MD\"])].index.values\n",
    "    availability = availability.sel(bus=buses)\n",
    "\n",
    "    # Save and plot for verification\n",
    "    availability.to_netcdf(snakemake.output[\"availability_matrix\"])\n",
    "    # availability.sum(dim=\"bus\").plot()\n",
    "    # plt.title(technology)\n",
    "    # plt.show()"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "",
   "language": "python",
   "name": ""
  },
  "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.8.12"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}