{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e61f2a4e",
   "metadata": {},
   "outputs": [],
   "source": [
    "import pandas as pd\n",
    "import pypsa\n",
    "\n",
    "import matplotlib.pyplot as plt\n",
    "import country_converter as coco"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "73333588",
   "metadata": {},
   "outputs": [],
   "source": [
    "cc = coco.CountryConverter()\n",
    "historic = pd.read_csv(\n",
    "    \"../resources/validation-test/historical_cross_border_flows.csv\",\n",
    "    index_col=0,\n",
    "    header=0,\n",
    "    parse_dates=True,\n",
    ")\n",
    "\n",
    "n = pypsa.Network(\"../results/validation-test/networks/elec_s_37_ec_lv1.0_Ept.nc\")\n",
    "n.loads.carrier = \"load\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b38a3232",
   "metadata": {},
   "outputs": [],
   "source": [
    "len(historic.index)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "489739b8",
   "metadata": {},
   "outputs": [],
   "source": [
    "len(n.snapshots)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a79b025a",
   "metadata": {},
   "outputs": [],
   "source": [
    "if len(historic.index) > len(n.snapshots):\n",
    "    print(\"yes\")\n",
    "    historic = historic.resample(n.snapshots.inferred_freq).mean().loc[n.snapshots]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8e5b5b18",
   "metadata": {},
   "outputs": [],
   "source": [
    "all_country = sorted(\n",
    "    list(\n",
    "        set(\n",
    "            [link[:2] for link in historic.columns[1:]]\n",
    "            + [link[5:] for link in historic.columns[1:]]\n",
    "        )\n",
    "    )\n",
    ")\n",
    "\n",
    "color_country = {\n",
    "    \"AL\": \"#440154\",\n",
    "    \"AT\": \"#482677\",\n",
    "    \"BA\": \"#43398e\",\n",
    "    \"BE\": \"#3953a4\",\n",
    "    \"BG\": \"#2c728e\",\n",
    "    \"CH\": \"#228b8d\",\n",
    "    \"CZ\": \"#1f9d8a\",\n",
    "    \"DE\": \"#29af7f\",\n",
    "    \"DK\": \"#3fbc73\",\n",
    "    \"EE\": \"#5ec962\",\n",
    "    \"ES\": \"#84d44b\",\n",
    "    \"FI\": \"#addc30\",\n",
    "    \"FR\": \"#d8e219\",\n",
    "    \"GB\": \"#fde725\",\n",
    "    \"GR\": \"#f0f921\",\n",
    "    \"HR\": \"#f1c25e\",\n",
    "    \"HU\": \"#f4a784\",\n",
    "    \"IE\": \"#f78f98\",\n",
    "    \"IT\": \"#f87ea0\",\n",
    "    \"LT\": \"#f87a9a\",\n",
    "    \"LU\": \"#f57694\",\n",
    "    \"LV\": \"#f3758d\",\n",
    "    \"ME\": \"#f37685\",\n",
    "    \"MK\": \"#f37b7c\",\n",
    "    \"NL\": \"#f28774\",\n",
    "    \"NO\": \"#f1976b\",\n",
    "    \"PL\": \"#efaa63\",\n",
    "    \"PT\": \"#ebb160\",\n",
    "    \"RO\": \"#e6c260\",\n",
    "    \"RS\": \"#e2d75e\",\n",
    "    \"SE\": \"#dedc5b\",\n",
    "    \"SI\": \"#d9e35a\",\n",
    "    \"SK\": \"#d3e75a\",\n",
    "}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0b3c8b87",
   "metadata": {},
   "outputs": [],
   "source": [
    "optimized_links = n.links_t.p0.rename(\n",
    "    columns=dict(n.links.bus0.str[:2] + \" - \" + n.links.bus1.str[:2])\n",
    ")\n",
    "optimized_lines = n.lines_t.p0.rename(\n",
    "    columns=dict(n.lines.bus0.str[:2] + \" - \" + n.lines.bus1.str[:2])\n",
    ")\n",
    "optimized = pd.concat([optimized_links, optimized_lines], axis=1)\n",
    "\n",
    "# Drop internal country connection\n",
    "optimized.drop([c for c in optimized.columns if c[:2] == c[5:]], axis=1, inplace=True)\n",
    "\n",
    "# align columns name\n",
    "for c1 in optimized.columns:\n",
    "    for c2 in optimized.columns:\n",
    "        if c1[:2] == c2[5:] and c2[:2] == c1[5:]:\n",
    "            optimized = optimized.rename(columns={c1: c2})\n",
    "\n",
    "optimized = optimized.groupby(lambda x: x, axis=1).sum()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "118dcec1",
   "metadata": {},
   "outputs": [],
   "source": [
    "n.links_t.p0"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8cdc309f",
   "metadata": {},
   "outputs": [],
   "source": [
    "optimized_lines[\"AL - GR\"].plot()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "10e575d8",
   "metadata": {},
   "outputs": [],
   "source": [
    "optimized"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "66fd3180",
   "metadata": {},
   "outputs": [],
   "source": [
    "def sort_one_country(country, df):\n",
    "    indices = [link for link in df.columns if country in link]\n",
    "    df_country = df[indices].copy()\n",
    "    for link in df_country.columns:\n",
    "        if country in link[5:]:\n",
    "            df_country[link] = -df_country[link]\n",
    "            link_reverse = str(link[5:] + \" - \" + link[:2])\n",
    "            df_country = df_country.rename(columns={link: link_reverse})\n",
    "\n",
    "    return df_country.reindex(sorted(df_country.columns), axis=1)\n",
    "\n",
    "\n",
    "def cross_border_time_series(data):\n",
    "    fig, ax = plt.subplots(2 * len(all_country), 1, figsize=(15, 10 * len(all_country)))\n",
    "    axis = 0\n",
    "\n",
    "    for country in all_country:\n",
    "        ymin = 0\n",
    "        ymax = 0\n",
    "        for df in data:\n",
    "            df_country = sort_one_country(country, df)\n",
    "            df_neg, df_pos = df_country.clip(upper=0), df_country.clip(lower=0)\n",
    "\n",
    "            color = [color_country[link[5:]] for link in df_country.columns]\n",
    "\n",
    "            df_pos.plot.area(\n",
    "                ax=ax[axis], stacked=True, linewidth=0.0, color=color, ylim=[-1, 1]\n",
    "            )\n",
    "\n",
    "            df_neg.plot.area(\n",
    "                ax=ax[axis], stacked=True, linewidth=0.0, color=color, ylim=[-1, 1]\n",
    "            )\n",
    "            if (axis % 2) == 0:\n",
    "                title = \"Historic\"\n",
    "            else:\n",
    "                title = \"Optimized\"\n",
    "\n",
    "            ax[axis].set_title(\n",
    "                title + \" Import / Export for \" + cc.convert(country, to=\"name_short\")\n",
    "            )\n",
    "\n",
    "            # Custom legend elements\n",
    "            legend_elements = []\n",
    "\n",
    "            for link in df_country.columns:\n",
    "                legend_elements = legend_elements + [\n",
    "                    plt.fill_between(\n",
    "                        [],\n",
    "                        [],\n",
    "                        color=color_country[link[5:]],\n",
    "                        label=cc.convert(link[5:], to=\"name_short\"),\n",
    "                    )\n",
    "                ]\n",
    "\n",
    "            # Create the legend\n",
    "            ax[axis].legend(handles=legend_elements, loc=\"upper right\")\n",
    "\n",
    "            # rescale the y axis\n",
    "            neg_min = df_neg.sum(axis=1).min() * 1.2\n",
    "            if neg_min < ymin:\n",
    "                ymin = neg_min\n",
    "\n",
    "            pos_max = df_pos.sum(axis=1).max() * 1.2\n",
    "            if pos_max < ymax:\n",
    "                ymax = pos_max\n",
    "\n",
    "            print(f\"{str(len(df_pos))} in {str(axis)}\")\n",
    "\n",
    "            axis = axis + 1\n",
    "\n",
    "        for x in range(axis - 2, axis):\n",
    "            ax[x].set_ylim([neg_min, pos_max])\n",
    "\n",
    "    fig.savefig(f\"trade_time_series/trade_time_series.png\", bbox_inches=\"tight\")\n",
    "\n",
    "\n",
    "def cross_border_bar(data):\n",
    "    df_positive = pd.DataFrame()\n",
    "    df_negative = pd.DataFrame()\n",
    "    color = []\n",
    "\n",
    "    for country in all_country:\n",
    "        order = 0\n",
    "        for df in data:\n",
    "            if (order % 2) == 0:\n",
    "                title = \"Historic\"\n",
    "            else:\n",
    "                title = \"Optimized\"\n",
    "\n",
    "            df_country = sort_one_country(country, df)\n",
    "            df_neg, df_pos = df_country.clip(upper=0), df_country.clip(lower=0)\n",
    "\n",
    "            color = [color_country[link[5:]] for link in df_country.columns] + color\n",
    "            df_positive_new = pd.DataFrame(data=df_pos.sum()).T.rename(\n",
    "                {0: title + \" \" + cc.convert(country, to=\"name_short\")}\n",
    "            )\n",
    "            df_negative_new = pd.DataFrame(data=df_neg.sum()).T.rename(\n",
    "                {0: title + \" \" + cc.convert(country, to=\"name_short\")}\n",
    "            )\n",
    "\n",
    "            df_positive = pd.concat([df_positive_new, df_positive])\n",
    "            df_negative = pd.concat([df_negative_new, df_negative])\n",
    "\n",
    "            order = order + 1\n",
    "\n",
    "    fig, ax = plt.subplots(figsize=(15, 60))\n",
    "\n",
    "    # for i in range(3):\n",
    "    #    color[i] = '#545454'\n",
    "\n",
    "    df_positive.plot.barh(ax=ax, stacked=True, color=color, zorder=2)\n",
    "    df_negative.plot.barh(ax=ax, stacked=True, color=color, zorder=2)\n",
    "\n",
    "    plt.grid(axis=\"x\", zorder=0)\n",
    "    plt.grid(axis=\"y\", zorder=0)\n",
    "\n",
    "    # Custom legend elements\n",
    "    legend_elements = []\n",
    "\n",
    "    for country in list(color_country.keys()):\n",
    "        legend_elements = legend_elements + [\n",
    "            plt.fill_between(\n",
    "                [],\n",
    "                [],\n",
    "                color=color_country[country],\n",
    "                label=cc.convert(country, to=\"name_short\"),\n",
    "            )\n",
    "        ]\n",
    "\n",
    "    # Create the legend\n",
    "    plt.legend(handles=legend_elements, loc=\"upper right\")\n",
    "\n",
    "    fig.savefig(\"cross_border_bar.png\", bbox_inches=\"tight\")\n",
    "\n",
    "    # Show the plot\n",
    "    plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b9684a69",
   "metadata": {
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "cross_border_bar([historic, optimized])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f8d67c3d",
   "metadata": {
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "cross_border_time_series([historic, optimized])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b5c36dbd",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "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.10.9"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}