"""C16 TED hires bitmap mode: 320x200, two colours per 8x8 cell.

Unlike the VIC-II, each of the two per-cell colours may be any of the TED's 128
colours.  The colours are stored across two 1K matrices (see MAME's mos7360
draw_bitmap):
    ch byte   (video base + $400):  high nibble = fg hue,  low nibble = bg hue
    attr byte (video base):         bits 0-2 = fg luminance, bits 4-6 = bg lum
A bitmap bit of 1 selects the foreground colour, 0 the background.
"""
from __future__ import annotations

import numpy as np

from ... import dither, palette as c64pal
from ...convert import base
from .. import palette as tedpal

WIDTH, HEIGHT = 320, 200
CELL_W, CELL_H = 8, 8
PIXEL_ASPECT = 1.0


def convert(img_rgb, palette_name="ted", dither_mode="floyd",
            intensive=False, base_color=None):
    plab = tedpal.palette_lab()
    img_lab = c64pal.srgb_to_lab(img_rgb)

    cells, rows, cols = base.cells_lab(img_lab, CELL_W, CELL_H)
    dist = base.cell_distance(cells, plab)
    sets = _select_pairs(dist)

    allowed = base.per_pixel_allowed(sets, rows, cols, CELL_W, CELL_H, HEIGHT, WIDTH)
    index_image = dither.quantize(img_lab, allowed, plab, dither_mode).astype(np.uint16)

    bitmap, attr, ch = _encode(index_image, sets, rows, cols)
    payload = bytes(bitmap) + bytes(attr) + bytes(ch)

    prgb = tedpal.get_palette().astype(np.uint8)
    return base.Conversion(
        mode="hires", width=WIDTH, height=HEIGHT, pixel_aspect=PIXEL_ASPECT,
        index_image=index_image, data=payload, viewer="c16",
        preview_rgb=prgb[index_image],
        error=base.perceptual_error(index_image, img_lab, plab),
        meta={"palette": "ted", "dither": dither_mode, "border": 0},
    )


def _select_pairs(dist, k=16):
    """Per-cell best 2 colours from the 128-colour TED palette.

    A full C(128,2) search is 8128 combos/cell; instead we restrict each cell to
    its ``k`` best single colours (the optimal pair is almost always among them)
    and search only C(k,2) pairs -- ~100x faster with effectively identical
    results.  dist: (n_cells, P, 128) squared CIELAB distances.
    """
    n_cells = dist.shape[0]
    single = dist.sum(1)                                    # (n_cells, 128)
    cand = np.argsort(single, axis=1)[:, :k]               # (n_cells, k)
    dist_c = np.take_along_axis(dist, cand[:, None, :], axis=2)   # (n_cells, P, k)
    best_err = np.full(n_cells, np.inf)
    best = np.zeros((n_cells, 2), dtype=np.int64)
    for i in range(k):
        for j in range(i + 1, k):
            m = np.minimum(dist_c[:, :, i], dist_c[:, :, j]).sum(1)
            upd = m < best_err
            best_err[upd] = m[upd]
            best[upd, 0] = cand[upd, i]
            best[upd, 1] = cand[upd, j]
    return best


def _encode(index_image, sets, rows, cols):
    """Return (bitmap 8000, attr 1000, ch 1000) for the TED hires layout."""
    bitmap = np.zeros(8000, dtype=np.uint8)
    attr = np.zeros(1000, dtype=np.uint8)
    ch = np.zeros(1000, dtype=np.uint8)
    for cr in range(rows):
        for cc in range(cols):
            ci = cr * cols + cc
            bg, fg = int(sets[ci, 0]), int(sets[ci, 1])
            fg_hue, fg_lum = fg & 0x0F, (fg >> 4) & 0x07
            bg_hue, bg_lum = bg & 0x0F, (bg >> 4) & 0x07
            ch[ci] = (fg_hue << 4) | bg_hue
            attr[ci] = (bg_lum << 4) | fg_lum
            base_addr = cr * 320 + cc * 8
            block = index_image[cr * 8:cr * 8 + 8, cc * 8:cc * 8 + 8]
            for r in range(8):
                row = block[r]
                byte = 0
                for x in range(8):
                    byte = (byte << 1) | (1 if int(row[x]) == fg else 0)
                bitmap[base_addr + r] = byte
    return bitmap, attr, ch