import {
    readFileSync,
    writeFileSync
} from 'fs';

import Jimp from 'Jimp';

function toLinear(val) {
    // use a 2.4 gamma approximation
    // this is BT.1886 compatible
    // and simpler than sRGB
    let unit = val / 255;
    unit **= 2.4;
    return unit * 255;
}

function fromLinear(val) {
    let unit = val / 255;
    unit **= (1 / 2.4);
    return unit * 255;
}

class RGB {
    constructor(r, g, b) {
        this.r = r;
        this.g = g;
        this.b = b;
    }

    clone() {
        return new RGB(this.r, this.g, this.b);
    }

    static fromHex(val) {
        let r = (val >>> 16) & 0xff;
        let g = (val >>> 8) & 0xff;
        let b = val & 0xff;
        return new RGB(r,g,b);
    }

    toLinear() {
        return new RGB(
            toLinear(this.r),
            toLinear(this.g),
            toLinear(this.b)
        );
    }

    fromLinear() {
        return new RGB(
            fromLinear(this.r),
            fromLinear(this.g),
            fromLinear(this.b)
        );
    }

    cap() {
        if (this.r < 0) {
            this.r = 0;
        }
        if (this.g < 0) {
            this.g = 0;
        }
        if (this.b < 0) {
            this.b = 0;
        }
        if (this.r > 255) {
            this.r = 255;
        }
        if (this.g > 255) {
            this.g = 255;
        }
        if (this.b > 255) {
            this.b = 255;
        }
    }

    inc(other) {
        this.r += other.r;
        this.g += other.g;
        this.b += other.b;
        return this;
    }

    static add(a, b) {
        return new RGB(
            a.r + b.r,
            a.g + b.g,
            a.b + b.b
        );
    }

    add(other) {
        return RGB.add(this, other);
    }

    difference(other) {
        return new RGB(
            this.r - other.r,
            this.g - other.g,
            this.b - other.b
        );
    }

    divide(scalar) {
        return new RGB(
            this.r / scalar,
            this.g / scalar,
            this.b / scalar,
        );
    }

    magnitude() {
        return Math.sqrt(this.magnitude2());
    }

    magnitude2() {
        return this.r * this.r +
            this.g * this.g +
            this.b * this.b;
    }
}

const maxDist = (new RGB(255, 255, 255)).magnitude();

