"""Atari 7800 MARIA 160A colour mode -- 160x192, 4 px/byte (2bpp).

MARIA gives 8 palettes of 3 colours each plus one shared background = up to 25
colours on screen.  The line is split into 8 objects of 20 px (5 bytes) each, and
every object may use a different palette, so each 20 px segment can pick the
3-colour palette (+ shared background) that fits it best -- chosen globally by
clustering the segments into 8 palettes.  The 2bpp bitmap packing is identical to
Atari GR.15, so the Atari encoder helpers are reused.

Conversion.data = bitmap(7680) + seg_palettes(192*8) + colours(25):
    bitmap        192 lines x 40 bytes, 2bpp, value 0..3 per pixel
    seg_palettes  one palette index (0..7) per 20px segment (8 per line)
    colours       [BACKGRND, P0C1,P0C2,P0C3, P1C1..P7C3]  (MARIA register order)
"""
from __future__ import annotations

import numpy as np

from ... import palette as c64pal
from ...convert.base import Conversion, perceptual_error, DIFFUSION_DITHERS
from ...atari import palette as apal
from ...atari.convert import _common

WIDTH, HEIGHT = 160, 192
PIXEL_ASPECT = 2.0
SEG_W = 40                       # pixels per object (10 bytes, 2bpp)
N_SEG = WIDTH // SEG_W           # 4 objects per line (MARIA DMA budget)
N_PAL = 8                        # MARIA palettes


def _pack_bitmap(val_image):
    return b"".join(bytes(b) for b in _common.pack_2bpp(val_image))


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

    bg, palettes, seg_pal, idx = _solve(img_lab, plab, dither_mode, intensive,
                                        candidates)

    # value image (0..3) per pixel: 0 = background, 1..3 = the segment palette's
    # three colours; build it from idx (palette indices) + the per-segment palette.
    val = np.zeros((HEIGHT, WIDTH), np.uint8)
    for row in range(HEIGHT):
        for s in range(N_SEG):
            x0 = s * SEG_W
            pal = [bg] + palettes[seg_pal[row, s]]
            lut = {c: v for v, c in enumerate(pal)}
            block = idx[row, x0:x0 + SEG_W]
            val[row, x0:x0 + SEG_W] = [lut[int(c)] for c in block]

    bitmap = _pack_bitmap(val)
    seg_bytes = bytes(seg_pal.reshape(-1).astype(np.uint8))
    colours = bytearray([bg])
    for p in range(N_PAL):
        colours += bytes(palettes[p])
    data = bitmap + seg_bytes + bytes(colours)

    err = perceptual_error(idx, img_lab, plab)
    return Conversion(
        mode=_mode, width=WIDTH, height=HEIGHT, pixel_aspect=PIXEL_ASPECT,
        index_image=idx.astype(np.uint16), data=data, data_addr=0,
        viewer="a7800", preview_rgb=np.repeat(prgb[idx], int(PIXEL_ASPECT), axis=1),
        error=err, meta={"palette": "ntsc", "dither": dither_mode, "bg": bg},
    )


def _pick(pixels_lab, plab, k, diff, candidates):
    """Pick k palette colours best representing the given pixels."""
    if diff:
        return _common.choose_palette_dither(pixels_lab, plab, k=k,
                                             candidates=candidates)
    if candidates is not None:
        sub = plab[candidates]
        return [candidates[i] for i in _common.choose_palette(pixels_lab, sub, k=k)]
    return _common.choose_palette(pixels_lab, plab, k=k)


def _solve(img_lab, plab, dither_mode, intensive, candidates=None):
    """Choose a shared background + 8 three-colour palettes by clustering the
    image's 40px segments by colour (so each palette is tuned to a group of
    similar segments), assign each segment its cluster's palette, and dither.
    Returns (bg, palettes[8][3], seg_pal, idx)."""
    from ... import dither as _dith
    diff = dither_mode in DIFFUSION_DITHERS
    H, W, _ = img_lab.shape

    # shared background = darkest of a small global palette
    gp = _pick(img_lab, plab, 8, diff, candidates)
    bg = min(gp, key=lambda c: plab[c, 0])

    # cluster the H*N_SEG segments by their mean colour into N_PAL groups
    seg_means = img_lab.reshape(H, N_SEG, SEG_W, 3).mean(axis=2)      # (H,N_SEG,3)
    feats = seg_means.reshape(-1, 3)
    rng = np.random.default_rng(0)
    cent = feats[rng.choice(len(feats), N_PAL, replace=False)].copy()
    labels = np.zeros(len(feats), np.int64)
    for _ in range(12):
        labels = ((feats[:, None] - cent[None]) ** 2).sum(-1).argmin(1)
        for g in range(N_PAL):
            m = labels == g
            if m.any():
                cent[g] = feats[m].mean(0)
    seg_pal = labels.reshape(H, N_SEG)

    # each palette = 3 colours tuned to the pixels of its cluster's segments
    seg_px = img_lab.reshape(H, N_SEG, SEG_W, 3)
    palettes = []
    for g in range(N_PAL):
        mask = seg_pal == g                                          # (H,N_SEG)
        if not mask.any():
            palettes.append([bg, bg, bg])
            continue
        px = seg_px[mask].reshape(-1, 1, 3)                          # (Npix,1,3)
        palettes.append(_pick(px, plab, 3, diff, candidates))

    # dither the whole image with each segment restricted to {bg} + its palette
    sets = [[bg] + palettes[g] for g in range(N_PAL)]
    allowed = np.zeros((H, W, 4), np.int64)
    for r in range(H):
        for s in range(N_SEG):
            allowed[r, s * SEG_W:s * SEG_W + SEG_W] = sets[seg_pal[r, s]]
    idx = _dith.quantize(img_lab, allowed, plab, dither_mode).astype(np.int64)
    return bg, palettes, seg_pal, idx