import matplotlib.pyplot as plt
import matplotlib as mpl
from matplotlib.colors import LinearSegmentedColormap
import numpy as np
import pandas as pd
from svgpathtools import svg2paths
from svgpath2mpl import parse_path

gradient_colors = [
    ["#CE5C46", "#DB9386"],
    ["#5678D7", "#8DADF2"],
    ["#303653", "#5C6077"],
]

color_dict = {
    (2004, "Norway"): "#303653",
    (2022, "Norway"): "#5C6077",
    (2004, "Denmark"): "#CE5C46",
    (2022, "Denmark"): "#DB9386",
    (2004, "Sweden"): "#5678D7",
    (2022, "Sweden"): "#8DADF2",
}

xy_ticklabel_color, xlabel_color, grand_totals_color, grid_color, datalabels_color = (
    "#101628",
    "#101628",
    "#101628",
    "#F2F3F4",
    "#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["year_lbl"] = "'" + df["year"].astype(str).str[-2:].astype(str)
df["sub_total"] = df.groupby("countries")["sites"].transform("sum")
df["pct_change"] = (
    df.groupby("countries", sort=False)["sites"]
    .apply(lambda x: x.pct_change())
    .to_numpy()
    .round(3)
    * 100
)
df = df.fillna(method="bfill")
# custom 
# custom sort
sort_order_dict = {"Denmark": 1, "Sweden": 2, "Norway": 3, 2004: 4, 2022: 5}
df = df.sort_values(
    by=[
        "year",
        "countries",
    ],
    key=lambda x: x.map(sort_order_dict),
)

# Add the color based on the color dictionary
df["color"] = df.set_index(["year", "countries"]).index.map(color_dict.get)
df

pcts = [int(num) if float(num).is_integer() else num for num in df["pct_change"]]
groups = df.groupby("countries", sort = False)
print(pcts[:3])

fig, axes = plt.subplots(ncols=df.countries.nunique(), nrows=1, sharey=True, figsize=(10, 6), facecolor="#FFFFFF")
fig.tight_layout(pad=0.1)

for (country, group), pct, ax in zip(groups, pcts[:3], axes.ravel()):

    ax.set(xlim = (-0.1, 1.1), ylim = (0, 16))
    year_labels = group.year_lbl
    sites = group.sites
    color = group.color.iloc[0]

    cm = LinearSegmentedColormap.from_list("Temperature Map", group.color.tolist())
    
    # fill between returns a polygon
    polygon = ax.fill_between(
        year_labels, 
        0, 
        sites, 
        lw = 0, 
        color = "none"
    ) 
    
    # get the shape of the poligon 
    verts = np.vstack(
        [p.vertices for p in polygon.get_paths()]
    )
    # create an imshow
    gradient = ax.imshow(
        np.linspace(0, 1, 256).reshape(-1, 1),
        cmap=cm,
        aspect="auto",
        origin="lower",
        extent=[
            verts[:, 0].min(),
            verts[:, 0].max(),
            verts[:, 1].min(),
            verts[:, 1].max(),
        ],
    )  
    # cut over the plot
    gradient.set_clip_path(
        polygon.get_paths()[0], transform=ax.transData
    )  

    #add the line and dots at the top
    ax.plot(
        year_labels,
        sites,
        "-o",
        ms=8,
        mec="w",
        clip_on=False,
        color=color,
        zorder=1,
    )

    #add data label
    for i, site in enumerate(sites):
        ax.annotate(
            site,
            xy=(i, site + 1),
            size=12,
            color=xlabel_color,
            ha="center",
            va="center",
        )

    #add country names
    ax.set_xlabel(
        f"{group.countries.iloc[0]}",
        size=20,
        color=xlabel_color,
        weight="bold",
        ha="center",
        labelpad =23
    )


    #add the heritage icon
    ax.plot(
        0.5,
        -5,
        marker=icon_marker,
        color=color,
        markersize=32,
        clip_on=False,
    )

    #add arrrow
    ax.annotate(
        f"  \u25b2\n\n{pct}%",
        xy=(0.5, 1),
        color=datalabels_color,
        size=28,
        weight="bold",
        ha="center",
    )

    #add lines between the plots
    ax.axvline(
        -0.2,
        -0.2, 
        1.1, 
        clip_on=False, 
        color=grid_color
        )

    #styling
    ax.tick_params(
        length=0, 
        labelsize=12, 
        colors=xy_ticklabel_color, 
        labelleft = False, 
        pad=23
        )

    ax.set_frame_on(False)

#add first line
  
ax.axvline(
    1.2,
    -0.2,
    1.1,
    clip_on=False,
    color=grid_color,
)