import {
    writeFileSync
} from 'fs';

import Jimp from 'Jimp';


function repeat(val, n) {
    let arr = [];
    for (let i = 0; i < n; i++) {
        arr.push(val);
    }
    return arr;
}

function zeroes(n) {
    return repeat(0, n);
}

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;
}

function fromSRGB(val) {
    val /= 255;
    if (val <= 0.04045) {
        val /= 12.92;
    } else {
        val = ((val + 0.055) / 1.055) ** 2.4;
    }
    val *= 255;
    return val;
}

function toSRGB(val) {
    val /= 255;
    if (val <= 0.0031308) {
        val *= 12.92;
    } else {
        val = (val * 1.055) ** (1.0 / 2.4) - 0.055;
    }
    val *= 255;
    return val;
}

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);
    }

    static fromGTIA(val) {
        // This seems off from what Atari800 does
        // https://forums.atariage.com/topic/107853-need-the-256-colors/page/2/#comment-1312467
        let cr = (val >> 4) & 15;
        let lm = val & 15;
        let crlv = cr ? 50 : 0;

        /*
        let phase = ((cr - 1) * 25 - 58) * (2 * Math.PI / 360);

        let y = 255 * (lm + 1) / 16;
        let i = crlv * Math.cos(phase);
        let q = crlv * Math.sin(phase);

        let r = y + 0.956 * i + 0.621 * q;
        let g = y - 0.272 * i - 0.647 * q;
        let b = y - 1.107 * i + 1.704 * q;
        */

        // PAL
        let phase = ((cr - 1) * 25.7 - 15) * (2 * Math.PI / 360);

        let y = 255 * (lm + 1) / 16;
        let i = crlv * Math.cos(phase);
        let q = crlv * Math.sin(phase);

        let r = y + 0.956 * i + 0.621 * q;
        let g = y - 0.272 * i - 0.647 * q;
        let b = y - 1.107 * i + 1.704 * q;

        return new RGB(r, g, b).clamp().fromSRGB();
    }

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

    fromNTSC() {
        return this.map(toLinear);
    }

    toNTSC() {
        return this.map(fromLinear);
    }

    fromSRGB() {
        return this.map(fromSRGB);
    }

    toSRGB() {
        return this.map(toSRGB);
    }

    clamp() {
        return this.map((val) => {
            if (val < 0) return 0;
            if (val > 255) return 255;
            return val;
        });
    }

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

    add(other) {
        return new RGB(
            this.r + other.r,
            this.g + other.g,
            this.b + other.b
        );
    }

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

    multiply(scalar) {
        return new RGB(
            this.r * scalar,
            this.g * scalar,
            this.b * scalar,
        );
    }

    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;
    }

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

    lumaScale() {
        return new RGB(
            this.r * 0.299,
            this.g * 0.586,
            this.b * 0.114
        );
    }

    luma() {
        return this.lumaScale().sum();
    }
}

let atariRGB = [];
for (let i = 0; i < 256; i++) {
    atariRGB[i] = RGB.fromGTIA(i);
}



/**
 * Dither RGB input data and pick a monochrome palette.
 * @param {RGB[]} input source scanline data, in linear RGB
 * @returns {{output: number[], palette: number[], error: RGB[]}}
 */
function colorize(input) {
    
    let width = input.length;

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

    // Apply dithering with given palette and collect color usage stats
    let dither = (palette) => {
        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 = zeroes(width);
        let distance = 0;
        let shortest = Infinity;

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

        // Try dithering with this palette.
        for (let x = 0; x < width; x++) {
            let rgb = inputPixel(x, error);
            let luma = Math.round(rgb.luma() * 15 / 255);
            let pick = luma;

            let diff = rgb.difference(atariRGB[palette[pick]]);
            let dist = diff.magnitude();
            nextError = diff;
            shortest = dist;
            //shortest = Math.min(shortest, dist);

            output[x] = pick;
            distance += shortest;
            //distance = shortest;

            let share = (n) => nextError.multiply(n / 16);

            error.cur[x + 1]?.inc(share(7));
            error.next[x - 1]?.inc(share(3));
            error.next[x    ]?.inc(share(5));
            error.next[x + 1]?.inc(share(1));
        }
        return {
            output,
            palette,
            distance,
            shortest,
            error: error.next
        };
    };

    let best = null;
    for (let hue = 0; hue < 16; hue++) {
        let palette = [];
        for (let i = 0; i < 16; i++) {
            palette[i] = (hue << 4) | i;
        }
        let variant = dither(palette);
        if (!best || variant.distance < best.distance) {
            best = variant;
        }
    }
    return best;
}

/**
 * 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;

    let aspect = width / height;
    let dar = 4 / 1.2;
    if (aspect > ((320 / 1.2) / 192)) {
        // wide
        width = 80;
        height = Math.round((width * image.bitmap.height / image.bitmap.width) * dar);
        if (height & 1) {
            height++;
        }
    } else {
        // tall
        height = 192;
        width = Math.round((height * image.bitmap.width / image.bitmap.height) / dar);
        if (width & 1) {
            width++;
        }
    }
    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]
        ).fromSRGB());
    }
    return input;
}

/**
 * Read an image file, squish to 80px if necessary,
 * and dither to 16 monochrome levels 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) {

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

    if (width > 80) {
        throw new Error(`expected <80px width, got ${width} pixels`);
    }

    if (height > 192) {
        throw new Error(`expected <192px 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');
    }

    let left = [], right = [];
    let padding = 0;
    if (width < 80) {
        padding = 80 - width;

        let black = new RGB(0, 0, 0);
        left = repeat(black, padding >> 1);
        right = repeat(black, padding + 1 >> 1);
    }

    let lines = [];
    for (let y = 0; y < height; y++) {
        let inputLine = input
            .slice(y * width, (y + 1) * width);

        if (padding) {
            inputLine = left.concat(inputLine, right);
        }
        if (y > 0) {
            let error = lines[y - 1].error;
            inputLine = inputLine.map((rgb, x) => rgb.add(error[x]).clamp());
        }
        let line = colorize(inputLine);
        lines.push(line);
    }
    return {
        width: width + padding,
        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),
        bitmap: new Uint8Array(stride * height),
    };
    for (let y = 0; y < height; y++) {
        let base = 0;
        frame.palette1[y] = lines[y + base].palette[1];
        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 displaylist

.segment "BUFFERS"

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

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

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

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




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

    ; ${height} lines graphics
    ; ANTIC mode f (320px 1bpp, 1 scan line per line)
    ; this is modified into GTIA grayscale mode
    .byte $4f
    .addr frame1_top
    .repeat ${height / 2 - 1}
        .byte $0f
    .endrep
    .byte $4f
    .addr frame1_bottom
    .repeat ${height / 2 - 1}
        .byte $0f
    .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]]].fromLinear();
            let rgb = atariRGB[palette[output[i]]].toSRGB();

            rgba[y * stride + x * 4 + 0] = rgb.r;
            rgba[y * stride + x * 4 + 1] = rgb.g;
            rgba[y * stride + x * 4 + 2] = 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 * 4), 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 = 4;

    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();