"""Amiga encoders: HAM6 (4096 colours), flat low-res (<=32 of 4096), greyscale.

HAM (Hold-And-Modify) is the showpiece: 6 bitplanes where the top 2 bits choose
whether a pixel is one of 16 base palette colours or modifies one R/G/B channel
of the pixel to its left -- giving up to 4096 colours on screen.  We pick 16 base
colours, then walk each scanline left-to-right choosing per pixel the option
(set / modify R / modify G / modify B) closest to the target.
"""
from __future__ import annotations

import numpy as np

from ... import dither, palette as c64pal
from ...convert.base import _box_blur
from .. import palette as apal

W, H = 320, 200


def _perr(final_rgb, img_rgb):
    a = _box_blur(c64pal.srgb_to_lab(final_rgb.astype(np.float64)))
    b = _box_blur(c64pal.srgb_to_lab(img_rgb.astype(np.float64)))
    return float(np.sqrt(((a - b) ** 2).sum(-1)).mean())


def planar_split(codes, nplanes):
    """codes (H,W) -> nplanes contiguous bitplanes (40 bytes/line, MSB left)."""
    out = bytearray()
    for p in range(nplanes):
        bit = ((codes >> p) & 1).astype(np.uint8)
        out += np.packbits(bit, axis=1).reshape(-1).tobytes()      # (H,40)
    return bytes(out)


def _kmeans_keys(img_lab, plab, k, iters=12):
    rng = np.random.default_rng(0)
    flat = img_lab.reshape(-1, 3)
    pts = flat[rng.choice(len(flat), min(6000, len(flat)), replace=False)]
    # k-means++ style seeding so every centroid is a real, distinct colour
    cen = pts[rng.integers(len(pts))][None].copy()
    for _ in range(k - 1):
        d = np.min(((pts[:, None, :] - cen[None]) ** 2).sum(2), 1)
        cen = np.vstack([cen, pts[int(d.argmax())]])      # farthest-point seed
    for _ in range(iters):
        lab = np.argmin(((pts[:, None, :] - cen[None]) ** 2).sum(2), 1)
        for j in range(k):
            m = pts[lab == j]
            cen[j] = m.mean(0) if len(m) else pts[rng.integers(len(pts))]
    # snap each centroid to the nearest 4096 palette key
    keys = [int(np.argmin(((plab - c) ** 2).sum(1))) for c in cen]
    return keys


def ham_encode(img_rgb, dither_mode):
    plab = apal.palette_lab()                         # (4096,3)
    img_lab = c64pal.srgb_to_lab(img_rgb)
    base_keys = _kmeans_keys(img_lab, plab, 16)
    base_lab = plab[base_keys]
    base_rgb4 = [(k >> 8 & 15, k >> 4 & 15, k & 15) for k in base_keys]

    t4 = np.clip(np.rint(img_rgb / 17.0), 0, 15).astype(np.int64)   # 4-bit targets
    # base option (independent of the held colour) precomputed for all pixels
    bd = ((img_lab[:, :, None, :] - base_lab[None, None]) ** 2).sum(3)   # (H,W,16)
    base_best = bd.argmin(2)
    base_cost = bd.min(2)

    codes = np.zeros((H, W), np.uint8)
    final = np.zeros((H, W, 3), np.uint8)
    P = plab
    for y in range(H):
        tl = img_lab[y]; t4y = t4[y]
        bb = base_best[y]; bc = base_cost[y]
        pr = pg = pb = 0                                # hardware holds black at line start
        for x in range(W):
            t0, t1, t2 = tl[x, 0], tl[x, 1], tl[x, 2]
            bi = int(bb[x]); best = bc[x]
            ctrl = 0; data = bi; nr, ng, nb = base_rgb4[bi]
            # force an absolute "set" on the first pixel so the line establishes a
            # colour regardless of the held-colour start (an all-modify run from
            # the wrong start would otherwise stay dark -- HAM left-edge bug).
            if x > 0:
                tr, tg, tb = int(t4y[x, 0]), int(t4y[x, 1]), int(t4y[x, 2])
                for c_ctrl, c_data, kr, kg, kb in (
                        (2, tr, tr, pg, pb), (3, tg, pr, tg, pb), (1, tb, pr, pg, tb)):
                    pk = P[(kr << 8) | (kg << 4) | kb]
                    dr = pk[0] - t0; dg = pk[1] - t1; db = pk[2] - t2
                    cc = dr * dr + dg * dg + db * db
                    if cc < best:
                        best = cc; ctrl = c_ctrl; data = c_data; nr, ng, nb = kr, kg, kb
            codes[y, x] = (ctrl << 4) | data
            pr, pg, pb = nr, ng, nb
            final[y, x, 0] = pr * 17; final[y, x, 1] = pg * 17; final[y, x, 2] = pb * 17

    planes = planar_split(codes, 6)
    colors = [apal.color_word(k) for k in base_keys]   # 16 base registers
    return planes, colors, final, _perr(final, img_rgb)


def flat_encode(img_rgb, n_colors, dither_mode, mono=False, base_color=None):
    plab = apal.palette_lab()
    prgb = apal.get_palette().astype(np.uint8)
    img_lab = c64pal.srgb_to_lab(img_rgb)
    nplanes = (n_colors - 1).bit_length()              # 32->5, 16->4

    if mono:
        keys = list(apal.GREYS)                        # 16 greys
        if base_color in range(4096):
            keys = sorted({keys[0], int(base_color), keys[-1]}, key=lambda i: plab[i, 0])
            keys = (keys + keys)[:n_colors]
        work = np.zeros_like(img_lab); work[..., 0] = img_lab[..., 0]
        pw = np.zeros_like(plab); pw[:, 0] = plab[:, 0]
    else:
        keys = _kmeans_keys(img_lab, plab, n_colors)
        work, pw = img_lab, plab

    allowed = np.tile(np.array(keys[:n_colors]), (H, W, 1))
    qidx = dither.quantize(work, allowed, pw, dither_mode).astype(np.int64)
    # map palette key -> pen index
    lut = {k: i for i, k in enumerate(keys[:n_colors])}
    pen = np.vectorize(lut.get)(qidx).astype(np.uint8)

    planes = planar_split(pen, nplanes)
    colors = [apal.color_word(k) for k in keys[:n_colors]]
    final = prgb[qidx]
    if mono:                              # measure greyscale against luminance
        g = img_rgb.mean(2, keepdims=True).repeat(3, 2)
        err = _perr(final, g.astype(np.uint8))
    else:
        err = _perr(final, img_rgb)
    return planes, colors, final, err