"""
Dithering algorithms for 6-color e-ink display testing.

Includes multiple error-diffusion and ordered dithering algorithms
for comparison testing.
"""

import numpy as np
from numba import jit
from PIL import Image

# 6-color ACeP palette (RGB format)
PALETTE_RGB = [
    (0, 0, 0),  # BLACK
    (255, 255, 255),  # WHITE
    (255, 255, 0),  # YELLOW
    (255, 0, 0),  # RED
    (0, 0, 255),  # BLUE
    (0, 255, 0),  # GREEN
]

PALETTE_NAMES = ["Black", "White", "Yellow", "Red", "Blue", "Green"]


def create_pil_palette_image() -> Image.Image:
    """Create a PIL palette image for quantization."""
    pal_image = Image.new("P", (1, 1))
    flat_palette = []
    for color in PALETTE_RGB:
        flat_palette.extend(color)
    # Pad to 256 colors (768 values)
    flat_palette.extend([0] * (768 - len(flat_palette)))
    pal_image.putpalette(flat_palette)
    return pal_image


def find_nearest_color(pixel: np.ndarray, palette: np.ndarray) -> tuple[int, np.ndarray]:
    """Find the nearest palette color using Euclidean distance."""
    distances = np.sqrt(np.sum((palette - pixel) ** 2, axis=1))
    idx = np.argmin(distances)
    return idx, palette[idx]


def find_nearest_color_weighted(pixel: np.ndarray, palette: np.ndarray) -> tuple[int, np.ndarray]:
    """Find nearest color using perceptually-weighted distance (human eye sensitivity)."""
    # Weights based on human perception: Green > Red > Blue
    weights = np.array([0.299, 0.587, 0.114])
    diff = palette - pixel
    weighted_diff = diff * weights
    distances = np.sqrt(np.sum(weighted_diff**2, axis=1))
    idx = np.argmin(distances)
    return idx, palette[idx]


# =============================================================================
# Error Diffusion Dithering Algorithms
# =============================================================================


def dither_floyd_steinberg(image: Image.Image, weighted: bool = False) -> Image.Image:
    """
    Floyd-Steinberg dithering (1976).

    Classic error diffusion algorithm with distribution:
        * 7/16
    3/16 5/16 1/16
    """
    img = np.array(image.convert("RGB"), dtype=np.float64)
    height, width = img.shape[:2]
    palette = np.array(PALETTE_RGB, dtype=np.float64)
    find_fn = find_nearest_color_weighted if weighted else find_nearest_color

    for y in range(height):
        for x in range(width):
            old_pixel = img[y, x].copy()
            _, new_pixel = find_fn(old_pixel, palette)
            img[y, x] = new_pixel
            error = old_pixel - new_pixel

            if x + 1 < width:
                img[y, x + 1] += error * 7 / 16
            if y + 1 < height:
                if x > 0:
                    img[y + 1, x - 1] += error * 3 / 16
                img[y + 1, x] += error * 5 / 16
                if x + 1 < width:
                    img[y + 1, x + 1] += error * 1 / 16

    img = np.clip(img, 0, 255).astype(np.uint8)
    return Image.fromarray(img, "RGB")


def dither_jarvis_judice_ninke(image: Image.Image, weighted: bool = False) -> Image.Image:
    """
    Jarvis, Judice, and Ninke dithering (1976).

    Spreads error over a larger area (12 pixels):
            * 7 5
    3 5 7 5 3
    1 3 5 3 1
    All divided by 48.
    """
    img = np.array(image.convert("RGB"), dtype=np.float64)
    height, width = img.shape[:2]
    palette = np.array(PALETTE_RGB, dtype=np.float64)
    find_fn = find_nearest_color_weighted if weighted else find_nearest_color

    for y in range(height):
        for x in range(width):
            old_pixel = img[y, x].copy()
            _, new_pixel = find_fn(old_pixel, palette)
            img[y, x] = new_pixel
            error = old_pixel - new_pixel

            # Row 0 (current row)
            if x + 1 < width:
                img[y, x + 1] += error * 7 / 48
            if x + 2 < width:
                img[y, x + 2] += error * 5 / 48

            # Row 1
            if y + 1 < height:
                for dx, w in [(-2, 3), (-1, 5), (0, 7), (1, 5), (2, 3)]:
                    if 0 <= x + dx < width:
                        img[y + 1, x + dx] += error * w / 48

            # Row 2
            if y + 2 < height:
                for dx, w in [(-2, 1), (-1, 3), (0, 5), (1, 3), (2, 1)]:
                    if 0 <= x + dx < width:
                        img[y + 2, x + dx] += error * w / 48

    img = np.clip(img, 0, 255).astype(np.uint8)
    return Image.fromarray(img, "RGB")


