"""Hires bitmap mode: 320x200, two colours per 8x8 cell.

Data file layout (PRG, load $2000), matched to viewer/hires.s:
    $2000  bitmap     8000 bytes  (VIC reads here directly)
    $3F40  screen RAM 1000 bytes  (viewer copies to $0400)
"""

from __future__ import annotations

import numpy as np

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

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


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)
    dist = base.cell_distance(cells, plab)
    sets, _ = base.select_cell_sets(dist, range(16), n_free=2)

    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 = _encode(index_image, sets, rows, cols)
    payload = bytes(bitmap) + bytes(screen)

    conv = base.Conversion(
        mode="hires", width=WIDTH, height=HEIGHT, pixel_aspect=PIXEL_ASPECT,
        index_image=index_image, data=payload, viewer="hires",
        error=base.mean_error(index_image, img_lab, plab),
        meta={"palette": palette_name, "dither": dither_mode},
    )
    # Standard OCP Art Studio hires file (load $2000): bitmap, screen, border.
    conv.extra_files = [("picture.art", base.prg(0x2000, payload + b"\x00"))]
    return conv


def _encode(index_image, sets, rows, cols):
    """Build the 8000-byte bitmap and 1000-byte screen RAM."""
    bitmap = np.zeros(8000, dtype=np.uint8)
    screen = np.zeros(1000, dtype=np.uint8)
    for cr in range(rows):
        for cc in range(cols):
            ci = cr * cols + cc
            bg_col, fg_col = int(sets[ci, 0]), int(sets[ci, 1])
            screen[ci] = ((fg_col & 0x0F) << 4) | (bg_col & 0x0F)
            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 row[x] == fg_col else 0)
                bitmap[base_addr + r] = byte
    return bitmap, screen