pqAutomationApp/app/tests/local_dimming.py

"""Local Dimming 测试逻辑（应用层）。

整合自原 drivers/local_dimming_test.py：窗口图片生成与测试主循环
直接落在本模块，UCD 通用操作通过 SignalService 完成。
"""

import atexit
import csv
import datetime
import os
import shutil
import sys
import threading
import time
import tkinter as tk
from tkinter import filedialog, messagebox

import numpy as np
from PIL import Image

from typing import TYPE_CHECKING

if TYPE_CHECKING:
    from pqAutomationApp import PQAutomationApp




# --------------------------------------------------------------------------
# 模块级常量与窗口图片缓存
# --------------------------------------------------------------------------

DEFAULT_WINDOW_PERCENTAGES = [1, 2, 5, 10, 18, 25, 50, 75, 100]
DEFAULT_CHESSBOARD_GRID = 5
INSTANT_PEAK_WINDOW_PERCENTAGE = 10
INSTANT_PEAK_CAPTURE_DELAY = 0.5

_TEMP_DIR = None
_IMAGE_CACHE = {}  # {(width, height, percentage): file_path}


def _cleanup_temp_dir():
    global _TEMP_DIR
    if _TEMP_DIR and os.path.exists(_TEMP_DIR):
        try:
            shutil.rmtree(_TEMP_DIR)
        except Exception:
            pass
    _TEMP_DIR = None
    _IMAGE_CACHE.clear()


def _get_temp_dir():
    global _TEMP_DIR
    if _TEMP_DIR is None:
        if getattr(sys, "frozen", False):
            base = os.path.dirname(sys.executable)
        else:
            base = os.getcwd()
        _TEMP_DIR = os.path.join(base, "temp_local_dimming")
        os.makedirs(_TEMP_DIR, exist_ok=True)
        atexit.register(_cleanup_temp_dir)
    return _TEMP_DIR


def _make_window_image_array(width, height, percentage):
    """生成黑底+居中白窗的 numpy 图像，保持屏幕比例。"""
    image = np.zeros((height, width, 3), dtype=np.uint8)
    if percentage >= 100:
        ww, wh = width, height
    else:
        scale = (percentage / 100.0) ** 0.5
        ww = int(width * scale)
        wh = int(height * scale)
    x1 = (width - ww) // 2
    y1 = (height - wh) // 2
    image[y1:y1 + wh, x1:x1 + ww] = 255
    return image


def _ensure_window_image(width, height, percentage):
    """生成或复用缓存的窗口 PNG 文件，返回路径。"""
    key = (width, height, percentage)
    cached = _IMAGE_CACHE.get(key)
    if cached and os.path.exists(cached):
        return cached
    arr = _make_window_image_array(width, height, percentage)
    fname = f"window_{width}x{height}_{percentage:03d}percent.png"
    path = os.path.join(_get_temp_dir(), fname)
    Image.fromarray(arr, mode="RGB").save(path, format="PNG")
    _IMAGE_CACHE[key] = path
    return path


def _ensure_solid_image(width, height, rgb, name):
    """生成或复用纯色 PNG 文件，返回路径。"""
    rgb = tuple(int(v) for v in rgb)
    key = ("solid", width, height, rgb)
    cached = _IMAGE_CACHE.get(key)
    if cached and os.path.exists(cached):
        return cached

    arr = np.zeros((height, width, 3), dtype=np.uint8)
    arr[:, :] = rgb
    path = os.path.join(_get_temp_dir(), f"{name}_{width}x{height}.png")
    Image.fromarray(arr, mode="RGB").save(path, format="PNG")
    _IMAGE_CACHE[key] = path
    return path


def _make_checkerboard_image_array(width, height, grid_size, center_white):
    """生成棋盘格图像，保证中心块可切换黑/白。"""
    image = np.zeros((height, width, 3), dtype=np.uint8)
    y_edges = np.linspace(0, height, grid_size + 1, dtype=int)
    x_edges = np.linspace(0, width, grid_size + 1, dtype=int)
    center_index = grid_size // 2

    for row in range(grid_size):
        for col in range(grid_size):
            block_is_white = (row + col) % 2 == 0
            if not center_white:
                block_is_white = not block_is_white
            value = 255 if block_is_white else 0
            image[
                y_edges[row]:y_edges[row + 1],
                x_edges[col]:x_edges[col + 1],
            ] = value

    center_value = 255 if center_white else 0
    image[
        y_edges[center_index]:y_edges[center_index + 1],
        x_edges[center_index]:x_edges[center_index + 1],
    ] = center_value
    return image