// snarfed from https://lospec.com/palette-list/atari-8-bit-family-gtia
// which was calculated with Retrospecs App's Atari 800 emulator
let atariRGB = [
    0x000000,
    0x111111,
    0x222222,
    0x333333,
    0x444444,
    0x555555,
    0x666666,
    0x777777,
    0x888888,
    0x999999,
    0xaaaaaa,
    0xbbbbbb,
    0xcccccc,
    0xdddddd,
    0xeeeeee,
    0xffffff,
    0x190700,
    0x2a1800,
    0x3b2900,
    0x4c3a00,
    0x5d4b00,
    0x6e5c00,
    0x7f6d00,
    0x907e09,
    0xa18f1a,
    0xb3a02b,
    0xc3b13c,
    0xd4c24d,
    0xe5d35e,
    0xf7e46f,
    0xfff582,
    0xffff96,
    0x310000,
    0x3f0000,
    0x531700,
    0x642800,
    0x753900,
    0x864a00,
    0x975b0a,
    0xa86c1b,
    0xb97d2c,
    0xca8e3d,
    0xdb9f4e,
    0xecb05f,
    0xfdc170,
    0xffd285,
    0xffe39c,
    0xfff4b2,
    0x420404,
    0x4f0000,
    0x600800,
    0x711900,
    0x822a0d,
    0x933b1e,
    0xa44c2f,
    0xb55d40,
    0xc66e51,
    0xd77f62,
    0xe89073,
    0xf9a183,
    0xffb298,
    0xffc3ae,
    0xffd4c4,
    0xffe5da,
    0x410103,
    0x50000f,
    0x61001b,
    0x720f2b,
    0x83203c,
    0x94314d,
    0xa5425e,
    0xb6536f,
    0xc76480,
    0xd87591,
    0xe986a2,
    0xfa97b3,
    0xffa8c8,
    0xffb9de,
    0xffcaef,
    0xfbdcf6,
    0x330035,
    0x440041,
    0x55004c,
    0x660c5c,
    0x771d6d,
    0x882e7e,
    0x993f8f,
    0xaa50a0,
    0xbb61b1,
    0xcc72c2,
    0xdd83d3,
    0xee94e4,
    0xffa5e4,
    0xffb6e9,
    0xffc7ee,
    0xffd8f3,
    0x1d005c,
    0x2e0068,
    0x400074,
    0x511084,
    0x622195,
    0x7332a6,
    0x8443b7,
    0x9554c8,
    0xa665d9,
    0xb776ea,
    0xc887eb,
    0xd998eb,
    0xe9a9ec,
    0xfbbaeb,
    0xffcbef,
    0xffdff9,
    0x020071,
    0x13007d,
    0x240b8c,
    0x351c9d,
    0x462dae,
    0x573ebf,
    0x684fd0,
    0x7960e1,
    0x8a71f2,
    0x9b82f7,
    0xac93f7,
    0xbda4f7,
    0xceb5f7,
    0xdfc6f7,
    0xf0d7f7,
    0xffe8f8,
    0x000068,
    0x000a7c,
    0x081b90,
    0x192ca1,
    0x2a3db2,
    0x3b4ec3,
    0x4c5fd4,
    0x5d70e5,
    0x6e81f6,
    0x7f92ff,
    0x90a3ff,
    0xa1b4ff,
    0xb2c5ff,
    0xc3d6ff,
    0xd4e7ff,
    0xe5f8ff,
    0x000a4d,
    0x001b63,
    0x002c79,
    0x023d8f,
    0x134ea0,
    0x245fb1,
    0x3570c2,
    0x4681d3,
    0x5792e4,
    0x68a3f5,
    0x79b4ff,
    0x8ac5ff,
    0x9bd6ff,
    0xace7ff,
    0xbdf8ff,
    0xceffff,
    0x001a26,
    0x002b3c,
    0x003c52,
    0x004d68,
    0x065e7c,
    0x176f8d,
    0x28809e,
    0x3991af,
    0x4aa2c0,
    0x5bb3d1,
    0x6cc4e2,
    0x7dd5f3,
    0x8ee6ff,
    0x9ff7ff,
    0xb0ffff,
    0xc1ffff,
    0x01250a,
    0x023610,
    0x004622,
    0x005738,
    0x05684d,
    0x16795e,
    0x278a6f,
    0x389b80,
    0x49ac91,
    0x5abda2,
    0x6bceb3,
    0x7cdfc4,
    0x8df0d5,
    0x9effe5,
    0xaffff1,
    0xc0fffd,
    0x04260d,
    0x043811,
    0x054713,
    0x005a1b,
    0x106b1b,
    0x217c2c,
    0x328d3d,
    0x439e4e,
    0x54af5f,
    0x65c070,
    0x76d181,
    0x87e292,
    0x98f3a3,
    0xa9ffb3,
    0xbaffbf,
    0xcbffcb,
    0x00230a,
    0x003510,
    0x044613,
    0x155613,
    0x266713,
    0x377813,
    0x488914,
    0x599a25,
    0x6aab36,
    0x7bbc47,
    0x8ccd58,
    0x9dde69,
    0xaeef7a,
    0xbfff8b,
    0xd0ff97,
    0xe1ffa3,
    0x001707,
    0x0e2808,
    0x1f3908,
    0x304a08,
    0x415b08,
    0x526c08,
    0x637d08,
    0x748e0d,
    0x859f1e,
    0x96b02f,
    0xa7c140,
    0xb8d251,
    0xc9e362,
    0xdaf473,
    0xebff82,
    0xfcff8e,
    0x1b0701,
    0x2c1801,
    0x3c2900,
    0x4d3b00,
    0x5f4c00,
    0x705e00,
    0x816f00,
    0x938009,
    0xa4921a,
    0xb2a02b,
    0xc7b43d,
    0xd8c64e,
    0xead760,
    0xf6e46f,
    0xfffa84,
    0xffff99,
].map((hex) => RGB.fromHex(hex).toLinear());

/**
 * Dither RGB input data with a target palette size.
 * If the number of used colors exceeds `n`, the
 * palette will be reduced until it fits.
 * @param {RGB[]} input source scanline data, in linear RGB
 * @param {number[]} palette - current working palette, as Atari 8-bit color values (low nybble luminance, high nybble hue)
 * @param {number} n - target color count
 * @param {RGB[]} inputError
 * @param {number} y
 * @returns {{output: Uint8Array, palette: number[], error: RGB[]}}
 */