def dither_stucki(image: Image.Image, weighted: bool = False) -> Image.Image:
    """
    Stucki dithering (1981).

    Similar to JJN but with different weights:
            * 8 4
    2 4 8 4 2
    1 2 4 2 1
    All divided by 42.
    """
    img = np.array(image.convert("RGB"), dtype=np.float64)
    height, width = img.shape[:2]
    palette = np.array(PALETTE_RGB, dtype=np.float64)
    find_fn = find_nearest_color_weighted if weighted else find_nearest_color

    for y in range(height):
        for x in range(width):
            old_pixel = img[y, x].copy()
            _, new_pixel = find_fn(old_pixel, palette)
            img[y, x] = new_pixel
            error = old_pixel - new_pixel

            if x + 1 < width:
                img[y, x + 1] += error * 8 / 42
            if x + 2 < width:
                img[y, x + 2] += error * 4 / 42

            if y + 1 < height:
                for dx, w in [(-2, 2), (-1, 4), (0, 8), (1, 4), (2, 2)]:
                    if 0 <= x + dx < width:
                        img[y + 1, x + dx] += error * w / 42

            if y + 2 < height:
                for dx, w in [(-2, 1), (-1, 2), (0, 4), (1, 2), (2, 1)]:
                    if 0 <= x + dx < width:
                        img[y + 2, x + dx] += error * w / 42

    img = np.clip(img, 0, 255).astype(np.uint8)
    return Image.fromarray(img, "RGB")


def dither_atkinson(image: Image.Image, weighted: bool = False) -> Image.Image:
    """
    Atkinson dithering (Bill Atkinson, Apple).

    Only diffuses 6/8 of the error (loses some detail but reduces noise):
        * 1 1
    1 1 1
      1
    All divided by 8 (but only 6/8 total error diffused).
    """
    img = np.array(image.convert("RGB"), dtype=np.float64)
    height, width = img.shape[:2]
    palette = np.array(PALETTE_RGB, dtype=np.float64)
    find_fn = find_nearest_color_weighted if weighted else find_nearest_color

    for y in range(height):
        for x in range(width):
            old_pixel = img[y, x].copy()
            _, new_pixel = find_fn(old_pixel, palette)
            img[y, x] = new_pixel
            error = old_pixel - new_pixel

            # Distribute 1/8 to each of 6 neighbors
            if x + 1 < width:
                img[y, x + 1] += error / 8
            if x + 2 < width:
                img[y, x + 2] += error / 8

            if y + 1 < height:
                if x > 0:
                    img[y + 1, x - 1] += error / 8
                img[y + 1, x] += error / 8
                if x + 1 < width:
                    img[y + 1, x + 1] += error / 8

            if y + 2 < height:
                img[y + 2, x] += error / 8

    img = np.clip(img, 0, 255).astype(np.uint8)
    return Image.fromarray(img, "RGB")


def dither_sierra(image: Image.Image, weighted: bool = False) -> Image.Image:
    """
    Sierra dithering (Frankie Sierra).

    Full Sierra (Sierra-3):
            * 5 3
    2 4 5 4 2
      2 3 2
    All divided by 32.
    """
    img = np.array(image.convert("RGB"), dtype=np.float64)
    height, width = img.shape[:2]
    palette = np.array(PALETTE_RGB, dtype=np.float64)
    find_fn = find_nearest_color_weighted if weighted else find_nearest_color

    for y in range(height):
        for x in range(width):
            old_pixel = img[y, x].copy()
            _, new_pixel = find_fn(old_pixel, palette)
            img[y, x] = new_pixel
            error = old_pixel - new_pixel

            if x + 1 < width:
                img[y, x + 1] += error * 5 / 32
            if x + 2 < width:
                img[y, x + 2] += error * 3 / 32

            if y + 1 < height:
                for dx, w in [(-2, 2), (-1, 4), (0, 5), (1, 4), (2, 2)]:
                    if 0 <= x + dx < width:
                        img[y + 1, x + dx] += error * w / 32

            if y + 2 < height:
                for dx, w in [(-1, 2), (0, 3), (1, 2)]:
                    if 0 <= x + dx < width:
                        img[y + 2, x + dx] += error * w / 32

    img = np.clip(img, 0, 255).astype(np.uint8)
    return Image.fromarray(img, "RGB")


