/*****************************************************************************
 *
 * This file is part of Mapnik (c++ mapping toolkit)
 *
 * Copyright (C) 2024 Artem Pavlenko
 *
 * This library is free software; you can redistribute it and/or
 * modify it under the terms of the GNU Lesser General Public
 * License as published by the Free Software Foundation; either
 * version 2.1 of the License, or (at your option) any later version.
 *
 * This library is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU
 * Lesser General Public License for more details.
 *
 * You should have received a copy of the GNU Lesser General Public
 * License along with this library; if not, write to the Free Software
 * Foundation, Inc., 51 Franklin St, Fifth Floor, Boston, MA  02110-1301  USA
 *
 *****************************************************************************/

#ifndef MAPNIK_OCTREE_HPP
#define MAPNIK_OCTREE_HPP

// mapnik
#include <mapnik/global.hpp>
#include <mapnik/palette.hpp>
#include <mapnik/util/noncopyable.hpp>

// stl
#include <algorithm>
#include <vector>
#include <deque>

namespace mapnik {

struct RGBPolicy
{
    const static unsigned MAX_LEVELS = 6;
    const static unsigned MIN_LEVELS = 2;
    inline static unsigned index_from_level(unsigned level, rgb const& c)
    {
        unsigned shift = 7 - level;
        return (((c.r >> shift) & 1) << 2) | (((c.g >> shift) & 1) << 1) | ((c.b >> shift) & 1);
    }
};

template<typename T, typename InsertPolicy = RGBPolicy>
class octree : private util::noncopyable
{
    struct node
    {
        node()
            : reds(0)
            , greens(0)
            , blues(0)
            , count(0)
            , count_cum(0)
            , children_count(0)
            , index(0)
        {
            std::fill(children_, children_ + 8, nullptr);
        }

        ~node()
        {
            for (unsigned i = 0; i < 8; ++i)
            {
                if (children_[i] != 0)
                {
                    delete children_[i];
                    children_[i] = 0;
                }
            }
        }

        bool is_leaf() const { return count == 0; }
        node* children_[8];
        std::uint64_t reds;
        std::uint64_t greens;
        std::uint64_t blues;
        unsigned count;
        double reduce_cost;
        unsigned count_cum;
        std::uint8_t children_count;
        std::uint8_t index;
    };
    struct node_cmp
    {
        bool operator()(const node* lhs, const node* rhs) const { return lhs->reduce_cost < rhs->reduce_cost; }
    };

    std::deque<node*> reducible_[InsertPolicy::MAX_LEVELS];
    unsigned max_colors_;
    unsigned colors_;
    unsigned offset_;
    unsigned leaf_level_;

  public:
    explicit octree(unsigned max_colors = 256)
        : max_colors_(max_colors)
        , colors_(0)
        , offset_(0)
        , leaf_level_(InsertPolicy::MAX_LEVELS)
        , root_(new node())
    {}

    ~octree() { delete root_; }

    unsigned colors() { return colors_; }

    void setMaxColors(unsigned max_colors) { max_colors_ = max_colors; }

    void setOffset(unsigned offset) { offset_ = offset; }

    unsigned getOffset() { return offset_; }

    void insert(T const& data)
    {
        unsigned level = 0;
        node* cur_node = root_;
        while (true)
        {
            cur_node->count_cum++;
            cur_node->reds += data.r;
            cur_node->greens += data.g;
            cur_node->blues += data.b;

            if (cur_node->count > 0 || level == leaf_level_)
            {
                cur_node->count += 1;
                if (cur_node->count == 1)
                    ++colors_;
                // if (colors_ >= max_colors_ - 1)
                // reduce();
                break;
            }

            unsigned idx = InsertPolicy::index_from_level(level, data);
            if (cur_node->children_[idx] == 0)
            {
                cur_node->children_count++;
                cur_node->children_[idx] = new node();
                if (level < leaf_level_ - 1)
                {
                    reducible_[level + 1].push_back(cur_node->children_[idx]);
                }
            }
            cur_node = cur_node->children_[idx];
            ++level;
        }
    }

