import java.awt.image.BufferedImage;
import java.io.File;
import java.util.*;
import javax.imageio.ImageIO;

/**
 * Resolves which decompiled generator body produces which official palette.
 *
 * The Ghidra listing shows 13 generator bodies (cases 0..10, 0xc, 0xd) but the
 * official app exposes 12 palettes, so the case->index mapping cannot be read
 * off the listing directly.  This tool identifies each body by matching it
 * against the palette preview image the official app ships
 * (res/mipmap-hdpi-v4/palette_<name>.png) and reporting the best assignment.
 *
 * Uses a 32x32x32 quantised lookup so the full 13 x 12 x 2 comparison is fast,
 * and reports both byte-order conventions.
 */
public class PalIdentify {
    static final int Q = 32;                  // quantisation levels per channel
    static final int STEP = 256 / Q;

    public static void main(String[] args) throws Exception {
        String apkRes = args[0];
        String[] names = {"white_hot","black_hot","iron_bow","rain_bow","glow_bow",
            "autumn","winter","hot_metal","jet","red_saturation","high_contrast","red_hot"};

        // build all 13 bodies in both conventions
        int[][] direct = new int[14][];
        int[][] swapped = new int[14][];
        for (int c : new int[]{0,1,2,3,4,5,6,7,8,9,10,12,13}) {
            byte[][] buf = PalBody.run(c);
            direct[c] = emit(buf, false);
            swapped[c] = emit(buf, true);
        }

        // preview histograms
        double[][] hist = new double[12][Q * Q * Q];
        for (int i = 0; i < 12; i++) {
            File f = new File(apkRes, "palette_" + names[i] + ".png");
            BufferedImage img = ImageIO.read(f);
            int n = 0;
            for (int y = 0; y < img.getHeight(); y++)
                for (int x = 0; x < img.getWidth(); x++) {
                    int a = img.getRGB(x, y);
                    if ((a >>> 24) != 0xFF) continue;
                    hist[i][qidx(a)]++;
                    n++;
                }
            for (int k = 0; k < hist[i].length; k++) hist[i][k] /= n;
        }

        System.out.println("=== best-matching generator body per official palette ===");
        System.out.printf("%-16s %-28s %-28s%n", "palette", "direct", "swapped");
        int[] bestBodyD = new int[12], bestBodyS = new int[12];
        for (int i = 0; i < 12; i++) {
            double[] dd = new double[14], ds = new double[14];
            int bd = -1, bs = -1;
            for (int c : new int[]{0,1,2,3,4,5,6,7,8,9,10,12,13}) {
                dd[c] = dist(hist[i], direct[c]);
                ds[c] = dist(hist[i], swapped[c]);
                if (bd < 0 || dd[c] < dd[bd]) bd = c;
                if (bs < 0 || ds[c] < ds[bs]) bs = c;
            }
            bestBodyD[i] = bd; bestBodyS[i] = bs;
            System.out.printf("%-16s %-28s %-28s%n", names[i],
                String.format("case %-4d d=%.2f", bd, dd[bd]),
                String.format("case %-4d d=%.2f", bs, ds[bs]));
        }

        System.out.println("\n=== body -> palettes it best explains (direct) ===");
        for (int c : new int[]{0,1,2,3,4,5,6,7,8,9,10,12,13}) {
            double best = Double.MAX_VALUE; String who = "-";
            for (int i = 0; i < 12; i++) {
                double d = dist(hist[i], direct[c]);
                if (d < best) { best = d; who = names[i]; }
            }
            System.out.printf("case %-3d -> %-16s (d=%.2f)%n", c, who, best);
        }
    }

    static int qidx(int argb) {
        int r = ((argb >> 16) & 0xFF) / STEP;
        int g = ((argb >> 8) & 0xFF) / STEP;
        int b = (argb & 0xFF) / STEP;
        return (r * Q + g) * Q + b;
    }

