/* === C R E D I T S  &  D I S C L A I M E R S ==============
 * Permission is given by the author to freely redistribute and include
 * this code in any program as long as this credit is given where due.
 *
 * CQuantizer (c)  1996-1997 Jeff Prosise
 *
 * 31/08/2003 Davide Pizzolato - www.xdp.it
 * - fixed minor bug in ProcessImage when bpp<=8
 * - better color reduction to less than 16 colors
 *
 * COVERED CODE IS PROVIDED UNDER THIS LICENSE ON AN "AS IS" BASIS, WITHOUT
 * WARRANTY OF ANY KIND, EITHER EXPRESSED OR IMPLIED, INCLUDING, WITHOUT
 * LIMITATION, WARRANTIES THAT THE COVERED CODE IS FREE OF DEFECTS,
 * MERCHANTABLE, FIT FOR A PARTICULAR PURPOSE OR NON-INFRINGING. THE ENTIRE
 * RISK AS TO THE QUALITY AND PERFORMANCE OF THE COVERED CODE IS WITH YOU.
 * SHOULD ANY COVERED CODE PROVE DEFECTIVE IN ANY RESPECT, YOU (NOT THE INITIAL
 * DEVELOPER OR ANY OTHER CONTRIBUTOR) ASSUME THE COST OF ANY NECESSARY
 * SERVICING, REPAIR OR CORRECTION. THIS DISCLAIMER OF WARRANTY CONSTITUTES AN
 * ESSENTIAL PART OF THIS LICENSE. NO USE OF ANY COVERED CODE IS AUTHORIZED
 * HEREUNDER EXCEPT UNDER THIS DISCLAIMER.
 *
 * Use at your own risk!
 * ==========================================================
 *
 * Modified for use with Haiku by David Powell & Stephan Aßmus.
 */
#include "ColorQuantizer.h"

#include <stddef.h>
#include <stdlib.h>
#include <string.h>


static inline uint8
clip(float component)
{
        if (component > 255.0)
                return 255;

        return (uint8)(component + 0.5f);
}

struct BColorQuantizer::Node {
        bool                    isLeaf;         // TRUE if node has no children
        uint32                  pixelCount;     // Number of pixels represented by this leaf
        uint32                  sumR;           // Sum of red components
        uint32                  sumG;           // Sum of green components
        uint32                  sumB;           // Sum of blue components
        uint32                  sumA;           // Sum of alpha components
        Node*                   child[8];       // Pointers to child nodes
        Node*                   next;           // Pointer to next reducible node
};


BColorQuantizer::BColorQuantizer(uint32 maxColors, uint32 bitsPerColor)
        : fTree(NULL),
          fLeafCount(0),
          fMaxColors(maxColors),
          fOutputMaxColors(maxColors),
          fBitsPerColor(bitsPerColor)
{
        // override parameters if out of range
        if (fBitsPerColor > 8)
                fBitsPerColor = 8;

        if (fMaxColors < 16)
                fMaxColors = 16;

        for (int i = 0; i <= (int)fBitsPerColor; i++)
                fReducibleNodes[i] = NULL;
}


BColorQuantizer::~BColorQuantizer()
{
        if (fTree != NULL)
                _DeleteTree(&fTree);
}


bool
BColorQuantizer::ProcessImage(const uint8* const * rowPtrs, int width,
        int height)
{
        for (int y = 0; y < height; y++) {
                for (int x = 0; x < width; x += 3) {
                        uint8 b = rowPtrs[y][x];
                        uint8 g = rowPtrs[y][x + 1];
                        uint8 r = rowPtrs[y][x + 2];
                        _AddColor(&fTree, r, g, b, 0, fBitsPerColor, 0, &fLeafCount,
                                fReducibleNodes);

                        while (fLeafCount > fMaxColors)
                                _ReduceTree(fBitsPerColor, &fLeafCount, fReducibleNodes);
                }
        }

        return true;
}


uint32
BColorQuantizer::GetColorCount() const
{
        return fLeafCount;
}


void
BColorQuantizer::GetColorTable(RGBA* table) const
{
        uint32 index = 0;
        if (fOutputMaxColors < 16) {
                uint32 sums[16];
                RGBA tmpPalette[16];
                _GetPaletteColors(fTree, tmpPalette, &index, sums);
                if (fLeafCount > fOutputMaxColors) {
                        for (uint32 j = 0; j < fOutputMaxColors; j++) {
                                uint32 a = (j * fLeafCount) / fOutputMaxColors;
                                uint32 b = ((j + 1) * fLeafCount) / fOutputMaxColors;
                                uint32 nr = 0;
                                uint32 ng = 0;
                                uint32 nb = 0;
                                uint32 na = 0;
                                uint32 ns = 0;
                                for (uint32 k = a; k < b; k++){
                                        nr += tmpPalette[k].r * sums[k];
                                        ng += tmpPalette[k].g * sums[k];
                                        nb += tmpPalette[k].b * sums[k];
                                        na += tmpPalette[k].a * sums[k];
                                        ns += sums[k];
                                }
                                table[j].r = clip((float)nr / ns);
                                table[j].g = clip((float)ng / ns);
                                table[j].b = clip((float)nb / ns);
                                table[j].a = clip((float)na / ns);
                        }
                } else {
                        memcpy(table, tmpPalette, fLeafCount * sizeof(RGBA));
                }
        } else {
                _GetPaletteColors(fTree, table, &index, NULL);
        }
}