def dither_sierra_lite(image: Image.Image, weighted: bool = False) -> Image.Image:
    """
    Sierra Lite (Sierra-2-4A) - faster variant.

        * 2
    1 1
    All divided by 4.
    """
    img = np.array(image.convert("RGB"), dtype=np.float64)
    height, width = img.shape[:2]
    palette = np.array(PALETTE_RGB, dtype=np.float64)
    find_fn = find_nearest_color_weighted if weighted else find_nearest_color

    for y in range(height):
        for x in range(width):
            old_pixel = img[y, x].copy()
            _, new_pixel = find_fn(old_pixel, palette)
            img[y, x] = new_pixel
            error = old_pixel - new_pixel

            if x + 1 < width:
                img[y, x + 1] += error * 2 / 4

            if y + 1 < height:
                if x > 0:
                    img[y + 1, x - 1] += error * 1 / 4
                img[y + 1, x] += error * 1 / 4

    img = np.clip(img, 0, 255).astype(np.uint8)
    return Image.fromarray(img, "RGB")


def dither_burkes(image: Image.Image, weighted: bool = False) -> Image.Image:
    """
    Burkes dithering.

        * 8 4
    2 4 8 4 2
    All divided by 32.
    """
    img = np.array(image.convert("RGB"), dtype=np.float64)
    height, width = img.shape[:2]
    palette = np.array(PALETTE_RGB, dtype=np.float64)
    find_fn = find_nearest_color_weighted if weighted else find_nearest_color

    for y in range(height):
        for x in range(width):
            old_pixel = img[y, x].copy()
            _, new_pixel = find_fn(old_pixel, palette)
            img[y, x] = new_pixel
            error = old_pixel - new_pixel

            if x + 1 < width:
                img[y, x + 1] += error * 8 / 32
            if x + 2 < width:
                img[y, x + 2] += error * 4 / 32

            if y + 1 < height:
                for dx, w in [(-2, 2), (-1, 4), (0, 8), (1, 4), (2, 2)]:
                    if 0 <= x + dx < width:
                        img[y + 1, x + dx] += error * w / 32

    img = np.clip(img, 0, 255).astype(np.uint8)
    return Image.fromarray(img, "RGB")


# =============================================================================
# Ordered Dithering Algorithms
# =============================================================================


def get_bayer_matrix(n: int) -> np.ndarray:
    """Generate a Bayer matrix of size 2^n x 2^n."""
    if n == 0:
        return np.array([[0]])
    smaller = get_bayer_matrix(n - 1)
    size = 2 ** (n - 1)
    result = np.zeros((2**n, 2**n))
    result[:size, :size] = 4 * smaller
    result[:size, size:] = 4 * smaller + 2
    result[size:, :size] = 4 * smaller + 3
    result[size:, size:] = 4 * smaller + 1
    return result


