"""Multicolor bitmap mode ("Koala"): 160x200, one shared background plus three
freely chosen colours per 4x8 cell.

Data file layout (PRG, load $2000), matched to viewer/multicolor.s:
    $2000  bitmap     8000 bytes  (VIC reads here directly)
    $3F40  screen RAM 1000 bytes  (viewer copies to $0400)
    $4328  colour RAM 1000 bytes  (viewer copies to $D800)
    $4710  background  1 byte     (viewer writes to $D021)
"""

from __future__ import annotations

import numpy as np

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

WIDTH, HEIGHT = 160, 200
CELL_W, CELL_H = 4, 8
PIXEL_ASPECT = 2.0
DATA_LOAD = 0x2000

# bit-pair -> colour source:  01 screen hi nibble, 10 screen lo nibble, 11 colour RAM
_SLOT_BITS = {1: 0b01, 2: 0b10, 3: 0b11}


def convert(img_rgb: np.ndarray, palette_name="colodore",
            dither_mode="bayer", intensive=False) -> base.Conversion:
    plab = pal.palette_lab(palette_name)
    img_lab = pal.srgb_to_lab(img_rgb)

    cells, rows, cols = base.cells_lab(img_lab, CELL_W, CELL_H)
    aware = dither_mode in base.DIFFUSION_DITHERS    # dither-aware colour selection

    if intensive:
        if aware:
            bg, sets, _ = base.optimize_background_dither(cells, plab, n_free=3)
        else:
            bg, sets, _ = base.optimize_background(base.cell_distance(cells, plab),
                                                   n_free=3)
    else:
        bg = base.best_global_color(img_lab, plab)
        avail = [i for i in range(16) if i != bg]
        if aware:
            sets, _ = base.select_cell_sets_dither(cells, plab, avail, n_free=3,
                                                   fixed=(bg,))
        else:
            dist = base.cell_distance(cells, plab)
            sets, _ = base.select_cell_sets(dist, avail, n_free=3, fixed=(bg,))

    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.uint8)

    bitmap, screen, colram = _encode(index_image, sets, bg, rows, cols)
    # This block also *is* a Koala body: bitmap, screen, colram, background.
    payload = bytes(bitmap) + bytes(screen) + bytes(colram) + bytes([bg])

    conv = base.Conversion(
        mode="multicolor", width=WIDTH, height=HEIGHT, pixel_aspect=PIXEL_ASPECT,
        index_image=index_image, data=payload, viewer="multicolor",
        error=base.perceptual_error(index_image, img_lab, plab),
        meta={"palette": palette_name, "dither": dither_mode, "background": bg},
    )
    # Standard "Koala Painter" file (load $6000) for use in other C64 art tools.
    conv.extra_files = [("picture.koa", base.prg(0x6000, payload))]
    return conv


def _encode(index_image, sets, bg, rows, cols):
    bitmap = np.zeros(8000, dtype=np.uint8)
    screen = np.zeros(1000, dtype=np.uint8)
    colram = np.zeros(1000, dtype=np.uint8)
    for cr in range(rows):
        for cc in range(cols):
            ci = cr * cols + cc
            # sets[ci] = [bg, c1, c2, c3]; assign the three non-bg colours to slots.
            c1, c2, c3 = int(sets[ci, 1]), int(sets[ci, 2]), int(sets[ci, 3])
            screen[ci] = ((c1 & 0x0F) << 4) | (c2 & 0x0F)
            colram[ci] = c3 & 0x0F
            color_to_bits = {bg: 0b00, c1: 0b01, c2: 0b10, c3: 0b11}
            base_addr = cr * 320 + cc * 8
            block = index_image[cr * 8:cr * 8 + 8, cc * 4:cc * 4 + 4]
            for r in range(8):
                row = block[r]
                byte = 0
                for x in range(4):
                    byte = (byte << 2) | color_to_bits.get(int(row[x]), 0b00)
                bitmap[base_addr + r] = byte
    return bitmap, screen, colram