/* osgEarth
* Copyright 2008-2012 Pelican Mapping
* MIT License
*/
#ifndef OSGEARTH_METATILE_H
#define OSGEARTH_METATILE_H

#include <osgEarth/Common>
#include <osgEarth/ImageUtils>
#include <osgEarth/GeoData>
#include <osgEarth/TileKey>
#include <osgEarth/Progress>
#include <osgEarth/Math>
#include <osg/Image>

namespace osgEarth { namespace Util
{
    using namespace osgEarth;

    /**
     * Metadata groups a collection of adjacent data tiles
     * together to facilitate operations that overlap multiple tiles.
     */
    class OSGEARTH_EXPORT MetaImage
    {
    public:
        //! Construct a new Metatiled image
        MetaImage() { }

        //! Sets the data at location (x,y), where (0,0) is the center.
        bool setImage(int x, int y, osg::Image* image, const osg::Matrix& scaleBias);

        //! Gets the image at the neighbor location (x,y).
        const osg::Image* getImage(int x, int y) const;

        //! Gets the positioning matrix for neightbor location (x,y).
        const osg::Matrix& getScaleBias(int x, int y) const;

        //! Reads the data from parametric location (u,v), where [u,v] in [-1, +2].
        //! Returns true upon success with the value in [output];
        //! false if there is no tile at the read location.
        virtual bool read(double u, double v, osg::Vec4f& output);

        void dump() const;

    protected:

        struct Tile {
            Tile();
            bool _failed;
            osg::ref_ptr<const osg::Image> _imageRef;
            ImageUtils::PixelReader _read;
            osg::Matrix _scaleBias;
        };

        Tile _tiles[3][3]; // col, row
    };


    class OSGEARTH_EXPORT TileKeyMetaImage : public MetaImage
    {
    public:
        using CreateImageFunction = std::function<
            GeoImage(const TileKey&, ProgressCallback*)>;

        //! Construct
        TileKeyMetaImage();

        //! Sets the center tilekey
        void setTileKey(const TileKey& value);

        //! CreateImage function
        void setCreateImageFunction(CreateImageFunction value);

    public: // MetaImage

        bool read(double u, double v, osg::Vec4f& output) override;

    private:
        TileKey _center;
        CreateImageFunction _createImage;
    };


    /**
     * MetaTile is a framework for treating a center tile and its 8 neighbors
     * as a single larger tile. As you sample the tile with the read() function,
     * it will automatically load the neighbors based on the u,v coordinates.
     */
    template<typename T>
    class MetaTile
    {
    public:
        using CreateTileFunction = std::function<T(const TileKey&, ProgressCallback*)>;

        MetaTile()
            : _progress(nullptr) { }

        //! Function that will create new neighbor tiles given a tile key
        inline void setCreateTileFunction(CreateTileFunction value);

        //! Center tile of this metatile. Must call this before attempting
        //! to read. Must call setCreateTileFunction before calling this.
        //! This will "fall back" on ancestor tiles until it find valid data.
        inline void setCenterTileKey(const TileKey& key, ProgressCallback* progress);

        //! Center tile of this metatile. Must call this before attempting
        //! to read. Must call setCreateTileFunction before calling this.
        inline void setCenterTileKey(const TileKey& key, const osg::Matrix& scale_bias);

        //! Read the value of a pixel of unit coordinates [u,v] relative to
        //! the center tile of the meta
        inline bool read(typename T::pixel_type& output, double u, double v);
        inline bool read(typename T::pixel_type& output, int s, int t);

        //! Read the value of a pixel of unit coordinates [u,v] relative to
        //! the center tile of the meta and return a pointer to the
        //! underlying data value instead of copying it to avoid expensive data copies.
        inline const typename T::pixel_type* read(double u, double v);
        inline const typename T::pixel_type* read(int s, int t);

        //! Read the interpolated value at the geospatial coordinate (x,y)
        //! (Bilinear interpolation is used)
        inline bool readAtCoord(typename T::pixel_type& output, double x, double y, RasterInterpolation = INTERP_BILINEAR);

        //! The scale&bias of the tile relative to the key originally passed
        //! to setCenterTileKey
        inline const osg::Matrix& getScaleBias() const {
            return _scale_bias;
        }

        inline bool valid() {
            return _tiles(0,0)._data.valid();
        }

        const T& getCenterTile() {
            return _tiles(0, 0)._data;
        }

