#include "monotonic_function_space.h"

#include <cmath>
#include <cassert>
#include <vector>

namespace toposmooth
{
// This function computes and returns the interpolating 't' value between 0 and 1 based on the
// theta value of the node given by 'node_i'.
inline real_t tree_ops_t::t_i( node_t node_i ) const
{
    const real_t t = ( .5 + .5 * cos( theta( node_i ) ) );
    return t;
}

// This function computes and returns the derivative of the
// interpolating 't' value between 0 and 1 based on the
// theta value of the node given by 'node_i'.
inline real_t tree_ops_t::dt_i( node_t node_i ) const
{
    const real_t dt = -.5 * sin( theta( node_i ) );
    return dt;
}

// This function computes and returns the second derivative of the
// interpolating 't' value between 0 and 1 based on the
// theta value of the node given by 'node_i'.
inline real_t tree_ops_t::d2t_i( node_t node_i ) const
{
    const real_t d2t = -.5 * cos( theta( node_i ) );
    return d2t;
}

// This function computes and returns the scalar value of the node given by 'node_i'.
inline real_t tree_ops_t::f_i( node_t node_i ) const
{
    if( is_fixed( node_i ) )
    {
        return fixed_value( node_i );
    }
    
    const real_t t = t_i( node_i );
    const real_t max_leaf_val = max_leaf_value( node_i );
    
    const node_t parent_i = parent_node( node_i );
    const real_t parent_val = f_i( parent_i );
    
    const real_t val = max_leaf_val + t * ( parent_val - max_leaf_val );
    return val;
}

// This function computes and returns the scalar value of the node given by 'node_i'.
// It does not use recursion, but instead uses dynamic memory allocation.
// The optional parameter 'depth' is a hint as to how many nodes deep 'node_i' is;
// If 'node_i' is the root, its depth is 0.
inline real_t tree_ops_t::f_i_nonrecursive( node_t node_i, unsigned int depth ) const
{
    if( is_fixed( node_i ) )
    {
        return fixed_value( node_i );
    }
    
    std::vector< node_t > node_i_and_lineage;
    node_i_and_lineage.reserve( depth + 1 );
    node_i_and_lineage.push_back( node_i );
    
    node_t parent = node_i;
    while( true )
    {
        parent = parent_node( parent );
        if( parent == node_i_and_lineage.back() ) break;
        
        node_i_and_lineage.push_back( parent );
    }
    
    assert( !node_i_and_lineage.empty() );
    return f_i( &node_i_and_lineage[0], node_i_and_lineage.size() );
}

// This function computes and returns the scalar value of the node given by 'node_i_and_lineage[0]',
// where 'node_i_and_lineage' is an array of length 'num_nodes', which must be at least one,
// and 'node_i_and_lineage[i+1]' is the parent of 'node_i_and_lineage[i]', all the way to the
// root.  The node given by 'node_i_and_lineage[num_nodes-1]' must always be fixed.
inline real_t tree_ops_t::f_i( node_t* node_i_and_lineage, int num_nodes ) const
{
    assert( node_i_and_lineage );
    assert( num_nodes >= 1 );
    assert( is_fixed( node_i_and_lineage[ num_nodes-1 ] ) );
    
    real_t result = fixed_value( node_i_and_lineage[ num_nodes-1 ] );
    for( int i = num_nodes-2; i >= 0; --i )
    {
        const node_t node_i = node_i_and_lineage[ i ];
        
        const real_t t = t_i( node_i );
        const real_t max_leaf_val = max_leaf_value( node_i );
        
        result = max_leaf_val + t * ( result - max_leaf_val );
    }
    return result;
}

// This function computes and returns the derivative of the scalar value of the node
// given by 'node_i' with respect to the theta of 'node_j'.
//
// tested (by finite differencing comparison)
real_t tree_ops_t::f_i_d_theta_j( node_t node_i, node_t node_j, bool* constant_out ) const
{
    // We can't take the derivative with respect to a fixed node.
    assert( !is_fixed( node_j ) );
    
    // Until we pass the constant-tests, assume we are constant.
    // (This simplifies the "return 0" code.)
    if( constant_out ) *constant_out = true;
    
    // The derivative of a fixed node is always 0.
    if( is_fixed( node_i ) ) return 0.;
    
    /// The derivative of a node with respect to a descendent is always 0.
    // if( is_ancestor( node_i, node_j ) ) return 0.;
    /// The derivative of a node with respect to a node that is *not* the same node
    /// or an ancestor --- perhaps from a different tree --- is always 0.
    if( !( node_i == node_j || is_ancestor( node_j, node_i ) ) ) return 0.;
    
    // If we made it here, then the result is not constantly zero.
    if( constant_out ) *constant_out = false;
    
    return f_i_d_theta_j_nonconstant( node_i, node_j );
}
// This version doesn't perform any constant checks, so the user must only call this version
// with the same inputs for which a call to the usual function would have return
// true in 'constant_out'.
real_t tree_ops_t::f_i_d_theta_j_nonconstant( node_t node_i, node_t node_j ) const
{
    real_t result = 1.;
    node_t parent = node_i;
    while( parent != node_j )
    {
        result *= t_i( parent );
        parent = parent_node( parent );
    }
    const real_t f_parent_parent = f_i( parent_node( parent ) );
    result *= ( f_parent_parent - max_leaf_value( parent ) );
    result *= dt_i( parent );
    
    return result;
}

// A version of f_i_d_theta_j_nonconstant() that doesn't make any recursive calls.
// It takes an additional, optional parameter estimating the depth of 'j' in the tree.
real_t tree_ops_t::f_i_d_theta_j_nonconstant_nonrecursive( node_t node_i, node_t node_j, unsigned int depth_j ) const
{
    real_t result = 1.;
    node_t parent = node_i;
    while( parent != node_j )
    {
        result *= t_i( parent );
        parent = parent_node( parent );
    }
    const real_t f_parent_parent = f_i_nonrecursive( parent_node( parent ), depth_j );
    result *= ( f_parent_parent - max_leaf_value( parent ) );
    result *= dt_i( parent );
    
    return result;
}

// This function computes and returns the second derivative of the scalar value of the node
// given by 'node_i' with respect to both the theta of 'node_j' and the theta of 'node_k'.
// Upon return, the optional output parameter 'constant_out' will be
// set to true if the result will always be zero (and false otherwise).
//
// untested
real_t tree_ops_t::f_i_d_theta_j_theta_k( node_t node_i, node_t node_j, node_t node_k, bool* constant_out ) const
{
    // We can't take the derivative with respect to a fixed node.
    assert( !is_fixed( node_j ) );
    assert( !is_fixed( node_k ) );
    
    // Until we pass the constant-tests, assume we are constant.
    // (This simplifies the "return 0" code.)
    if( constant_out ) *constant_out = true;
    
    // The derivative of a fixed node is always 0.
    if( is_fixed( node_i ) ) return 0.;
    
    /// The derivative of a node with respect to a descendent is always 0.
    // if( is_ancestor( node_i, node_j ) ) return 0.;
    // if( is_ancestor( node_i, node_k ) ) return 0.;
    /// The derivative of a node with respect to a node that is *not* the same node
    /// or an ancestor --- perhaps from a different tree --- is always 0.
    if( !( node_i == node_j || is_ancestor( node_j, node_i ) ) ) return 0.;
    if( !( node_i == node_k || is_ancestor( node_k, node_i ) ) ) return 0.;
    
    
    // If we made it here, then the result is not constantly zero.
    if( constant_out ) *constant_out = false;
    
    return f_i_d_theta_j_theta_k_nonconstant( node_i, node_j, node_k );
}
// This version doesn't perform any constant checks, so the user must only call this version
// with the same inputs for which a call to the usual function would have return
// true in 'constant_out'.
real_t tree_ops_t::f_i_d_theta_j_theta_k_nonconstant( node_t node_i, node_t node_j, node_t node_k ) const
#if 0
{
    // Make sure k is older than j.
    if( is_ancestor( node_j, node_k ) )
    {
        const node_t temp = node_j;
        node_j = node_k;
        node_k = temp;
    }
    
    real_t result = 1.;
    node_t parent = node_i;
    // node_k is oldest, so terminate when we see it.
    while( parent != node_k )
    {
        if( parent == node_j )
        {
            result *= dt_i( parent );
        }
        else
        {
            result *= t_i( parent );
        }
        
        parent = parent_node( parent );
    }
    const real_t f_parent_parent = f_i( parent_node( parent ) );
    result *= ( f_parent_parent - max_leaf_value( parent ) );
    // If node_k and node_j are the same, compute the second derivative.
    if( node_k == node_j )
    {
        result *= d2t_i( parent );
    }
    else
    {
        result *= dt_i( parent );
    }
    
    return result;
}
#else
{
    // Make sure k is older than j.
    // UPDATE: is_ancestor() is slow, so let's work around it.
    /*
    if( is_ancestor( node_j, node_k ) )
    {
        const node_t temp = node_j;
        node_j = node_k;
        node_k = temp;
    }
    */
    
    real_t result = 1.;
    node_t parent = node_i;
    // Terminate when we see the older node.
    // We will implement this by keeping a flag indicating whether we have already seen one,
    // and, once we have, terminate when we see either.
    // UPDATE: The first one should terminate us if node_k == node_j.
    bool seen_one = false;
    if( node_j == node_k ) seen_one = true;
    while( !seen_one || !(parent == node_k || parent == node_j) )
    {
        if( parent == node_j || parent == node_k )
        {
            result *= dt_i( parent );
            seen_one = true;
        }
        else
        {
            result *= t_i( parent );
        }
        
        parent = parent_node( parent );
    }
    const real_t f_parent_parent = f_i( parent_node( parent ) );
    result *= ( f_parent_parent - max_leaf_value( parent ) );
    // If node_k and node_j are the same, compute the second derivative.
    if( node_k == node_j )
    {
        result *= d2t_i( parent );
    }
    else
    {
        result *= dt_i( parent );
    }
    
    return result;
}
#endif

// A version of f_i_d_theta_j_theta_k_nonconstant() that doesn't make any recursive calls.
// It takes an additional, optional parameter estimating the depth of 'j' or 'k' in the tree,
// whichever is smaller.
real_t tree_ops_t::f_i_d_theta_j_theta_k_nonconstant_nonrecursive( node_t node_i, node_t node_j, node_t node_k, unsigned int depth_j ) const
{
    // Make sure k is older than j.
    // UPDATE: is_ancestor() is slow, so let's work around it.
    /*
    if( is_ancestor( node_j, node_k ) )
    {
        const node_t temp = node_j;
        node_j = node_k;
        node_k = temp;
    }
    */
    
    real_t result = 1.;
    node_t parent = node_i;
    // Terminate when we see the older node.
    // We will implement this by keeping a flag indicating whether we have already seen one,
    // and, once we have, terminate when we see either.
    // UPDATE: The first one should terminate us if node_k == node_j.
    bool seen_one = false;
    if( node_j == node_k ) seen_one = true;
    while( !seen_one || !(parent == node_k || parent == node_j) )
    {
        if( parent == node_j || parent == node_k )
        {
            result *= dt_i( parent );
            seen_one = true;
        }
        else
        {
            result *= t_i( parent );
        }
        
        parent = parent_node( parent );
    }
    const real_t f_parent_parent = f_i_nonrecursive( parent_node( parent ), depth_j );
    result *= ( f_parent_parent - max_leaf_value( parent ) );
    // If node_k and node_j are the same, compute the second derivative.
    if( node_k == node_j )
    {
        result *= d2t_i( parent );
    }
    else
    {
        result *= dt_i( parent );
    }
    
    return result;
}

}
// ~namespace toposmooth
