#!/usr/bin/env python3
"""Generate 1x5 figure: data collection strategy schematics.

Colors: blue = sampled / active, yellow = unsampled / inactive.
Blue–yellow is generally colorblind-friendly for the common red–green deficiencies
(protanopia/deuteranopia); tritan (blue–yellow) deficiency is much rarer.
Palette inspired by Wong / Okabe–Ito style contrasts on white.
"""
from pathlib import Path

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

# Colorblind-friendly contrast on white: dark blue vs lighter yellow/gold
ACTIVE = "#0072B2"  # Wong blue (sampled)
INACTIVE = "#F0DE8A"  # lighter yellow (still distinct from white at dpi)

N = 10
FIG_DIR = Path(__file__).resolve().parents[1] / "figures"


def _configure_font():
    """Sans-serif (distinct from previous Times/Liberation Serif): Helvetica / Arial / DejaVu Sans."""
    available = {f.name for f in fm.fontManager.ttflist}
    mpl.rcParams["font.family"] = "sans-serif"
    preferred = ["Helvetica", "Arial", "Liberation Sans", "DejaVu Sans"]
    chosen = [n for n in preferred if n in available]
    mpl.rcParams["font.sans-serif"] = chosen + ["DejaVu Sans", "sans-serif"]
    mpl.rcParams["axes.labelsize"] = 11
    mpl.rcParams["axes.titlesize"] = 14
    return chosen[0] if chosen else "sans-serif"


def grid_points():
    x, y = np.meshgrid(np.arange(N), np.arange(N), indexing="xy")
    return x.ravel(), y.ravel()


def plot_dots(ax, mask_active, title):
    x, y = grid_points()
    c = np.where(mask_active, ACTIVE, INACTIVE)
    ax.scatter(x, y, c=c, s=118, edgecolors="none", zorder=2, linewidths=0)
    ax.set_title(title, fontsize=14, fontweight="bold", pad=8)
    ax.set_xlabel("Factor 1 Values", fontsize=11, fontweight="bold")
    ax.set_ylabel("Factor 2 Values", fontsize=11, fontweight="bold")
    ax.set_xlim(-0.5, N - 0.5)
    ax.set_ylim(-0.5, N - 0.5)
    ax.set_xticks(np.arange(N))
    ax.set_yticks(np.arange(N))
    ax.set_xticklabels([])
    ax.set_yticklabels([])
    ax.tick_params(length=4)
    ax.set_aspect("equal")


def main():
    FIG_DIR.mkdir(parents=True, exist_ok=True)
    _configure_font()

    fig, axes = plt.subplots(1, 5, figsize=(17.8, 3.6), constrained_layout=True)

    x_idx, y_idx = np.meshgrid(np.arange(N), np.arange(N), indexing="xy")
    flat_x = x_idx.ravel()
    flat_y = y_idx.ravel()

    # 1) Complete
    mask = np.ones(N * N, dtype=bool)
    plot_dots(axes[0], mask, "Complete")

    # 2) Random (~9 points, reproducible)
    rng = np.random.default_rng(0)
    k = 9
    choice = rng.choice(N * N, size=k, replace=False)
    mask = np.zeros(N * N, dtype=bool)
    mask[choice] = True
    plot_dots(axes[1], mask, "Random")

    # 3) Diagonal (Factor1 == Factor2)
    mask = flat_x == flat_y
    plot_dots(axes[2], mask, "Diagonal")

    # 4) L: left column + bottom row only
    mask = (flat_x == 0) | (flat_y == 0)
    plot_dots(axes[3], mask, "L")

    # 5) Left column + main diagonal (no bottom row except where it meets the column/diagonal)
    mask = (flat_x == 0) | (flat_x == flat_y)
    plot_dots(axes[4], mask, "Ours")

    out = FIG_DIR / "data_collection_strategies_1x5.png"
    fig.savefig(out, dpi=200, bbox_inches="tight", facecolor="white")
    plt.close(fig)
    print(f"Saved {out}")


if __name__ == "__main__":
    main()