chain.h

#pragma once

#include <funcy/concepts.h>
#include <funcy/util/compute_chain.h>
#include <funcy/util/compute_sum.h>
#include <funcy/util/derivative_wrappers.h>
#include <funcy/util/evaluate_if_present.h>
#include <funcy/util/type_traits.h>

#include <utility>

namespace funcy
{
    template < class >
    struct Chainer;

    namespace mathop
    {
        template < Function F, Function G >
        struct Chain : Chainer< Chain< F, G > >
        {
        private:
            using FArg = decltype( std::declval< G >()() );

            template < class IndexedArgX, class IndexedArgY, class IndexedFArgX,
                       class IndexedFArgY >
            using D2LazyType =
                ComputeSum< ComputeChainD2< F, D1< G, IndexedArgX >, D1< G, IndexedArgY >,
                                            IndexedFArgX, IndexedFArgY >,
                            ComputeChainD1< F, D2< G, IndexedArgX, IndexedArgY >, IndexedFArgX > >;

            template < class IndexedArgX, class IndexedArgY, class IndexedArgZ, class IndexedFArgX,
                       class IndexedFArgY, class IndexedFArgZ >
            using D3LazyType = ComputeSum<
                ComputeChainD3< F, D1< G, IndexedArgX >, D1< G, IndexedArgY >, D1< G, IndexedArgZ >,
                                IndexedFArgX, IndexedFArgY, IndexedFArgZ >,
                ComputeChainD2< F, D2< G, IndexedArgX, IndexedArgZ >, D1< G, IndexedArgY >,
                                IndexedFArgX, IndexedFArgY >,
                ComputeChainD2< F, D1< G, IndexedArgX >, D2< G, IndexedArgY, IndexedArgZ >,
                                IndexedFArgX, IndexedFArgY >,
                ComputeChainD2< F, D2< G, IndexedArgX, IndexedArgY >, D1< G, IndexedArgZ >,
                                IndexedFArgX, IndexedFArgZ >,
                ComputeChainD1< F, D3< G, IndexedArgX, IndexedArgY, IndexedArgZ >, IndexedFArgX > >;

        public:
            constexpr Chain( const F& f_, const G& g_ ) : g( g_ ), f( f_ )
            {
                update_if_present( f, g() );
            }

            constexpr Chain( const F& f_, G&& g_ ) : g( std::move( g_ ) ), f( f_ )
            {
                update_if_present( f, g() );
            }

            constexpr Chain( F&& f_, const G& g_ ) : g( g_ ), f( std::move( f_ ) )
            {
                update_if_present( f, g() );
            }

            constexpr Chain( F&& f_, G&& g_ ) : g( std::move( g_ ) ), f( std::move( f_ ) )
            {
                update_if_present( f, g() );
            }

            template < class Arg >
            void update( const Arg& x )
            {
                update_if_present( g, x );
                update_if_present( f, g() );
            }

            template < int index, class Arg >
            void update( const Arg& x )
            {
                update_if_present< index >( g, x );
                update_if_present( f, g() );
            }

            template < class... IndexedArgs >
            void bulk_update( IndexedArgs&&... args )
            {
                bulk_update_if_present( g, std::forward< IndexedArgs >( args )... );
                update_if_present( f, g() );
            }

            constexpr decltype( auto ) d0() const noexcept
            {
                return f();
            }

            template < int id, class Arg, class IndexedArg = IndexedType< Arg, id >,
                       class IndexedFArg = IndexedType< FArg, id > >
            auto d1( Arg const& dx ) const
                requires( ComputeChainD1< F, D1< G, IndexedArg >, IndexedFArg >::present )
            {
                return chain< IndexedFArg >( f, D1< G, IndexedArg >( g, dx ) )();
            }

            template < int idx, int idy, class ArgX, class ArgY,
                       class IndexedArgX = IndexedType< ArgX, idx >,
                       class IndexedArgY = IndexedType< ArgY, idy >,
                       class IndexedFArgX = IndexedType< FArg, idx >,
                       class IndexedFArgY = IndexedType< FArg, idy > >
            auto d2( ArgX const& dx, ArgY const& dy ) const requires(
                D2LazyType< IndexedArgX, IndexedArgY, IndexedFArgX, IndexedFArgY >::present )
            {
                return sum(
                    chain< IndexedFArgX, IndexedFArgY >( f, D1< G, IndexedArgX >( g, dx ),
                                                         D1< G, IndexedArgY >( g, dy ) ),
                    chain< IndexedFArgX >( f, D2< G, IndexedArgX, IndexedArgY >( g, dx, dy ) ) )();
            }

            template < int idx, int idy, int idz, class ArgX, class ArgY, class ArgZ,
                       class IndexedArgX = IndexedType< ArgX, idx >,
                       class IndexedArgY = IndexedType< ArgY, idy >,
                       class IndexedArgZ = IndexedType< ArgZ, idz >,
                       class IndexedFArgX = IndexedType< FArg, idx >,
                       class IndexedFArgY = IndexedType< FArg, idy >,
                       class IndexedFArgZ = IndexedType< FArg, idz > >
            auto d3( ArgX const& dx, ArgY const& dy, ArgZ const& dz ) const
                requires( D3LazyType< IndexedArgX, IndexedArgY, IndexedArgZ, IndexedFArgX,
                                      IndexedFArgY, IndexedFArgZ >::present )
            {
                D1< G, IndexedArgX > dGdx( g, dx );
                D1< G, IndexedArgY > dGdy( g, dy );
                D1< G, IndexedArgZ > dGdz( g, dz );
                return sum(
                    chain< IndexedFArgX, IndexedFArgY, IndexedFArgZ >( f, dGdx, dGdy, dGdz ),
                    chain< IndexedFArgX, IndexedFArgY >(
                        f, D2< G, IndexedArgX, IndexedArgZ >( g, dx, dz ), dGdy ),
                    chain< IndexedFArgX, IndexedFArgY >(
                        f, dGdx, D2< G, IndexedArgY, IndexedArgZ >( g, dy, dz ) ),
                    chain< IndexedFArgX, IndexedFArgZ >(
                        f, D2< G, IndexedArgX, IndexedArgY >( g, dx, dy ), dGdz ),
                    chain< IndexedFArgX >(
                        f, D3< G, IndexedArgX, IndexedArgY, IndexedArgZ >( g, dx, dy, dz ) ) )();
            }

        private:
            G g;
            F f;
        };
    } // namespace mathop
} // namespace funcy