//////////////////////////////////////////////////////////////////////////////
//
// Asura_KDTree.h
//
// Created by NJD - 12 September 2006
//
// Copyright (C) 2006 Rebellion, All Rights Reserved.
//
//////////////////////////////////////////////////////////////////////////////

// B E G I N   F I L E   W R A P P E R ///////////////////////////////////////

#ifdef ASURA_USE_PRAGMA_ONCE
#pragma once
#endif

#ifndef __ASURA_KDTREE_H__
#define __ASURA_KDTREE_H__ 1

// E N D   F I L E   W R A P P E R ///////////////////////////////////////////

// I N C L U D E S ///////////////////////////////////////////////////////////

#include "Core/Asura_Maths.h"

//////////////////////////////////////////////////////////////////////////////
//
// Asura_KDTree_DataType
//
// A suggested type for use with the Asura_KDTree. Any other type can be used,
// however, provided it supports the interface detailed below.
//
//////////////////////////////////////////////////////////////////////////////

class Asura_KDTree_DataType
{
public:
    Asura_KDTree_DataType() : m_uDummyValueX( 0 ), m_uPlaneSplitFlag( 0 ), m_uDummyValueZ( 0 ) {}

    u_int GetSplittingPlane() const { return m_uPlaneSplitFlag & 3; }
    void  SetSplittingPlane( u_int uSplittingPlane ) { m_uPlaneSplitFlag &= uSplittingPlane | (~3); }

    void GetPosition(Asura_Vector_3& xPosition) const { xPosition.SetX(m_afPosition[0]); xPosition.SetY(m_afPosition[1]); xPosition.SetZ(m_afPosition[2]); }
    void SetPosition( const Asura_Vector_3& xPosition ) { m_afPosition[ 0 ] = xPosition.x; m_afPosition[ 1 ] = xPosition.y; m_afPosition[ 2 ] = xPosition.z; }

    float GetElement( const u_int uAxis ) const { return m_afPosition[ uAxis ]; }

    union
    {
        float	  m_afPosition[3];
        struct
        {
            u_int m_uDummyValueX;
            u_int m_uPlaneSplitFlag;
            u_int m_uDummyValueZ;
        };
    };
};

//////////////////////////////////////////////////////////////////////////////
//
// Asura_KDTree
//
// Pass the tree an array of DataTypes and set the size of the array.
// The data array should already be sorted into a format that the algorithm can use.
// Use ConstructLeftBalancedTree for this, during tools.
// After calling CalculateNearestElements the elements can be accessed, from the nearest
// to the furthest away, using GetNearestElement, up to MAX_PROXIMAL_RESULTS.
//
// Types used with Asura_KDTree must support the following:
//  u_int GetSplittingPlane( void ) const;
//  void  SetSplittingPlane( u_int uSplittingPlane );
//  const Asura_Vector_3& GetPosition() const;
//  float GetElement( const u_int uAxis ) const;
//
// The splitting plane will be filled in when ConstructLeftBalancedTree is called.
// The data must be preserved between calling ConstructLeftBalancedTree and CalculateNearestElements.
// The data array must be freed by the caller after use.
//
//////////////////////////////////////////////////////////////////////////////

template < class DataType >
class Asura_KDTree
{
public:
    Asura_KDTree() : m_uNumElements( 0 ), m_pxDataArray( NULL ), m_axResultIndexDistances(), m_uNumResults( 0 ) {}
    ~Asura_KDTree() {}

    // Set/get the number of elements and the data itself.
    void SetNumElements( const u_int uNumElements ) { m_uNumElements = uNumElements; }
    void SetDataArray( DataType* pxDataArray )      { m_pxDataArray = pxDataArray; }
    u_int     GetNumElements() const { return m_uNumElements; }
    DataType* GetDataArray()   const { return m_pxDataArray; }

    // The number of results that will be calculated.
    enum { MAX_PROXIMAL_RESULTS = 6 };
    u_int           CalculateNearestElements( const class Asura_Vector_3& xPosition );
    const DataType& GetNearestElement( const int iResultIndex ) const { return m_pxDataArray[ m_axResultIndexDistances[ iResultIndex ].m_uIndex ]; }

    void ConstructLeftBalancedTree();

	void Reset(void)	{ m_uNumElements = 0; m_pxDataArray = NULL; m_uNumResults = 0; }

private:
    void BalanceSegment( DataType* pxBalancedArray, DataType** ppxOriginalArray, const u_int uIndex, const u_int uStart, const u_int uEnd );
    void MedianSplit( DataType** ppxArray, const u_int uStart, const u_int uEnd, const u_int uMedian, const u_int uAxis );

    static int CompareRIDs( const void* pOne, const void* pTwo )
    {
        const ResultIndexDistanceSq* pxRIDOne = static_cast< const ResultIndexDistanceSq* >( pOne );
        const ResultIndexDistanceSq* pxRIDTwo = static_cast< const ResultIndexDistanceSq* >( pTwo );
        if ( pxRIDOne->m_fDistanceSq < pxRIDTwo->m_fDistanceSq )
        {
            return -1;
        }
        else if ( pxRIDOne->m_fDistanceSq > pxRIDTwo->m_fDistanceSq )
        {
            return 1;
        }
        else
        {
            return 0;
        }
    }

