{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "422a5b1d-46a3-4359-a2fd-0295141e01a0",
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import glob\n",
    "import csv\n",
    "import requests\n",
    "from pathlib import Path\n",
    "import xarray as xr\n",
    "import rioxarray as rio\n",
    "import pandas as pd\n",
    "import geopandas as gpd\n",
    "import numpy as np\n",
    "import rasterio"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "20745ca3-cc15-422d-b365-1c2f9131b33f",
   "metadata": {},
   "outputs": [],
   "source": [
    "city = 'Luanda'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "0d28adac-232b-4e5b-8cb2-a3446322fb7f",
   "metadata": {},
   "outputs": [],
   "source": [
    "data_folder = Path(r\"D:\\World Bank\\CRP\\data\\WaterGAP\")\n",
    "output_folder = Path('output')\n",
    "int_output_folder = Path('intermediate_output')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "5ab27697-98be-44d8-9c79-1369977078ad",
   "metadata": {},
   "outputs": [],
   "source": [
    "# read catchment AOI\n",
    "aoi = gpd.read_file('AOI/luanda_catchment_level4.shp').to_crs(epsg = 4326)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "b53ab321-8c6c-4173-84cf-8b79edde76d2",
   "metadata": {},
   "outputs": [],
   "source": [
    "var_list = ['ncrun', 'qs']\n",
    "year_list = range(1987, 2017)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "79b07184-b927-45a8-a5b7-240c205d7796",
   "metadata": {},
   "source": [
    "### Convert netcdf to raster"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "6506fa64-d1ac-4535-9976-4c68c1c29d53",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "for var in var_list:\n",
    "    if not os.path.exists(data_folder / f'{var}_raster2d'):\n",
    "        os.mkdir(data_folder / f'{var}_raster2d')\n",
    "    \n",
    "    # open netcdf\n",
    "    nc = rio.open_rasterio(data_folder / f'watergap_22d_WFDEI-GPCC_histsoc_{var}_monthly_1901_2016.nc4',\n",
    "                           decode_times = False)\n",
    "    # set time dimension\n",
    "    units, reference_date = nc.time.attrs['units'].split('since')\n",
    "    nc['time'] = pd.date_range(start = reference_date, periods = nc.sizes['time'], freq = 'MS')\n",
    "    \n",
    "    # output raster 2D if it does not already exist\n",
    "    for year in year_list:\n",
    "        for month in range(1, 13):\n",
    "            output_raster = data_folder / f'{var}_raster2d' / f'{var}_{str(year)}{str(month).zfill(2)}.tif'\n",
    "            if not os.path.exists(output_raster):\n",
    "                nc.rio.write_crs(\"epsg:4326\", inplace = True).sel(time = f'{str(year)}-{str(month).zfill(2)}-01').rio.to_raster(output_raster)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "id": "fd027ce5-3760-490f-9bb4-d4ddbaf61200",
   "metadata": {},
   "outputs": [],
   "source": [
    "# clip each raster\n",
    "for var in var_list:\n",
    "    for year in year_list:\n",
    "        for month in range(1, 13):\n",
    "            input_raster = data_folder / f'{var}_raster2d' / f'{var}_{str(year)}{str(month).zfill(2)}.tif'\n",
    "            with rasterio.open(input_raster) as src:\n",
    "                # shapely presumes all operations on two or more features exist in the same Cartesian plane.\n",
    "                out_image, out_transform = rasterio.mask.mask(\n",
    "                    src, aoi.geometry, crop = True)\n",
    "                if year == year_list[0] and month == 1:\n",
    "                    out_meta = src.meta.copy()\n",
    "\n",
    "            out_meta.update({\"driver\": \"GTiff\",\n",
    "                             \"height\": out_image.shape[1],\n",
    "                             \"width\": out_image.shape[2],\n",
    "                             \"transform\": out_transform})\n",
    "\n",
    "            output_file = f'{var}_{str(year)}{str(month).zfill(2)}.tif'\n",
    "            with rasterio.open(int_output_folder / output_file, \"w\", **out_meta) as dest:\n",
    "                dest.write(out_image)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "id": "7bb594ac-5243-46d2-82f0-55848c2d64b1",
   "metadata": {},
   "outputs": [],
   "source": [
    "# average raster (or other summary stats) and save\n",
    "for var in var_list:\n",
    "    raster_list = []\n",
    "    \n",
    "    for year in year_list:\n",
    "        for month in range(1, 13):\n",
    "            with rasterio.open(int_output_folder / f'{var}_{str(year)}{str(month).zfill(2)}.tif') as src:\n",
    "                raster_list.append(src.read(1))\n",
    "                if year == year_list[0] and month == 1:\n",
    "                    meta = src.meta\n",
    "    \n",
    "    output_raster_mean = np.nanmean(raster_list, axis = 0)\n",
    "    output_raster_std = np.nanstd(raster_list, axis = 0)\n",
    "    \n",
    "    with rasterio.open(output_folder / f'{city.lower()}_{var}_mean.tif', 'w', **meta) as dst:\n",
    "        dst.write(output_raster_mean, 1)\n",
    "    with rasterio.open(output_folder / f'{city.lower()}_{var}_std.tif', 'w', **meta) as dst:\n",
    "        dst.write(output_raster_std, 1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 30,
   "id": "a7abfb1b-d525-4733-b007-29ca0da67717",
   "metadata": {},
   "outputs": [],
   "source": [
    "# generate time series stats\n",
    "for var in var_list:\n",
    "    time_series = {'mean': {}, 'std': {},\n",
    "                   'max' : {}, 'min': {}}\n",
    "    \n",
    "    for year in year_list:\n",
    "        for month in range(1, 13):\n",
    "            with rasterio.open(int_output_folder / f'{var}_{str(year)}{str(month).zfill(2)}.tif') as src:\n",
    "                src_array = src.read(1)\n",
    "                src_array = src_array[src_array != src.meta['nodata']]\n",
    "                time_series['mean'][f'{str(year)}{str(month).zfill(2)}'] = np.nanmean(src_array)\n",
    "                time_series['std'][f'{str(year)}{str(month).zfill(2)}'] = np.nanstd(src_array)\n",
    "                time_series['min'][f'{str(year)}{str(month).zfill(2)}'] = np.nanmin(src_array)\n",
    "                time_series['max'][f'{str(year)}{str(month).zfill(2)}'] = np.nanmax(src_array)\n",
    "\n",
    "    with open(output_folder / f'{city.lower()}_{var}.csv', 'w') as f:\n",
    "        f.write('year,month,mean,std,min,max\\n')\n",
    "        for year in year_list:\n",
    "            for month in range(1, 13):\n",
    "                f.write('%s,%s,%s,%s,%s,%s\\n' % (str(year), str(month).zfill(2),\n",
    "                                                 time_series['mean'][f'{str(year)}{str(month).zfill(2)}'],\n",
    "                                                 time_series['std'][f'{str(year)}{str(month).zfill(2)}'],\n",
    "                                                 time_series['min'][f'{str(year)}{str(month).zfill(2)}'],\n",
    "                                                 time_series['max'][f'{str(year)}{str(month).zfill(2)}']))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "074f54d9-e384-4b4e-8087-acdb406857cd",
   "metadata": {},
   "outputs": [],
   "source": [
    "# delete intermediate outputs\n",
    "for var in var_list:\n",
    "    for year in year_list:\n",
    "        for month in range(1, 13):\n",
    "            os.remove(int_output_folder / f'{var}_{str(year)}{str(month).zfill(2)}.tif')"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "crp",
   "language": "python",
   "name": "crp"
  },
  "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.12"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}