#!/usr/bin/env python
# coding: utf-8

# # Creating Plots of Temporal Changes using Matplotlib

# [![Binder](https://mybinder.org/badge_logo.svg)](https://mybinder.org/v2/gh/ai4er-cdt/gtc-biodiversity/main?filepath=notebooks%2F5-demo-nodediff.ipynb)

# This tutorial shows how to create plots of dynamics and changes in GeoGraphs over time using Matplotlib.

# ## 1. Setup and Loading package

# In[1]:


# %load ../jupyter_setup.txt
# Convenient jupyter setup
get_ipython().run_line_magic('load_ext', 'autoreload')
get_ipython().run_line_magic('autoreload', '2')
get_ipython().run_line_magic('config', 'IPCompleter.greedy=True')
get_ipython().run_line_magic('config', 'IPCompleter.use_jedi=False')


# In[2]:


import matplotlib.pyplot as plt
import matplotlib as mpl
import seaborn as sns
import numpy as np
import pandas as pd
import rioxarray as rxr
import geopandas as gpd
import pylandstats as pls
from geograph import GeoGraph
from geograph.geotimeline import GeoGraphTimeline
from geograph.constants import UTM35N
from geograph.demo.binder_constants import DATA_DIR, ROIS
from geograph.demo.plot_settings import ps_defaults, set_dim

ps_defaults(use_tex=True)

# Parse geotif landcover data
chernobyl_path = (
    lambda year: DATA_DIR / "chernobyl" / "esa_cci" / f"esa_cci_{year}_chernobyl.tif"
)

# Parse ROIS
rois = gpd.read_file(ROIS)
cez = rois[rois["name"] == "Chernobyl Exclusion Zone"]
ez = rois[rois.name.str.contains("Exclusion")]


# ## 2. Load Chernobyl Exclusion Zone data

# In[3]:


def clip_and_reproject(xrdata, clip_geometry=None, to_crs=UTM35N, x_res=300, y_res=300):

    if clip_geometry is not None:
        clipped_data = xrdata.rio.clip(clip_geometry)
    else:
        clipped_data = xrdata

    if to_crs is not None:
        reprojected_data = clipped_data.rio.reproject(to_crs, resolution=(x_res, y_res))
    else:
        reprojected_data = clipped_data

    return reprojected_data


# In[4]:


# Loading raster data
years = list(range(2000, 2015))
cez_rasters = {
    year: clip_and_reproject(
        rxr.open_rasterio(chernobyl_path(year)), clip_geometry=cez.geometry
    )
    for year in years
}


# In[5]:


## NOTE: For faster loading you can load the graphs from memory.
#        The demo path includes pre-loaded graphs for faster loading. Simply uncomment.
# demo_path = DATA_DIR / "chernobyl" / "graphs"
# years = list(range(2000,2015))
# cez_graphs = {year: GeoGraph(chernobyl_path(yera))
#              for year in years}

# Polygonising raster and transforming into graph
cez_graphs = {}
for year, raster in cez_rasters.items():
    print(f"Analysing year {year}")
    # Load geograph from the raster data (construction takes ~5s)
    cez_graphs[year] = GeoGraph(
        data=raster.data.squeeze(),
        transform=raster.rio.transform(),
        mask=raster.data.squeeze() > 0,
        crs=UTM35N,
        connectivity=8,
    )


# In[6]:


cez_timeline = GeoGraphTimeline(cez_graphs)


# In[7]:


# Perform node identification between adjacent time slices (takes ~10s)
cez_timeline.timestack()


# ## 3. Plots

# Let us now visualise the ecosystem dynamics from 2013 to 2014

# In[8]:


# Identify node dynamics for the year 2014
cez_timeline.calculate_node_dynamics(2014)


# In[9]:


cez_timeline[2014].df.node_dynamic.unique()


# In[10]:


graph = cez_timeline[2014]

plot_scale_factor = 1
dynamic_to_int = {
    "split": 0,
    "shrank": 1,
    "unchanged": 2,
    "complex": 3,
    "grew": 4,
    "merged": 5,
    "birth": 6,
}

colors = sns.color_palette("Paired").as_hex()
dynamic = lambda x: dynamic_to_int[x]
graph.df["dynamic_class"] = graph.df.node_dynamic.map(dynamic)

fig, ax = plt.subplots(1)
plt.title(
    "Chernobyl Exclusion Zone 2013 to 2014\n(ESA CCI 300m resolution)",
    fontsize=9 * plot_scale_factor,
)
set_dim(fig, fraction_of_line_width=plot_scale_factor)
vmin, vmax = 0, 7
cmap = mpl.colors.ListedColormap(
    [colors[7], colors[6], "lightgrey", colors[0], colors[2], colors[3], colors[9]]
)
graph.df.plot(column="dynamic_class", cmap=cmap, vmin=vmin, vmax=vmax, ax=ax)
ax.set_xticks([])
ax.set_yticks([])

inset_text = (
    "Node dynamics (#):\n"
    f"  Splits: {np.sum(graph.df['node_dynamic'] == 'split')}\n"
    f"  Shrinking: {np.sum(graph.df['node_dynamic'] == 'shrank')}\n"
    f"  Unchanged: {np.sum(graph.df['node_dynamic'] == 'unchanged')}\n"
    f"  Complex: {np.sum(graph.df['node_dynamic'] == 'complex')}\n"
    f"  Merges: {np.sum(graph.df['node_dynamic'] == 'merged')}\n"
    f"  Growth: {np.sum(graph.df['node_dynamic'] == 'grew')}\n"
    f"  Births: {np.sum(graph.df['node_dynamic'] == 'birth')}"
)

# these are matplotlib.patch.Patch properties
props = dict(boxstyle="round", facecolor="white", alpha=0.8)