    u_int     m_uNumElements;
    DataType* m_pxDataArray;
    struct ResultIndexDistanceSq
    {
        ResultIndexDistanceSq() : m_uIndex( 0 ), m_fDistanceSq( Asura_Maths::LargeFloat ) {}
        u_int m_uIndex;
        float m_fDistanceSq;
    }         m_axResultIndexDistances[ MAX_PROXIMAL_RESULTS ];
    u_int     m_uNumResults;
    // Not sure what the stack size should be. This was the value used in code elsewhere.
    enum { STACK_SIZE = 32 };
    u_int     m_auStack[ STACK_SIZE ];

    Asura_Vector_3 g_xBBMin;
    Asura_Vector_3 g_xBBMax;
};

//////////////////////////////////////////////////////////////////////////////
//
// 
//
//////////////////////////////////////////////////////////////////////////////

template< class DataType >
u_int Asura_KDTree< DataType >::CalculateNearestElements( const class Asura_Vector_3& xPosition )
{
    if ( m_uNumElements == 0 || m_pxDataArray == NULL )
    {
        Asura_Assert( false, "Asura_KDTree::CalculateNearestElements - no elements" );
        return(0);
    }

    float fMaxDistanceSquared = 0;
    m_uNumResults = 0;
    u_int uFurthestSample = 0;
    u_int* puStack = m_auStack;

    *puStack++ = 0;

    u_int uIndex = 1;

    while(true)
    {
        const DataType& xData = m_pxDataArray[ uIndex - 1 ];
        const u_int uCurrentDataIndex = uIndex;

        const u_int uAxis = xData.GetSplittingPlane();

        const float fPlaneDist = xPosition[uAxis] - xData.GetElement( uAxis );

        // implicit children: 2*n is left, 2*n+1 is right
        u_int uLeftIndex = uIndex * 2;
        u_int uRightIndex = uLeftIndex;	// we'll add one to whichever one is needed in a minute

        if (uLeftIndex < m_uNumElements - 1)
        {
            // effectively this code swaps the two values in the negative case, so we no longer have to proces two cases !
            if (fPlaneDist>0)
            {
                ++uRightIndex;
            }
            else
            {
                ++uLeftIndex;	// note we can only get away with this because the value is not reused, but recreated on next loop
            }

            if ( (m_uNumResults < MAX_PROXIMAL_RESULTS) || (fPlaneDist * fPlaneDist < fMaxDistanceSquared) )
            {
                *puStack++ = uLeftIndex;
            }
            *puStack++ = uRightIndex;
        }

        uIndex = *(--puStack);

        if (uIndex == 0) break;

        const float fDistanceSquared = ( xData.GetPosition() - xPosition ).MagnitudeSquared();

        // if we haven't reached the maximum number of samples, just add it
        if (m_uNumResults < MAX_PROXIMAL_RESULTS)
        {
            // keep track of the max distance found, and which sample it was
            if (fDistanceSquared > fMaxDistanceSquared)
            {
                fMaxDistanceSquared = fDistanceSquared;
                uFurthestSample = m_uNumResults;
            }

            // set the data to record the proximal sample
            m_axResultIndexDistances[ m_uNumResults ].m_uIndex      = uCurrentDataIndex - 1;
            m_axResultIndexDistances[ m_uNumResults ].m_fDistanceSq = fDistanceSquared;

            ++m_uNumResults;
        }
        // otherwise only add it if it's closer than the samples we already have
        else if (fDistanceSquared < fMaxDistanceSquared)
        {
            // replace the furthest out sample
            m_axResultIndexDistances[ uFurthestSample ].m_uIndex      = uCurrentDataIndex - 1;
            m_axResultIndexDistances[ uFurthestSample ].m_fDistanceSq = fDistanceSquared;

            // And find the new furthest out sample.
            // Since number of proximal samples is likely to always be small (<=5 say) we'll
            // just test the distances directly.
            uFurthestSample = 0;
            fMaxDistanceSquared = m_axResultIndexDistances[ 0 ].m_fDistanceSq;

            for (u_int uSample=1; uSample<MAX_PROXIMAL_RESULTS; ++uSample)
            {
                if (fMaxDistanceSquared < m_axResultIndexDistances[ uSample ].m_fDistanceSq)
                {
                    fMaxDistanceSquared = m_axResultIndexDistances[ uSample ].m_fDistanceSq;
                    uFurthestSample = uSample;
                }
            }
        }
    }

    // we now have a list of most proximal samples
    // For ease of access we should sort the list now, from nearest to furthest.
    qsort( m_axResultIndexDistances, m_uNumResults, sizeof( ResultIndexDistanceSq ), CompareRIDs );
	return(m_uNumResults);
}