def _ensure_checkerboard_image(width, height, grid_size, center_white):
    """生成或复用棋盘格 PNG 文件，返回路径。"""
    key = ("checkerboard", width, height, grid_size, center_white)
    cached = _IMAGE_CACHE.get(key)
    if cached and os.path.exists(cached):
        return cached

    arr = _make_checkerboard_image_array(width, height, grid_size, center_white)
    center_name = "white_center" if center_white else "black_center"
    path = os.path.join(
        _get_temp_dir(),
        f"checkerboard_{grid_size}x{grid_size}_{center_name}_{width}x{height}.png",
    )
    Image.fromarray(arr, mode="RGB").save(path, format="PNG")
    _IMAGE_CACHE[key] = path
    return path


def _build_ld_result_row(test_item, pattern_label, value, x="--", y="--"):
    timestamp = datetime.datetime.now().strftime("%H:%M:%S")
    if isinstance(value, (int, float, np.floating)):
        display_value = f"{float(value):.4f}"
    else:
        display_value = str(value)

    return {
        "test_item": test_item,
        "pattern": pattern_label,
        "value": display_value,
        "x": x if isinstance(x, str) else f"{x:.4f}",
        "y": y if isinstance(y, str) else f"{y:.4f}",
        "time": timestamp,
    }


def _measure_ld_row(self: "PQAutomationApp", test_item, pattern_label):
    """读取一次 CA410 数据并包装为表格行。"""
    x, y, lv, _X, _Y, _Z = self.ca.readAllDisplay()
    if lv is None:
        raise RuntimeError(f"{pattern_label} 采集失败")
    return _build_ld_result_row(test_item, pattern_label, lv, x, y), lv


def _send_ld_image(self: "PQAutomationApp", image_path):
    self.signal_service.send_image(image_path)


def _apply_ld_ucd_params(self: "PQAutomationApp") -> bool:
    """发送 Local Dimming 图案前，按当前测试类型写入 UCD 参数。"""
    test_type = getattr(self.config, "current_test_type", "screen_module")
    cfg = self.config.current_test_types.get(test_type, {})

    try:
        self.signal_service.apply_config(self.config)

        if test_type == "screen_module":
            ok = self.signal_service.update_signal_format(
                color_space=(
                    self.screen_module_color_space_var.get()
                    if hasattr(self, "screen_module_color_space_var")
                    else cfg.get("colorimetry", "sRGB")
                ),
                data_range=(
                    self.screen_module_data_range_var.get()
                    if hasattr(self, "screen_module_data_range_var")
                    else cfg.get("data_range", "Full")
                ),
                bit_depth=(
                    self.screen_module_bit_depth_var.get()
                    if hasattr(self, "screen_module_bit_depth_var")
                    else f"{int(cfg.get('bpc', 8))}bit"
                ),
                output_format=(
                    self.screen_module_output_format_var.get()
                    if hasattr(self, "screen_module_output_format_var")
                    else cfg.get("color_format", "RGB")
                ),
            )
        elif test_type == "sdr_movie":
            ok = self.signal_service.update_signal_format(
                color_space=(
                    self.sdr_color_space_var.get()
                    if hasattr(self, "sdr_color_space_var")
                    else cfg.get("colorimetry", "sRGB")
                ),
                data_range=(
                    self.sdr_data_range_var.get()
                    if hasattr(self, "sdr_data_range_var")
                    else cfg.get("data_range", "Full")
                ),
                bit_depth=(
                    self.sdr_bit_depth_var.get()
                    if hasattr(self, "sdr_bit_depth_var")
                    else f"{int(cfg.get('bpc', 8))}bit"
                ),
                output_format=(
                    self.sdr_output_format_var.get()
                    if hasattr(self, "sdr_output_format_var")
                    else cfg.get("color_format", "RGB")
                ),
            )
        elif test_type == "hdr_movie":
            ok = self.signal_service.update_signal_format(
                color_space=(
                    self.hdr_color_space_var.get()
                    if hasattr(self, "hdr_color_space_var")
                    else cfg.get("colorimetry", "sRGB")
                ),
                data_range=(
                    self.hdr_data_range_var.get()
                    if hasattr(self, "hdr_data_range_var")
                    else cfg.get("data_range", "Full")
                ),
                bit_depth=(
                    self.hdr_bit_depth_var.get()
                    if hasattr(self, "hdr_bit_depth_var")
                    else f"{int(cfg.get('bpc', 8))}bit"
                ),
                output_format=(
                    self.hdr_output_format_var.get()
                    if hasattr(self, "hdr_output_format_var")
                    else cfg.get("color_format", "RGB")
                ),
                max_cll=(
                    self.hdr_maxcll_var.get()
                    if hasattr(self, "hdr_maxcll_var")
                    else None
                ),
                max_fall=(
                    self.hdr_maxfall_var.get()
                    if hasattr(self, "hdr_maxfall_var")
                    else None
                ),
            )
        elif test_type == "local_dimming":
            ok = self.signal_service.update_signal_format(
                color_space=(
                    self.local_dimming_color_space_var.get()
                    if hasattr(self, "local_dimming_color_space_var")
                    else cfg.get("colorimetry", "sRGB")
                ),
                data_range=(
                    self.local_dimming_data_range_var.get()
                    if hasattr(self, "local_dimming_data_range_var")
                    else cfg.get("data_range", "Full")
                ),
                bit_depth=(
                    self.local_dimming_bit_depth_var.get()
                    if hasattr(self, "local_dimming_bit_depth_var")
                    else f"{int(cfg.get('bpc', 8))}bit"
                ),
                output_format=(
                    self.local_dimming_output_format_var.get()
                    if hasattr(self, "local_dimming_output_format_var")
                    else cfg.get("color_format", "RGB")
                ),
            )
        else:
            self._dispatch_ui(
                self.log_gui.log,
                f"Local Dimming 不支持的测试类型: {test_type}",
                "error",
            )
            return False

        if not ok:
            self._dispatch_ui(self.log_gui.log, "Local Dimming UCD 参数设置失败", "error")
            return False

        return True
    except Exception as e:
        self._dispatch_ui(self.log_gui.log, f"Local Dimming UCD 参数设置异常: {e}", "error")
        return False


