"""TI-99/4A Graphics Mode 2: 256x192, 2 colours per 8x8 cell (15-colour palette).

Like C64 hires but on the TMS9918A.  Produces the bitmap pattern table (6144 B)
and one colour byte per cell (768 B); the cartridge viewer expands each cell's
colour across its 8 rows of the VDP colour table.
"""

from __future__ import annotations

import numpy as np

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

WIDTH, HEIGHT = 256, 192
CELL_W, CELL_H = 8, 8
PIXEL_ASPECT = 1.0
N_COLS, N_ROWS = 32, 24


def convert(img_rgb, palette_name="tms9918", dither_mode="floyd",
            intensive=False, base_color=None):
    plab = tpal.palette_lab()
    prgb = tpal.get_palette().astype(np.uint8)
    img_lab = c64pal.srgb_to_lab(img_rgb)

    cells, rows, cols = base.cells_lab(img_lab, CELL_W, CELL_H)
    # Dither-aware colour selection for error-diffusion modes -- each cell's two
    # colours are chosen so the segment between them brackets the cell, letting
    # dithering blend to the true shade (far smoother than nearest-colour, which
    # bands).  Ordered/none keep plain nearest-colour selection.
    if dither_mode in base.DIFFUSION_DITHERS:
        sets, _ = base.select_cell_sets_dither(cells, plab, tpal.USABLE, n_free=2)
    else:
        dist = base.cell_distance(cells, plab)
        sets, _ = base.select_cell_sets(dist, tpal.USABLE, n_free=2)

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

    pattern, colors = _encode(idx, sets, rows, cols)
    data = bytes(pattern) + bytes(colors)            # 6144 + 768

    return base.Conversion(
        mode="gm2", width=WIDTH, height=HEIGHT, pixel_aspect=PIXEL_ASPECT,
        index_image=idx.astype(np.uint16), data=data, data_addr=0,
        viewer="gm2", preview_rgb=prgb[idx],
        error=base.perceptual_error(idx, img_lab, plab),
        meta={"palette": "tms9918", "dither": dither_mode},
    )


def _encode(idx, sets, rows, cols):
    pattern = np.zeros(6144, dtype=np.uint8)         # 768 cells x 8 rows
    colors = np.zeros(768, dtype=np.uint8)           # 1 colour byte per cell
    for cr in range(rows):
        for cc in range(cols):
            ci = cr * cols + cc
            c0, c1 = int(sets[ci, 0]), int(sets[ci, 1])
            # brighter colour = foreground (bit 1); store fg in high nibble.
            bg, fg = (c0, c1)
            colors[ci] = ((fg & 0x0F) << 4) | (bg & 0x0F)
            block = idx[cr * 8:cr * 8 + 8, cc * 8:cc * 8 + 8]
            base = ci * 8
            for r in range(8):
                row = block[r]
                byte = 0
                for x in range(8):
                    byte = (byte << 1) | (1 if row[x] == fg else 0)
                pattern[base + r] = byte
    return pattern, colors