function decimate(input, palette, n, inputError, y) {
    // to brute-force, the possible palettes are:
    // 255 * 254 * 253 = 16,386,810
    //
    // we could brute force it but that's a lot :D
    // but can do some bisection :D
    //
    // need a fitness metric.
    // each pixel in the dithered line gives a distance
    // sum/average them? median? maximum?
    // summing evens out the ups/downs from dithering
    // but doesn't distinguish between two close and two distant options
    // consider median, 90th-percentile, and max of abs(distance)
    // consider doing the distance for each channel?

    let width = input.length;

    let inputPixel = (x, error) => {
        let rgb = input[x].clone();
        if (error) {
            rgb.inc(error.cur[x]);
        }
        if (inputError) {
            rgb.inc(inputError[x]);
        }
        rgb.cap();
        return rgb;
    };

    // Apply dithering with given palette and collect color usage stats
    let dither = (palette) => {
        let fitness = new Float64Array(width);
        let error = {
            cur: [],
            next: [],
        };
        for (let i = 0; i < width; i++) {
            error.cur[i] = new RGB(0, 0, 0);
            error.next[i] = new RGB(0, 0, 0);
        }

        let output = new Uint8Array(width);
        let popularity = new Int32Array(palette.length);
        let distance2 = 0;

        let nextError = new RGB(0, 0, 0);

        // Try dithering with this palette.
        for (let x = 0; x < width; x++) {
            let rgb = inputPixel(x, error);

            // find the closest possible color
            // @todo consider doing the difference scoring in luminance and hue spaces
            let shortest = Infinity;
            let pick = 1;

            for (let i = 0; i < palette.length; i++) {
                let diff = rgb.difference(atariRGB[palette[i]]);
                let dist = diff.magnitude2();
                if (dist < shortest) {
                    nextError = diff;
                    shortest = dist;
                    pick = i;
                }
            }

            output[x] = pick;
            popularity[pick]++;

            let shares = 8;
            let single = nextError.divide(shares);
            let double = nextError.divide(shares / 2);
            error.cur[x + 1]?.inc(double);
            error.cur[x + 2]?.inc(single);

            error.next[x - 1]?.inc(single);
            error.next[x]?.inc(double);
            error.next[x + 1]?.inc(double);

            let mag = nextError.magnitude();
            fitness[x] = maxDist / mag;
            // 442 is the 3d distance across the rgb cube
            //fitness[x] = 442 - (nextError.magnitude());
            //fitness[x] = 442 / (442 - nextError.magnitude());
            fitness[x] = 255 / (256 - Math.max(0, nextError.r, nextError.g, nextError.b));

            let mag2 = nextError.magnitude2();
            distance2 += mag2;
        }
        return {
            output,
            palette,
            fitness,
            distance2,
            popularity,
            error: error.next
        };
    };


    let decimated = palette.slice();
    
    // force to grayscale
    //decimated = [0, 5, 10, 15];

    // force to rgb
    //decimated = [0, 0x36, 0xb6, 0x86];

    // force to rWb
    //decimated = [0, 0x36, 0x0f, 0x86];
    
    let reserved = [0]; // black
    //reserved = [0, 15]; // black, white
    //reserved = [0, 5, 10, 15]; // grayscale
    //reserved = [0, 0x48, 0x78, 15]; // vaporwave
    //reserved = [0, 0x3c, 0x78, 15]; // red/blue/white

    /*
    if (( y & 1 ) === 0) {
        reserved = [0, 0x3c, 0x1e, 15]; // red/yellow/white
    } else {
        reserved = [0, 0x76, 0x9e, 0xb8]; // blue/cyan/green
    }
    */

    let keepers = new Uint8Array(256);
    for (let i of reserved) {
        keepers[i & 0xfe] = 1; // drop that 0 luminance bit!
    }

    let zeros = (n) => {
        let arr = [];
        for (let i = 0; i < n; i++) {
            arr.push(0);
        }
        return arr;
    };

    // this takes the top hues, and uses the brightest of each hue
    // needs tuning
    /*
    // first, dither to the total atari palette
    while (decimated.length > n) {
        let {popularity} = dither(decimated);
        let hues = zeros(16);
        let lumas = zeros(16);
        for (let i = 0; i < decimated.length; i++) {
            let color = decimated[i];
            let hue = color >> 4;
            let luma = color & 0xe;
            hues[hue] += popularity[i];
            if (luma > lumas[hue]) {
                lumas[hue] = luma;
            }
        }

        let a = hues;
        a = a.map((count, hue) => { return {count, hue} });
        a = a.sort((a, b) => b.count - a.count);
        a = a.map(({hue}) => hue);
        a = a.filter((color) => !keepers[color]);
        //a = a.slice(0, n - reserved.length);
        a = a.slice(0, decimated.length - 1 - reserved.length);
        a = a.map((hue) => (hue << 4) | lumas[hue]);
        a = a.sort((a, b) => a - b);

        decimated = reserved.concat(a);
    }
    console.log('end', decimated);
    */

    // popularity? not really working right
    // first, dither to the total atari palette
    /*
    while (decimated.length > n) {
        //console.log(y);
        let {popularity} = dither(decimated);
        let a = decimated;
        // temporarily strip the reserved items
        a = a.filter((color) => !keepers[color]);
        a = a.sort((a, b) => popularity[b] - popularity[a]);
        a = reserved.concat(a);
        decimated = a.slice(0, n);
        //console.log(decimated);
    }
    //console.log('end', decimated);
    */

    while (decimated.length > n) {
        let {popularity, fitness, output} = dither(decimated);

        // Try dropping least used color on each iteration
        let least = Infinity;
        let pick = -1;
        for (let i = 1; i < decimated.length; i++) {

            //let coolFactor = popularity[i];

            let coolFactor = 0;
            if (popularity[i]) {
                for (let x = 0; x < width; x++) {
                    if (output[x] == i) {
                        // Scale up the scoring for close matches to prioritize
                        // color accuracy over raw linear usage.
                        //coolFactor += (fitness[x] ** 2);
                        coolFactor += (fitness[x] ** 4);
                    }
                }
            }
            if (coolFactor < least) {
                pick = i;
                least = coolFactor;
            }
            decimated = decimated.filter((color, i) => {
                if (i == 0) {
                    return true;
                }
                if (i == pick) {
                    return false;
                }
                // Also drop any non-black unused colors to save trouble.
                // However -- this may change dither results.
                // Try this with/without later.
                if (popularity[i] == 0) {
                    return false;
                }
                return true;
            });
        }
    }
                                                    
    // Palette fits
    return dither(decimated);
}

