# tutorial
import matplotlib.pyplot as plt
import math
import numpy as np
import pandas as pd


color_dict = {"Norway": "#2B314D", "Denmark": "#A54836", "Sweden": "#5375D4"}

xy_ticklabel_color, grid_color, datalabels_color = "#757C85", "#ECEFF1", "#FFFFFF"

data = {
    "year": [2004, 2022, 2004, 2022, 2004, 2022],
    "countries": ["Denmark", "Denmark", "Norway", "Norway", "Sweden", "Sweden"],
    "sites": [4, 10, 5, 8, 13, 15],
}
df = pd.DataFrame(data)
df = df.sort_values(["countries", "year"], ascending=False).reset_index(drop=True)
# map the colors of a dict to a dataframe
df["color"] = df.countries.map(color_dict)
df

fig, axes = plt.subplots(ncols=len(df.year.unique()), nrows=len(df.countries.unique()), sharex=True, figsize=(5, 5))

for row, shift_ax, ax in zip(df.itertuples(), shift_axes, axes.ravel()):
    ax.set(xlim=(-2, 2),ylim=(-0, 1))

    #add the triangles
    ax.scatter(
        X,
        Y,
        s=row.sites * 400,
        marker=10,
        color=row.color,
        zorder=1,
    )

    #add the vertial lines
    ax.axvline(
        x=0,
        ymin=0,
        ymax=1,
        color=grid_color,
        linestyle="-",
        clip_on=False,
        zorder=0,
    )

    #add site labels
    ax.annotate(
        row.sites,
        (X, Y + math.sqrt(row.sites) / 8),
        ha="center",
        va="top",
        color=datalabels_color,
    )

    box = ax.get_position()
    box.x0 = box.x0 + shift_ax  # x0 first coordinate of the box
    box.x1 = box.x1 + shift_ax  # x1 last coordinate of the box
    ax.set_position(box)
    ax.patch.set_alpha(0.1)  # transparent axis

    ax.set_frame_on(False)
    ax.tick_params(length=0, labelleft = False, labelbottom = False)


for ax, country in zip(axes[:, 0], df.countries.unique()):
    ax.set_ylabel(
        country, 
        rotation = 0,
        size = 10, 
        color = xy_ticklabel_color,
        labelpad =23
        )


for ax, year in zip(axes[0], df.year.unique()):
    ax.set_xlabel(
        year, 
        size = 10, 
        color = xy_ticklabel_color
        )
    ax.xaxis.set_label_position("top")