package ip.color;

public class Octree {
    private static final int MAXDEPTH = 7;
    private int numNodes = 0;
    private int maxNodes = 0;
    private int size, level, leafLevel;
    private RGB color_reg[];
    private Node tree = null;
    private Node reduceList[] = new Node[MAXDEPTH + 1];
    private int k;
    private short r[][], g[][], b[][];

    public Octree() {

    }

    public void octreeQuantization(short ra[][], short ga[][], short ba[][], int ki) {
        r = ra;
        g = ga;
        b = ba;
        k = ki;

        setColor();
        reMap(r, g, b);
    }

    public void reMap(short r[][], short g[][], short b[][]) {
        for (int x = 0; x < r.length; x++)
            for (int y = 0; y < r[0].length; y++) {
                RGB c = new RGB();
                c.red = r[x][y];
                c.green = g[x][y];
                c.blue = b[x][y];
                int id = findColor(tree, c);
                c = color_reg[id];
                r[x][y] = c.red;
                g[x][y] = c.blue;
                b[x][y] = c.green;
            }
    }

    public void setColor() {
        RGB color = new RGB();

        color_reg = new RGB[k];
        tree = null;
        size = 0;
        level = MAXDEPTH;
        leafLevel = level + 1;
        for (int y = 0; y < r[0].length; y++) {
            for (int x = 0; x < r.length; x++) {
                color.red = r[x][y];
                color.green = g[x][y];
                color.blue = b[x][y];
                tree = insertNode(tree, color, 0);
                if (size > k)
                    reduceTree();
            }
        }
        int index[] = new int[1];
        index[0] = 0;
        initVGAPalette(tree, index);

    }

    public int findColor(Node tree, RGB color) {
        if (tree.leaf)
            return tree.colorIndex;
        else
            return findColor(tree.link[
                    ((color.red >> (MAXDEPTH - tree.level)) & 1) << 2 |
                    ((color.green >> (MAXDEPTH - tree.level)) & 1) << 1 |
                    (color.blue >> (MAXDEPTH - tree.level)) & 1], color);
    }

    public Node insertNode(Node node, RGB color, int depth) {
        int branch;

        if (node == null) // create new node
        {
            node = new Node();
            numNodes++;
            if (numNodes > maxNodes)
                maxNodes = numNodes;
            node.level = depth;
            node.leaf = (depth >= leafLevel)?true:false;
            if (node.leaf)
                size++;
        }
        node.colorCount++;
        node.RGBSum.r += color.red;
        node.RGBSum.g += color.green;
        node.RGBSum.b += color.blue;
        if (!(node.leaf) && (depth < leafLevel)) {
            branch = ((color.red >> (MAXDEPTH - depth)) & 1) << 2 |
                    ((color.green >> (MAXDEPTH - depth)) & 1) << 1 |
                    (color.blue >> (MAXDEPTH - depth)) & 1;
            if (node.link[branch] == null) {
                node.child++;
                if (node.child == 2) {
                    node.nextReduceable = reduceList[depth];
                    reduceList[depth] = node;
                }
            }
            node.link[branch] = insertNode(node.link[branch], color, depth + 1);
        }

        return node;
    }

    public Node killTree(Node tree) {
        if (tree == null)
            return null;
        for (int i = 0; i < 8; i++)
            tree.link[i] = killTree(tree.link[i]);

        numNodes--;

        return null;
    }

    public void reduceTree() {
        Node node;
        int new_Level;
        int depth;

        new_Level = level;
        while (reduceList[new_Level] == null)
            new_Level--;
        node = reduceList[new_Level];
        reduceList[new_Level] = reduceList[new_Level].nextReduceable;
        node.leaf = true;
        size = size - node.child + 1;
        depth = node.level;
        for (int i = 0; i < 8; i++)
            node.link[i] = killTree(node.link[i]);
        if (depth < level) {
            level = depth;
            leafLevel = level + 1;
        }
    }

    public void initVGAPalette(Node tree, int index[]) {
        if (tree != null) {
            if (tree.leaf || tree.level == leafLevel) {
                if (color_reg[index[0]] == null)
                    color_reg[index[0]] = new RGB();
                color_reg[index[0]].red = (short) (tree.RGBSum.r / tree.colorCount);
                color_reg[index[0]].green = (short) (tree.RGBSum.g / tree.colorCount);
                color_reg[index[0]].blue = (short) (tree.RGBSum.b / tree.colorCount);
                tree.colorIndex = index[0]++;
                tree.leaf = true;
            } else
                for (int octant = 0; octant < 8; octant++)
                    initVGAPalette(tree.link[octant], index);
        }
    }

    public static int getMAXDEPTH() {
        return MAXDEPTH;
    }

    public int getNumNodes() {
        return numNodes;
    }

    public void setNumNodes(int numNodes) {
        this.numNodes = numNodes;
    }

    public int getMaxNodes() {
        return maxNodes;
    }

    public void setMaxNodes(int maxNodes) {
        this.maxNodes = maxNodes;
    }

    public int getSize() {
        return size;
    }

    public void setSize(int size) {
        this.size = size;
    }

    public int getLevel() {
        return level;
    }

    public void setLevel(int level) {
        this.level = level;
    }

    public int getLeafLevel() {
        return leafLevel;
    }

    public void setLeafLevel(int leafLevel) {
        this.leafLevel = leafLevel;
    }

    public Node getTree() {
        return tree;
    }

    public int getK() {
        return k;
    }

    public short[][] getR() {
        return r;
    }

    public short[][] getG() {
        return g;
    }

    public short[][] getB() {
        return b;
    }

    public RGB[] getColor_reg() {
        return color_reg;
    }

    public Node[] getReduceList() {
        return reduceList;
    }
}

class ColorSum {
    public long r,g,b;

    public ColorSum() {
        r = g = b = 0;
    }
}

class RGB {
    public short red, green, blue;
}

class Node {
    boolean leaf = false;
    int level = 0;
    int colorIndex = 0;
    int child = 0;
    long colorCount = 0;
    ColorSum RGBSum = new ColorSum();
    Node nextReduceable = null;
    Node link[] = new Node[8];

    public Node() {
        for (int i = 0; i < 8; i++)
            link[i] = null;
    }
}