//
//  main.cpp
//  unique_colors
//
//  Created by Larry Bank on 7/23/25.
//

#include "PNGdec.h"
//
// A table to accelerate the testing of 2-bit images for the number
// of unique colors. Each entry sets bits 0-3 depending on the presence
// of colors 0-3 in each 2-bit pixel
//
const uint8_t ucTwoBitFlags[256] = {
0x01,0x03,0x05,0x09,0x03,0x03,0x07,0x0b,0x05,0x07,0x05,0x0d,0x09,0x0b,0x0d,0x09,
0x03,0x03,0x07,0x0b,0x03,0x03,0x07,0x0b,0x07,0x07,0x07,0x0f,0x0b,0x0b,0x0f,0x0b,
0x05,0x07,0x05,0x0d,0x07,0x07,0x07,0x0f,0x05,0x07,0x05,0x0d,0x0d,0x0f,0x0d,0x0d,
0x09,0x0b,0x0d,0x09,0x0b,0x0b,0x0f,0x0b,0x0d,0x0f,0x0d,0x0d,0x09,0x0b,0x0d,0x09,
0x03,0x03,0x07,0x0b,0x03,0x03,0x07,0x0b,0x07,0x07,0x07,0x0f,0x0b,0x0b,0x0f,0x0b,
0x03,0x03,0x07,0x0b,0x03,0x02,0x06,0x0a,0x07,0x06,0x06,0x0e,0x0b,0x0a,0x0e,0x0a,
0x07,0x07,0x07,0x0f,0x07,0x06,0x06,0x0e,0x07,0x06,0x06,0x0e,0x0f,0x0e,0x0e,0x0e,
0x0b,0x0b,0x0f,0x0b,0x0b,0x0a,0x0e,0x0a,0x0f,0x0e,0x0e,0x0e,0x0b,0x0a,0x0e,0x0a,
0x05,0x07,0x05,0x0d,0x07,0x07,0x07,0x0f,0x05,0x07,0x05,0x0d,0x0d,0x0f,0x0d,0x0d,
0x07,0x07,0x07,0x0f,0x07,0x06,0x06,0x0e,0x07,0x06,0x06,0x0e,0x0f,0x0e,0x0e,0x0e,
0x05,0x07,0x05,0x0d,0x07,0x06,0x06,0x0e,0x05,0x06,0x04,0x0c,0x0d,0x0e,0x0c,0x0c,
0x0d,0x0f,0x0d,0x0d,0x0f,0x0e,0x0e,0x0e,0x0d,0x0e,0x0c,0x0c,0x0d,0x0e,0x0c,0x0c,
0x09,0x0b,0x0d,0x09,0x0b,0x0b,0x0f,0x0b,0x0d,0x0f,0x0d,0x0d,0x09,0x0b,0x0d,0x09,
0x0b,0x0b,0x0f,0x0b,0x0b,0x0a,0x0e,0x0a,0x0f,0x0e,0x0e,0x0e,0x0b,0x0a,0x0e,0x0a,
0x0d,0x0f,0x0d,0x0d,0x0f,0x0e,0x0e,0x0e,0x0d,0x0e,0x0c,0x0c,0x0d,0x0e,0x0c,0x0c,
0x09,0x0b,0x0d,0x09,0x0b,0x0a,0x0e,0x0a,0x0d,0x0e,0x0c,0x0c,0x09,0x0a,0x0c,0x08
};

PNG png; // static instance of class

//
// This function gets called for each line decoded in the image
// It returns 1 to continue or 0 to abort the decode.
// For 2-bpp images, as soon as the number of unique colors
// passes 2, it aborts the decoding to save time
// For 8-bpp images, as soon as the number of unique colors
// passes 16, it aborts
//
int PNGDraw(PNGDRAW *pDraw)
{
uint8_t c, *s;
int x;
uint64_t u64;

    u64 = *(uint64_t *)pDraw->pUser;
    s = (uint8_t *)pDraw->pPixels;
    if (pDraw->iBpp == 2) { // This case needs to be the fastest for TRMNL
        for (x=0; x<pDraw->iWidth; x+=4) {
            u64 |= ucTwoBitFlags[*s++]; // do 4 pixels at a time
        }
        *(uint64_t *)pDraw->pUser = u64;
        if (__builtin_popcount(u64) >= 3) printf("Aborting early at line %d\n", pDraw->y);
        return (__builtin_popcount(u64) < 3); // Abort early if we pass 2 unique colors
    } else if (pDraw->iBpp == 8) { // check if more than 16 colors are used
        for (x=0; x<pDraw->iWidth; x++) {
            c = *s++;
            if (c > 63) c = 63; // we only have a 64-bit vector
            u64 |= (1 << c);
        }
        *(uint64_t *)pDraw->pUser = u64;
        return (__builtin_popcount(u64) < 17); // Abort early if we pass 16 unique colors
    }
    return 0; // error
} /* PNGDraw() */

int main(int argc, const char * argv[]) {
    uint64_t iColors = 0;
    int rc = 0;
    uint8_t *pData;
    int iDataSize;
    FILE *ihandle;
   
    if (argc != 2) {
       printf("Usage: unique_colors <infile.png>\n");
       return 0;
    }
    ihandle = fopen(argv[1],"rb"); // open input file
    if (ihandle == NULL)
    {
        fprintf(stderr, "Unable to open file: %s\n", argv[1]);
        return 0; // bad filename passed
    }
    fseek(ihandle, 0L, SEEK_END); // get the file size
    iDataSize = (int)ftell(ihandle);
    fseek(ihandle, 0, SEEK_SET);
    pData = (uint8_t *)malloc(iDataSize);
    fread(pData, 1, iDataSize, ihandle);
    fclose(ihandle);
    
    rc = png.openRAM(pData, iDataSize, PNGDraw);
    if (rc == PNG_SUCCESS) {
//        printf("image specs: (%d x %d), %d bpp, pixel type: %d\n", png.getWidth(), png.getHeight(), png.getBpp(), png.getPixelType());
        rc = png.decode((void *)&iColors, 0);
        png.close();
    }
    free(pData);
    iColors = __builtin_popcount(iColors);
    printf("Unique colors: %d\n", (int)iColors); // DEBUG
    return iColors; // return the number of unique colors found
} /* main() */