const { PNG } = require('pngjs');
const fs = require('fs');

async function segmentationMask(input1, input2, output, threshold = 0.02) {
    let img1 = await loadPngFile(input1);
    let img2 = await loadPngFile(input2);

    if (img1.width !== img2.width || img1.height !== img2.height) {
        throw new Error("Cannot compute mask on images with different sizes");
    }

    for (let i = 0; i < img1.data.length; i += 4) {
        let r1 = img1.data[i + 0] / 255;
        let g1 = img1.data[i + 1] / 255;
        let b1 = img1.data[i + 2] / 255;

        let r2 = img2.data[i + 0] / 255;
        let g2 = img2.data[i + 1] / 255;
        let b2 = img2.data[i + 2] / 255;

        // Test difference
        let difference = Math.sqrt((r1 - r2) * (r1 - r2) + (g1 - g2) * (g1 - g2) + (b1 - b2) * (b1 - b2));
        let pixelsAreDifferent = difference > threshold;

        if (pixelsAreDifferent) {
            img1.data[i + 0] = 255;
            img1.data[i + 1] = 255;
            img1.data[i + 2] = 255;
        } else {
            img1.data[i + 0] = 0;
            img1.data[i + 1] = 0;
            img1.data[i + 2] = 0;
        }
    }

    let outputStream = fs.createWriteStream(output);
    await img1.pack().pipe(outputStream);
}

// The following is fixed from https://www.npmjs.com/package/png-quality

async function loadPngFile(pathOrBuffer) {
    // Load buffer of path
    if (!(pathOrBuffer instanceof Buffer)) {
        pathOrBuffer = await new Promise((resolve, reject) => {
            fs.readFile(pathOrBuffer, (err, data) => {
                if (err) return reject(err)
                resolve(data)
            });
        })
    }

    // Load PNG from buffer
    return await new Promise((resolve, reject) => {
        const png = new PNG()
        png.parse(pathOrBuffer, err => {
            if (err) return reject(err)
            resolve(png)
        })
    })
}

async function mse(png1, png2) {
    const pngs = [png1, png2];
    for (let i in pngs) {
        if (!(pngs[i] instanceof PNG)) pngs[i] = await loadPngFile(pngs[i])
    }

    if (pngs[0].width !== pngs[1].width || pngs[0].height !== pngs[1].height) {
        throw new Error('Width or height does not equal')
    }

    if (pngs[0].data.length !== pngs[1].data.length) {
        throw new Error('Data buffer length does not equal')
    }

    const square = (a) => a * a,
        channelIndex = [0, 1, 2],
        channelMax = 255 * 255,
        area = pngs[0].width * pngs[1].height

    let mse = 0
    for (let i = 0; i < pngs[0].data.length; i += 4) {
        const rgbas = pngs.map(png => png.data.slice(i, i + 4))
        const rgbs = rgbas.map(rgba => channelIndex.map(i => rgba[i] * rgba[3]))
        channelIndex.forEach(i => mse += square(rgbs[0][i] - rgbs[1][i]))
    }

    return mse / 3.0 / (channelMax * channelMax) / area
}

async function psnr(png1, png2) {
    const m = await mse(png1, png2)
    return 10 * Math.log10(1 / m)
}

module.exports = {
    loadPngFile,
    mse,
    psnr,
    segmentationMask,
};