/* osgEarth
 * Copyright 2025 Pelican Mapping
 * MIT License
 */
#ifndef OSGEARTH_COVERAGE_H
#define OSGEARTH_COVERAGE_H 1

#include <osgEarth/GeoData>
#include <osgEarth/ImageUtils>
#include <array>
#include <bitset>
#include <cstdint>
#include <limits>
#include <unordered_map>
#include <vector>
#include <map>

namespace osgEarth
{
    /**
     * Grid of templated values.
     *
     * The parameter <T> must follow implement the interface:
     *
     *    // default constructor, which will result in valid()==false
     *    T() { }
     *
     *    // return true for valid data, false for NO DATA.
     *    bool valid() const { return ...has a valid data value...; }
     *
     *    // "less than" operator for sorting
     *    bool operator < (const T& rhs) const
     *
     * Note: This is a sparse grid for memory conservation purposes,
     * but includes some optimizations for fast mass-insertion.
     */
    template<typename T>
    class Coverage
    {
    public:
        using Ptr = std::shared_ptr<Coverage<T>>;
        
        using Index = std::uint8_t;
        static constexpr Index NoData = 0;
        static constexpr std::size_t MaxIndex = std::numeric_limits<Index>::max();

        // maps T to index for duplicate values (std::map so that T doesn't need std::hash)
        using LUT = std::map<T, Index>;                 // value -> palette index

        // reverse lookup for reading; maps index to value
        using RLUT = std::array<T, MaxIndex + 1>;       // palette index -> value

        // raster of indices
        using Grid = std::vector<Index>;

        using pixel_type = T;

        //! create a new coverage (must be a shared_ptr)
        static Ptr create()
        {
            return Ptr(new Coverage<T>());
        }

        //! copy constructor
        Coverage(const Coverage<T>& rhs) :
            _width(rhs._width),
            _height(rhs._height),
            _valid_count(rhs._valid_count),
            _pixels(rhs._pixels),
            _lut(rhs._lut),
            _rlut(rhs._rlut)
        {
            //nop
        }

        //! move constructor
        Coverage(Coverage<T>&& rhs) :
            _width(rhs._width),
            _height(rhs._height),
            _valid_count(rhs._valid_count),
            _pixels(std::move(rhs._pixels)),
            _lut(std::move(rhs._lut)),
            _rlut(std::move(rhs._rlut))
        {
            //nop
        }

        //! allocate width*height pixels
        void allocate(unsigned w, unsigned h)
        {
            _width = w;
            _height = h;
            _lut.clear();
            _rlutValid.reset();
            _pixels.assign(static_cast<std::size_t>(w) * static_cast<std::size_t>(h), NoData);
            _valid_count = 0;
        }

        //! raster width
        unsigned s() const
        {
            return _width;
        }

        //! raster height
        unsigned t() const
        {
            return _height;
        }

        //! Number of pixels containing "NO DATA" value
        unsigned nodataCount() const
        {
            return (_width*_height) - _valid_count;
        }

        //! Is this coverage empty?
        bool empty() const
        {
            return _valid_count == 0;
        }

        //! Write a value to a pixel.
        //! Optional optimization: This method returns a value "k". If you
        //! are writing the same value again on the next call, pass "k" back
        //! in as your final argument, and this will dramatically speed up 
        //! performance. This helps a lot when writing the same value many
        //! times in a row.
        int write(const T& value, unsigned s, unsigned t, int k = -1)
        {
            const std::size_t ptr = static_cast<std::size_t>(t) * _width + s;
            const Index oldIndex = _pixels[ptr];
            const bool oldValid = oldIndex != NoData;
            const bool newValid = value.valid();

            Index newIndex = NoData;

            if (newValid)
            {
                if (k > 0)
                {
                    newIndex = static_cast<Index>(k);
                }
                else
                {
                    auto it = _lut.find(value);
                    if (it != _lut.end())
                    {
                        newIndex = it->second;
                    }
                    else
                    {
                        if (_lut.size() >= MaxIndex)
                            return 0; // or assert/throw: palette overflow

                        newIndex = static_cast<Index>(_lut.size() + 1); // reserve 0 for NoData
                        _lut.emplace(value, newIndex);
                        _rlut[newIndex] = value;
                        _rlutValid.set(newIndex);
                    }
                }
            }

            if (!oldValid && newValid) ++_valid_count;
            else if (oldValid && !newValid) --_valid_count;

            _pixels[ptr] = newIndex;
            return static_cast<int>(newIndex);
        }

        bool read(T& output, unsigned s, unsigned t) const
        {
            const T* value = read(s, t);
            if (value) output = *value;
            return value != nullptr;
        }

        bool read(T& output, double u, double v) const
        {
            const T* value = read(u, v);
            if (value) output = *value;
            return value != nullptr;
        }

        const T* read(double u, double v) const
        {
            unsigned s, t;
            ImageUtils::nnUVtoST(u, v, s, t, _width, _height);
            return read(s, t);
        }

        const T* read(unsigned s, unsigned t) const
        {
            OE_SOFT_ASSERT_AND_RETURN(s < _width && t < _height, nullptr);
            const std::size_t ptr = static_cast<std::size_t>(t) * _width + s;
            const Index i = _pixels[ptr];
            if (i != NoData && _rlutValid.test(i))
                return &_rlut[i];
            return nullptr;
        }

        bool hasDataAt(unsigned s, unsigned t) const
        {
            OE_SOFT_ASSERT_AND_RETURN(s < _width && t < _height, false);
            int ptr = t * _width + s;
            return _pixels[ptr] > 0;
        }

