etude_lego_jurassic_world/lib/plots/part_rarity.py

"""Visualisation des pièces les plus rares observées dans les sets filtrés."""

import csv
from pathlib import Path
from typing import List

import matplotlib.pyplot as plt
from matplotlib.offsetbox import AnnotationBbox, OffsetImage
from PIL import Image, ImageDraw, ImageFont

from lib.filesystem import ensure_parent_dir


def load_part_rarity(path: Path) -> List[dict]:
    """Charge le CSV des pièces rares."""
    rows: List[dict] = []
    with path.open() as csv_file:
        reader = csv.DictReader(csv_file)
        for row in reader:
            rows.append(row)
    return rows


def select_printed_exclusive(rows: List[dict], resources_dir: Path) -> List[dict]:
    """Filtre les pièces imprimées exclusives aux sets filtrés disposant d'une image locale."""
    filtered: List[dict] = []
    for row in rows:
        if row.get("other_sets_quantity", "0") != "0":
            continue
        if "print" not in row["part_name"].lower():
            continue
        image_path = resources_dir / row.get("sample_set_id", "") / "rare_parts" / f"{row['part_num']}.jpg"
        if not image_path.exists():
            continue
        filtered.append(row)
    filtered.sort(key=lambda r: (r["part_name"], r["part_num"]))
    return filtered


def format_label(row: dict) -> str:
    """Formate l’étiquette de l’axe vertical."""
    return f"{row['part_num']} — {row['part_name']}"


def load_part_image(row: dict, resources_dir: Path) -> Image.Image | None:
    """Charge l'image associée à une pièce si elle est disponible."""
    path = resources_dir / row["sample_set_id"] / "rare_parts" / f"{row['part_num']}.jpg"
    if not path.exists():
        return None
    return Image.open(path)


def plot_part_rarity(
    path: Path,
    destination_path: Path,
    resources_dir: Path = Path("figures/rebrickable"),
    show_images: bool = True,
) -> None:
    """Trace un bar chart horizontal des pièces les plus rares avec leurs visuels."""
    rows = load_part_rarity(path)
    selected = rows
    labels = [format_label(row) for row in selected]
    filtered_counts = [int(row["filtered_quantity"]) for row in selected]
    other_counts = [int(row["other_sets_quantity"]) for row in selected]
    positions = list(range(len(selected)))

    fig, ax = plt.subplots(figsize=(13, 0.55 * len(selected) + 1.4))
    ax.barh(positions, filtered_counts, color="#1f78b4", label="Sets filtrés")
    ax.barh(positions, other_counts, left=filtered_counts, color="#b2df8a", label="Autres sets")
    ax.set_yticks(positions)
    ax.set_yticklabels(labels)
    ax.set_xlabel("Occurrences de la pièce (rechanges incluses)")
    ax.grid(axis="x", linestyle="--", alpha=0.4)
    ax.legend()

    if show_images:
        max_count = max((f + o) for f, o in zip(filtered_counts, other_counts)) if selected else 0
        pad = max_count * 0.15 if max_count > 0 else 1.0
        ax.set_xlim(left=-pad, right=max_count + pad * 0.3)
        for row, pos in zip(selected, positions):
            image = load_part_image(row, resources_dir)
            if image is None:
                continue
            target_height = 28
            ratio = target_height / image.height
            resized = image.resize((int(image.width * ratio), target_height))
            imagebox = OffsetImage(resized)
            ab = AnnotationBbox(
                imagebox,
                (-pad * 0.45, pos),
                xycoords=("data", "data"),
                box_alignment=(0.5, 0.5),
                frameon=False,
            )
            ax.add_artist(ab)

    fig.subplots_adjust(left=0.42)
    fig.tight_layout()

    ensure_parent_dir(destination_path)
    fig.savefig(destination_path, dpi=150)
    plt.close(fig)


def plot_printed_exclusive_parts(
    path: Path,
    destination_path: Path,
    resources_dir: Path = Path("figures/rebrickable"),
    columns: int = 5,
) -> None:
    """Assemble les images des pièces imprimées exclusives aux sets filtrés."""
    rows = load_part_rarity(path)
    selected = select_printed_exclusive(rows, resources_dir)
    selected.sort(key=lambda r: (int(r.get("sample_set_year", "9999") or 9999), r["sample_set_num"], r["part_num"]))
    if not selected:
        return
    images: List[Image.Image] = []
    labels: List[str] = []
    for row in selected:
        image_path = resources_dir / row["sample_set_id"] / "rare_parts" / f"{row['part_num']}.jpg"
        img = Image.open(image_path).convert("RGBA")
        max_side = 180
        ratio = min(max_side / img.width, max_side / img.height, 1.0)
        if ratio < 1.0:
            img = img.resize((int(img.width * ratio), int(img.height * ratio)))
        images.append(img)
        labels.append(f"{row.get('sample_set_year', '')} • {row['sample_set_num']}")

    columns = max(1, columns)
    rows_count = (len(images) + columns - 1) // columns
    cell_width = 220
    font = ImageFont.load_default()
    draw_temp = ImageDraw.Draw(Image.new("RGB", (10, 10)))
    def measure(text: str) -> tuple[int, int]:
        bbox = draw_temp.textbbox((0, 0), text, font=font)
        return bbox[2] - bbox[0], bbox[3] - bbox[1]
    text_height = max(measure(label)[1] for label in labels)
    cell_height = 190 + text_height + 14
    width = columns * cell_width
    height = rows_count * cell_height
    canvas = Image.new("RGBA", (width, height), (255, 255, 255, 255))
    draw = ImageDraw.Draw(canvas)
    for index, (img, label) in enumerate(zip(images, labels)):
        col = index % columns
        row_idx = index // columns
        x = col * cell_width + (cell_width - img.width) // 2
        y = row_idx * cell_height + 8
        canvas.paste(img, (x, y), img)
        text_width, _ = measure(label)
        text_x = col * cell_width + (cell_width - text_width) // 2
        text_y = y + img.height + 6
        draw.text((text_x, text_y), label, fill="#111111", font=font)

    ensure_parent_dir(destination_path)
    canvas.convert("RGB").save(destination_path, "PNG")