    int quantize(unsigned val) const
    {
        unsigned level = 0;
        rgb c(val);
        node* cur_node = root_;
        while (cur_node)
        {
            if (cur_node->children_count == 0)
            {
                return cur_node->index + offset_;
            }
            unsigned idx = InsertPolicy::index_from_level(level, c);
            cur_node = cur_node->children_[idx];
            ++level;
        }
        return -1;
    }

    void create_palette(std::vector<rgb>& palette)
    {
        reduce();
        palette.reserve(colors_);
        create_palette(palette, root_);
    }

    void computeCost(node* r)
    {
        r->reduce_cost = 0;
        if (r->children_count == 0)
        {
            return;
        }

        double mean_r = static_cast<double>(r->reds) / r->count_cum;
        double mean_g = static_cast<double>(r->greens) / r->count_cum;
        double mean_b = static_cast<double>(r->blues) / r->count_cum;
        for (unsigned idx = 0; idx < 8; ++idx)
        {
            if (r->children_[idx] != 0)
            {
                double dr, dg, db;
                computeCost(r->children_[idx]);
                dr = r->children_[idx]->reds / r->children_[idx]->count_cum - mean_r;
                dg = r->children_[idx]->greens / r->children_[idx]->count_cum - mean_g;
                db = r->children_[idx]->blues / r->children_[idx]->count_cum - mean_b;
                r->reduce_cost += r->children_[idx]->reduce_cost;
                r->reduce_cost += (dr * dr + dg * dg + db * db) * r->children_[idx]->count_cum;
            }
        }
    }

    void reduce()
    {
        computeCost(root_);
        reducible_[0].push_back(root_);

        // sort reducible by reduce_cost
        for (unsigned i = 0; i < InsertPolicy::MAX_LEVELS; ++i)
        {
            std::sort(reducible_[i].begin(), reducible_[i].end(), node_cmp());
        }
        while (colors_ > max_colors_ && colors_ > 1)
        {
            while (leaf_level_ > 0 && reducible_[leaf_level_ - 1].size() == 0)
            {
                --leaf_level_;
            }

            if (leaf_level_ <= 0)
            {
                return;
            }

            // select best of all reducible:
            unsigned red_idx = leaf_level_ - 1;
            unsigned bestv = static_cast<unsigned>((*reducible_[red_idx].begin())->reduce_cost);
            for (unsigned i = red_idx; i >= InsertPolicy::MIN_LEVELS; i--)
            {
                if (!reducible_[i].empty())
                {
                    node* nd = *reducible_[i].begin();
                    unsigned gch = 0;
                    for (unsigned idx = 0; idx < 8; idx++)
                    {
                        if (nd->children_[idx])
                            gch += nd->children_[idx]->children_count;
                    }
                    if (gch == 0 && nd->reduce_cost < bestv)
                    {
                        bestv = static_cast<unsigned>(nd->reduce_cost);
                        red_idx = i;
                    }
                }
            }

            typename std::deque<node*>::iterator pos = reducible_[red_idx].begin();
            node* cur_node = *pos;
            unsigned num_children = 0;
            for (unsigned idx = 0; idx < 8; ++idx)
            {
                if (cur_node->children_[idx] != 0)
                {
                    cur_node->children_count--;
                    ++num_children;
                    cur_node->count += cur_node->children_[idx]->count;
                    // todo: case of nonleaf children, if someday sorting by reduce_cost doesn't handle it
                    delete cur_node->children_[idx], cur_node->children_[idx] = 0;
                }
            }

            reducible_[red_idx].erase(pos);
            if (num_children > 0)
            {
                colors_ -= (num_children - 1);
            }
        }
    }

    void create_palette(std::vector<rgb>& palette, node* itr) const
    {
        if (itr->count != 0)
        {
            unsigned count = itr->count;
            palette.push_back(rgb(std::uint8_t(itr->reds / float(count)),
                                  std::uint8_t(itr->greens / float(count)),
                                  std::uint8_t(itr->blues / float(count))));
            itr->index = static_cast<unsigned>(palette.size()) - 1;
        }
        for (unsigned i = 0; i < 8; ++i)
        {
            if (itr->children_[i] != 0)
            {
                create_palette(palette, itr->children_[i]);
            }
        }
    }

  private:
    node* root_;
};

} // namespace mapnik

#endif // MAPNIK_OCTREE_HPP
