import matplotlib.pyplot as plt
import numpy as np
import pandas as pd


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

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)
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))
df['ctry_code'] = df.countries.map(code_dict)
df['colors']= df.countries.map(color_dict)
df['year_lbl'] ="'"+df['year'].astype(str).str[-2:].astype(str)
#df = df[df.countries=="Sweden"]
df

# Number of dots per side
n = df.sites.max()

bottom_white, left_white, right_white = df.sites[:3].values
print(bottom_white, left_white, right_white)

# define the triangle vertices
A = np.array([0, 0])
B = np.array([1, 0])
C = np.array([0.5, np.sqrt(3) / 2])

# Calculate middle point of the triangle for year labels
centroid = (A + B + C) / 3

# add even spacing
t = np.linspace(0, 1, n)

# define the triangle sides
bottom = np.outer(1 - t, A) + np.outer(t, B)  # B to A
left   = np.outer(1 - t, A) + np.outer(t, C)  # A to C
right  = np.outer(1 - t, B) + np.outer(t, C)  # B to C


line_params =  {
    "solid_capstyle": "round",
    "linewidth": 8,    
    "zorder": 1
}

scatter_params = {
    "s": 4,
    "zorder": 2
}

label_params = {
    'ha': 'center',
    'va': 'center',
    'fontsize': 12
}

fig, axes = plt.subplots(nrows=2, figsize=(6, 10))

for ax, (year, group) in zip(axes, df.groupby("year", sort=False)):
    # Extract all data at once
    white_counts = group.sites.values
    codes = group.ctry_code.values
    colors = group.colors.values
    
    # Plot sides
    for i, side_data in enumerate([bottom, left, right]):
        white_count = white_counts[i]
        color = colors[i]
        
        # if side is left i =1 start at the beginning of the array, otherwise start at the end.
        line_segment = slice(None, white_count) if i == 1 else slice(-white_count, None)
        
        # Plot line
        ax.plot(
            side_data[line_segment, 0],
            side_data[line_segment, 1],
            color=color,
            **line_params
        )
        
        # Plot points
        for i, p in enumerate(side_data):
            is_white = (i < white_count) if i == 1 else (i >= len(side_data) - white_count)
            ax.scatter(
                p[0], p[1],
                color="white" if is_white else "black",
                **scatter_params
            )

    # Add labels
    offset = 0.05
    ax.text(A[0]-offset, A[1]-offset, codes[1], color=colors[1], **label_params)
    ax.text(B[0]+offset, B[1]-offset, codes[0], color=colors[0], **label_params)
    ax.text(C[0], C[1]+offset, codes[2], color=colors[2], **label_params)
    ax.text(centroid[0], centroid[1], year, ha='center', va='center', fontsize=22, weight="light")

    ax.set_aspect("equal")
    ax.axis("off")