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

#include <frantic/channels/channel_map_adaptor.hpp>
#include <frantic/geometry/trimesh3.hpp>
#include <frantic/particles/streams/particle_istream.hpp>

#include <boost/random.hpp>
#include <boost/scoped_array.hpp>

namespace frantic {
namespace particles {
namespace streams {

struct steal {}; // Empty class for generating overloaded constructors.

/**
 * @tparam SurfaceCollection Adapter for geometry representing a surface. Must implement:
 *            void get_native_map( frantic::channels::channel_map& outNativeMap );
 *            void set_channel_map( const frantic::channels::channel_map& seedMap );
 *            std::size_t surface_count();
 *            std::size_t element_count( std::size_t surfaceIndex );
 *            float element_area( std::size_t surfaceIndex, std::size_t elementIndex );
 *            template <class RandomGen> void seed_particle( char* outParticle, std::size_t surfaceIndex, std::size_t
 * elementIndex, RandomGen& randomnessGenerator );
 */
template <class SurfaceCollection>
class surface_particle_istream : public particle_istream {
  public:
    surface_particle_istream( const SurfaceCollection& surfaceCollection );

    /**
     * @overload If SurfaceCollection is non-copyable, this alternative constructor uses the 'steal' member function to
     * steal data from the passed reference.
     * @param surfaceCollection The surface data. The referenced object will not be valid after passing it to this
     * constructor.
     */
    surface_particle_istream( SurfaceCollection& surfaceCollection, steal );

#if _MSC_VER >= 1600
    surface_particle_istream( SurfaceCollection&& surfaceCollection );
#endif

    virtual ~surface_particle_istream();

    /**
     * Sets the total number of particles generated to a specific value.
     * @param particleCount The explicit number of particles to generate.
     */
    void set_particle_count( boost::int64_t particleCount );

    /**
     * Sets the total number of particle generated by using the surface collection's total area divided by the average
     * spacing
     * @param averageSpacing The desired average spacing between particles.
     */
    void set_particle_spacing( float averageSpacing );

    /**
     * Sets the value used to generate randomness, so that differerent or predicatable randomness can be produced.
     * @param seed The new seed value for the random number generator.
     */
    void set_random_seed( boost::uint32_t seed );

    /**
     * Enable or disable element area-based weighting when seeding particles. If disabled, all elements are equally
     * probable despite variations in area.
     * @param enabled Whether area-based weighting should be used.
     */
    // void set_weighting_enabled( bool enabled );

    virtual void close();

    // The stream can return its filename or other identifier for better error messages.
    virtual frantic::tstring name() const;

    // This is the size of the particle structure which will be loaded, in bytes.
    virtual std::size_t particle_size() const;

    // Returns the number of particles, or -1 if unknown
    virtual boost::int64_t particle_count() const;
    virtual boost::int64_t particle_index() const;
    virtual boost::int64_t particle_count_left() const;

    virtual boost::int64_t particle_progress_count() const;
    virtual boost::int64_t particle_progress_index() const;

    virtual boost::int64_t particle_count_guess() const;

    // This allows you to change the particle layout that's being loaded on the fly, in case it couldn't
    // be set correctly at creation time.
    virtual void set_channel_map( const frantic::channels::channel_map& pcm );

    // This is the particle channel map which specifies the byte layout of the particle structure that is being used.
    virtual const frantic::channels::channel_map& get_channel_map() const;

    // This is the particle channel map which specifies the byte layout of the input to this stream.
    virtual const frantic::channels::channel_map& get_native_channel_map() const;

    /** This provides a default particle which should be used to fill in channels of the requested channel map
     *	which are not supplied by the native channel map.
     *	IMPORTANT: Make sure the buffer you pass in is at least as big as particle_size() bytes.
     */
    virtual void set_default_particle( char* rawParticleBuffer );

    // This reads a particle into a buffer matching the channel_map.
    // It returns true if a particle was read, false otherwise.
    // IMPORTANT: Make sure the buffer you pass in is at least as big as particle_size() bytes.
    virtual bool get_particle( char* rawParticleBuffer );