/**
 * Read an image file into a buffer
 * @param {string} src 
 * @returns {{width: number, height: number, rgba: Uint8Array}}
 */
async function loadImage(src) {
    let image = await Jimp.read(src);

    let width = image.bitmap.width;
    let height = image.bitmap.height;
    if (width != 160 || height != 160) {
        width = 160;
        height = 160;
        image = image.resize(width, height);
    }

    let rgba = image.bitmap.data.slice();
    return {
        width,
        height,
        rgba,
    };
}

function imageToLinearRGB(rgba) {
    let input = [];
    for (let i = 0; i < rgba.length; i += 4) {
        input.push(new RGB(
            rgba[i + 0],
            rgba[i + 1],
            rgba[i + 2]
        ).toLinear());
    }
    return input;
}

/**
 * Read an image file, squish to 160px if necessary,
 * and dither to 4 colors per scan line.
 *
 * @param {string} source path to source image file
 * @returns {{width: number, height: number, lines: {palette: Array, output: Uint8Array}[]}}
 */
async function convert(source, nbits) {

    let {
        width,
        height,
        rgba
    } = await loadImage(source);

    if (width !== 160) {
        throw new Error(`expected 160px-compatible width, got ${width} pixels`);
    }

    if (height !== 160) {
        // @fixme support up to 240px
        throw new Error(`expected 160px height, got ${height} pixels`);
    }

    if (rgba.length != width * 4 * height) {
        console.log(`
        width: ${width}
        height: ${height}
        rgba.length: ${rgba.length}`)
        throw new Error('inconsistent data size');
    }

    let input = imageToLinearRGB(rgba);

    if (input.length != width * height) {
        console.log(`
        width: ${width}
        height: ${height}
        rgba.length: ${input.length}`)
        throw new Error('inconsistent data size on input');
    }

    // Start with all colors usable with regular CTIA modes
    // (not the 16-luminance special mode on GTIA)
    let allColors = [];
    for (let i = 0; i < 256; i += 2) {
        allColors.push(i);
    }

    let lines = [];
    for (let y = 0; y < height; y++) {
        let error = lines[y - 1]?.error;
        let inputLine = input.slice(y * width, (y + 1) * width);
        let line = decimate(inputLine, allColors, 4, error, y);
        lines.push(line);
    }
    return {
        width,
        height,
        lines
    };
}


function indexedToBitmap(width, nbits, src, dest) {
    let nbytes = width * nbits / 8;
    let x = 0;
    for (let i = 0; i < nbytes; i++) {
        let a = 0;
        for (let b = 0; b < 8; b += nbits) {
            a <<= nbits;
            a |= src[x++];
        }
        dest[i] = a;
    }
}