def _run_ld_measurement_step(self: "PQAutomationApp", width, height, wait_time, step, log):
    label = step["label"]
    test_item = step["test_item"]
    kind = step["kind"]

    if kind == "window":
        percentage = step["percentage"]
        image_path = _ensure_window_image(width, height, percentage)
        _send_ld_image(self, image_path)
        settle_time = wait_time
    elif kind == "black":
        image_path = _ensure_solid_image(width, height, (0, 0, 0), "black")
        _send_ld_image(self, image_path)
        settle_time = wait_time
    elif kind == "checkerboard":
        image_path = _ensure_checkerboard_image(
            width,
            height,
            DEFAULT_CHESSBOARD_GRID,
            step["center_white"],
        )
        _send_ld_image(self, image_path)
        settle_time = wait_time
    elif kind == "instant_peak":
        black_image = _ensure_solid_image(width, height, (0, 0, 0), "black")
        peak_image = _ensure_window_image(
            width,
            height,
            step["percentage"],
        )
        _send_ld_image(self, black_image)
        log(f"  黑场预置 {wait_time:.1f} 秒", level="info")
        time.sleep(wait_time)
        _send_ld_image(self, peak_image)
        settle_time = min(wait_time, INSTANT_PEAK_CAPTURE_DELAY)
    else:
        raise ValueError(f"未知 Local Dimming 测试步骤: {kind}")

    log(f"  等待 {settle_time:.1f} 秒后采集...", level="info")
    time.sleep(settle_time)
    return _measure_ld_row(self, test_item, label)


def _set_current_ld_pattern(self: "PQAutomationApp", test_item, pattern_label, percentage=None):
    self.current_ld_test_item = test_item
    self.current_ld_pattern_label = pattern_label
    self.current_ld_percentage = percentage


# --------------------------------------------------------------------------
# GUI 入口（绑定为 PQAutomationApp 方法）
# --------------------------------------------------------------------------

def start_local_dimming_test(self: "PQAutomationApp"):
    """Local Dimming 不再提供自动测试，保留接口仅提示用户使用手动模式。"""
    messagebox.showinfo("提示", "Local Dimming 请使用手动发送图案后再采集亮度")


def update_ld_results(self: "PQAutomationApp", results):
    """把批量测试结果填入 Treeview。"""
    for row in results:
        self.ld_tree.insert(
            "", tk.END,
            values=(
                row["test_item"],
                row["pattern"],
                row["value"],
                row["x"],
                row["y"],
                row["time"],
            ),
        )


def stop_local_dimming_test(self: "PQAutomationApp"):
    """兼容旧接口，无操作。"""
    return


