// SPDX-FileCopyrightText: © 2026 Jeffrey C. Ollie <jeff@ocjtech.us>
// SPDX-License-Identifier: MIT

//! Choosing the colors a sixel image is written in.
//!
//! A sixel image draws with a limited set of color registers, 256 on most
//! terminals that support it, so a picture with more colors than that has to
//! be reduced to a palette first. An image that already fits is written with
//! exactly its own colors and nothing is lost but the rounding to whole
//! percentages, which every sixel color goes through.
//!
//! Otherwise the palette is chosen by median cut (Heckbert, 1982) over a
//! histogram of the image at five bits per channel: the box of colors that
//! is most populous for its size is split at its median along its longest
//! side, until there are as many boxes as colors wanted, and each box
//! becomes the average of the pixels in it. Floyd–Steinberg error diffusion
//! then spreads what each pixel loses to its neighbors, which turns banding
//! into a fine grain the eye averages away. Both are deterministic: the
//! same surface always gives the same output.

const std = @import("std");
const Allocator = std.mem.Allocator;
const palette = @import("palette.zig");
const Rgb = palette.Rgb;

pub const Dither = enum { none, floyd_steinberg };

/// One pixel of the image being written: straight (not premultiplied) color,
/// and whether it is drawn at all.
pub const Pixel = struct {
    color: Rgb,
    opaque_: bool,
};

/// The index of a pixel that is not drawn.
pub const transparent: u16 = std.math.maxInt(u16);

pub const Result = struct {
    /// The registers to define, each already rounded to what sixel can say.
    colors: []Rgb,
    /// One register per pixel, row by row, or `transparent`.
    indices: []u16,

    pub fn deinit(self: *Result, gpa: Allocator) void {
        gpa.free(self.colors);
        gpa.free(self.indices);
        self.* = undefined;
    }
};

/// Assigns each pixel of a `width`-wide image a register, using at most
/// `max_colors` registers.
pub fn quantize(
    gpa: Allocator,
    pixels: []const Pixel,
    width: usize,
    max_colors: u16,
    dither: Dither,
) Allocator.Error!Result {
    std.debug.assert(max_colors >= 1);
    const indices = try gpa.alloc(u16, pixels.len);
    errdefer gpa.free(indices);

    if (try exact(gpa, pixels, indices, max_colors)) |colors| {
        return .{ .colors = colors, .indices = indices };
    }

    const colors = try medianCut(gpa, pixels, max_colors);
    errdefer gpa.free(colors);
    try map(gpa, pixels, width, colors, indices, dither);
    return .{ .colors = colors, .indices = indices };
}

/// The image's own colors, if there are few enough, with `indices` filled
/// in. Colors that round to the same percentages are the same color here,
/// since they would be once written.
fn exact(gpa: Allocator, pixels: []const Pixel, indices: []u16, max_colors: u16) Allocator.Error!?[]Rgb {
    var seen: std.AutoArrayHashMapUnmanaged(u24, void) = .empty;
    defer seen.deinit(gpa);
    for (pixels, indices) |p, *index| {
        if (!p.opaque_) {
            index.* = transparent;
            continue;
        }
        const gop = try seen.getOrPut(gpa, key(p.color.snapped()));
        if (gop.index >= max_colors) return null;
        index.* = @intCast(gop.index);
    }
    const colors = try gpa.alloc(Rgb, seen.count());
    for (seen.keys(), colors) |k, *c| c.* = unkey(k);
    return colors;
}

fn key(c: Rgb) u24 {
    return @as(u24, c.r) << 16 | @as(u24, c.g) << 8 | c.b;
}

fn unkey(k: u24) Rgb {
    return .{ .r = @truncate(k >> 16), .g = @truncate(k >> 8), .b = @truncate(k) };
}

const bins = 1 << 15;

fn bin(c: Rgb) u15 {
    return @as(u15, c.r >> 3) << 10 | @as(u15, c.g >> 3) << 5 | c.b >> 3;
}

const Bin = struct {
    count: u64 = 0,
    sum: [3]u64 = .{ 0, 0, 0 },

    fn mean(self: Bin, axis: usize) u8 {
        return @intCast(self.sum[axis] / self.count);
    }
};

/// A box of histogram bins, as a range of the sorted bin list.
const Box = struct {
    start: usize,
    end: usize,
    count: u64,
    /// The longest side, and how long it is.
    axis: usize,
    range: u8,
};

