{
 "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",
    "        snakemake = mock_snakemake('determine_availability_matrix_MD_UA', technology='solar')\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 = gpd.read_file(snakemake.input.regions).set_index('name').rename_axis('bus')\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 = [20, 30, 40, 60, 100, 111, 112, 113, 114, 115, 116, 121, 122, 123, 124, 125, 126]\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(snakemake.input.copernicus, codes=codes, invert=True, crs=\"EPSG:4326\")\n",
    "    if \"distance\" in corine and corine.get(\"distance\", 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(snakemake.input.copernicus, codes=codes, buffer=buffer, crs=\"EPSG:4326\")\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(snakemake.input.country_shapes, buffer=buffer, invert=True)\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": "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.8.12"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}