def send_ld_window(self: "PQAutomationApp", percentage):
    """发送指定百分比的白色窗口（手动模式）。"""
    if not self.signal_service.is_connected:
        messagebox.showwarning("警告", "请先连接 UCD323 设备")
        return

    self.log_gui.log(f"🔆 发送 {percentage}% 窗口...", level="info")
    _set_current_ld_pattern(self, "峰值亮度", f"{percentage}%窗口", percentage)

    def send():
        if not _apply_ld_ucd_params(self):
            return
        width, height = self.signal_service.current_resolution()
        try:
            image_path = _ensure_window_image(width, height, percentage)
        except Exception as e:
            self._dispatch_ui(self.log_gui.log, f"图像生成失败: {e}")
            return
        try:
            self.signal_service.send_image(image_path)
            ok = True
        except Exception:
            ok = False
        msg = (
            f"{percentage}% 窗口已发送" if ok
            else f"{percentage}% 窗口发送失败"
        )
        self._dispatch_ui(self.log_gui.log, msg)

    threading.Thread(target=send, daemon=True).start()


def send_ld_checkerboard(self: "PQAutomationApp", center_white):
    """发送棋盘格图案（手动模式）。"""
    if not self.signal_service.is_connected:
        messagebox.showwarning("警告", "请先连接 UCD323 设备")
        return

    pattern_label = "棋盘格(中心白)" if center_white else "棋盘格(中心黑)"
    self.log_gui.log(f"🔲 发送 {pattern_label}...", level="info")
    _set_current_ld_pattern(self, "棋盘格对比度", pattern_label)

    def send():
        if not _apply_ld_ucd_params(self):
            return
        width, height = self.signal_service.current_resolution()
        try:
            image_path = _ensure_checkerboard_image(
                width,
                height,
                DEFAULT_CHESSBOARD_GRID,
                center_white,
            )
        except Exception as e:
            self._dispatch_ui(self.log_gui.log, f"图像生成失败: {e}")
            return

        try:
            self.signal_service.send_image(image_path)
            ok = True
        except Exception:
            ok = False

        msg = f"{pattern_label} 已发送" if ok else f"{pattern_label} 发送失败"
        self._dispatch_ui(self.log_gui.log, msg)

    threading.Thread(target=send, daemon=True).start()


def send_ld_black_pattern(self: "PQAutomationApp"):
    """发送全黑图案（手动模式）。"""
    if not self.signal_service.is_connected:
        messagebox.showwarning("警告", "请先连接 UCD323 设备")
        return

    self.log_gui.log("⚫ 发送全黑画面...", level="info")
    _set_current_ld_pattern(self, "黑电平", "全黑画面")

    def send():
        if not _apply_ld_ucd_params(self):
            return
        width, height = self.signal_service.current_resolution()
        try:
            image_path = _ensure_solid_image(width, height, (0, 0, 0), "black")
        except Exception as e:
            self._dispatch_ui(self.log_gui.log, f"图像生成失败: {e}")
            return

        try:
            self.signal_service.send_image(image_path)
            ok = True
        except Exception:
            ok = False

        msg = "全黑画面已发送" if ok else "全黑画面发送失败"
        self._dispatch_ui(self.log_gui.log, msg)

    threading.Thread(target=send, daemon=True).start()


def send_ld_instant_peak(self: "PQAutomationApp"):
    """发送瞬时峰值亮度图案：先黑场，再切到 10% 窗口并保持。"""
    if not self.signal_service.is_connected:
        messagebox.showwarning("警告", "请先连接 UCD323 设备")
        return

    pattern_label = f"黑场后切 {INSTANT_PEAK_WINDOW_PERCENTAGE}%窗口"
    self.log_gui.log(f"⚡ 发送瞬时峰值图案: {pattern_label}", level="info")
    _set_current_ld_pattern(
        self,
        "瞬时峰值亮度",
        pattern_label,
        INSTANT_PEAK_WINDOW_PERCENTAGE,
    )

    def send():
        if not _apply_ld_ucd_params(self):
            return
        width, height = self.signal_service.current_resolution()
        try:
            black_image = _ensure_solid_image(width, height, (0, 0, 0), "black")
            peak_image = _ensure_window_image(
                width,
                height,
                INSTANT_PEAK_WINDOW_PERCENTAGE,
            )
        except Exception as e:
            self._dispatch_ui(self.log_gui.log, f"图像生成失败: {e}")
            return

        try:
            self.signal_service.send_image(black_image)
            time.sleep(INSTANT_PEAK_CAPTURE_DELAY)
            self.signal_service.send_image(peak_image)
            ok = True
        except Exception:
            ok = False

        msg = (
            f"瞬时峰值图案已发送，当前保持 {INSTANT_PEAK_WINDOW_PERCENTAGE}%窗口"
            if ok else
            "瞬时峰值图案发送失败"
        )
        self._dispatch_ui(self.log_gui.log, msg)

    threading.Thread(target=send, daemon=True).start()