fn medianCut(gpa: Allocator, pixels: []const Pixel, max_colors: u16) Allocator.Error![]Rgb {
    const hist = try gpa.alloc(Bin, bins);
    defer gpa.free(hist);
    @memset(hist, .{});
    for (pixels) |p| {
        if (!p.opaque_) continue;
        const b = &hist[bin(p.color)];
        b.count += 1;
        b.sum[0] += p.color.r;
        b.sum[1] += p.color.g;
        b.sum[2] += p.color.b;
    }

    var used: std.ArrayList(Bin) = .empty;
    defer used.deinit(gpa);
    for (hist) |b| if (b.count > 0) try used.append(gpa, b);
    if (used.items.len == 0) return gpa.alloc(Rgb, 0);

    var boxes: std.ArrayList(Box) = .empty;
    defer boxes.deinit(gpa);
    try boxes.append(gpa, measure(used.items, 0, used.items.len));

    while (boxes.items.len < max_colors) {
        // The box most worth splitting: many pixels spread over a long side.
        var best: ?usize = null;
        var best_score: u64 = 0;
        for (boxes.items, 0..) |b, i| {
            if (b.end - b.start < 2) continue;
            const score = b.count * (@as(u64, b.range) + 1);
            if (best == null or score > best_score) {
                best = i;
                best_score = score;
            }
        }
        const i = best orelse break;
        const b = boxes.items[i];
        const slice = used.items[b.start..b.end];
        std.mem.sortUnstable(Bin, slice, b.axis, struct {
            fn lessThan(axis: usize, l: Bin, r: Bin) bool {
                return l.mean(axis) < r.mean(axis);
            }
        }.lessThan);
        // The population median, kept off both ends so that neither half
        // is empty.
        var acc: u64 = 0;
        var cut: usize = b.start + 1;
        for (slice, b.start..) |bn, j| {
            acc += bn.count;
            if (acc * 2 >= b.count) {
                cut = std.math.clamp(j + 1, b.start + 1, b.end - 1);
                break;
            }
        }
        boxes.items[i] = measure(used.items, b.start, cut);
        try boxes.append(gpa, measure(used.items, cut, b.end));
    }

    const colors = try gpa.alloc(Rgb, boxes.items.len);
    for (boxes.items, colors) |b, *c| {
        var total: Bin = .{};
        for (used.items[b.start..b.end]) |bn| {
            total.count += bn.count;
            for (0..3) |axis| total.sum[axis] += bn.sum[axis];
        }
        c.* = (Rgb{ .r = total.mean(0), .g = total.mean(1), .b = total.mean(2) }).snapped();
    }
    return colors;
}

fn measure(used: []const Bin, start: usize, end: usize) Box {
    var lo: [3]u8 = .{ 255, 255, 255 };
    var hi: [3]u8 = .{ 0, 0, 0 };
    var count: u64 = 0;
    for (used[start..end]) |b| {
        count += b.count;
        for (0..3) |axis| {
            lo[axis] = @min(lo[axis], b.mean(axis));
            hi[axis] = @max(hi[axis], b.mean(axis));
        }
    }
    var axis: usize = 0;
    for (1..3) |a| {
        if (hi[a] - lo[a] > hi[axis] - lo[axis]) axis = a;
    }
    return .{ .start = start, .end = end, .count = count, .axis = axis, .range = hi[axis] - lo[axis] };
}

