from __future__ import annotations

import numpy as np
import matplotlib as mpl
import matplotlib.pyplot as plt

mpl.rcParams.update({
    "pdf.fonttype": 42,
    "ps.fonttype": 42,
    "figure.dpi": 170,
    "savefig.dpi": 240,
    "font.family": "DejaVu Sans",
    "font.size": 8.2,
    "axes.titlesize": 9.6,
    "axes.labelsize": 8.4,
    "xtick.labelsize": 7.2,
    "ytick.labelsize": 7.2,
    "legend.fontsize": 7.2,
    "axes.linewidth": 0.75,
})

COL_INK = "#172033"
COL_MUTED = "#667085"
COL_GRID = "#d9dee7"
COL_BLUE = "#356ca5"
COL_TEAL = "#2f9c95"
COL_ORANGE = "#d8863b"
COL_RED = "#c75756"
COL_PURPLE = "#8066a8"
COL_GREEN = "#5f9b68"


def polish_axes(ax, grid_axis="y"):
    ax.spines["top"].set_visible(False)
    ax.spines["right"].set_visible(False)
    ax.spines["left"].set_color("#303642")
    ax.spines["bottom"].set_color("#303642")
    ax.tick_params(length=0, colors=COL_MUTED, pad=3)
    if grid_axis:
        ax.grid(True, axis=grid_axis, color=COL_GRID, linewidth=0.55, alpha=0.85)
        ax.set_axisbelow(True)


def save(fig):
    fig.savefig("augmented.png", bbox_inches="tight", facecolor="white")
    fig.savefig("augmented.pdf", bbox_inches="tight", facecolor="white")

# DATA SECTOR: same countries, annual business-registration days, and populations as original.py.
countries = ['Algeria', 'Antigua & Barbuda', 'Ghana', 'Latvia', 'Vietnam', 'Kenya', 'Nigeria', 'Ethiopia', 'South Africa', 'Morocco', 'Egypt', 'Tunisia', 'Gambia', 'Mozambique', 'Rwanda', 'Botswana', 'Namibia', 'Zambia', 'Eritrea', 'Somalia', 'Senegal', 'Sudan', 'South Sudan', 'Lesotho', 'Niger']
years = [2019, 2020, 2021, 2022, 2023, 2024, 2025, 2026, 2027]
day_series = {
    2019: [32, 28, 19, 20, 27, 29, 31, 28, 33, 34, 31, 22, 24, 25, 23, 27, 29, 28, 26, 29, 22, 30, 30, 24, 33],
    2020: [31, 26, 18, 19, 26, 28, 30, 27, 30, 31, 29, 21, 23, 24, 22, 26, 28, 27, 25, 28, 21, 29, 29, 22, 32],
    2021: [32, 28, 19, 20, 27, 29, 31, 28, 32, 33, 30, 22, 24, 25, 23, 27, 30, 29, 27, 30, 22, 31, 31, 23, 31],
    2022: [33, 28, 20, 21, 28, 30, 32, 29, 33, 34, 31, 22, 25, 26, 24, 28, 31, 30, 28, 31, 23, 32, 32, 25, 32],
    2023: [34, 29, 21, 22, 29, 31, 33, 30, 34, 35, 32, 23, 26, 27, 25, 29, 32, 31, 29, 33, 24, 34, 34, 26, 33],
    2024: [35, 30, 22, 23, 30, 32, 34, 31, 35, 36, 33, 24, 27, 28, 26, 30, 33, 32, 30, 34, 25, 35, 35, 27, 34],
    2025: [22, 21, 13, 15, 20, 22, 24, 21, 24, 25, 23, 16, 18, 19, 17, 21, 22, 21, 17, 19, 18, 20, 21, 15, 25],
    2026: [23, 22, 14, 16, 21, 23, 25, 22, 25, 26, 24, 17, 19, 20, 18, 22, 23, 22, 18, 20, 19, 21, 22, 16, 26],
    2027: [24, 23, 15, 17, 22, 24, 26, 23, 26, 27, 25, 18, 20, 21, 19, 23, 24, 23, 19, 21, 20, 22, 23, 17, 27],
}
population = [44.2, 0.12, 32.5, 2.1, 98.3, 55.1, 216.4, 123.2, 60.5, 37.2, 106.1, 12.3, 2.6, 32.2, 13.1, 2.5, 2.6, 20.3, 3.8, 16.5, 17.0, 45.1, 11.2, 2.1, 24.2]
avg_days = np.array([[day_series[yr][idx] for yr in years] for idx in range(len(countries))]).mean(axis=1)
angles = np.linspace(0.0, 2 * np.pi, len(countries), endpoint=False)
width = 2 * np.pi / len(countries) * 0.86
norm = mpl.colors.Normalize(min(population), max(population))
cmap = mpl.cm.viridis
colors = cmap(norm(population))
fig, ax = plt.subplots(figsize=(6.3, 5.25), subplot_kw=dict(polar=True))
bars = ax.bar(angles, avg_days, width=width, bottom=0.0, color=colors, edgecolor="white", linewidth=0.55)
ax.set_theta_zero_location("N")
ax.set_theta_direction(-1)
ax.set_ylim(0, 36)
ax.set_yticks([10, 20, 30])
ax.set_yticklabels(["10", "20", "30"], color=COL_MUTED, fontsize=7)
ax.grid(color=COL_GRID, linewidth=0.55, alpha=0.85)
ax.spines["polar"].set_color("#303642")
ax.spines["polar"].set_linewidth(0.75)
ax.set_xticks(angles)
short = [c.replace("Antigua & Barbuda", "Antigua").replace("South Africa", "S. Africa").replace("South Sudan", "S. Sudan") for c in countries]
ax.set_xticklabels(short, fontsize=6.2, color=COL_MUTED)
ax.tick_params(pad=2)
ax.set_title("Average days to register a business (2019-2027)", va="bottom", color=COL_INK, fontweight=600, pad=12)
sm = mpl.cm.ScalarMappable(cmap=cmap, norm=norm)
sm.set_array([])
cbar = fig.colorbar(sm, ax=ax, pad=0.10, fraction=0.042)
cbar.set_label("Population (millions)", color=COL_MUTED)
cbar.ax.tick_params(length=0, colors=COL_MUTED, labelsize=7)
cbar.outline.set_linewidth(0.55)
fig.subplots_adjust(left=0.05, right=0.88, top=0.91, bottom=0.07)
save(fig)