# place a text box in upper left in axes coords
ax.text(
    0.03,
    0.97,
    inset_text,
    transform=ax.transAxes,
    fontsize=9 * plot_scale_factor,
    verticalalignment="top",
    bbox=props,
)

from mpl_toolkits.axes_grid1.anchored_artists import AnchoredSizeBar
import matplotlib.font_manager as fm

fontprops = fm.FontProperties(size=6 * plot_scale_factor)
scalebar = AnchoredSizeBar(
    ax.transData,
    1e4,
    "10 km",
    "lower left",
    pad=0,
    borderpad=0.8,
    color="black",
    frameon=False,
    size_vertical=250,
    fontproperties=fontprops,
)

ax.add_artist(scalebar)

# add colorbar
fig = ax.get_figure()
cax = fig.add_axes(
    [0.12, 0.08, 0.78, 0.05]
)  # left-offset, # bottom offset # width, # height
sm = plt.cm.ScalarMappable(cmap=cmap, norm=plt.Normalize(vmin=vmin, vmax=vmax))
sm._A = []
cbar = fig.colorbar(sm, cax=cax, orientation="horizontal")
cbar.set_ticks(np.arange(vmin + 0.5, vmax + 1))
cbar.set_ticklabels(
    ["Split", "Shrank", "Unchanged", "Complex", "Grew", "Merged", "Birth"]
)
cbar.ax.tick_params(labelsize=9 * plot_scale_factor)
ax.spines["bottom"].set_visible(False)
ax.spines["top"].set_visible(False)
ax.spines["right"].set_visible(False)
ax.spines["left"].set_visible(False)


# plt.savefig("CEZ_node_dynamics.svg")
# plt.savefig("CEZ_node_dynamics.png")


# In[11]:


plot_scale_factor = 1
fig, ax = plt.subplots(1)
set_dim(fig, fraction_of_line_width=plot_scale_factor)

cmap = sns.diverging_palette(6, 120, s=360, l=55, as_cmap=True)
norm = mpl.colors.TwoSlopeNorm(vcenter=0, vmin=-10e5, vmax=10e5)

graph.df.plot(
    "absolute_growth", ax=ax, cmap=cmap, norm=norm, edgecolor="grey", linewidth=0.1
)
ax.set_xticks([])
ax.set_yticks([])
from mpl_toolkits.axes_grid1.anchored_artists import AnchoredSizeBar
import matplotlib.font_manager as fm

fontprops = fm.FontProperties(size=6 * plot_scale_factor)
scalebar = AnchoredSizeBar(
    ax.transData,
    1e4,
    "10 km",
    "lower left",
    pad=0,
    borderpad=0.3,
    color="black",
    frameon=False,
    size_vertical=250,
    fontproperties=fontprops,
)

ax.add_artist(scalebar)

cbar = fig.colorbar(
    mpl.cm.ScalarMappable(norm=norm, cmap=cmap),
    orientation="horizontal",
    label="Absolute growth rate [ha / yr]",
    # aspect=9,
    shrink=0.88,
    pad=0.04,
)
cbar.set_ticks(
    [-10e5, -7.5 * 1e5, -5e5, -2.5 * 1e5, 0, 2.5 * 1e5, 5e5, 7.5 * 1e5, 10e5]
)
cbar.set_ticklabels([-100, -75, -50, -25, 0, 25, 50, 75, 100])
cbar.ax.tick_params(labelsize=9 * plot_scale_factor)
cbar.set_label("Absolute growth rate [ha/yr]", fontsize=9 * plot_scale_factor)
plt.title(
    "Chernobyl Exclusion Zone 2013 to 2014\n(ESA CCI 300m resolution)",
    fontsize=9 * plot_scale_factor,
)
ax.spines["bottom"].set_visible(False)
ax.spines["top"].set_visible(False)
ax.spines["right"].set_visible(False)
ax.spines["left"].set_visible(False)
plt.savefig("CEZ_node_growth_rates.svg")
plt.savefig("CEZ_node_growth_rates.pdf")
plt.savefig("CEZ_node_growth_rates.png")


# In[12]:


cez_timeline[2014].get_patch_metrics()


# In[13]:


fig, ax = plt.subplots(1)
plot_scale_factor = 2
set_dim(fig, fraction_of_line_width=plot_scale_factor)

cez_timeline[2013].df.loc[1629:1629].plot(ax=ax, color="red", alpha=0.6)
cez_timeline[2014].df.loc[1187:1187].plot(ax=ax, color="blue", alpha=0.6)
cez_timeline[2014].df.loc[1625:1625].plot(ax=ax, color="green", alpha=0.8)

inset_text = (
    f"Evergreen forest (70)\n\n"
    f"Total area: {graph.df.area.loc[1625]/1e4:.0f} ha\n"
    f"Perimeter: {graph.df.perimeter.loc[1625]/1e3:.0f} km\n"
    f"Fractal dimension: {graph.df.fractal_dimension.loc[1625]:.2f}\n"
    f"Shape index: {graph.df.shape_index.loc[1625]:.2f}"
)

# these are matplotlib.patch.Patch properties
props = dict(boxstyle="round", facecolor="white", alpha=0.8)

# place a text box in upper left in axes coords
ax.text(
    0.6,
    0.95,
    inset_text,
    transform=ax.transAxes,
    fontsize=8 * plot_scale_factor,
    verticalalignment="top",
    bbox=props,
)

ax.set_xticks([])
ax.set_yticks([])
ax.spines["bottom"].set_visible(False)
ax.spines["top"].set_visible(False)
ax.spines["right"].set_visible(False)
ax.spines["left"].set_visible(False)
plt.savefig("CEZ-nodediff-example.svg")
plt.savefig("CEZ-nodediff-example.pdf")