#/usr/bin/env python3
# usage: uv run owid-energy.py
# 
# requirements: matplotlib
import csv
import urllib.request
import matplotlib.pyplot as plt
import matplotlib.patches as mpatches
from collections import defaultdict

url = "https://raw.githubusercontent.com/owid/energy-data/master/owid-energy-data.csv"
startYear = 1990
regions = ["Africa", "Asia", "Europe", "North America", "South America", "Oceania"]

# label → (csv_column, color)
sources = {
    "Coal":     ("coal_electricity",            "#3b3b3b"),
    "Oil":      ("oil_electricity",             "#8b4513"),
    "Gas":      ("gas_electricity",             "#cd853f"),
    "Nuclear":  ("nuclear_electricity",         "#9400d3"),
    "Hydro":    ("hydro_electricity",           "#1e90ff"),
    "Wind":     ("wind_electricity",            "#00ced1"),
    "Solar":    ("solar_electricity",           "#ffd700"),
    "Other RE": ("other_renewable_electricity", "#32cd32"),
}

try:
    urllib.request.urlretrieve(url, "owid-energy-data.csv")
except Exception as e:
    print(f"Download failed ({e}), using local file.")

data = defaultdict(lambda: defaultdict(lambda: defaultdict(float)))
with open("owid-energy-data.csv", newline="", encoding="utf-8") as f:
    for row in csv.DictReader(f):
        if row["country"] not in regions or not row["year"].isdigit():
            continue
        if int(row["year"]) < startYear:
            continue
        for label, (col, _) in sources.items():
            try:
                data[row["country"]][int(row["year"])][label] = float(row.get(col) or 0)
            except ValueError:
                pass

fig, axes = plt.subplots(3, 2, figsize=(12, 16))
for ax, region in zip(axes.flatten(), regions):
    years = sorted(data[region])
    rows  = {label: [] for label in sources}
    valid = []
    for y in years:
        vals  = {l: data[region][y][l] for l in sources}
        total = sum(vals.values())
        if total <= 0:
            continue
        valid.append(y)
        for l in sources:
            rows[l].append(vals[l] / total * 100)

    ax.stackplot(valid, rows.values(), colors=[c for _, c in sources.values()], alpha=0.88)
    ax.set(title=region, xlim=(startYear, valid[-1]), ylim=(0, 100),
           xlabel="Year", ylabel="Share (%)")
    ax.grid(axis="y", linestyle="--", linewidth=0.5, alpha=0.5)

# handles = [mpatches.Patch(color=c, label=l) for l, (_, c) in sources.items()]
# fig.legend(handles=handles, loc="lower center", ncol=len(sources),
#            bbox_to_anchor=(0.5, 0.0), fontsize=10)
fig.suptitle(f"Electricity Generation Mix by Region ({startYear}–present)\n"
             "Source: Our World in Data", fontsize=14, fontweight="bold")

plt.tight_layout(rect=[0, 0.06, 1, 0.97])
fig.savefig("energy_mix.png", dpi=150, bbox_inches="tight")
plt.show()
