"""MFD palette LUT helpers for VideoReceiverSFP.

Provides a small, self-contained implementation of the MFD LUT
generation and application compatible with the ControlPanel defaults.

`build_mfd_lut()` returns an (256,3) uint8 array in RGB order suitable
for `Image.fromarray(..., mode='RGB')`.
"""
from typing import Optional
import numpy as np


def build_mfd_lut(raw_map_intensity: int = 255) -> np.ndarray:
    """Build a 256x3 RGB LUT for MFD indices using the same defaults
    ControlPanel uses, but return values in RGB order suitable for PIL.

    The ControlPanel `update_mfd_lut` builds an array in BGR order from
    category definitions. We replicate the category/pixel mapping and
    intensity calculation, then convert the resulting BGR -> RGB so the
    LUT can be used directly with PIL `Image.fromarray(..., mode='RGB')`.
    """
    # Build in BGR order following ControlPanel logic
    bgr_lut = np.zeros((256, 3), dtype=np.uint8)

    # Categories as defined in controlpanel.app_state._initialize_mfd_params
    categories = {
        "Occlusion": {"color": (0, 0, 0), "intensity": 255, "pixels": [0, 1]},
        "Cat A": {"color": (255, 255, 255), "intensity": 255, "pixels": [2]},
        "Cat B": {"color": (255, 255, 255), "intensity": 255, "pixels": [3, 18]},
        "Cat C": {"color": (255, 255, 255), "intensity": 255, "pixels": [4, 5, 6, 16]},
        "Cat C1": {"color": (255, 255, 255), "intensity": 255, "pixels": [7, 8, 9]},
        "Cat C2": {"color": (255, 255, 255), "intensity": 255, "pixels": [10, 11, 12]},
        "Cat C3": {"color": (255, 255, 255), "intensity": 255, "pixels": [13, 14, 15]},
        "Reserved": {"color": (0, 0, 0), "intensity": 255, "pixels": list(range(17, 32))},
    }

    pixel_to_category = {p: cat for cat, data in categories.items() for p in data["pixels"]}

    raw_map_factor = float(raw_map_intensity) / 255.0

    for idx in range(256):
        cat_name = pixel_to_category.get(idx)
        if cat_name:
            cat_data = categories[cat_name]
            bgr = cat_data["color"]
            intensity_factor = cat_data["intensity"] / 255.0
            b = int(np.clip(int(round(float(bgr[0]) * intensity_factor)), 0, 255))
            g = int(np.clip(int(round(float(bgr[1]) * intensity_factor)), 0, 255))
            r = int(np.clip(int(round(float(bgr[2]) * intensity_factor)), 0, 255))
            bgr_lut[idx, 0] = b
            bgr_lut[idx, 1] = g
            bgr_lut[idx, 2] = r
        elif 32 <= idx <= 255:
            raw_intensity = (float(idx) - 32.0) * (255.0 / 223.0)
            final_gray = int(round(np.clip(raw_intensity * raw_map_factor, 0, 255)))
            bgr_lut[idx, :] = (final_gray, final_gray, final_gray)

    # Convert BGR -> RGB for PIL consumers
    rgb_lut = bgr_lut[:, ::-1].copy()
    return rgb_lut


def apply_mfd_lut(indices: np.ndarray, lut: Optional[np.ndarray]) -> Optional[np.ndarray]:
    """Map an indices array (H,W) to an RGB image (H,W,3) using the LUT.

    Returns a uint8 NumPy array or None on error.
    """
    if lut is None:
        return None
    try:
        # Use NumPy advanced indexing to map (H,W) -> (H,W,3)
        rgb = lut[indices]
        return rgb
    except Exception:
        return None