"""Shared helpers for the BBC Micro converters."""

from __future__ import annotations

import itertools

import numpy as np

from ... import dither
from ...convert.base import Conversion, perceptual_error, _box_blur
from ...palette import srgb_to_lab
from .. import palette as bpal


def _choose_physicals(img_lab, n, dither_mode):
    """Pick the n physical colours (of 8) that best reproduce the image, then
    dither once with that palette.  Candidates are ranked by a fast vectorised
    proxy -- the perceptual error of the nearest-colour (un-dithered)
    reconstruction -- so we do NOT run the slow Floyd dither for all C(8,n)
    combinations (doing so made the 4-colour modes hang for a minute+); Floyd runs
    only on the winning palette.  Returns (physical_indices, logical_idx, error)."""
    plab_all = bpal.phys_lab()
    H, W, _ = img_lab.shape
    target_blur = _box_blur(img_lab)
    best_combo, best_score = None, np.inf
    for combo in itertools.combinations(range(8), n):
        sub = plab_all[list(combo)]
        nidx = ((img_lab[:, :, None, :] - sub[None, None]) ** 2).sum(-1).argmin(-1)
        diff = _box_blur(sub[nidx]) - target_blur
        score = float(np.sqrt((diff ** 2).sum(-1)).mean())
        if score < best_score:
            best_score, best_combo = score, list(combo)
    sub = plab_all[best_combo]
    allowed = np.tile(np.arange(n), (H, W, 1))
    idx = dither.quantize(img_lab, allowed, sub, dither_mode).astype(np.uint8)
    return best_combo, idx, perceptual_error(idx, img_lab, sub)


def build(img_rgb, *, mode, bbc_mode, ncol, bpp, width, height, base,
          dither_mode, mono=False):
    if mono:
        L = srgb_to_lab(img_rgb)[..., 0]
        img_lab = np.zeros((height, width, 3))
        img_lab[..., 0] = L
        plab = np.zeros((2, 3)); plab[:, 0] = bpal.mono_lab()[:, 0]
        allowed = np.tile(np.array([0, 1]), (height, width, 1))
        idx = dither.quantize(img_lab, allowed, plab, dither_mode).astype(np.uint8)
        physicals = [0, 7]                     # black, white
        err = perceptual_error(idx, img_lab, plab)
        prgb = bpal.PHYS[[0, 7]].astype(np.uint8)
    else:
        img_lab = srgb_to_lab(img_rgb)
        if ncol >= 8:
            physicals = list(range(8))
            plab = bpal.phys_lab()
            allowed = np.tile(np.arange(8), (height, width, 1))
            idx = dither.quantize(img_lab, allowed, plab, dither_mode).astype(np.uint8)
            err = perceptual_error(idx, img_lab, plab)
        else:
            physicals, idx, err = _choose_physicals(img_lab, ncol, dither_mode)
        prgb = bpal.PHYS[physicals].astype(np.uint8)

    data = bpal.pack(idx, width, bpp)
    preview = prgb[idx]
    return Conversion(
        mode=mode, width=width, height=height,
        pixel_aspect=(4 / 3) / (width / height),
        index_image=idx.astype(np.uint16), data=data, data_addr=base,
        viewer="bbc", preview_rgb=preview, error=err,
        meta={"palette": "bbc", "dither": dither_mode, "bbc_mode": bbc_mode,
              "ncol": ncol, "physicals": physicals, "base": base},
    )