function byte2byte(arr) {
    let lines = [];
    for (let i=0; i < arr.length; i++) {
        lines.push(`.byte ${arr[i]}`);
    }
    return lines.join('\n');
}

//let [even, odd] = [0, 1].map((bit) => (arr) => arr.filter((_item, index) => (index & 1) === bit));
function even(arr) {
    return arr.filter((_item, index) => !(index & 1));
}
function odd(arr) {
    return arr.filter((_item, index) => (index & 1));
}

function genAssembly(width, height, nbits, lines) {
    let stride = width * nbits / 8;
    let half = stride * height / 2;
    let frame = {
        palette1: new Uint8Array(height),
        palette2: new Uint8Array(height),
        palette3: new Uint8Array(height),
        bitmap: new Uint8Array(stride * height),
    };
    for (let y = 0; y < height; y++) {
        let base = 0;
        frame.palette1[y] = lines[y + base].palette[1];
        frame.palette2[y] = lines[y + base].palette[2];
        frame.palette3[y] = lines[y + base].palette[3];
        indexedToBitmap(
            width,
            nbits,
            lines[y + base].output,
            frame.bitmap.subarray(y * stride, (y + 1) * stride)
        );
    }

    return `.data
.export frame1_top
.export frame1_bottom
.export frame1_palette1_even
.export frame1_palette1_odd
.export frame1_palette2_even
.export frame1_palette2_odd
.export frame1_palette3_even
.export frame1_palette3_odd
.export displaylist

.segment "BUFFERS"

.align 4096
frame1_top:
${byte2byte(frame.bitmap.slice(0, half))}

.align 128
frame1_palette1_even:
${byte2byte(even(frame.palette1))}

.align 128
frame1_palette1_odd:
${byte2byte(odd(frame.palette1))}

.align 128
frame1_palette2_even:
${byte2byte(even(frame.palette2))}

.align 128
frame1_palette2_odd:
${byte2byte(odd(frame.palette2))}

.align 128
frame1_palette3_even:
${byte2byte(even(frame.palette3))}

.align 128
frame1_palette3_odd:
${byte2byte(odd(frame.palette3))}

.align 4096
frame1_bottom:
${byte2byte(frame.bitmap.slice(half))}



.align 1024
displaylist:
    ; 40 lines overscan
    .repeat 4
        .byte $70 ; 8 blank lines
    .endrep
    ; include a DLI to mark us as frame 0
    .byte $f0 ; 8 blank lines

    ; 160 lines graphics
    ; ANTIC mode e (160px 2bpp, 1 scan line per line)
    .byte $4e
    .addr frame1_top
    .repeat ${height / 2 - 1}
        .byte $0e
    .endrep
    .byte $4e
    .addr frame1_bottom
    .repeat ${height / 2 - 1}
        .byte $0e
    .endrep

    .byte $41 ; jump and blank
    .addr displaylist
`;
}

/**
 * Double width and save as an image file
 * @param {number} width 
 * @param {number} height 
 * @param {{output: number[], palette: number[]}} lines 
 * @param {string} dest 
 */
async function saveImage(width, height, lines, dest) {
    let width2 = width * 2;
    let stride = width2 * 4;
    let rgba = new Uint8Array(stride * height);

    for (let y = 0; y < height; y++) {
        let {output, palette} = lines[y];
        for (let x = 0; x < width2; x++) {
            let i = x >> 1;
            if (i >= width) {
                throw new Error('i >= width');
            }
            let rgb = atariRGB[palette[output[i]]];

            rgba[y * stride + x * 4 + 0] = fromLinear(rgb.r);
            rgba[y * stride + x * 4 + 1] = fromLinear(rgb.g);
            rgba[y * stride + x * 4 + 2] = fromLinear(rgb.b);
            rgba[y * stride + x * 4 + 3] = 255;
        }
    }

    let image = await new Promise((resolve, reject) => {
        new Jimp({
            data: rgba,
            width: width2,
            height,
        }, (err, image) => {
            if (err) reject(err);
            resolve(image);
        });
    });
    await image.resize(Math.round(width2 * 2 / 1.2), height * 2);
    await image.writeAsync(dest);
}

async function main() {
    if (process.argv.length < 3) {
        console.error("Usage: node dither-image.js source-image.jpg dest-asm.s");
        process.exit(1);
    }

    let nbits = 2;

    let {width, height, lines} = await convert(process.argv[2], nbits);

    let asm = genAssembly(width, height, nbits, lines);
    writeFileSync(process.argv[3], asm, "utf-8");

    await saveImage(width, height, lines, `${process.argv[3]}.png`);

    process.exit(0);
}

main();