"""Atari GR.15 + DLI: 160x192, a fresh set of 4 colours every 2 scanlines.

A display-list interrupt rewrites the four colour registers for each 2-line band
(96 bands).  Every *single* scanline is impossible -- four register writes don't
fit the inter-DLI window -- but every 2 lines leaves a comfortable budget, and
96x4 colours is still far beyond flat GR.15.  The display list (with DLI bits on
the right lines) is generated here and shipped in the data block.
"""

from __future__ import annotations

import numpy as np

from ... import dither, palette as c64pal
from ...convert.base import (Conversion, mean_error, perceptual_error,
                             DIFFUSION_DITHERS)
from .. import palette as apal
from . import _common

WIDTH, HEIGHT = 160, 192
PIXEL_ASPECT = 2.0
BAND_H = 2
N_BANDS = HEIGHT // BAND_H            # 96
COLOR_ADDR = 0x6000
DL_ADDR = 0x6400                      # display list, after the colour table


def make_dlist() -> bytes:
    """ANTIC mode-E display list, 4K-split, DLI bit on the last line of each
    2-line band (odd lines 1..189) so the handler sets up the next band."""
    dl = bytearray([0x70, 0x70, 0x70])           # 24 blank lines
    dl += bytes([0x4e, 0x00, 0x40])              # line 0: LMS $4000 (no DLI)
    for ln in range(1, 102):                     # lines 1..101
        dl.append(0x8e if ln % 2 == 1 else 0x0e)
    dl += bytes([0x4e, 0x00, 0x50])              # line 102: LMS $5000 (no DLI)
    for ln in range(103, 192):                   # lines 103..191
        dl.append(0x8e if (ln % 2 == 1 and ln != 191) else 0x0e)
    dl += bytes([0x41, DL_ADDR & 0xFF, DL_ADDR >> 8])   # JVB -> start
    return bytes(dl)


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

    band_sets = np.zeros((N_BANDS, 4), dtype=np.int64)
    aware = dither_mode in DIFFUSION_DITHERS
    iters = 10 if intensive else 5
    # restrict the per-band dither-aware search to the image's own gamut (fast).
    cand = _common.relevant_candidates(img_lab, plab) if aware else None
    for b in range(N_BANDS):
        block = img_lab[b * BAND_H:(b + 1) * BAND_H].reshape(-1, 3)
        cols = _common.choose_palette(block, plab, k=4, iters=iters)
        if aware:   # span each band's gamut so dithering blends to the true shade
            cols = _common.choose_palette_dither(block, plab, k=4, init=cols,
                                                 iters=4 if intensive else 3,
                                                 candidates=cand)
        cols.sort(key=lambda c: plab[c, 0])
        band_sets[b] = cols

    allowed = np.repeat((band_sets[np.arange(HEIGHT) // BAND_H])[:, None, :],
                        WIDTH, axis=1)
    idx = dither.quantize(img_lab, allowed, plab, dither_mode).astype(np.int64)

    val = np.zeros((HEIGHT, WIDTH), dtype=np.uint8)
    for y in range(HEIGHT):
        lut = {int(c): v for v, c in enumerate(band_sets[y // BAND_H])}
        val[y] = [lut.get(int(p), 0) for p in idx[y]]

    bitmap = _common.split_screen(_common.pack_2bpp(val))     # 8192 -> $4000..$5FFF
    coltab = band_sets.astype(np.uint8).tobytes()             # 384 -> $6000
    region = coltab + bytes((DL_ADDR - COLOR_ADDR) - len(coltab)) + make_dlist()
    data = bitmap + region

    preview = np.repeat(prgb[idx], int(PIXEL_ASPECT), axis=1)
    return Conversion(
        mode="gr15dli", width=WIDTH, height=HEIGHT, pixel_aspect=PIXEL_ASPECT,
        index_image=idx.astype(np.uint16), data=data, data_addr=_common.DATA_ADDR,
        viewer="gr15dli", preview_rgb=preview,
        error=(perceptual_error if aware else mean_error)(idx, img_lab, plab),
        meta={"palette": palette_name, "dither": dither_mode},
    )