// #pragma mark - private


void
BColorQuantizer::_AddColor(Node** _node, uint8 r, uint8 g, uint8 b, uint8 a,
        uint32 bitsPerColor, uint32 level, uint32* _leafCount,
        Node** reducibleNodes)
{
        static const uint8 kMask[8]
                = {0x80, 0x40, 0x20, 0x10, 0x08, 0x04, 0x02, 0x01};

        // If the node doesn't exist, create it.
        if (*_node == NULL)
                *_node = _CreateNode(level, bitsPerColor, _leafCount, reducibleNodes);

        // Update color information if it's a leaf node.
        if ((*_node)->isLeaf) {
                (*_node)->pixelCount++;
                (*_node)->sumR += r;
                (*_node)->sumG += g;
                (*_node)->sumB += b;
                (*_node)->sumA += a;
        } else {
                // Recurse a level deeper if the node is not a leaf.
                int shift = 7 - level;
                int index = (((r & kMask[level]) >> shift) << 2) |
                                         (((g & kMask[level]) >> shift) << 1) |
                                         (( b & kMask[level]) >> shift);
                _AddColor(&((*_node)->child[index]), r, g, b, a, bitsPerColor,
                        level + 1, _leafCount, reducibleNodes);
        }
}


BColorQuantizer::Node*
BColorQuantizer::_CreateNode(uint32 level, uint32 bitsPerColor,
        uint32* _leafCount, Node** reducibleNodes)
{
        Node* node = (Node*)calloc(1, sizeof(Node));

        if (node == NULL)
                return NULL;

        node->isLeaf = (level == bitsPerColor) ? true : false;
        if (node->isLeaf)
                (*_leafCount)++;
        else {
                node->next = reducibleNodes[level];
                reducibleNodes[level] = node;
        }
        return node;
}


void
BColorQuantizer::_ReduceTree(uint32 bitsPerColor, uint32* _leafCount,
        Node** reducibleNodes)
{
        int i = bitsPerColor - 1;
        // Find the deepest level containing at least one reducible node.
        for (; i > 0 && reducibleNodes[i] == NULL; i--)
                ;

        // Reduce the node most recently added to the list at level i.
        Node* node = reducibleNodes[i];
        reducibleNodes[i] = node->next;

        uint32 sumR = 0;
        uint32 sumG = 0;
        uint32 sumB = 0;
        uint32 sumA = 0;
        uint32 childCount = 0;

        for (i = 0; i < 8; i++) {
                if (node->child[i] != NULL) {
                        sumR += node->child[i]->sumR;
                        sumG += node->child[i]->sumG;
                        sumB += node->child[i]->sumB;
                        sumA += node->child[i]->sumA;
                        node->pixelCount += node->child[i]->pixelCount;

                        free(node->child[i]);
                        node->child[i] = NULL;

                        childCount++;
                }
        }

        node->isLeaf = true;
        node->sumR = sumR;
        node->sumG = sumG;
        node->sumB = sumB;
        node->sumA = sumA;

        *_leafCount -= (childCount - 1);
}


void
BColorQuantizer::_DeleteTree(Node** _node)
{
        for (int i = 0; i < 8; i++) {
                if ((*_node)->child[i] != NULL)
                        _DeleteTree(&((*_node)->child[i]));
        }
        free(*_node);
        *_node = NULL;
}


void
BColorQuantizer::_GetPaletteColors(Node* node, RGBA* table, uint32* _index,
        uint32* sums) const
{
        if (node == NULL)
                return;

        if (node->isLeaf) {
                table[*_index].r = clip((float)node->sumR / node->pixelCount);
                table[*_index].g = clip((float)node->sumG / node->pixelCount);
                table[*_index].b = clip((float)node->sumB / node->pixelCount);
                table[*_index].a = clip((float)node->sumA / node->pixelCount);
                if (sums)
                        sums[*_index] = node->pixelCount;
                (*_index)++;
        } else {
                for (int i = 0; i < 8; i++) {
                        if (node->child[i] != NULL)
                                _GetPaletteColors(node->child[i], table, _index, sums);
                }
        }
}