"""Apple IIGS Super Hi-Res encoder: 320x200, 16 palettes x 16 colours (of 4096),
one palette per scanline (the SCB).

Lines are clustered into 16 palette groups; each group's 16 colours come from a
k-means of its pixels (quantised to 12-bit); each line is then assigned its best
palette and dithered to that palette's 16 colours.  Mono uses a single 16-grey
palette for all lines.
"""
from __future__ import annotations

import numpy as np

from ... import dither, palette as c64pal
from ...convert.base import perceptual_error
from .. import palette as gs

W, H = 320, 200
NPAL = 16
PALSIZE = 16
# SHR RAM block ($2000-$9FFF) offsets, relative to the 32K block base ($2000):
PIX_OFF, SCB_OFF, PAL_OFF = 0x0000, 0x7D00, 0x7E00
BLOCK_LEN = 0x8000


def _kmeans(pts, k, iters=12, seed=0):
    rng = np.random.default_rng(seed)
    if len(pts) <= k:
        c = np.zeros((k, pts.shape[1])); c[:len(pts)] = pts; return c
    cen = pts[rng.choice(len(pts), k, replace=False)].copy()
    for _ in range(iters):
        lab = np.argmin(((pts[:, None, :] - cen[None]) ** 2).sum(2), 1)
        for j in range(k):
            m = pts[lab == j]
            if len(m):
                cen[j] = m.mean(0)
    return cen


def _build_palettes(img_lab, img_rgb, mono, base_color):
    """Return (pal_rgb (NPAL,16,3) float 0-255, pal_q4 (NPAL,16,3) 4-bit,
    line_pal (H,) palette index per scanline)."""
    if mono:
        rgb, q4 = gs.greys12(PALSIZE)
        if base_color is not None:
            # tint: black -> base hue -> white ramp
            base = np.array(base_color, float)
            t = np.linspace(0, 1, PALSIZE)[:, None]
            ramp = np.where(t < 0.5, base * (t / 0.5),
                            base + (255 - base) * ((t - 0.5) / 0.5))
            rgb, q4 = gs.quantize12(ramp)
        pal_rgb = np.repeat(rgb[None], NPAL, 0)
        pal_q4 = np.repeat(q4[None], NPAL, 0)
        return pal_rgb, pal_q4, np.zeros(H, np.int64)

    # cluster the 200 lines into NPAL groups by mean colour
    line_mean = img_lab.mean(1)                            # (H,3)
    gc = _kmeans(line_mean, NPAL)
    grp = np.argmin(((line_mean[:, None, :] - gc[None]) ** 2).sum(2), 1)

    pal_rgb = np.zeros((NPAL, PALSIZE, 3))
    pal_q4 = np.zeros((NPAL, PALSIZE, 3), np.int64)
    rng = np.random.default_rng(0)
    for g in range(NPAL):
        lines = np.where(grp == g)[0]
        if len(lines) == 0:
            lines = np.array([g % H])
        pool = img_rgb[lines].reshape(-1, 3)
        if len(pool) > 4000:
            pool = pool[rng.choice(len(pool), 4000, replace=False)]
        cen = _kmeans(c64pal.srgb_to_lab(pool.reshape(1, -1, 3))[0], PALSIZE)
        # k-means in Lab -> map centroids back to RGB via nearest pool pixel
        plab_pool = c64pal.srgb_to_lab(pool.reshape(1, -1, 3))[0]
        rgb16 = np.zeros((PALSIZE, 3))
        for i in range(PALSIZE):
            rgb16[i] = pool[np.argmin(((plab_pool - cen[i]) ** 2).sum(1))]
        rgb_q, q4 = gs.quantize12(rgb16)
        pal_rgb[g], pal_q4[g] = rgb_q, q4

    # assign each line its best palette (min nearest-colour error)
    pal_lab = c64pal.srgb_to_lab(pal_rgb.reshape(1, -1, 3))[0].reshape(NPAL, PALSIZE, 3)
    line_pal = np.zeros(H, np.int64)
    for y in range(H):
        row = img_lab[y]                                   # (W,3)
        best, berr = 0, np.inf
        for g in range(NPAL):
            e = np.min(((row[:, None, :] - pal_lab[g][None]) ** 2).sum(2), 1).sum()
            if e < berr:
                berr, best = e, g
        line_pal[y] = best
    return pal_rgb, pal_q4, line_pal


def encode(img_rgb, dither_mode, mono=False, base_color=None):
    img_lab = c64pal.srgb_to_lab(img_rgb)
    pal_rgb, pal_q4, line_pal = _build_palettes(img_lab, img_rgb, mono, base_color)

    # combined colour table (NPAL*16) for the dither; per-pixel allowed = its
    # line's 16 palette colours
    rgb_all = pal_rgb.reshape(-1, 3)
    plab_all = c64pal.srgb_to_lab(rgb_all.reshape(1, -1, 3))[0]
    allowed = np.zeros((H, W, PALSIZE), np.int64)
    for y in range(H):
        base = line_pal[y] * PALSIZE
        allowed[y, :, :] = np.arange(base, base + PALSIZE)
    idx = dither.quantize(img_lab, allowed, plab_all, dither_mode).astype(np.int64)
    nib = (idx - line_pal[:, None] * PALSIZE).astype(np.uint8)   # 0-15 within palette

    # ---- emit the 32K SHR block ($2000-$9FFF) ----
    block = bytearray(BLOCK_LEN)
    for y in range(H):
        row = nib[y]
        off = PIX_OFF + y * 160
        for bx in range(160):
            block[off + bx] = (int(row[bx * 2]) << 4) | (int(row[bx * 2 + 1]) & 0x0F)
        block[SCB_OFF + y] = line_pal[y] & 0x0F             # 320 mode, palette n
    for p in range(NPAL):
        for c in range(PALSIZE):
            r4, g4, b4 = pal_q4[p, c]
            block[PAL_OFF + p * 32 + c * 2:PAL_OFF + p * 32 + c * 2 + 2] = \
                gs.color_word(int(r4), int(g4), int(b4))

    preview = rgb_all[idx].astype(np.uint8)
    if mono:
        # measure against luminance (the greys have no hue), like other mono modes
        lum = img_lab.copy(); lum[..., 1:] = 0.0
        plum = plab_all.copy(); plum[:, 1:] = 0.0
        err = perceptual_error(idx, lum, plum)
    else:
        err = perceptual_error(idx, img_lab, plab_all)
    return bytes(block), preview, err