template< class DataType >
void Asura_KDTree< DataType >::ConstructLeftBalancedTree()
{
    if (m_uNumElements==0) return;

    DataType* pxSortArray1 = new DataType[m_uNumElements+1];
    DataType** ppxSortArray2 = new DataType*[m_uNumElements+1];

    Asura_Bounding_Box g_xBB;
    g_xBB.Reset();

    ppxSortArray2[0] = 0;
    for (u_int u=0; u<m_uNumElements; u++)
    {
        ppxSortArray2[u+1] = &m_pxDataArray[u];
        g_xBB.Set(ppxSortArray2[u+1]->GetPosition());
        ppxSortArray2[u+1]->SetSplittingPlane(0);
    }

    g_xBBMax = Asura_Vector_3(g_xBB.MaxX,g_xBB.MaxY, g_xBB.MaxZ);
    g_xBBMin = Asura_Vector_3(g_xBB.MinX,g_xBB.MinY, g_xBB.MinZ);

    BalanceSegment(pxSortArray1, ppxSortArray2, 1, 1, m_uNumElements);

    delete [] ppxSortArray2;

    for ( u_int u = 0; u < m_uNumElements; ++u )
    {
        m_pxDataArray[ u ] = pxSortArray1[ u + 1 ];
    }

    delete [] pxSortArray1;
}

template< class DataType >
void Asura_KDTree< DataType >::BalanceSegment( DataType* pxBalancedArray, DataType** ppxOriginalArray, const u_int uIndex, const u_int uStart, const u_int uEnd )
{
    u_int uMedian = 1;

    while ((4*uMedian)<=(uEnd-uStart+1))
    {
        uMedian += uMedian;
    }

    if ((3*uMedian)<=(uEnd-uStart+1))
    {
        uMedian += uMedian;
        uMedian += uStart-1;
    }
    else
    {
        uMedian = uEnd - uMedian + 1;
    }

    float fLengthAlongX = g_xBBMax.x-g_xBBMin.x;
    float fLengthAlongY = g_xBBMax.y-g_xBBMin.y;
    float fLengthAlongZ = g_xBBMax.z-g_xBBMin.z;


    u_int uAxis;

    if (fLengthAlongX > fLengthAlongY)
    {
        if (fLengthAlongX > fLengthAlongZ)
        {
            uAxis = 0;
        }
        else
        {
            uAxis = 2;
        }
    }
    else if (fLengthAlongY > fLengthAlongZ)
    {
        uAxis = 1;
    }
    else
    {
        uAxis = 2;
    }


    MedianSplit(ppxOriginalArray, uStart, uEnd, uMedian, uAxis);

    pxBalancedArray[uIndex] = *ppxOriginalArray[uMedian];
    pxBalancedArray[uIndex].SetSplittingPlane(uAxis);

    if (uMedian>uStart)
    {
        if (uStart<uMedian-1)
        {
            const float tmp = g_xBBMax[uAxis];

            g_xBBMax[uAxis] = pxBalancedArray[uIndex].GetPosition()[uAxis];
            BalanceSegment(pxBalancedArray,ppxOriginalArray,2*uIndex,uStart, uMedian-1);
            g_xBBMax[uAxis] = tmp;
        }
        else
        {
            pxBalancedArray[2*uIndex] = *ppxOriginalArray[uStart];
        }
    }

    if (uMedian<uEnd)
    {
        if (uMedian+1 < uEnd)
        {
            const float tmp = g_xBBMin[uAxis];

            g_xBBMin[uAxis] = pxBalancedArray[uIndex].GetPosition()[uAxis];
            BalanceSegment(pxBalancedArray,ppxOriginalArray,2*uIndex+1,uMedian+1, uEnd);
            g_xBBMin[uAxis] = tmp;
        }
        else
        {
            pxBalancedArray[2*uIndex+1] = *ppxOriginalArray[uEnd];
        }
    }
}

template< class DataType >
void Asura_KDTree< DataType >::MedianSplit( DataType** ppxArray, const u_int uStart, const u_int uEnd, const u_int uMedian, const u_int uAxis )
{
    u_int uLeft = uStart;
    u_int uRight = uEnd;

    while (uRight>uLeft)
    {
        const float f = ppxArray[uRight]->GetPosition()[uAxis];

        u_int i = uLeft-1;
        u_int j = uRight;

        for (;;)
        {
            while ( ppxArray[++i]->GetPosition()[uAxis] < f );
            while ( ppxArray[--j]->GetPosition()[uAxis] > f && j>uLeft);

            if (i>=j) break;

            DataType* pxTemp = ppxArray[i];
            ppxArray[i] = ppxArray[j];
            ppxArray[j] = pxTemp;
        }

        DataType* pxTemp = ppxArray[i];
        ppxArray[i] = ppxArray[uRight];
        ppxArray[uRight] = pxTemp;

        if (i>=uMedian)
        {
            uRight = i-1;
        }
        if (i<=uMedian)
        {
            uLeft =  i+1;
        }
    }
}

//////////////////////////////////////////////////////////////////////////////
//
// 
//
//////////////////////////////////////////////////////////////////////////////

#endif //__ASURA_KDTREE_H__
