"""PET image encoder: monochrome quadrant-block pseudo-bitmap.

40-column models -> 80x50 pixels (40x25 chars); 80-column -> 160x50.  The image
is dithered to one bit, then each 2x2 pixel block becomes the PETSCII quadrant
character with that pattern.  Output is the screen-RAM byte array ($8000): one
screen code per character cell, row-major.
"""
from __future__ import annotations

import numpy as np

from .. import dither, palette as c64pal, imageprep
from ..convert.base import Conversion, perceptual_error
from . import palette as petpal

ROWS = 25                      # character rows (both 40- and 80-col PETs)


def _dims(cols):
    return cols * 2, ROWS * 2, cols   # pixel W, pixel H, char cols


def convert(img_rgb, cols=40, dither_mode="floyd", base_color=None):
    W, H, _ = _dims(cols)
    plab = petpal.palette_lab()
    prgb = petpal.get_palette().astype(np.uint8)

    # one-bit luminance dither (1 = lit phosphor)
    L = c64pal.srgb_to_lab(img_rgb)[..., 0]
    mono = np.zeros((H, W, 3)); mono[..., 0] = L
    pmono = np.zeros_like(plab); pmono[:, 0] = plab[:, 0]
    allowed = np.tile(np.array([0, 1]), (H, W, 1))
    idx = dither.quantize(mono, allowed, pmono, dither_mode).astype(np.uint8)

    # each 2x2 block -> quadrant screen code
    screen = bytearray(cols * ROWS)
    for r in range(ROWS):
        for c in range(cols):
            tl = idx[r * 2, c * 2]; tr = idx[r * 2, c * 2 + 1]
            bl = idx[r * 2 + 1, c * 2]; br = idx[r * 2 + 1, c * 2 + 1]
            key = (tl << 3) | (tr << 2) | (bl << 1) | br
            screen[r * cols + c] = petpal.QUAD[key]

    return Conversion(
        mode="mono", width=W, height=H,
        pixel_aspect=0.83 if cols == 40 else 0.42,
        index_image=idx.astype(np.uint16), data=bytes(screen), data_addr=0x8000,
        viewer="pet", preview_rgb=prgb[idx],
        error=perceptual_error(idx, mono, pmono),
        meta={"palette": "pet", "dither": dither_mode, "cols": cols},
    )


def convert_image(path_or_img, cols=40, dither_mode="floyd", intensive=False,
                  prep_opt=None, base_color=None):
    prep_opt = prep_opt or imageprep.PrepOptions()
    W, H, _ = _dims(cols)
    aspect = 0.83 if cols == 40 else 0.42
    img_rgb = imageprep.prepare(path_or_img, W, H, aspect, prep_opt,
                                border_rgb=(0, 0, 0))
    return convert(img_rgb, cols, dither_mode, base_color=base_color)