def measure_ld_luminance(self: "PQAutomationApp"):
    """测量当前显示的亮度并追加一行到 Treeview。"""
    if not self.ca:
        messagebox.showwarning("警告", "请先连接 CA410 色度计")
        return
    if getattr(self, "current_ld_pattern_label", None) is None:
        messagebox.showinfo("提示", "请先发送一个窗口图案")
        return

    self.log_gui.log("📏 正在采集亮度...", level="info")

    def measure():
        try:
            x, y, lv, _X, _Y, _Z = self.ca.readAllDisplay()
        except Exception as e:
            self._dispatch_ui(self.log_gui.log, f"采集异常: {str(e)}")
            return
        if lv is None:
            self._dispatch_ui(self.log_gui.log, "采集失败")
            return
        timestamp = datetime.datetime.now().strftime("%H:%M:%S")
        self._dispatch_ui(
            self.ld_result_label.config,
            text=f"亮度: {lv:.2f} cd/m² | x: {x:.4f} | y: {y:.4f}",
        )
        self._dispatch_ui(
            self.ld_tree.insert, "", tk.END,
            values=(
                getattr(self, "current_ld_test_item", "手动采集"),
                self.current_ld_pattern_label,
                f"{lv:.4f}",
                f"{x:.4f}",
                f"{y:.4f}",
                timestamp,
            ),
        )
        self._dispatch_ui(self.log_gui.log, f"采集完成: {lv:.2f} cd/m²")

    threading.Thread(target=measure, daemon=True).start()


def clear_ld_records(self: "PQAutomationApp"):
    """清空 Treeview 中的测试记录。"""
    for item in self.ld_tree.get_children():
        self.ld_tree.delete(item)
    self.ld_result_label.config(text="亮度: -- cd/m² | x: -- | y: --")
    self.current_ld_percentage = None
    self.current_ld_test_item = None
    self.current_ld_pattern_label = None
    self.log_gui.log("测试记录已清空", level="info")


def save_local_dimming_results(self: "PQAutomationApp"):
    """把 Treeview 中的全部记录导出为 CSV。"""
    if len(self.ld_tree.get_children()) == 0:
        messagebox.showinfo("提示", "没有可保存的数据")
        return

    default_name = (
        f"LocalDimming_{datetime.datetime.now().strftime('%Y%m%d_%H%M%S')}.csv"
    )
    save_path = filedialog.asksaveasfilename(
        title="保存测试结果",
        initialfile=default_name,
        defaultextension=".csv",
        filetypes=[("CSV 文件", "*.csv"), ("所有文件", "*.*")],
    )
    if not save_path:
        return

    try:
        with open(save_path, "w", newline="", encoding="utf-8-sig") as f:
            writer = csv.writer(f)
            writer.writerow(["测试项目", "图案", "亮度/结果", "x", "y", "时间"])
            for item in self.ld_tree.get_children():
                writer.writerow(self.ld_tree.item(item, "values"))
        self.log_gui.log(f"测试结果已保存: {save_path}", level="success")
        messagebox.showinfo("成功", f"测试结果已保存到:\n{save_path}")
    except Exception as e:
        self.log_gui.log(f"保存失败: {str(e)}", level="error")
        messagebox.showerror("错误", f"保存失败: {str(e)}")


class LocalDimmingMixin:
    """由 tools/refactor_to_mixins.py 自动生成。
    把本模块的自由函数挂到 PQAutomationApp 上，便于 F12 跳转与类型推断。
    """
    start_local_dimming_test = start_local_dimming_test
    update_ld_results = update_ld_results
    stop_local_dimming_test = stop_local_dimming_test
    send_ld_window = send_ld_window
    send_ld_checkerboard = send_ld_checkerboard
    send_ld_black_pattern = send_ld_black_pattern
    send_ld_instant_peak = send_ld_instant_peak
    measure_ld_luminance = measure_ld_luminance
    clear_ld_records = clear_ld_records
    save_local_dimming_results = save_local_dimming_results