        Config getConfig() const
        {
            Config root("coverage");
            root.set("width", _width);
            root.set("height", _height);

            Config values("values");
            for (auto& entry : _lut)
            {
                Config value_conf = entry.first.getConfig();
                value_conf.set("__coverage_index", static_cast<int>(entry.second));
                values.add(value_conf);
            }
            root.add(values);

            std::ostringstream pixels_string;
            int count = 0;
            int last_value = -1;
            for(auto& pixel : _pixels)
            {
                if (last_value >= 0 && pixel != last_value) {
                    pixels_string << (int)count << ' ' << (int)last_value << ' ';
                    count = 0;
                }
                last_value = pixel;
                ++count;
            }
            if (count > 0) {
                pixels_string << (int)count << ' ' << (int)last_value;
            }
            root.add(Config("pixels", pixels_string.str()));

            return root;
        }

        void setConfig(const Config& root)
        {
            unsigned w = 256, h = 256;
            root.get("width", w);
            root.get("height", h);
            allocate(w, h);

            Config values = root.child("values");
            for (auto& value_conf : values.children())
            {
                T value(value_conf);
                int index = 0;
                if (!value_conf.get("__coverage_index", index))
                    continue;
                if (index > 0 && index <= static_cast<int>(MaxIndex))
                {
                    Index i = static_cast<Index>(index);
                    _lut[value] = i;
                    _rlut[i] = value;
                    _rlutValid.set(static_cast<std::size_t>(i));
                }
            }

            std::string pixels_string = root.value("pixels");
            std::istringstream pixels_stream(pixels_string);
            int count, pixel;
            int ptr = 0;
            while (!pixels_stream.eof() && ptr < (int)(w*h))
            {
                pixels_stream >> count >> pixel;
                for (int i = 0; i < count; ++i) {
                    _pixels[ptr++] = pixel;
                    if (pixel > 0)
                        ++_valid_count;
                }
            }
        }

    private:
        Coverage()
        {
            allocate(0, 0);
        }

    private:
        unsigned _width, _height;
        unsigned _valid_count;
        Grid _pixels;
        LUT _lut;
        mutable RLUT _rlut;
        std::bitset<MaxIndex + 1> _rlutValid;
    };

    template<typename T>
    class GeoCoverage
    {
    public:
        using pixel_type = T;

        //! Construct a GeoCoverage
        //! @param coverage Coverage to include in the geocoverage
        //! @param extent geospatial extent of the coverage data
        GeoCoverage(typename Coverage<T>::Ptr coverage, const GeoExtent& extent)
            : _coverage(coverage), _extent(extent)
        {
            //nop
        }

        GeoCoverage() = default;

        inline bool valid() const {
            return _coverage != nullptr;
        }

        inline unsigned s() const {
            return _coverage ? _coverage->s() : 0u;
        }

        inline unsigned t() const {
            return _coverage ? _coverage->t() : 0u;
        }

        inline unsigned nodataCount() const {
            return _coverage ? _coverage->nodataCount() : 0u;
        }

        inline unsigned empty() const {
            return _coverage ? _coverage->empty() : true;
        }

        void write(const T& value, unsigned s, unsigned t)
        {
            OE_SOFT_ASSERT_AND_RETURN(_coverage != nullptr, void());
            _coverage->write(value, s, t);
        }

        bool read(T& value, unsigned s, unsigned t) const
        {
            OE_SOFT_ASSERT_AND_RETURN(_coverage != nullptr, false);
            return _coverage->read(value, s, t);
        }

        bool read(T& value, double u, double v) const
        {
            OE_SOFT_ASSERT_AND_RETURN(_coverage != nullptr, false);
            return _coverage->read(value, u, v);
        }

        const T* read(unsigned s, unsigned t) const
        {
            OE_SOFT_ASSERT_AND_RETURN(_coverage != nullptr, nullptr);
            return _coverage->read(s, t);
        }

        const T* read(double u, double v) const
        {
            OE_SOFT_ASSERT_AND_RETURN(_coverage != nullptr, nullptr);
            return _coverage->read(u, v);
        }

        bool readAtCoords(T& value, double x, double y) const
        {
            if (_extent.contains(x, y))
            {
                unsigned xs = (unsigned)((double)(s() - 1) * (x - _extent.xMin()) / _extent.width());
                unsigned yt = (unsigned)((double)(t() - 1) * (y - _extent.yMin()) / _extent.height());
                if (xs >= 0 && xs < s() && yt >= 0 && yt < t())
                {
                    read(value, xs, yt);
                    return true;
                }
            }
            return false;
        }
        
        const T* readAtCoords(double x, double y) const
        {
            if (_extent.contains(x, y))
            {
                unsigned xs = (unsigned)((double)(s() - 1) * (x - _extent.xMin()) / _extent.width());
                unsigned yt = (unsigned)((double)(t() - 1) * (y - _extent.yMin()) / _extent.height());
                if (xs >= 0 && xs < s() && yt >= 0 && yt < t())
                {
                    return read(xs, yt);
                }
            }
            return nullptr;
        }

        const GeoExtent& getExtent() const
        {
            return _extent;
        }

        Config getConfig() const
        {
            Config root("geocoverage");
            root.set("extent", _extent.getConfig());
            root.add(_coverage->getConfig());
            return root;
        }

        void setConfig(const Config& in)
        {
            _extent.fromConfig(in.child("extent"));
            _coverage->setConfig(in.child("coverage"));
        }

    private:
        GeoExtent _extent;
        bool _valid;
        typename Coverage<T>::Ptr _coverage;
    };

    // required in C++14
    template<typename T>
    constexpr typename Coverage<T>::Index Coverage<T>::NoData;
}

#endif // OSGEARTH_COVERAGE_H
