#!/usr/bin/env python3
"""3D lattice plots (8³ nodes): five standalone PNGs plus one 1×5 row (data_collection_strategies_3d_1x5.png).

Blue = active / sampled, yellow = inactive.

Logical lattice uses origin (0,0,0) for L / Ours joints; points are mapped to plot space with
    (x', y', z') = (N-1-ix, N-1-iy, iz)
so the common intersection maps to corner (N-1, N-1, 0). Camera is chosen so that corner reads
as the lower-right of the cube on the page (tune VIEW_ELEV / VIEW_AZIM if needed).
"""
from pathlib import Path

import matplotlib as mpl
import matplotlib.font_manager as fm
import matplotlib.pyplot as plt
import numpy as np
from mpl_toolkits.mplot3d import Axes3D  # noqa: F401 — registers 3d projection

N = 8
# Y arm for "Ours": fewer samples along logical y than full N
OURS_Y_LEN = 5

ACTIVE = "#0072B2"
INACTIVE = "#D4B84A"  # slightly deeper yellow/gold (was very pale)
INACTIVE_ALPHA = 0.42
FIG_DIR = Path(__file__).resolve().parents[1] / "figures"

# Camera: tune so the joint corner (plot coords N-1, N-1, 0) projects near the figure's lower-right
VIEW_ELEV = 9
VIEW_AZIM = 22


def _configure_font():
    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"]


def lattice_coords():
    ix, iy, iz = np.meshgrid(
        np.arange(N), np.arange(N), np.arange(N), indexing="ij"
    )
    return ix.ravel(), iy.ravel(), iz.ravel()


def logical_to_plot(IX, IY, IZ):
    """Map logical indices so joint (0,0,0) → plot (N-1,N-1,0)."""
    return (N - 1 - IX, N - 1 - IY, IZ)


def wireframe_cube(ax, m, color="#bfbfbf", lw=0.9):
    z0 = [(0, 0, 0), (m, 0, 0), (m, m, 0), (0, m, 0)]
    z1 = [(0, 0, m), (m, 0, m), (m, m, m), (0, m, m)]
    for i in range(4):
        p0, p1 = z0[i], z0[(i + 1) % 4]
        ax.plot([p0[0], p1[0]], [p0[1], p1[1]], [p0[2], p1[2]], c=color, lw=lw)
    for i in range(4):
        p0, p1 = z1[i], z1[(i + 1) % 4]
        ax.plot([p0[0], p1[0]], [p0[1], p1[1]], [p0[2], p1[2]], c=color, lw=lw)
    for a, b in zip(z0, z1):
        ax.plot([a[0], b[0]], [a[1], b[1]], [a[2], b[2]], c=color, lw=lw)


