"""Monochrome / grayscale mode -- the highest-resolution path.

Renders at hires (320x200) but matches the image by *luminance* to a small ramp
of palette colours, so detail is carried entirely by spatial dithering.  With the
grayscale ramp (black -> dark grey -> grey -> light grey -> white) this gives a
proper greyscale photo; pick any base colour and the ramp becomes that hue's
shades (e.g. black -> blue -> light blue -> white) for a tinted monochrome.

Output is ordinary hires-format data, so it reuses the hires viewer.
"""

from __future__ import annotations

import numpy as np

from .. import dither, palette as pal
from . import base, hires

WIDTH, HEIGHT = 320, 200
CELL_W, CELL_H = 8, 8
PIXEL_ASPECT = 1.0
DATA_LOAD = 0x2000

# Luminance-ordered grey ramp: black, dark grey, grey, light grey, white.
GRAY_RAMP = [0, 11, 12, 15, 1]
# A few palette colours have a lighter sibling, giving a richer tinted ramp.
SIBLINGS = {2: 10, 10: 2, 5: 13, 13: 5, 6: 14, 14: 6, 8: 9, 9: 8}


def build_ramp(base_color, plab):
    """Return palette indices (luminance-sorted) used to render the image."""
    if base_color is None or base_color in (0, 1, 11, 12, 15):
        ramp = list(GRAY_RAMP)
    else:
        ramp = {0, 1, base_color}
        if base_color in SIBLINGS:
            ramp.add(SIBLINGS[base_color])
        ramp = list(ramp)
    ramp.sort(key=lambda i: plab[i, 0])      # by Lab lightness
    return ramp


def convert(img_rgb, palette_name="colodore", dither_mode="floyd",
            intensive=False, base_color=None):
    plab = pal.palette_lab(palette_name)

    # Work purely in luminance: collapse image and palette to (L, 0, 0).
    L_pix = pal.srgb_to_lab(img_rgb)[..., 0]
    img_mono = np.zeros((HEIGHT, WIDTH, 3))
    img_mono[..., 0] = L_pix
    plab_mono = np.zeros((16, 3))
    plab_mono[:, 0] = plab[:, 0]

    ramp = build_ramp(base_color, plab)
    n_free = min(2, len(ramp))

    cells, rows, cols = base.cells_lab(img_mono, CELL_W, CELL_H)
    dist = base.cell_distance(cells, plab_mono)
    sets, _ = base.select_cell_sets(dist, ramp, n_free=n_free)
    if n_free == 1:                          # pad to 2 colours per cell for hires
        sets = np.concatenate([sets, sets], axis=1)

    allowed = base.per_pixel_allowed(sets, rows, cols, CELL_W, CELL_H, HEIGHT, WIDTH)
    index_image = dither.quantize(img_mono, allowed, plab_mono, dither_mode).astype(np.uint8)

    bitmap, screen = hires._encode(index_image, sets, rows, cols)
    payload = bytes(bitmap) + bytes(screen)

    conv = base.Conversion(
        mode="mono", width=WIDTH, height=HEIGHT, pixel_aspect=PIXEL_ASPECT,
        index_image=index_image, data=payload, data_addr=DATA_LOAD, viewer="hires",
        error=base.mean_error(index_image, img_mono, plab_mono),
        meta={"palette": palette_name, "dither": dither_mode,
              "base_color": base_color, "ramp": ramp},
    )
    conv.extra_files = [("picture.art", base.prg(0x2000, payload + b"\x00"))]
    return conv