/// Fills in `indices` with the nearest of `colors` to each pixel, diffusing
/// the error if asked to.
fn map(
    gpa: Allocator,
    pixels: []const Pixel,
    width: usize,
    colors: []const Rgb,
    indices: []u16,
    dither: Dither,
) Allocator.Error!void {
    // Nearest-color searches remembered by five-bit bin. Two colors in one
    // bin can have different nearest entries, so this is an approximation,
    // but one finer than the palette it searches.
    const cache = try gpa.alloc(u16, bins);
    defer gpa.free(cache);
    @memset(cache, transparent);

    // The error carried into this row and the next, in 1/16ths, with a
    // column of slack at each end so the edges need no special case.
    const err = try gpa.alloc([3]i32, 2 * (width + 2));
    defer gpa.free(err);
    @memset(err, .{ 0, 0, 0 });
    var cur = err[0 .. width + 2];
    var next = err[width + 2 ..];

    for (pixels, indices, 0..) |p, *index, n| {
        const x = n % width;
        if (x == 0 and n > 0) {
            std.mem.swap([][3]i32, &cur, &next);
            @memset(next, .{ 0, 0, 0 });
        }
        if (!p.opaque_) {
            index.* = transparent;
            continue;
        }
        var want: [3]i32 = .{ p.color.r, p.color.g, p.color.b };
        if (dither == .floyd_steinberg) {
            for (0..3) |a| want[a] = std.math.clamp(want[a] + @divTrunc(cur[x + 1][a], 16), 0, 255);
        }
        const target: Rgb = .{ .r = @intCast(want[0]), .g = @intCast(want[1]), .b = @intCast(want[2]) };
        const slot = &cache[bin(target)];
        if (slot.* == transparent) slot.* = nearest(colors, target);
        index.* = slot.*;

        if (dither == .floyd_steinberg) {
            const got = colors[slot.*];
            const e: [3]i32 = .{ want[0] - got.r, want[1] - got.g, want[2] - got.b };
            for (0..3) |a| {
                cur[x + 2][a] += e[a] * 7;
                next[x][a] += e[a] * 3;
                next[x + 1][a] += e[a] * 5;
                next[x + 2][a] += e[a] * 1;
            }
        }
    }
}

fn nearest(colors: []const Rgb, c: Rgb) u16 {
    var best: u16 = 0;
    var best_d: u32 = std.math.maxInt(u32);
    for (colors, 0..) |p, i| {
        const dr = @as(i32, p.r) - c.r;
        const dg = @as(i32, p.g) - c.g;
        const db = @as(i32, p.b) - c.b;
        const d: u32 = @intCast(dr * dr + dg * dg + db * db);
        if (d < best_d) {
            best_d = d;
            best = @intCast(i);
        }
    }
    return best;
}

const testing = std.testing;

fn solid(r: u8, g: u8, b: u8) Pixel {
    return .{ .color = .{ .r = r, .g = g, .b = b }, .opaque_ = true };
}

test "few colors are kept exactly" {
    const pixels = [_]Pixel{
        solid(255, 0, 0), solid(0, 255, 0), .{ .color = .{ .r = 1, .g = 2, .b = 3 }, .opaque_ = false },
        solid(255, 0, 0), solid(0, 0, 255), solid(0, 255, 0),
    };
    var r = try quantize(testing.allocator, &pixels, 3, 256, .floyd_steinberg);
    defer r.deinit(testing.allocator);
    try testing.expectEqual(@as(usize, 3), r.colors.len);
    try testing.expectEqualSlices(u16, &.{ 0, 1, transparent, 0, 2, 1 }, r.indices);
    try testing.expect(r.colors[2].eql(.{ .r = 0, .g = 0, .b = 255 }));
}

test "colors that round alike are one register" {
    const pixels = [_]Pixel{ solid(128, 0, 0), solid(127, 0, 0) };
    var r = try quantize(testing.allocator, &pixels, 2, 256, .none);
    defer r.deinit(testing.allocator);
    try testing.expectEqual(@as(usize, 1), r.colors.len);
}

test "many colors are reduced, deterministically" {
    var pixels: [64 * 64]Pixel = undefined;
    for (&pixels, 0..) |*p, i| p.* = solid(@intCast(i % 64 * 4), @intCast(i / 64 * 4), @intCast((i * 7) % 256));
    var a = try quantize(testing.allocator, &pixels, 64, 16, .floyd_steinberg);
    defer a.deinit(testing.allocator);
    var b = try quantize(testing.allocator, &pixels, 64, 16, .floyd_steinberg);
    defer b.deinit(testing.allocator);
    try testing.expectEqual(@as(usize, 16), a.colors.len);
    try testing.expectEqualSlices(u16, a.indices, b.indices);
    for (a.indices) |i| try testing.expect(i < 16);
    for (a.colors) |c| try testing.expect(c.eql(c.snapped()));

    var c = try quantize(testing.allocator, &pixels, 64, 1, .none);
    defer c.deinit(testing.allocator);
    try testing.expectEqual(@as(usize, 1), c.colors.len);
}

test "an image with nothing drawn has no colors" {
    const pixels: [4]Pixel = @splat(.{ .color = .{ .r = 0, .g = 0, .b = 0 }, .opaque_ = false });
    var r = try quantize(testing.allocator, &pixels, 2, 1, .none);
    defer r.deinit(testing.allocator);
    try testing.expectEqual(@as(usize, 0), r.colors.len);
}