def plot_strategy(
    ax,
    mask_active: np.ndarray,
    title: str,
    show_inactive: bool,
    *,
    compact: bool = False,
):
    IX, IY, IZ = lattice_coords()
    PX, PY, PZ = logical_to_plot(IX, IY, IZ)
    m = N - 1

    if compact:
        s_in, s_act = 10, 18
        title_fs, label_fs = 14, 10
        title_pad = -22  # pull subplot titles (Complete, …) down toward the cube
        wire_lw = 0.55
        lw_scatter = 0.22
        labelpad_xy, labelpad_z = -10, -7
    else:
        s_in, s_act = 24, 42
        title_fs, label_fs = 14, 11
        title_pad = 2
        wire_lw = 0.9
        lw_scatter = 0.35
        labelpad_xy, labelpad_z = -14, -10

    wireframe_cube(ax, m, lw=wire_lw)

    if show_inactive:
        inactive = ~mask_active
        if np.any(inactive):
            ax.scatter(
                PX[inactive],
                PY[inactive],
                PZ[inactive],
                c=INACTIVE,
                s=s_in,
                alpha=INACTIVE_ALPHA,
                edgecolors="none",
                depthshade=True,
            )

    if np.any(mask_active):
        ax.scatter(
            PX[mask_active],
            PY[mask_active],
            PZ[mask_active],
            c=ACTIVE,
            s=s_act,
            alpha=1.0,
            edgecolors="black",
            linewidths=lw_scatter,
            depthshade=True,
        )

    ax.set_title(title, fontsize=title_fs, fontweight="bold", pad=title_pad)

    ticks = np.arange(N)
    ax.set_xticks(ticks)
    ax.set_yticks(ticks)
    ax.set_zticks(ticks)
    ax.set_xticklabels([])
    ax.set_yticklabels([])
    ax.set_zticklabels([])

    lim_lo, lim_hi = -0.4, m + 0.4
    ax.set_xlim(lim_lo, lim_hi)
    ax.set_ylim(lim_lo, lim_hi)
    ax.set_zlim(lim_lo, lim_hi)
    ax.view_init(elev=VIEW_ELEV, azim=VIEW_AZIM)
    ax.set_box_aspect((1, 1, 1))

    ax.set_xlabel("Verb", fontsize=label_fs, fontweight="bold", labelpad=labelpad_xy)
    ax.set_ylabel("Color", fontsize=label_fs, fontweight="bold", labelpad=labelpad_xy)
    ax.set_zlabel("Object", fontsize=label_fs, fontweight="bold", labelpad=labelpad_z)


def build_masks():
    IX, IY, IZ = lattice_coords()

    complete = np.ones(IX.shape[0], dtype=bool)

    rng = np.random.default_rng(0)
    n_random = max(1, int(0.09 * N**3))
    choice = rng.choice(N**3, size=n_random, replace=False)
    random_m = np.zeros(N**3, dtype=bool)
    random_m[choice] = True

    diagonal = (IX == IY) & (IY == IZ)

    L_mask = ((IY == 0) & (IZ == 0)) | ((IX == 0) & (IZ == 0)) | ((IX == 0) & (IY == 0))

    x_arm = (IY == 0) & (IZ == 0)
    y_arm = (IX == 0) & (IZ == 0) & (IY < OURS_Y_LEN)
    diag_arm = (IX == IY) & (IY == IZ)
    ours = x_arm | y_arm | diag_arm

    return {
        "complete": complete,
        "random": random_m,
        "diagonal": diagonal,
        "L": L_mask,
        "ours": ours,
    }


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

    specs = [
        ("complete", "Complete", False),
        ("random", "Random", True),
        ("diagonal", "Diagonal", True),
        ("L", "L", True),
        ("ours", "Ours", True),
    ]

    for key, title, show_inactive in specs:
        fig = plt.figure(figsize=(7.2, 6.4))
        ax = fig.add_subplot(111, projection="3d")
        plot_strategy(ax, masks[key], title, show_inactive=show_inactive)
        out = FIG_DIR / f"data_collection_strategies_3d_{key}.png"
        fig.savefig(
            out, dpi=200, bbox_inches="tight", pad_inches=0.03, facecolor="white"
        )
        plt.close(fig)
        print(f"Saved {out}")

    # Single row of five 3D panels (same layout idea as data_collection_strategies_1x5.png)
    fig_row = plt.figure(figsize=(19.0, 3.95))
    for i, (key, title, show_inactive) in enumerate(specs):
        ax = fig_row.add_subplot(1, 5, i + 1, projection="3d")
        plot_strategy(
            ax, masks[key], title, show_inactive=show_inactive, compact=True
        )
    fig_row.subplots_adjust(left=0.02, right=0.98, top=0.88, bottom=0.06, wspace=-0.18)
    out_row = FIG_DIR / "data_collection_strategies_3d_1x5.png"
    fig_row.savefig(
        out_row, dpi=200, bbox_inches="tight", pad_inches=0.06, facecolor="white"
    )
    plt.close(fig_row)
    print(f"Saved {out_row}")


if __name__ == "__main__":
    main()