#!/usr/bin/env python3
import csv
from pathlib import Path

import matplotlib.pyplot as plt

HERE = Path(__file__).parent
CSV_PATH = HERE / "distribution.csv"
OUT_PATH = HERE / "img" / "distribution.png"

BAR_COLORS = {
    "red": "#d94f4f",
    "yellow": "#e8c14a",
    "gray": "#a3a3a3",
    "green": "#5cad6b",
    "white": "#f2ede4",
    "black": "#3a332c",
    "blue": "#5b8ac4",
    "purple": "#9b7fc7",
}

INK = "#f2e9dc"
EDGE = "#f2e9dc"

def load_data():
    with CSV_PATH.open() as f:
        rows = list(csv.DictReader(f))
    rows.sort(key=lambda r: int(r["count"]), reverse=True)
    return rows

def plot(rows):
    labels = [r["color"] for r in rows]
    counts = [int(r["count"]) for r in rows]
    colors = [BAR_COLORS[c] for c in labels]

    fig, ax = plt.subplots(figsize=(8, 5), dpi=200)
    fig.patch.set_alpha(0)
    ax.patch.set_alpha(0)

    bars = ax.bar(
        labels,
        counts,
        color=colors,
        edgecolor=EDGE,
        linewidth=1.1,
        width=0.65,
        zorder=3,
    )

    for bar, count in zip(bars, counts):
        ax.text(
            bar.get_x() + bar.get_width() / 2,
            bar.get_height() + max(counts) * 0.02,
            str(count),
            ha="center",
            va="bottom",
            color=INK,
            fontsize=11,
            fontweight="bold",
        )

    ax.set_title("Bottle Cap Colors", color=INK, fontsize=15, fontweight="bold", pad=16)
    ax.set_ylabel("Count", color=INK, fontsize=11)

    ax.set_ylim(0, max(counts) * 1.15)
    ax.tick_params(colors=INK, labelsize=11)
    ax.set_xticks(range(len(labels)))
    ax.set_xticklabels([l.capitalize() for l in labels])

    for spine in ("top", "right", "left"):
        ax.spines[spine].set_visible(False)
    ax.spines["bottom"].set_color(INK)
    ax.spines["bottom"].set_alpha(0.4)

    ax.yaxis.grid(True, color=INK, alpha=0.15, linewidth=0.8, zorder=0)
    ax.set_axisbelow(True)

    fig.tight_layout()
    OUT_PATH.parent.mkdir(exist_ok=True)
    fig.savefig(OUT_PATH, transparent=True)
    print(f"Saved {OUT_PATH}")


if __name__ == "__main__":
    plot(load_data())