    private:
        TileKey _centerKey;
        CreateTileFunction _createTile;
        ProgressCallback* _progress;

        // one component tile
        struct Tile {
            Tile() : _failed(false) { }
            bool _failed;
            T _data;
        };

        // sparse grid for metatile components
        struct Grid : public std::unordered_map<int, Tile> {
            inline Tile& operator()(int x, int y) {
                return this->operator[](y * 100 + x);
            }
        };
        Grid _tiles;
        osg::Matrix _scale_bias; // scale/bias matrix of _centerKey
        unsigned _width, _height;

        Tile* getTile(int tx, int ty);
        bool getTileAndPixelCoords(double u, double v, Tile*& tile, double& s, double& t);
    };

    template<typename T>
    void MetaTile<T>::setCenterTileKey(const TileKey& original_key, ProgressCallback* progress)
    {
        OE_HARD_ASSERT(_createTile != nullptr, "Must call setCreateTileFunction() before calling setCenterTileKey()");
        // Fall back on parent keys until we get real data
        TileKey key;
        for (key = original_key;
            !_tiles(0,0)._data.valid() && key.valid();
            key.makeParent())
        {
            _tiles(0,0)._data = _createTile(key, progress);
            _centerKey = key;
            _width = _tiles(0, 0)._data.s();
            _height = _tiles(0, 0)._data.t();
        }
        original_key.getExtent().createScaleBias(_centerKey.getExtent(), _scale_bias);
        _progress = progress;
    }

    template<typename T>
    void MetaTile<T>::setCenterTileKey(
        const TileKey& original_key,
        const osg::Matrix& scale_bias)
    {
        OE_HARD_ASSERT(_createTile != nullptr, "Must call setCreateTileFunction() before calling setCenterTileKey()");
        // Fall back on parent keys until we get real data
        TileKey key;
        for (key = original_key;
            !_tiles(0, 0)._data.valid() && key.valid();
            key.makeParent())
        {
            _tiles(0, 0)._data = _createTile(key, nullptr);
            _centerKey = key;
            _width = _tiles(0, 0)._data.s();
            _height = _tiles(0, 0)._data.t();
        }
        original_key.getExtent().createScaleBias(_centerKey.getExtent(), _scale_bias);
        _scale_bias.preMult(scale_bias);
    }

    template<typename T>
    void MetaTile<T>::setCreateTileFunction(typename MetaTile<T>::CreateTileFunction value)
    {
        _createTile = value;
    }

    template<typename T>
    bool MetaTile<T>::read(typename T::pixel_type& output, double u, double v)
    {
        // scale and bias the u,v to the real center key
        u = u * _scale_bias(0, 0) + _scale_bias(3, 0);
        v = v * _scale_bias(1, 1) + _scale_bias(3, 1);

        // tile number:
        // TODO: when this hits an exact boundary (i.e. 1.0) which is the
        // correct tile to choose? The actual answer is that when using 
        // metatile, you should probably always use forEachPixelOnCenter
        // which will never land on a boundary exactly. -gw
        int x = (int)::floor(u);
        int y = (int)::floor(v);

        Tile& tile = _tiles(x, y);

        // if we already tried to load this tile and failed, bail out
        if (tile._failed)
            return false;

        // if we still need to load this tile, do so
        if (!tile._data.valid() && _createTile != nullptr)
        {
            TileKey key = _centerKey.createNeighborKey(x, -y);
            tile._data = _createTile(key, nullptr);
            if (!tile._data.valid())
            {
                tile._failed = true;
            }
        }

        if (tile._failed)
            return false;

        // tile-local coordinates for sampling:
        u = u - (double)x;
        v = v - (double)y;

        return tile._data.read(output, u, v);
    }

    template<typename T>
    const typename T::pixel_type* MetaTile<T>::read(double u, double v)
    {
        // scale and bias the u,v to the real center key
        u = u * _scale_bias(0, 0) + _scale_bias(3, 0);
        v = v * _scale_bias(1, 1) + _scale_bias(3, 1);

        // tile number:
        // TODO: when this hits an exact boundary (i.e. 1.0) which is the
        // correct tile to choose? The actual answer is that when using 
        // metatile, you should probably always use forEachPixelOnCenter
        // which will never land on a boundary exactly. -gw
        int x = (int)::floor(u);
        int y = (int)::floor(v);

        Tile& tile = _tiles(x, y);

        // if we already tried to load this tile and failed, bail out
        if (tile._failed)
            return nullptr;

        // if we still need to load this tile, do so
        if (!tile._data.valid() && _createTile != nullptr)
        {
            TileKey key = _centerKey.createNeighborKey(x, -y);
            tile._data = _createTile(key, nullptr);
            if (!tile._data.valid())
            {
                tile._failed = true;
            }
        }

        if (tile._failed)
            return nullptr;

        // tile-local coordinates for sampling:
        u = u - (double)x;
        v = v - (double)y;

        return tile._data.read(u, v);
    }