    // This reads a group of particles. Returns false if the end of the source
    // was reached during the read.
    virtual bool get_particles( char* rawParticleBuffer, std::size_t& numParticles );

  private:
    void init();

  private:
    SurfaceCollection m_surfaceData;

    std::vector<float> m_surfaceAreas;                     // The total area of each of the N surfaces.
    std::vector<std::vector<float>> m_surfaceElementAreas; // There are N surfaces, each with M(N) elements with area.

    float m_totalArea;
    float m_densityScale;
    bool m_useWeighting;

    boost::mt19937 m_rndGen;

    boost::int64_t m_particleIndex, m_particleCount;

    boost::scoped_array<char> m_defaultParticle;

    frantic::channels::channel_map m_outMap, m_nativeMap;
    frantic::channels::channel_cvt_accessor<float> m_densityAccessor;
};

template <class SurfaceCollection>
inline surface_particle_istream<SurfaceCollection>::surface_particle_istream(
    const SurfaceCollection& surfaceCollection )
    : m_surfaceData( surfaceCollection ) {
    init();
}

template <class SurfaceCollection>
inline surface_particle_istream<SurfaceCollection>::surface_particle_istream( SurfaceCollection& surfaceCollection,
                                                                              steal ) {
    m_surfaceData.steal( surfaceCollection );

    init();
}

#if _MSC_VER >= 1600
template <class SurfaceCollection>
inline surface_particle_istream<SurfaceCollection>::surface_particle_istream( SurfaceCollection&& surfaceCollection )
    : m_surfaceData( std::move( surfaceCollection ) ) {
    init();
}
#endif

template <class SurfaceCollection>
inline surface_particle_istream<SurfaceCollection>::~surface_particle_istream() {}

template <class SurfaceCollection>
inline void surface_particle_istream<SurfaceCollection>::init() {
    m_particleIndex = -1;
    m_particleCount = 0;
    m_densityScale = 1.f;
    m_useWeighting = true;

    m_surfaceData.get_native_map( m_nativeMap );
    m_nativeMap.define_channel<float>( _T("Density") );
    m_nativeMap.end_channel_definition();

    frantic::channels::channel_map tempMap;
    tempMap.define_channel<frantic::graphics::vector3f>( _T("Position") );
    tempMap.end_channel_definition();

    set_channel_map( tempMap );

    // Collect each surface element's area.
    std::size_t numSurfaces = m_surfaceData.surface_count();

    /*if( numSurfaces == 0 )
      throw std::runtime_error( "surface_particle_istream cannot handle an empty surface collection" );*/

    m_surfaceAreas.assign( numSurfaces, 0.f );
    m_surfaceElementAreas.clear();
    m_surfaceElementAreas.resize( numSurfaces );

    m_totalArea = 0.f;

    bool hasElements = false;

    // There are multiple surfaces (ie. distinct objects with potentially different object to world transforms) to
    // accumulate areas for.
    for( std::size_t i = 0; i < numSurfaces; ++i ) {
        std::size_t numElements = m_surfaceData.element_count( i );

        if( numElements > 0 )
            hasElements = true;

        std::vector<float>& elementAreas = m_surfaceElementAreas[i];

        elementAreas.reserve( numElements );

        // Calculate the surface area for each element of the current surface, and store the running total at each.
        float surfaceArea = 0.f;
        for( std::size_t j = 0; j < numElements; ++j ) {
            float elementArea = m_surfaceData.element_area( i, j );

            surfaceArea += elementArea;

            elementAreas.push_back( surfaceArea ); // We want to store the running total up to and include the current
                                                   // element for this surface.
        }

        m_totalArea += surfaceArea;

        m_surfaceAreas[i] = m_totalArea;
    }

    if( !hasElements || m_totalArea == 0.0 )
        m_surfaceAreas.clear();
    // throw std::runtime_error( "surface_particle_istream cannot handle an empty surface collection" );
}

template <class SurfaceCollection>
inline void surface_particle_istream<SurfaceCollection>::set_particle_count( boost::int64_t particleCount ) {
    if( m_totalArea > 0.f ) {
        m_particleCount = particleCount;
        m_densityScale = static_cast<float>(
            static_cast<double>( 100000ll ) /
            static_cast<double>( std::max( static_cast<boost::int64_t>( 100000ll ),
                                           m_particleCount ) ) ); // We make 100000 particles be the baseline
    }

    if( m_surfaceAreas.empty() )
        m_particleCount = 0;
}

template <class SurfaceCollection>
inline void surface_particle_istream<SurfaceCollection>::set_particle_spacing( float averageSpacing ) {
    if( m_totalArea > 0.f ) {
        // If we assume each particle is an (averageSpacing)^2 area square, we can divide the total area by that to get
        // the number of particles that fit.
        double averageArea = static_cast<double>( averageSpacing );
        averageArea = averageArea * averageArea;

        double averageCount = static_cast<double>( m_totalArea ) / averageArea;

        m_particleCount = static_cast<boost::int64_t>( std::ceil( averageCount ) );
        m_densityScale = static_cast<float>(
            static_cast<double>( 100000ll ) /
            static_cast<double>( std::max( static_cast<boost::int64_t>( 100000ll ),
                                           m_particleCount ) ) ); // We make 100000 particles be the baseline
    }

    if( m_surfaceAreas.empty() )
        m_particleCount = 0;
}

template <class SurfaceCollection>
inline void surface_particle_istream<SurfaceCollection>::set_random_seed( boost::uint32_t seed ) {
    m_rndGen.seed( seed );
}

// template <class SurfaceCollection>
// inline void surface_particle_istream<SurfaceCollection>::set_weighting_enabled( bool enabled ){
//	m_useWeighting = enabled;
// }

template <class SurfaceCollection>
inline void surface_particle_istream<SurfaceCollection>::close() {}

template <class SurfaceCollection>
inline frantic::tstring surface_particle_istream<SurfaceCollection>::name() const {
    return _T("surface_stream");
}

template <class SurfaceCollection>
inline std::size_t surface_particle_istream<SurfaceCollection>::particle_size() const {
    return m_outMap.structure_size();
}

template <class SurfaceCollection>
inline boost::int64_t surface_particle_istream<SurfaceCollection>::particle_count() const {
    return m_particleCount;
}

template <class SurfaceCollection>
inline boost::int64_t surface_particle_istream<SurfaceCollection>::particle_index() const {
    return m_particleIndex;
}

template <class SurfaceCollection>
inline boost::int64_t surface_particle_istream<SurfaceCollection>::particle_count_left() const {
    return m_particleCount - m_particleIndex - 1;
}

template <class SurfaceCollection>
inline boost::int64_t surface_particle_istream<SurfaceCollection>::particle_progress_count() const {
    return this->particle_count();
}

template <class SurfaceCollection>
inline boost::int64_t surface_particle_istream<SurfaceCollection>::particle_progress_index() const {
    return this->particle_index();
}

template <class SurfaceCollection>
inline boost::int64_t surface_particle_istream<SurfaceCollection>::particle_count_guess() const {
    return this->particle_count();
}

template <class SurfaceCollection>
inline void surface_particle_istream<SurfaceCollection>::set_channel_map( const frantic::channels::channel_map& pcm ) {
    { // Swap in a new default particle.
        boost::scoped_array<char> newDefault( new char[pcm.structure_size()] );

        pcm.construct_structure( newDefault.get() );

        if( m_defaultParticle ) {
            frantic::channels::channel_map_adaptor tempAdaptor( pcm, m_outMap );
            tempAdaptor.copy_structure( newDefault.get(), m_defaultParticle.get() );
        }

        m_defaultParticle.swap( newDefault );
    }

    m_outMap = pcm;

    m_surfaceData.set_channel_map( m_outMap );

    // Assign new accessors, etc.
    m_densityAccessor.reset();

    if( m_outMap.has_channel( _T("Density") ) )
        m_densityAccessor = m_outMap.get_cvt_accessor<float>( _T("Density") );
}

template <class SurfaceCollection>
inline const frantic::channels::channel_map& surface_particle_istream<SurfaceCollection>::get_channel_map() const {
    return m_outMap;
}

template <class SurfaceCollection>
inline const frantic::channels::channel_map&
surface_particle_istream<SurfaceCollection>::get_native_channel_map() const {
    return m_nativeMap;
}

template <class SurfaceCollection>
inline void surface_particle_istream<SurfaceCollection>::set_default_particle( char* rawParticleBuffer ) {
    boost::scoped_array<char> newDefault( new char[m_outMap.structure_size()] );
    m_outMap.copy_structure( newDefault.get(), rawParticleBuffer );
    m_defaultParticle.swap( newDefault );
}

template <class SurfaceCollection>
inline bool surface_particle_istream<SurfaceCollection>::get_particle( char* rawParticleBuffer ) {
    if( ++m_particleIndex >= m_particleCount )
        return false;

    // std::size_t i = 0, iEnd = m_surfaceData.surface_count(), j = 0, jEnd = 0;
    std::size_t surfaceIndex = 0, elementIndex = 0;

    // if( m_useWeighting ){
    boost::variate_generator<boost::mt19937&, boost::uniform_real<float>> rnd(
        m_rndGen, boost::uniform_real<float>( 0, m_totalArea ) );

    float target = rnd();

    // Find the index of the first surface which is greater than 'target'.
    surfaceIndex = static_cast<std::size_t>( std::upper_bound( m_surfaceAreas.begin(), m_surfaceAreas.end(), target ) -
                                             m_surfaceAreas.begin() );

    float elemTarget = target;
    if( surfaceIndex > 0 )
        elemTarget -= m_surfaceAreas[surfaceIndex - 1];

    std::vector<float>& elementAreas = m_surfaceElementAreas[surfaceIndex];

    // Find the index of the first element in the surface that is greater than 'elemTarget'. This is where we will plant
    // the new particle.
    elementIndex = static_cast<std::size_t>( std::upper_bound( elementAreas.begin(), elementAreas.end(), elemTarget ) -
                                             elementAreas.begin() );

    /*}else{
      boost::variate_generator< boost::mt19937&, boost::uniform_int<std::size_t> > rndSurf( m_rndGen,
    boost::uniform_int<std::size_t>(0, m_surfaceData.surface_count() - 1) );

      i = rndSurf();

      jEnd = m_surfaceData.element_count( i );

      boost::variate_generator< boost::mt19937&, boost::uniform_int<std::size_t> > rndElem( m_rndGen,
    boost::uniform_int<std::size_t>(0, jEnd - 1) );

      j = rndElem();
    }*/

    m_outMap.copy_structure( rawParticleBuffer, m_defaultParticle.get() );

    if( elementIndex < m_surfaceData.element_count( surfaceIndex ) ) {
        m_surfaceData.seed_particle( rawParticleBuffer, surfaceIndex, elementIndex, m_rndGen );

        if( m_densityAccessor.is_valid() )
            m_densityAccessor.set( rawParticleBuffer, m_densityScale );
    }

    return true;
}

template <class SurfaceCollection>
inline bool surface_particle_istream<SurfaceCollection>::get_particles( char* rawParticleBuffer,
                                                                        std::size_t& numParticles ) {
    std::size_t particleSize = particle_size();
    for( std::size_t i = 0; i < numParticles; ++i ) {
        if( !get_particle( rawParticleBuffer ) ) {
            numParticles = i;
            return false;
        }
        rawParticleBuffer += particleSize;
    }
    return true;
}

particle_istream_ptr create_surface_particle_istream_using_count( const frantic::geometry::trimesh3& mesh,
                                                                  const boost::uint64_t particleCount,
                                                                  const boost::uint32_t randomSeed = 5489 );
particle_istream_ptr create_surface_particle_istream_using_spacing( const frantic::geometry::trimesh3& mesh,
                                                                    const float particleSpacing,
                                                                    const boost::uint32_t randomSeed = 5489 );

} // namespace streams
} // namespace particles
} // namespace frantic