def dither_ordered_bayer(
    image: Image.Image, matrix_size: int = 4, strength: float = 1.0
) -> Image.Image:
    """
    Ordered dithering using Bayer matrix.

    Args:
        image: Input image
        matrix_size: Size of Bayer matrix (2, 4, 8, or 16)
        strength: Dithering strength multiplier (0.0-2.0)
    """
    img = np.array(image.convert("RGB"), dtype=np.float64)
    height, width = img.shape[:2]
    palette = np.array(PALETTE_RGB, dtype=np.float64)

    # Get appropriate Bayer matrix
    n = {2: 1, 4: 2, 8: 3, 16: 4}.get(matrix_size, 2)
    bayer = get_bayer_matrix(n)
    bayer_size = bayer.shape[0]

    # Normalize Bayer matrix to -0.5 to 0.5 range, then scale
    bayer_normalized = (bayer / (bayer_size**2) - 0.5) * strength * 128

    result = np.zeros_like(img)

    for y in range(height):
        for x in range(width):
            threshold = bayer_normalized[y % bayer_size, x % bayer_size]
            adjusted_pixel = img[y, x] + threshold
            adjusted_pixel = np.clip(adjusted_pixel, 0, 255)
            _, new_pixel = find_nearest_color(adjusted_pixel, palette)
            result[y, x] = new_pixel

    return Image.fromarray(result.astype(np.uint8), "RGB")


_NUMBA_PALETTE = np.array(
    [
        [0, 0, 0],
        [255, 255, 255],
        [255, 255, 0],
        [255, 0, 0],
        [0, 0, 255],
        [0, 255, 0],
    ],
    dtype=np.float64,
)
_NUMBA_WEIGHTS = np.array([0.299, 0.587, 0.114], dtype=np.float64)


@jit(nopython=True)
def _numba_find_nearest(r, g, b, palette, weights):
    best_idx = 0
    best_dist = 1e10
    for i in range(palette.shape[0]):
        dr = (palette[i, 0] - r) * weights[0]
        dg = (palette[i, 1] - g) * weights[1]
        db = (palette[i, 2] - b) * weights[2]
        dist = dr * dr + dg * dg + db * db
        if dist < best_dist:
            best_dist = dist
            best_idx = i
    return best_idx


@jit(nopython=True)
def _numba_atkinson(img, palette, weights):
    height, width = img.shape[0], img.shape[1]
    for y in range(height):
        for x in range(width):
            old_r, old_g, old_b = img[y, x, 0], img[y, x, 1], img[y, x, 2]
            idx = _numba_find_nearest(old_r, old_g, old_b, palette, weights)
            new_r, new_g, new_b = palette[idx, 0], palette[idx, 1], palette[idx, 2]
            img[y, x, 0], img[y, x, 1], img[y, x, 2] = new_r, new_g, new_b
            err_r = (old_r - new_r) / 8.0
            err_g = (old_g - new_g) / 8.0
            err_b = (old_b - new_b) / 8.0
            if x + 1 < width:
                img[y, x + 1, 0] += err_r
                img[y, x + 1, 1] += err_g
                img[y, x + 1, 2] += err_b
            if x + 2 < width:
                img[y, x + 2, 0] += err_r
                img[y, x + 2, 1] += err_g
                img[y, x + 2, 2] += err_b
            if y + 1 < height:
                if x > 0:
                    img[y + 1, x - 1, 0] += err_r
                    img[y + 1, x - 1, 1] += err_g
                    img[y + 1, x - 1, 2] += err_b
                img[y + 1, x, 0] += err_r
                img[y + 1, x, 1] += err_g
                img[y + 1, x, 2] += err_b
                if x + 1 < width:
                    img[y + 1, x + 1, 0] += err_r
                    img[y + 1, x + 1, 1] += err_g
                    img[y + 1, x + 1, 2] += err_b
            if y + 2 < height:
                img[y + 2, x, 0] += err_r
                img[y + 2, x, 1] += err_g
                img[y + 2, x, 2] += err_b
    return img


def dither_atkinson_numba(image: Image.Image) -> Image.Image:
    """Numba-accelerated Atkinson dithering with perceptual weighting (~150x faster)."""
    img = np.array(image.convert("RGB"), dtype=np.float64)
    img = _numba_atkinson(img, _NUMBA_PALETTE, _NUMBA_WEIGHTS)
    img = np.clip(img, 0, 255).astype(np.uint8)
    return Image.fromarray(img, "RGB")


# =============================================================================
# PIL Built-in (for comparison)
# =============================================================================


def dither_pil_floyd_steinberg(image: Image.Image) -> Image.Image:
    """PIL's built-in Floyd-Steinberg dithering for comparison."""
    pal_image = create_pil_palette_image()
    img = image.convert("RGB")
    quantized = img.quantize(dither=Image.Dither.FLOYDSTEINBERG, palette=pal_image)
    return quantized.convert("RGB")