    template<typename T>
    bool MetaTile<T>::read(typename T::pixel_type& output, int s, int t)
    {
        if (_tiles(0, 0)._failed)
            return false;

        // scale and bias the u,v to the real center key
        s = (int)floor((double)s * _scale_bias(0, 0) + _scale_bias(3, 0)*(double)_width);
        t = (int)floor((double)t * _scale_bias(1, 1) + _scale_bias(3, 1)*(double)_height);

        // tile number:
        int x = (int)::floor((double)s / (double)_width);
        int y = (int)::floor((double)t / (double)_height);

        Tile& tile = _tiles(x, y);

        // if we already tried to load this tile and failed, bail out
        if (tile._failed)
            return false;

        // if we still need to load this tile, do so
        if (!tile._data.valid() && _createTile != nullptr)
        {
            TileKey key = _centerKey.createNeighborKey(x, -y);
            tile._data = _createTile(key, nullptr);
            if (!tile._data.valid())
            {
                tile._failed = true;
            }
        }

        if (tile._failed)
            return false;

        // tile-local coordinates for sampling:
        if (s < 0) {
            s = s % _width;
            if (s < 0) s += _width;
        }
        else s = s % _width;

        if (t < 0) {
            t = t % _height;
            if (t < 0) t += _height;
        }
        else t = t % _height;

        return tile._data.read(output, (unsigned)s, (unsigned)t);
    }

    template<typename T>
    const typename T::pixel_type* MetaTile<T>::read(int s, int t)
    {
        if (_tiles(0, 0)._failed)
            return nullptr;

        // scale and bias the u,v to the real center key
        s = (int)floor((double)s * _scale_bias(0, 0) + _scale_bias(3, 0) * (double)_width);
        t = (int)floor((double)t * _scale_bias(1, 1) + _scale_bias(3, 1) * (double)_height);

        // tile number:
        int x = (int)::floor((double)s / (double)_width);
        int y = (int)::floor((double)t / (double)_height);

        Tile& tile = _tiles(x, y);

        // if we already tried to load this tile and failed, bail out
        if (tile._failed)
            return nullptr;

        // if we still need to load this tile, do so
        if (!tile._data.valid() && _createTile != nullptr)
        {
            TileKey key = _centerKey.createNeighborKey(x, -y);
            tile._data = _createTile(key, nullptr);
            if (!tile._data.valid())
            {
                tile._failed = true;
            }
        }

        if (tile._failed)
            return nullptr;

        // tile-local coordinates for sampling:
        if (s < 0) {
            s = s % _width;
            if (s < 0) s += _width;
        }
        else s = s % _width;

        if (t < 0) {
            t = t % _height;
            if (t < 0) t += _height;
        }
        else t = t % _height;

        return tile._data.read((unsigned)s, (unsigned)t);
    }

    template<typename T>
    typename MetaTile<T>::Tile* MetaTile<T>::getTile(int tx, int ty)
    {
        Tile& tile = _tiles(tx, ty);

        // if we already tried to load this tile and failed, bail out
        if (tile._failed)
            return nullptr;

        // if we still need to load this tile, do so
        if (!tile._data.valid() && _createTile != nullptr)
        {
            TileKey key = _centerKey.createNeighborKey(tx, -ty);
            tile._data = _createTile(key, nullptr);
            if (!tile._data.valid())
            {
                tile._failed = true;
            }
        }

        if (tile._failed)
            return nullptr;

        return &tile;
    }

