// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
// SPDX-License-Identifier: Apache-2.0
#pragma once

#include <boost/math/constants/constants.hpp>

#include <frantic/math/utils.hpp>

#include <frantic/graphics2d/image_channel.hpp>
#include <frantic/graphics2d/size2.hpp>
#include <frantic/graphics2d/vector2.hpp>

namespace frantic {
namespace graphics2d {

/**
 * Generate a mipmap for an image_channel.
 * The image channel must be square and have a resolution which is a power of 2.
 * The DataType template parameter must have operator+( DataType, DataType ), and be divisible by a scalar.
 * The Combiner is the strategy for combining pixels and must implement the following:
 *	Combiner( const image_t& past, image_t& current )
 *	void operator()( std::size_t currX, std::size_t currY )
 *
 * Note that past is the image one resolution larger, while current is the current image being generated.
 * currX and currY are the coordinates of the pixel in current that should be written by the operator().
 * An example of a Combiner can be found below in the form of four_pixel_mean.
 *
 * @tparam    DataType    The data type used by the image_channel.
 * @tparam    Combiner    The combiner to be used to generate pixels for a resolution based on the previous resolution's
 *                        image_channel.
 * @param     original    The full resolution image_channel to generate a mipmap for.
 * @return                A vector containing the mipmap, starting with the 1x1 texture and ending with the original
 *                        texture at the index log2 the width of the original image.
 */
template <typename DataType, typename Combiner>
std::vector<frantic::graphics2d::image_channel<DataType>>
generate_mipmap( const frantic::graphics2d::image_channel<DataType>& original ) {
    if( original.size().get_xsize() != original.size().get_ysize() ||
        !frantic::math::is_power_of_two( original.size().get_xsize() ) ) {
        throw std::runtime_error( "generate_mipmap(): original texture must be square with a power of 2 for a width." );
    }

    const std::size_t logged = static_cast<std::size_t>( frantic::math::log2_uint32( original.size().get_xsize() ) );
    std::vector<frantic::graphics2d::image_channel<DataType>> result( logged + 1 );
    std::size_t index = logged - 1;
    result[logged] = original;

    for( int width = original.size().get_xsize() / 2; width > 0; width /= 2, --index ) {
        const frantic::graphics2d::image_channel<DataType>& past = result[index + 1];
        frantic::graphics2d::image_channel<DataType> current( frantic::graphics2d::size2( width, width ) );

        Combiner cmb( past, current );

        for( int currX = 0; currX < width; ++currX ) {
            for( int currY = 0; currY < width; ++currY ) {
                cmb( currX, currY );
            }
        }

        result[index] = current;
    }

    return result;
}

/**
 * Find the texture in a mipmap of the resolution closest to the resolution being utilized.
 *
 * @tparam    DataType    The data type used by the mipmap.
 * @param     mipmap      A vector representing a mipmap, as generated by generate_mipmap().
 * @param     resolution  The target resolution to find a texture closest to.
 * @return                The image_channel in the mipmap closest to resolution.
 */
template <typename DataType>
inline const frantic::graphics2d::image_channel<DataType>&
nearest_mipmap_resolution( const std::vector<frantic::graphics2d::image_channel<DataType>>& mipmap,
                           std::size_t resolution ) {
    const std::size_t size = mipmap.size();
    if( resolution <= 1 || size <= 1 ) {
        if( size < 1 ) {
            throw std::runtime_error( "nearest_mipmap_resolution(): Empty mipmap." );
        }

        return mipmap.front();
    } else if( resolution >= mipmap.back().width() || mipmap.size() == 2 ) {
        return mipmap.back();
    } else {
        const float logged = std::log( static_cast<float>( resolution ) ) / boost::math::constants::ln_two<float>();
        const float loggedFloor = std::floor( logged );
        const std::size_t incremented = static_cast<std::size_t>( loggedFloor + 1.0f );

        if( size <= incremented ) {
            throw std::runtime_error( "nearest_resolution(): Malformed mipmap." );
        }

        if( logged - loggedFloor > 0.5f ) {
            return mipmap[incremented];
        }
        return mipmap[static_cast<std::size_t>( loggedFloor )];
    }
}

/**
 * A Combiner which generates mipmaps based on a four pixel mean
 */
template <typename DataType>
class four_pixel_mean {
  protected:
    typedef image_channel<DataType> image_t;
    const image_t& m_past;
    image_t& m_current;

  public:
    four_pixel_mean( const image_t& past, image_t& current )
        : m_past( past )
        , m_current( current ) {}

    virtual void operator()( int currX, int currY ) {
        const int pastX = currX * 2;
        const int pastY = currY * 2;

        const DataType mean = ( m_past.at( pastX, pastY ) + m_past.at( pastX + 1, pastY ) +
                                m_past.at( pastX, pastY + 1 ) + m_past.at( pastX + 1, pastY + 1 ) ) /
                              4;
        m_current.set_pixel( frantic::graphics2d::vector2( currX, currY ), mean );
    }
};
} // namespace graphics2d
} // namespace frantic