"""Shared CoCo 3 GIME encoder helpers: global palette selection + linear packing.

GIME native graphics modes have no per-cell colour limit: a flat palette of pens
(16/4/2), each pen any of the 64 colours.  Pick the best N-colour sub-palette,
dither to it, then pack pens into the LINEAR 80-byte x 192-row screen.
"""
from __future__ import annotations

import numpy as np

from ... import dither, palette as c64pal
from .. import palette as g

BYTES_PER_ROW, ROWS = 80, 192


def choose_inks(img_lab, plab, n):
    """Greedy forward selection of the ``n`` palette colours (0-63) minimising
    nearest-colour error over the image."""
    flat = img_lab.reshape(-1, 3)
    d = np.sum((flat[:, None, :] - plab[None, :, :]) ** 2, axis=-1)
    chosen, best = [], np.full(flat.shape[0], np.inf)
    for _ in range(n):
        cand = np.minimum(best[:, None], d).sum(0)
        for c in chosen:
            cand[c] = np.inf
        c = int(cand.argmin())
        chosen.append(c)
        best = np.minimum(best, d[:, c])
    return sorted(chosen)


def render(img_rgb, n, dither_mode, ramp=None):
    """Return (pen (H,W) 0..n-1, inks (n 6-bit colours), idx (H,W) palette
    indices, img_lab, plab, prgb)."""
    plab = g.palette_lab()
    prgb = g.get_palette().astype(np.uint8)
    img_lab = c64pal.srgb_to_lab(img_rgb)

    pal_idx = list(ramp) if ramp is not None else choose_inks(img_lab, plab, n)
    allowed = np.tile(np.array(pal_idx), (*img_lab.shape[:2], 1))
    idx = dither.quantize(img_lab, allowed, plab, dither_mode).astype(np.int64)
    lut = {p: k for k, p in enumerate(pal_idx)}
    pen = np.vectorize(lut.get)(idx).astype(np.uint8)
    return pen, list(pal_idx), idx.astype(np.uint16), img_lab, plab, prgb


def _pack_16(pen):
    """160x192 pens (0-15) -> 15360 bytes.  2 px/byte, high nibble = left."""
    scr = bytearray(BYTES_PER_ROW * ROWS)
    for y in range(ROWS):
        row = pen[y]
        base = y * BYTES_PER_ROW
        for bx in range(BYTES_PER_ROW):
            scr[base + bx] = (int(row[bx * 2]) << 4) | (int(row[bx * 2 + 1]) & 0x0F)
    return bytes(scr)


def _pack_4(pen):
    """320x192 pens (0-3) -> 15360 bytes.  4 px/byte, MSB = left."""
    scr = bytearray(BYTES_PER_ROW * ROWS)
    for y in range(ROWS):
        row = pen[y]
        base = y * BYTES_PER_ROW
        for bx in range(BYTES_PER_ROW):
            scr[base + bx] = ((int(row[bx * 4]) << 6) | (int(row[bx * 4 + 1]) << 4) |
                              (int(row[bx * 4 + 2]) << 2) | int(row[bx * 4 + 3]))
    return bytes(scr)


def _pack_2(pen):
    """640x192 pens (0-1) -> 15360 bytes.  8 px/byte, MSB = left."""
    scr = bytearray(BYTES_PER_ROW * ROWS)
    for y in range(ROWS):
        row = pen[y]
        base = y * BYTES_PER_ROW
        for bx in range(BYTES_PER_ROW):
            byte = 0
            for k in range(8):
                byte |= (int(row[bx * 8 + k]) & 1) << (7 - k)
            scr[base + bx] = byte
    return bytes(scr)


PACK = {0: _pack_2, 1: _pack_4, 2: _pack_16}    # keyed by CRES