pqAutomationApp/app/plots/gamut_background.py

"""CIE 色度图底图渲染与缓存。

将"重型图像渲染"（colour-science 的谱迹颜色填充）与"轻量框架数据层"
（参考/实测三角形、标签、覆盖率）解耦。

底图：
- 仅在首次调用或缓存失效时通过 colour-science 渲染一次；
- 渲染结果保存为 numpy RGBA 数组，同时落盘到 settings/cache/，
  下次启动直接 imread 加载，避免重新跑色彩科学计算。

调用方在每次绘图时只需 `ax.imshow(bg, extent=bbox)`，再叠加自己的矢量层。
"""

from __future__ import annotations

import hashlib
import os
import threading
from typing import Tuple

import numpy as np

# 谱迹底图分辨率（边长，单位像素）。1024 对于 14 inch 画布足够细腻，
# 文件大小 ~1-2MB，单次渲染 ~0.5-1 s，缓存后毫秒级加载。
_DIAGRAM_RES = 1024

# 缓存版本号：当渲染参数或风格调整时递增，强制重新生成。
_CACHE_VERSION = "v1"

_BBox = Tuple[float, float, float, float]  # (xmin, xmax, ymin, ymax)

_CIE1931_BBOX: _BBox = (0.0, 0.8, 0.0, 0.9)
_CIE1976_BBOX: _BBox = (0.0, 0.65, 0.0, 0.6)


_memory_cache: dict[str, np.ndarray] = {}
_lock = threading.Lock()


def _cache_dir() -> str:
    # 项目根目录通过本文件位置反推：app/plots/ -> 项目根
    here = os.path.dirname(os.path.abspath(__file__))
    root = os.path.abspath(os.path.join(here, "..", ".."))
    d = os.path.join(root, "settings", "cache")
    os.makedirs(d, exist_ok=True)
    return d


def _cache_key(kind: str, bbox: _BBox) -> str:
    sig = f"{kind}|{bbox}|{_DIAGRAM_RES}|{_CACHE_VERSION}"
    h = hashlib.md5(sig.encode("utf-8")).hexdigest()[:10]
    return f"chromaticity_{kind}_{h}.npy"


def _cache_path(kind: str, bbox: _BBox) -> str:
    return os.path.join(_cache_dir(), _cache_key(kind, bbox))


def _render_chromaticity(kind: str, bbox: _BBox) -> np.ndarray:
    """通过 colour-science 离屏渲染谱迹底图，返回 RGBA float 数组。"""
    # 延迟导入：仅在缓存未命中时支付 colour.plotting 的加载开销。
    import matplotlib
    prev_backend = matplotlib.get_backend()
    try:
        matplotlib.use("Agg", force=True)
    except Exception:
        pass

    import matplotlib.pyplot as plt
    from colour.plotting import (
        plot_chromaticity_diagram_CIE1931,
        plot_chromaticity_diagram_CIE1976UCS,
    )

    xmin, xmax, ymin, ymax = bbox
    aspect = (xmax - xmin) / (ymax - ymin)
    height = _DIAGRAM_RES
    width = int(round(height * aspect))

    fig = plt.figure(figsize=(width / 100.0, height / 100.0), dpi=100)
    ax = fig.add_axes([0.0, 0.0, 1.0, 1.0])

    if kind == "cie1931":
        plot_chromaticity_diagram_CIE1931(
            axes=ax, show=False, title=False,
            tight_layout=False, transparent_background=True,
            bounding_box=bbox,
        )
    elif kind == "cie1976":
        plot_chromaticity_diagram_CIE1976UCS(
            axes=ax, show=False, title=False,
            tight_layout=False, transparent_background=True,
            bounding_box=bbox,
        )
    else:
        plt.close(fig)
        raise ValueError(f"unknown diagram kind: {kind!r}")

    ax.set_xlim(xmin, xmax)
    ax.set_ylim(ymin, ymax)
    ax.set_axis_off()
    ax.set_position([0.0, 0.0, 1.0, 1.0])

    fig.canvas.draw()
    # 从 canvas 抓取 RGBA 数组
    buf = np.asarray(fig.canvas.buffer_rgba()).copy()
    plt.close(fig)

    try:
        matplotlib.use(prev_backend, force=True)
    except Exception:
        pass

    return buf


def _load_or_render(kind: str, bbox: _BBox) -> np.ndarray:
    key = _cache_key(kind, bbox)
    with _lock:
        if key in _memory_cache:
            return _memory_cache[key]

        disk = _cache_path(kind, bbox)
        if os.path.isfile(disk):
            try:
                arr = np.load(disk)
                _memory_cache[key] = arr
                return arr
            except Exception:
                # 缓存损坏则重新渲染
                try:
                    os.remove(disk)
                except OSError:
                    pass

        arr = _render_chromaticity(kind, bbox)
        _memory_cache[key] = arr
        try:
            np.save(disk, arr)
        except Exception:
            pass
        return arr


def get_cie1931_background() -> Tuple[np.ndarray, _BBox]:
    """返回 (RGBA 数组, bbox)，可直接 ax.imshow(arr, extent=[*bbox])。"""
    return _load_or_render("cie1931", _CIE1931_BBOX), _CIE1931_BBOX


def get_cie1976_background() -> Tuple[np.ndarray, _BBox]:
    return _load_or_render("cie1976", _CIE1976_BBOX), _CIE1976_BBOX


def clear_cache(*, disk: bool = False) -> None:
    """清空内存缓存（可选连同磁盘）。供调试/样式调整时使用。"""
    with _lock:
        _memory_cache.clear()
        if disk:
            d = _cache_dir()
            for name in os.listdir(d):
                if name.startswith("chromaticity_") and name.endswith(".npy"):
                    try:
                        os.remove(os.path.join(d, name))
                    except OSError:
                        pass