    static double dist(double[] hist, int[] pal) {
        double[] lut = new double[Q * Q * Q];
        for (int r = 0; r < Q; r++)
            for (int g = 0; g < Q; g++)
                for (int b = 0; b < Q; b++) {
                    int rr = r * STEP, gg = g * STEP, bb = b * STEP;
                    int best = Integer.MAX_VALUE;
                    for (int v : pal) {
                        int dr = rr - ((v >> 16) & 0xFF);
                        int dg = gg - ((v >> 8) & 0xFF);
                        int db = bb - (v & 0xFF);
                        int d2 = dr * dr + dg * dg + db * db;
                        if (d2 < best) best = d2;
                    }
                    lut[(r * Q + g) * Q + b] = Math.sqrt(best);
                }
        double sum = 0;
        for (int k = 0; k < hist.length; k++) if (hist[k] != 0) sum += hist[k] * lut[k];
        return sum;
    }

    static int[] emit(byte[][] buf, boolean swap) {
        int[] out = new int[256];
        for (int i = 0; i < 256; i++) {
            int b0 = buf[i][swap ? 2 : 0] & 0xFF;
            int b1 = buf[i][1] & 0xFF;
            int b2 = buf[i][swap ? 0 : 2] & 0xFF;
            out[i] = 0xFF000000 | (b2 << 16) | (b1 << 8) | b0;
        }
        return out;
    }
}

/** Raw generator bodies: returns the 256x4 byte buffer each case produces. */
class PalBody {
    static final int BASE = 0x20cc;
    static byte[][] buf;
    static int[] writes;