def dither_pil_none(image: Image.Image) -> Image.Image:
    """PIL quantization with no dithering (nearest color only)."""
    pal_image = create_pil_palette_image()
    img = image.convert("RGB")
    quantized = img.quantize(dither=Image.Dither.NONE, palette=pal_image)
    return quantized.convert("RGB")


# =============================================================================
# Algorithm Registry
# =============================================================================

DITHER_ALGORITHMS = {
    "none": {
        "name": "No Dithering (PIL)",
        "func": dither_pil_none,
        "description": "Simple nearest-color quantization without error diffusion",
    },
    "pil_fs": {
        "name": "Floyd-Steinberg (PIL)",
        "func": dither_pil_floyd_steinberg,
        "description": "PIL built-in Floyd-Steinberg implementation",
    },
    "floyd_steinberg": {
        "name": "Floyd-Steinberg",
        "func": dither_floyd_steinberg,
        "description": "Classic error diffusion (1976), good balance of speed and quality",
    },
    "floyd_steinberg_weighted": {
        "name": "Floyd-Steinberg (Weighted)",
        "func": lambda img: dither_floyd_steinberg(img, weighted=True),
        "description": "Floyd-Steinberg with perceptual color weighting",
    },
    "atkinson": {
        "name": "Atkinson",
        "func": dither_atkinson,
        "description": "Bill Atkinson (Apple), diffuses only 75% of error for cleaner results",
    },
    "atkinson_weighted": {
        "name": "Atkinson (Weighted)",
        "func": lambda img: dither_atkinson(img, weighted=True),
        "description": "Atkinson with perceptual color weighting",
    },
    "atkinson_fast": {
        "name": "Atkinson (Numba Fast)",
        "func": dither_atkinson_numba,
        "description": "Numba-accelerated Atkinson (~150x faster, requires numba)",
    },
    "jarvis": {
        "name": "Jarvis-Judice-Ninke",
        "func": dither_jarvis_judice_ninke,
        "description": "Larger diffusion kernel (1976), smoother gradients but slower",
    },
    "stucki": {
        "name": "Stucki",
        "func": dither_stucki,
        "description": "Similar to JJN with modified weights (1981)",
    },
    "sierra": {
        "name": "Sierra",
        "func": dither_sierra,
        "description": "Full Sierra dithering, balanced results",
    },
    "sierra_lite": {
        "name": "Sierra Lite",
        "func": dither_sierra_lite,
        "description": "Faster Sierra variant with smaller kernel",
    },
    "burkes": {
        "name": "Burkes",
        "func": dither_burkes,
        "description": "Simplified two-row error diffusion",
    },
    "bayer2": {
        "name": "Ordered (Bayer 2x2)",
        "func": lambda img: dither_ordered_bayer(img, matrix_size=2),
        "description": "Ordered dithering with 2x2 Bayer matrix",
    },
    "bayer4": {
        "name": "Ordered (Bayer 4x4)",
        "func": lambda img: dither_ordered_bayer(img, matrix_size=4),
        "description": "Ordered dithering with 4x4 Bayer matrix",
    },
    "bayer8": {
        "name": "Ordered (Bayer 8x8)",
        "func": lambda img: dither_ordered_bayer(img, matrix_size=8),
        "description": "Ordered dithering with 8x8 Bayer matrix",
    },
    "bayer4_strong": {
        "name": "Ordered (Bayer 4x4 Strong)",
        "func": lambda img: dither_ordered_bayer(img, matrix_size=4, strength=1.5),
        "description": "Bayer 4x4 with increased dithering strength",
    },
}


def get_algorithm_names() -> list[str]:
    """Return list of available algorithm names."""
    return list(DITHER_ALGORITHMS.keys())


def apply_dithering(image: Image.Image, algorithm: str) -> Image.Image:
    """Apply the specified dithering algorithm to an image."""
    if algorithm not in DITHER_ALGORITHMS:
        raise ValueError(f"Unknown algorithm: {algorithm}. Available: {get_algorithm_names()}")
    return DITHER_ALGORITHMS[algorithm]["func"](image)