    template<typename T>
    bool MetaTile<T>::readAtCoord(typename T::pixel_type& output, double x, double y, RasterInterpolation interp)
    {
        // geoextents of center tile:
        double xmin, ymin, xmax, ymax;
        _centerKey.getExtent().getBounds(xmin, ymin, xmax, ymax);

        // (u, v) ndc relative to center tile:
        double u = (x - xmin) / (xmax - xmin);
        double v = (y - ymin) / (ymax - ymin);

        // tile index (offset from (0,0) center tile)
        int tx = (int)u;
        int ty = (int)v;

        // Anchor:
        Tile* p00_tile = getTile(tx, ty);
        if (!p00_tile)
            return false;

        double p00_s = fract(u) * (double)_width, p00_t = fract(v) * (double)_height;

        // nearest-neighbor pixel:
        unsigned s00 = (unsigned)floor(p00_s);
        unsigned t00 = (unsigned)floor(p00_t);

        if (interp == INTERP_NEAREST)
        {
            return p00_tile->_data.read(output, s00, t00);
        }

        // intepolation weights:
        double left_weight = 1.0 - (p00_s + 0.5 - (double)floor(p00_s + 0.5));
        double bottom_weight = 1.0 - (p00_t + 0.5 - (double)floor(p00_t + 0.5));

        // prepare for interpolation.
        int delta_s = 1, delta_t = 1;
        if (fract(p00_s) < 0.5)
        {
            delta_s = -1;
            left_weight = 1.0 - left_weight;
        }
        if (fract(p00_t) < 0.5)
        {
            delta_t = -1;
            bottom_weight = 1.0 - bottom_weight;
        }

        // Neighbors:
        Tile* p10_tile = p00_tile;
        double p10_s = p00_s + delta_s, p10_t = p00_t;
        if (p10_s >= (double)_width || p10_s < 0)
        {
            p10_tile = getTile(tx + delta_s, ty);
            p10_s -= (double)_width * delta_s;
        }

        Tile* p01_tile = p00_tile;
        double p01_s = p00_s, p01_t = p00_t + delta_t;
        if (p01_t >= _height || p01_t < 0)
        {
            p01_tile = getTile(tx, ty + delta_t);
            p01_t -= (double)_height * (double)delta_t;
        }

        Tile* p11_tile = p00_tile;
        double p11_s = p00_s + delta_s, p11_t = p00_t + delta_t;
        int ntx = tx, nty = ty;
        if (p11_s >= (double)_width || p11_s < 0)
            ntx = tx + delta_s;
        if (p11_t >= (double)_height || p11_t < 0)
            nty = ty + delta_t;
        if (ntx != tx || nty != ty)
        {
            p11_tile = getTile(ntx, nty);
            if (p11_s > (double)_width || p11_s < 0) p11_s -= (double)_width * (double)delta_s;
            if (p11_t > (double)_height || p11_t < 0) p11_t -= (double)_height * (double)delta_t;
        }

        if (!p01_tile || !p10_tile || !p11_tile)
        {
            return p00_tile->_data.read(output, s00, t00);
        }
        else
        {
            typename T::pixel_type p00, p10, p01, p11;

            unsigned s01 = clamp((unsigned)floor(p01_s), 0u, _width-1);
            unsigned t01 = clamp((unsigned)floor(p01_t), 0u, _height-1);
            unsigned s10 = clamp((unsigned)floor(p10_s), 0u, _width-1);
            unsigned t10 = clamp((unsigned)floor(p10_t), 0u, _height-1);
            unsigned s11 = clamp((unsigned)floor(p11_s), 0u, _width-1);
            unsigned t11 = clamp((unsigned)floor(p11_t), 0u, _height-1);

            if (!p00_tile->_data.read(p00, s00, t00) ||
                !p01_tile->_data.read(p01, s01, t01) ||
                !p10_tile->_data.read(p10, s10, t10) ||
                !p11_tile->_data.read(p11, s11, t11))
            {
                return false;
            }

            if (interp == INTERP_AVERAGE)
            {
                output =
                    mult(p00, left_weight * bottom_weight) +
                    mult(p10, (1.0 - left_weight) * bottom_weight) +
                    mult(p01, left_weight * (1.0 - bottom_weight)) +
                    mult(p11, (1.0 - left_weight) * (1.0 - bottom_weight));
            }

            else // INTERP_BILINEAR
            {
                typename T::pixel_type p0, p1;
                p0 = mult(p00, left_weight) + mult(p10, 1.0 - left_weight);
                p1 = mult(p01, left_weight) + mult(p11, 1.0 - left_weight);
                output = mult(p0, bottom_weight) + mult(p1, 1.0 - bottom_weight);
            }

            return true;
        }
    }
       

} }

#endif // OSGEARTH_METATILE_H