    static byte[][] run(int c) {
        buf = new byte[256][4];
        writes = new int[256];
        switch (c) {
            case 0: {
                int v = 0, a = BASE;
                do { st(a+1,v); st(a+2,v); st(a,v); st(a+3,0); a+=4; v++; } while (v != 0x100);
                break;
            }
            case 1: {
                int v = 0x100, a = BASE;
                do { st(a+3,0); v--; st(a+1,v); st(a+2,v); st(a,v); a+=4; } while (v != 0);
                break;
            }
            case 2: {
                for (int k = 0; k < 22; k++) { int a=BASE+k*4; st(a+1,0); st(a+2,0); st(a,u(0x75*k,0x16)); st(a+3,0); }
                for (int k = 0; k < 71; k++) { int a=BASE+(22+k)*4; st(a+1,0); st(a+3,0); st(a+2,u(0xbc*k,0x47)); st(a,u(u(0x25*k,0x47)+0x75)); }
                for (int k = 0; k < 29; k++) { int a=BASE+(93+k)*4; st(a+1,u(0xff*k,0x94)); st(a+3,0); st(a+2,u(u(0x27*k,0x1d)-0x44)); st(a,u(s(-0x3b*k,0x1d)-0x66)); }
                for (int k = 0; k < 18; k++) { int a=BASE+(122+k)*4; st(a+1,u(0x1ce3+0xff*k,0x94)); st(a+3,0); st(a+2,u(u(0x1c*k,0x1e)-0x1d)); st(a,u(s(-0x50*k,0x12)+0x5f)); }
                for (int k = 0; k < 12; k++) { int a=BASE+(140+k)*4; st(a+1,u(0x2ed1+0xff*k,0x94)); st(a+3,0); st(a+2,u(u(0x1f8+0x1c*k,0x1e)-0x1d)); st(a,0x0f); }
                for (int k = 0; k < 71; k++) { int a=BASE+(152+k)*4; st(a+1,u(0x3ac5+0xff*k,0x94)); st(a+3,0); st(a+2,0xff); st(a,0x0f); }
                for (int k = 0; k < 18; k++) { int a=BASE+(223+k)*4; st(a+1,u(0x817e+0xff*k,0x94)); st(a+3,0); st(a+2,0xff); st(a,u((0xf0*k>>5)+0x0f)); }
                for (int k = 0; k < 15; k++) { int a=BASE+(241+k)*4; st(a+1,0xff); st(a+3,0); st(a+2,0xff); st(a,u(((0xf0e0+0xf0*k)>>5)+0x0f)); }
                break;
            }
            case 3: {
                int a=BASE,i,u4=0,u10=0;
                for (i=0x40;i!=0;i--){ st(a+2,0); st(a+3,0); st(a+1,u(u10>>7)); st(a,u((u4>>6)+0x32)); u4+=0xcd; u10+=0xff; a+=4; }
                a=BASE+0x100; u4=0; int i5=0;
                for (i=0x40;i!=0;i--){ st(a+3,0); st(a+2,u(u4>>6)); st(a+1,u((u4+0x3fc0)>>7)); st(a,u(s(i5,0x40)-1)); a+=4; u4+=0xff; i5+=-0xff; }
                a=BASE+0x200; i5=0; u4=0;
                for (i=0x40;i!=0;i--){ st(a+2,0xff); st(a+3,0); st(a,u(u4>>7)); st(a+1,u(s(i5,0x40)-1)); i5+=-0xff; u4+=0xff; a+=4; }
                int t=BASE+0x300; u4=0xbfc0;
                for (i=0x40;i!=0;i--){ st(t+2,0xff); st(t+3,0); st(t,u(u4>>7)); st(t+1,u((u4+0x40)>>6)); u4+=0xff; t+=4; }
                break;
            }
            case 4: {
                int a=BASE,i;
                for (i=-0x32;i!=0;i++){ st(a+2,0x32); st(a,0); st(a+3,0); st(a+1,u(i+0x32)); a+=4; }
                a=BASE+0xc8; int i5=-0x96, i11=0;
                for (;i5!=0;i5++){ st(a,0); st(a+3,0); st(a+1,u(i5-0x38)); st(a+2,u(u(i11,0x96)+0x32)); a+=4; i11+=0xcd; }
                int c2=-0xe0, cv=0xc8;
                do { st(c2+0x24cd,cv); st(c2+0x24ce,0xff); st(c2+0x24cc,0); st(c2+0x24cf,0); c2+=4; cv++; } while (c2 != 0);
                break;
            }
            case 5: {
                int a=BASE,i,u4=0,u10=0,u13=0;
                for (i=0x40;i!=0;i--){ st(a+3,0); st(a+1,u(u10>>6)); st(a+2,u(u13>>6)); st(a,u(u4>>6)); u4+=0x24; u10+=0x16; u13+=0xcb; a+=4; }
                a=BASE+0x100; u4=0; u10=0; int i5=0;
                for (i=0x40;i!=0;i--){ st(a+3,0); st(a+1,u((u10>>6)+0x16)); st(a+2,u(u(u4>>6)-0x35)); st(a,u(s(i5,0x40)+0x24)); a+=4; u4+=0x34; u10+=0x44; i5+=-0x24; }
                a=BASE+0x200; u4=0;
                for (i=0x40;i!=0;i--){ st(a+2,0xff); st(a,0); st(a+3,0); st(a+1,u((u4>>6)+0x5a)); u4+=0x7c; a+=4; }
                a=BASE+0x300; u4=0; u10=0;
                for (i=0x40;i!=0;i--){ st(a+2,0xff); st(a+3,0); st(a,u(u10>>6)); st(a+1,u(u(u4>>6)-0x2a)); u4+=0x29; u10+=0xff; a+=4; }
                break;
            }
            case 6: {
                int a=BASE, i5=0, i11=0xff;
                do { st(a+1,i11); st(a,i11); st(a+3,0); st(a+2,u(s(i5,100)-1)); a+=4; i5+=-0xff; i11--; } while (i11 != 0x9b);
                a=BASE+0x190; i5=0x9b; i11=0;
                do { st(a+1,i5); st(a,i5); st(a+3,0); st(a+2,u(u(i11,0x69))); a+=4; i11+=0xff; i5--; } while (i5 != 0x32);
                int t=BASE+0x334; i11=0x33;
                do { st(t+2,0xff); st(t+3,0); i11--; st(t+1,i11); st(t,i11); t+=4; } while (i11 != 0);
                break;
            }
            case 7: {
                int a=BASE,i,cv=0;
                for (i=0x80;i!=0;i--){ st(a+1,0); st(a+2,cv); st(a,0); st(a+3,0); a+=4; cv+=2; }
                int t=BASE+0x200; int i5=0x80, i11=0;
                do { st(t+2,0xff); st(t+3,0); st(t+1,u(u(i11,0x7f))); st(t,u(u(i11,0x7f))); i11+=0xff; t+=4; i5--; } while (i5 != 0);
                break;
            }
            case 8: {
                int i11=-0x80;
                do { st(i11+0x214d,0); st(i11+0x214e,0); st(i11+0x214c,u(i11-1)); st(i11+0x214f,0); i11+=4; } while (i11 != 0);
                int a=BASE+0x80, u4=0, i;
                for (i=0x40;i!=0;i--){ st(a+2,0); st(a,0xff); st(a+3,0); st(a+1,u(u4>>6)); u4+=0xff; a+=4; }
                a=BASE+0x180; u4=0; int i1=0;
                for (i=0x40;i!=0;i--){ st(a+1,0xff); st(a+3,0); st(a+2,u(u4>>6)); st(a,u(s(i1,0x40)-1)); a+=4; u4+=0xff; i1+=-0xff; }
                a=BASE+0x280; int i5j=0x40; int j=0;
                do { st(a+2,0xff); st(a,0); st(a+3,0); st(a+1,u(s(j,0x40)-1)); j+=-0xff; a+=4; i5j--; } while (i5j != 0);
                int k=-0x80, m=0;
                do { st(k+0x24cd,0); st(k+0x24cc,0); st(k+0x24cf,0); st(k+0x24ce,u(s(m,0x20)-1)); m+=-0x80; k+=4; } while (k != 0);
                break;
            }
            case 9: {
                int a=BASE, i11=0, i5=0x99, cv;
                do { st(a+3,0); cv=u(u(i11,200)); st(a+1,cv); st(a+2,cv); st(a,cv); a+=4; i11+=0xff; i5--; } while (i5 != 0);
                int t=-0xbc, uu=0x9867;
                do { st(t+0x23ed,0xc3); st(t+0x23ee,u(u(uu,200))); st(t+0x23ec,0xc3); st(t+0x23ef,0); uu+=0xff; t+=4; } while (t != 0);
                int c2=-0xe0, v=0;
                do { st(c2+0x24ce,0xff); st(c2+0x24cf,0); int g=u(s(v,0x37)-0x3d); st(c2+0x24cd,g); st(c2+0x24cc,g); v+=-0xc3; c2+=4; } while (c2 != 0);
                break;
            }
            case 10: {
                int a=BASE, i11=0;
                do { st(a+1,0); st(a+2,i11); st(a,i11); st(a+3,0); a+=4; i11+=7; } while (i11 != 0xfc);
                int cv=0xfc, o=-0x94, u4=0;
                do { st(o+0x21f1,0); st(o+0x21f2,cv); st(o+0x21f3,0); st(o+0x21f0,u(-4-(u4>>2))); u4+=7; cv+=-7; o+=4; } while (o != 0);
                int p=BASE+0x124; int cv1=0; int u4b=0xf5; boolean more;
                do { cv1+=7; st(p+1,cv1); st(p+2,0); st(p+3,0); st(p,u(-4-(u4b>>2))); p+=4; more=u4b!=0; u4b-=7; } while (more);
                int q=BASE+0x1b4; int u4c=0; int i11b=0xfc;
                do { st(q+2,0); st(q,i11b); st(q+3,0); st(q+1,u(-4-(u4c>>1))); u4c+=7; q+=4; more=i11b!=0; i11b-=7; } while (more);
                int r=BASE+0x248; int cvr=0x7; int u4d=0xf5;
                do { st(r+2,cvr); st(r,0); st(r+3,0); st(r+1,u(-4-(u4d>>1))); r+=4; cvr+=7; more=u4d!=0; u4d-=7; } while (more);
                int u4e=0; int s=-0x94; int cvs=0xfc;
                do { st(s+0x2439,cvs); st(s+0x2438,0); st(s+0x243b,0); st(s+0x243a,u(-4-(u4e>>2))); u4e+=7; cvs+=-7; s+=4; } while (s != 0);
                int d=BASE+0x36c; int u4f=0xfc; int cv6=0;
                do { st(d+1,cv6); st(d,cv6); st(d+3,0); st(d+2,u(-4-(u4f>>2))); d+=4; cv6+=7; more=u4f!=0; u4f-=7; } while (more);
                break;
            }
            case 12: {   // case 0xc
                int a=BASE, i11=0;
                do { st(a+1,0); st(a+2,0); st(a,i11); st(a+3,0); a+=4; i11+=7; } while (i11 != 0xfc);
                int o=-0x94, cv=0;
                do { st(o+0x21f1,cv); st(o+0x21f2,0); st(o+0x21f0,0xff); st(o+0x21f3,0); cv+=7; o+=4; } while (o != 0);
                int p=BASE+0x124; int i11c=0;
                do { st(p+1,0xff); st(p+2,0); st(p+3,0); st(p,u(i11c-0xb)); i11c+=-7; p+=4; } while (i11c != -0xfc);
                int o2=-0x94, cv2=0;
                do { st(o2+0x2315,0xff); st(o2+0x2316,cv2); st(o2+0x2314,0); st(o2+0x2317,0); cv2+=7; o2+=4; } while (o2 != 0);
                int p2=BASE+0x248; int i11d=0;
                do { st(p2+2,0xff); st(p2,0); st(p2+3,0); st(p2+1,u(i11d-0xb)); i11d+=-7; p2+=4; } while (i11d != -0xfc);
                int o3=-0x94, cv3=0;
                do { st(o3+0x2439,0); st(o3+0x243a,0xff); st(o3+0x2438,cv3); st(o3+0x243b,0); cv3+=7; o3+=4; } while (o3 != 0);
                int o4=-0x94, cv4=0;
                do { st(o4+0x24cd,cv4); st(o4+0x24ce,0xff); st(o4+0x24cc,0xff); st(o4+0x24cf,0); cv4+=7; o4+=4; } while (o4 != 0);
                break;
            }
            case 13: {   // case 0xd
                int a=BASE, i5=0, i;
                for (i=0x40;i!=0;i--){ st(a,0xff); st(a+1,u(i5)); st(a+2,0); st(a+3,0); i5+=4; a+=4; }
                int b=BASE+0x100; int cv=0xfc;
                for (i=0x40;i!=0;i--){ st(b+1,0xff); st(b+2,0); st(b,cv); st(b+3,0); cv-=4; b+=4; }
                int c2=BASE+0x200; int cv1=0, cv6=0xfc;
                for (i=0x40;i!=0;i--){ st(c2+1,0xff); st(c2+2,cv1); st(c2,0); st(c2+3,0); cv1+=4; c2+=4; }
                int d=BASE+0x300;
                for (i=0x40;i!=0;i--){ st(d+1,cv6); st(d+2,0xff); st(d,0); st(d+3,0); cv6-=4; d+=4; }
                break;
            }
        }
        for (int i = 0; i < 256; i++) {
            if (writes[i] != 1) throw new IllegalStateException(
                "body " + c + " entry " + i + " written " + writes[i] + " times");
        }
        return buf;
    }

    static void st(int addr, int v) {
        int rel = addr - BASE, idx = rel / 4, off = rel % 4;
        if (idx < 0 || idx > 255) throw new IllegalStateException("oob " + Integer.toHexString(addr));
        buf[idx][off] = (byte) v;
        if (off == 0) writes[idx]++;
    }
    static int u(int v) { return v & 0xFF; }
    static int u(int a, int b) { return Integer.divideUnsigned(a, b); }
    static int s(int a, int b) { return a / b; }
}