//   Copyright 2015 Patrick Putnam
//
//   Licensed under the Apache License, Version 2.0 (the "License");
//   you may not use this file except in compliance with the License.
//   You may obtain a copy of the License at
//
//       http://www.apache.org/licenses/LICENSE-2.0
//
//   Unless required by applicable law or agreed to in writing, software
//   distributed under the License is distributed on an "AS IS" BASIS,
//   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
//   See the License for the specific language governing permissions and
//   limitations under the License.

#ifndef ASSORTATIVE_SELECTOR_HPP_
#define ASSORTATIVE_SELECTOR_HPP_

#include <boost/random/discrete_distribution.hpp>
#include <boost/random/uniform_real_distribution.hpp>
#include <boost/math/constants/constants.hpp>
#include <cmath>

#include <iterator>
#include <cmath>
#include <cassert>
#include <iostream>

template < class URNG >
struct assortative_selector {
    typedef URNG rng_type;
    typedef unsigned int result_type;

    typedef boost::random::discrete_distribution< result_type, double > dist_type;

    typedef std::vector< std::pair< double, unsigned int > > population_fitness_type;
    typedef typename population_fitness_type::iterator      fitness_iterator;

    static const unsigned int BIN_SIZE = 100;
    static const unsigned int BIN_COUNT = 100;

    rng_type * m_rng;
    dist_type m_dist;
    boost::random::uniform_real_distribution< double > m_uniform;

    unsigned int m_next_bin;

    population_fitness_type m_pop_fit;
    std::vector< size_t >   m_bin_sizes;

    struct simple_comp {
        bool operator()( const std::pair< double, unsigned int > & lhs, const std::pair< double, unsigned int > & rhs ) {
            return (lhs.first < rhs.first);
        }
    };

    assortative_selector( const assortative_selector< URNG > & other ) :
        m_rng( other.m_rng)
        , m_dist( other.m_dist )
        , m_uniform()
        , m_next_bin( other.m_next_bin )
        , m_pop_fit( other.m_pop_fit )
        , m_bin_sizes( other.m_bin_sizes )
    {}

    template < class Iter >
    assortative_selector( rng_type & rng, Iter first, Iter last ) :
        m_rng( &rng )
        , m_dist()
        , m_uniform()
        , m_next_bin( -1 )
        , m_pop_fit()
    {
        init_method2( first, last );
    }

    /// allows selfing
    unsigned int operator()() {
        unsigned int bstart = 0, bsize = 0;
        if( m_next_bin == -1 ) {
            unsigned int bin_idx = m_dist( *m_rng );
            unsigned int bin_y = (int) ((sqrt( 8 * bin_idx + 1) - 1) / 2);
            m_next_bin = bin_idx - (( bin_y + 1 ) * (bin_y) / 2);

            bstart = (( bin_y == 0 ) ? 0 : m_bin_sizes[ bin_y - 1] );
            bsize  = m_bin_sizes[ bin_y ] - bstart;
//            if( bsize == 0 ) {
//                std::cerr << bin_y << ", " << m_bin_sizes[bin_y] << std::endl;
//                std::cerr << "Bin Start: " << bstart << "; Bin Size: " << bsize << std::endl;
//                std::cerr << m_bin_sizes[bin_y] << std::endl;
//                std::cerr << m_dist.probabilities()[ bin_y ] << std::endl;
//                assert( false );
//            }
        } else {
            bstart = (( m_next_bin == 0 ) ? 0 : m_bin_sizes[ m_next_bin - 1] );
            bsize =  m_bin_sizes[m_next_bin] - bstart;
            m_next_bin = -1;
        }

        if( bsize == 0 ) {
            std::cerr << "Bin Start: " << bstart << "; Bin Size: " << bsize << std::endl;
            assert( false );
        }

        double r = m_uniform( *m_rng );
        return m_pop_fit[ bstart + r * bsize ].second;
    }

    virtual ~assortative_selector() {}
protected:

    template < class Iter >
    void init_method1( Iter first, Iter last ) {
        unsigned int N = 0;
        while( first != last ) {
            m_pop_fit.push_back( std::make_pair( *first, N++ ) );
            ++first;
        }

        simple_comp comp;

        std::sort( m_pop_fit.begin(), m_pop_fit.end(), comp );

        std::vector< double > bin_mean;
        if( m_pop_fit.size() % BIN_SIZE ) {
            bin_mean.reserve( m_pop_fit.size() / BIN_SIZE + 1 );
        } else {
            bin_mean.reserve( m_pop_fit.size() / BIN_SIZE );
        }

        fitness_iterator pit = m_pop_fit.begin();
        double accum = 0.0;
        unsigned int b = BIN_SIZE;

        unsigned int bin_size_accum = 0;
        while( pit != m_pop_fit.end() ) {
            accum += pit->first;
            if( ! --b ) {
                bin_mean.push_back( accum / (double) BIN_SIZE );
                bin_size_accum += BIN_SIZE;
                m_bin_sizes.push_back( bin_size_accum );
                accum = 0.0;
                b = BIN_SIZE;
            }
            ++pit;
        }

        if( b != BIN_SIZE ) {
            bin_mean.push_back( accum / (BIN_SIZE - b) );
            bin_size_accum += (BIN_SIZE - b);
            m_bin_sizes.push_back( bin_size_accum );
        }
        assert(m_bin_sizes.back() == N );

        init_discrete_distribution( bin_mean );
    }

    template < class Iter >
    void init_method2( Iter first, Iter last ) {
//        std::cerr << "Init method 2 " << std::endl;
        unsigned int N = 0;
        double min_fit = 2.0, max_fit = 0.0;
        while( first != last ) {
            m_pop_fit.push_back( std::make_pair( *first, N++ ) );

            if( *first > max_fit ) {
                max_fit = *first;
            }

            if( *first < min_fit ) {
                min_fit = *first;
            }
            ++first;
        }

        simple_comp comp;
        std::sort( m_pop_fit.begin(), m_pop_fit.end(), comp );

        std::vector< double > bin_mean;
        bin_mean.reserve( BIN_COUNT );
        m_bin_sizes.reserve( BIN_COUNT );

//        double bin_size = (max_fit - min_fit) / (double) BIN_COUNT;
        if( max_fit == min_fit) {
            // all individuals have the same fitness
            // therefore all bins have same average fitness
            // divide population into BIN_COUNT bins
                                                            // if N = 101 && BIN_COUNT = 100 then
            unsigned int ind_per_bin = (N / BIN_COUNT);     // ind_per_bin = 1
            unsigned int b = BIN_COUNT;                     // b = 100
            unsigned int n = N % BIN_COUNT;                 // n = 1
            unsigned int baccum = 0;

            while( n-- ) {                                  // single round
                // N - (BIN_COUNT - 1) * ind_per_bin
                bin_mean.push_back( max_fit );
                baccum += (ind_per_bin + 1);
                m_bin_sizes.push_back( baccum ); // push_back(1 + 1)
                --b;                                        // b = 99
            }

            while( b-- ) {                                  // 99 rounds
                bin_mean.push_back( max_fit );              //
                baccum += (ind_per_bin);                    //
                m_bin_sizes.push_back( baccum );            // push_back( )
            }
                                                            // | m_bin_sizes | = | < 2, 3, 4, 5, ..., 100, 101 > | = 100
        } else {
            double bin_step = (1.0000001 * max_fit - min_fit) / (double)(BIN_COUNT);
            fitness_iterator pit = m_pop_fit.begin();
            double accum = 0.0, bin_max = min_fit + bin_step;

//            std::cerr << "Bin Range: <" << min_fit << ", " << max_fit << ", "  << bin_step << " >" << std::endl;

            unsigned int count = 0, bin_size_accum = 0;
            while( pit != m_pop_fit.end() ) {
                if( pit->first >= bin_max ) {

                    bin_mean.push_back( accum / (double) count );
                    bin_size_accum += count;
                    m_bin_sizes.push_back( bin_size_accum );

                    accum = 0.0;
                    count = 0;
                    bin_max += bin_step;

                    // advance to next containing bin
                    while( pit->first >= bin_max ) {
                        bin_mean.push_back( 0.0 );
                        m_bin_sizes.push_back( bin_size_accum );
                        bin_max += bin_step;
                    }
                }
                accum += pit->first;
                ++count;

                ++pit;
            }
            if( bin_size_accum < (double) N ) {
                bin_mean.push_back( accum / (double) count );
                bin_size_accum += count;
                assert( bin_size_accum == (double) N );

                m_bin_sizes.push_back( bin_size_accum );
            }
        }

//        std::cerr << "Last two bin sizes: " << m_bin_sizes[ m_bin_sizes.size() - 2] << ", " << m_bin_sizes.back() << std::endl;
        assert( m_bin_sizes.back() == N);
//        std::cerr << "bin count: " << bin_mean.size() << std::endl;
        assert( bin_mean.size() == BIN_COUNT );

        init_discrete_distribution( bin_mean );
    }

    void init_discrete_distribution( std::vector< double > & bin_mean ) {
        // N * (N + 1) / 2
//        std::cerr << "Bin Mean Size: " << bin_mean.size() << std::endl;
        unsigned int bin_pairs = bin_mean.size() * (bin_mean.size() + 1)  / 2;
        assert( bin_pairs != -1 );

        std::vector< double > pairs( bin_pairs, 0.0);
        std::vector< double >::iterator it = pairs.begin();
        std::vector< double >::iterator y = bin_mean.begin();
        while( y != bin_mean.end() ) {
            double fit = *y;
            ++y;
            std::vector<double>::iterator x = bin_mean.begin();
            while( x != y ) {
                (*it) = prob2( fit, (*x) );
                ++it;
                ++x;
            }
        }

        assert( it == pairs.end() );

        typename dist_type::param_type p(pairs.begin(), pairs.end());
        m_dist.param(p);
    }

    double prob1( double x, double y ) {
        return x * y;
    }

    double prob2( double x, double y, double mu = 0.0, double sigma = 1.0 ) {
        double res = x * y;
        if( res == 0.0 ) return 0.0;

        double coeff = (1.0/sigma) * boost::math::double_constants::one_div_root_two_pi;
        double var = sigma * sigma;

        res = ( (1.0 - res) - mu); // (1 - xy) - mu
        res *= res;
        res /= (2.0 * var);

        return coeff * exp(-res);
    }
};

#endif  // ASSORTATIVE_SELECTOR_HPP_


#ifndef __BESIQ_MODELS_H__
#define __BESIQ_MODELS_H__

#include <armadillo>

#include <plink/snp_row.hpp>
#include <besiq/stats/dirichlet.hpp>
#include <besiq/stats/log_scale.hpp>
#include <besiq/stats/snp_count.hpp>

/**
 * This class represents a general model that can compute a
 * model likelihood for two snps.
 */
class model
{
public:
    /**
     * Constructor.
     *
     * @param prior The prior for the model.
     * @param alpha The prior parameters.
     */
    model(log_double prior, const arma::vec &alpha)
    : m_prior( prior ),
      m_alpha( alpha )
    {
    }

    /**
     * Destructor.
     */
    virtual ~model(){ };

    /**
     * Returns the prior of the model.
     *
     * @return The prior probability.
     */
    log_double prior( )
    {
        return m_prior;
    }

    /**
     * Returns the prior parameters.
     *
     * @return the prior parameters.
     */
    arma::vec get_alpha()
    {
        return m_alpha;
    }

    /**
     * Computes the likelihood of the phenotype under this
     * model.
     *
     * @param row1 The first snp.
     * @param row2 The §second snp.
     * @param phenotype The phenotype, discrete 0.0 and 1.0.
     * @param weight A weight for each sample, this will be used instead of 1.0 as a count.
     *
     * @return The likelihood of the snps.
     */
    virtual log_double prob(const snp_row &row1, const snp_row &row2, const arma::vec &phenotype, const arma::vec &weight) = 0;

private:
    /**
     * The prior probability for the model.
     */
    log_double m_prior;

    /**
     * The beta prior parameters.
     */
    arma::vec m_alpha;
};

/**
 * This model represents a saturated model in the sense
 * that each cell in the penetrance table has it is own
 * parameter.
 */
class saturated
: public model
{
public:
    /**
     * Constructor.
     *
     * @param prior The prior probability for the model.
     */
    saturated(log_double prior, const arma::vec &alpha);

    /**
     * @see model::prob.
     */
    virtual log_double prob(const snp_row &row1, const snp_row &row2, const arma::vec &phenotype, const arma::vec &weight);
};

/**
 * This class represents a model where no snp is associated with
 * the phenotype. The snps may however be in ld.
 */
class null
: public model
{
public:
    /**
     * Constructor.
     *
     * @param prior The prior probability for the model.
     */
    null(log_double prior, const arma::vec &alpha);

    /**
     * Computes a sensible prior alpha based on the minor allele frequencies
     * of both SNPs.
     *
     * @param row1 The first snp.
     * @param row2 The second snp.
     *
     * @return A vector of alphas that reflects the estiamted mafs.
     */
    static arma::vec compute_alpha(const snp_row &row1, const snp_row &row2);

    /**
     * Computes the probability of the snps under the null model.
     *
     * Note: Convenient to use in other models, as they always compute the
     *       null or prior probability of the snps.
     *
     * @param row1 The first snp.
     * @param row2 The second snp.
     * @param phenotype The phenotype, discrete 0.0 and 1.0.
     * @param weight A weight for each sample, this will be used instead of 1.0 as a count.
     *
     * @return the probability of the snps under the null model.
     */
    static log_double snp_prob(const snp_row &row1, const snp_row &row2, const arma::vec &phenotype, const arma::vec &weight);

    /**
     * Computes the probability of the phenotype under the null model.
     *
     * @param row1 The first snp.
     * @param row2 The second snp.
     * @param phenotype The phenotype, discrete 0.0 and 1.0.
     * @param weight A weight for each sample, this will be used instead of 1.0 as a count.
     *
     * @return the probability of the phenotype under the null model.
     */
    static log_double pheno_prob(const snp_row &row1, const snp_row &row2, const arma::vec &phenotype, const arma::vec &weight);

    /**
     * @see model::prob.
     */
    virtual log_double prob(const snp_row &row1, const snp_row &row2, const arma::vec &phenotype, const arma::vec &weight);
};

/**
 * This class represents a model where one snp is associated with
 * the phenotype, and the other possibly in ld with the first.
 */
class ld_assoc
: public model
{
public:
    /**
     * Constructor.
     *
     * @param prior The prior probability for the model.
     * @param is_first If true the first snp is associated, otherwise the second.
     */
    ld_assoc(log_double prior, const arma::vec &alpha, bool is_first);

    /**
     * @see model::prob.
     */
    virtual log_double prob(const snp_row &row1, const snp_row &row2, const arma::vec &phenotype, const arma::vec &weight);

private:
    /**
     * Indicates which snp is associated with the phenotype, if true
     * the first, false the second.
     */
    bool m_is_first;
};

/**
 * This class represents a model where two snps are independently
 * associated to the phenotype.
 *
 * The posterior of this model takes significantly longer than the
 * other models and should only be used on a small subset of
 * snp pairs.
 */
class sindependent
: public model
{
public:
    /**
     * Constructor.
     *
     * @param prior The prior probability for the model.
     * @param num_mc_iterations The number of monte carlo iterations.
     */
    sindependent(log_double prior, const arma::vec &alpha, int num_mc_iterations);

    /**
     * @see model::prob.
     */
    virtual log_double prob(const snp_row &row1, const snp_row &row2, const arma::vec &phenotype, const arma::vec &weight);

private:
    dir_generator m_rdir;

    /**
     * Number of monte carlo iterations.
     */
    int m_num_mc_iterations;
};

#endif /* End of __BESIQ_MODELS_H__ */


/* boost random/mersenne_twister.hpp header file
 *
 * Copyright Jens Maurer 2000-2001
 * Copyright Steven Watanabe 2010
 * Distributed under the Boost Software License, Version 1.0. (See
 * accompanying file LICENSE_1_0.txt or copy at
 * http://www.boost.org/LICENSE_1_0.txt)
 *
 * See http://www.boost.org for most recent version including documentation.
 *
 * $Id$
 *
 * Revision history
 *  2013-10-14  fixed some warnings with Wshadow (mgaunard)
 *  2001-02-18  moved to individual header files
 */

#ifndef BOOST_RANDOM_MERSENNE_TWISTER_HPP
#define BOOST_RANDOM_MERSENNE_TWISTER_HPP

#include <iosfwd>
#include <istream>
#include <stdexcept>
#include <boost/config.hpp>
#include <boost/cstdint.hpp>
#include <boost/integer/integer_mask.hpp>
#include <boost/random/detail/config.hpp>
#include <boost/random/detail/ptr_helper.hpp>
#include <boost/random/detail/seed.hpp>
#include <boost/random/detail/seed_impl.hpp>
#include <boost/random/detail/generator_seed_seq.hpp>
#include <boost/random/detail/polynomial.hpp>

#include <boost/random/detail/disable_warnings.hpp>

namespace boost {
namespace random {

/**
 * Instantiations of class template mersenne_twister_engine model a
 * \pseudo_random_number_generator. It uses the algorithm described in
 *
 *  @blockquote
 *  "Mersenne Twister: A 623-dimensionally equidistributed uniform
 *  pseudo-random number generator", Makoto Matsumoto and Takuji Nishimura,
 *  ACM Transactions on Modeling and Computer Simulation: Special Issue on
 *  Uniform Random Number Generation, Vol. 8, No. 1, January 1998, pp. 3-30.
 *  @endblockquote
 *
 * @xmlnote
 * The boost variant has been implemented from scratch and does not
 * derive from or use mt19937.c provided on the above WWW site. However, it
 * was verified that both produce identical output.
 * @endxmlnote
 *
 * The seeding from an integer was changed in April 2005 to address a
 * weakness.
 *
 * The quality of the generator crucially depends on the choice of the
 * parameters.  User code should employ one of the sensibly parameterized
 * generators such as \mt19937 instead.
 *
 * The generator requires considerable amounts of memory for the storage of
 * its state array. For example, \mt11213b requires about 1408 bytes and
 * \mt19937 requires about 2496 bytes.
 */
template<class UIntType,
         std::size_t w, std::size_t n, std::size_t m, std::size_t r,
         UIntType a, std::size_t u, UIntType d, std::size_t s,
         UIntType b, std::size_t t,
         UIntType c, std::size_t l, UIntType f>
class mersenne_twister_engine
{
public:
    typedef UIntType result_type;
    BOOST_STATIC_CONSTANT(std::size_t, word_size = w);
    BOOST_STATIC_CONSTANT(std::size_t, state_size = n);
    BOOST_STATIC_CONSTANT(std::size_t, shift_size = m);
    BOOST_STATIC_CONSTANT(std::size_t, mask_bits = r);
    BOOST_STATIC_CONSTANT(UIntType, xor_mask = a);
    BOOST_STATIC_CONSTANT(std::size_t, tempering_u = u);
    BOOST_STATIC_CONSTANT(UIntType, tempering_d = d);
    BOOST_STATIC_CONSTANT(std::size_t, tempering_s = s);
    BOOST_STATIC_CONSTANT(UIntType, tempering_b = b);
    BOOST_STATIC_CONSTANT(std::size_t, tempering_t = t);
    BOOST_STATIC_CONSTANT(UIntType, tempering_c = c);
    BOOST_STATIC_CONSTANT(std::size_t, tempering_l = l);
    BOOST_STATIC_CONSTANT(UIntType, initialization_multiplier = f);
    BOOST_STATIC_CONSTANT(UIntType, default_seed = 5489u);

    // backwards compatibility
    BOOST_STATIC_CONSTANT(UIntType, parameter_a = a);
    BOOST_STATIC_CONSTANT(std::size_t, output_u = u);
    BOOST_STATIC_CONSTANT(std::size_t, output_s = s);
    BOOST_STATIC_CONSTANT(UIntType, output_b = b);
    BOOST_STATIC_CONSTANT(std::size_t, output_t = t);
    BOOST_STATIC_CONSTANT(UIntType, output_c = c);
    BOOST_STATIC_CONSTANT(std::size_t, output_l = l);

    // old Boost.Random concept requirements
    BOOST_STATIC_CONSTANT(bool, has_fixed_range = false);

    /**
     * Constructs a @c mersenne_twister_engine and calls @c seed().
     */
    mersenne_twister_engine() { seed(); }

    /**
     * Constructs a @c mersenne_twister_engine and calls @c seed(value).
     */
    BOOST_RANDOM_DETAIL_ARITHMETIC_CONSTRUCTOR(mersenne_twister_engine,
                                               UIntType, value)
    { seed(value); }
    template<class It> mersenne_twister_engine(It& first, It last)
    { seed(first,last); }

    /**
     * Constructs a mersenne_twister_engine and calls @c seed(gen).
     *
     * @xmlnote
     * The copy constructor will always be preferred over
     * the templated constructor.
     * @endxmlnote
     */
    BOOST_RANDOM_DETAIL_SEED_SEQ_CONSTRUCTOR(mersenne_twister_engine,
                                             SeedSeq, seq)
    { seed(seq); }

    // compiler-generated copy ctor and assignment operator are fine

    /** Calls @c seed(default_seed). */
    void seed() { seed(default_seed); }

    /**
     * Sets the state x(0) to v mod 2w. Then, iteratively,
     * sets x(i) to
     * (i + f * (x(i-1) xor (x(i-1) rshift w-2))) mod 2w
     * for i = 1 .. n-1. x(n) is the first value to be returned by operator().
     */
    BOOST_RANDOM_DETAIL_ARITHMETIC_SEED(mersenne_twister_engine, UIntType, value)
    {
        // New seeding algorithm from
        // http://www.math.sci.hiroshima-u.ac.jp/~m-mat/MT/MT2002/emt19937ar.html
        // In the previous versions, MSBs of the seed affected only MSBs of the
        // state x[].
        const UIntType mask = (max)();
        x[0] = value & mask;
        for (i = 1; i < n; i++) {
            // See Knuth "The Art of Computer Programming"
            // Vol. 2, 3rd ed., page 106
            x[i] = (f * (x[i-1] ^ (x[i-1] >> (w-2))) + i) & mask;
        }

        normalize_state();
    }

    /**
     * Seeds a mersenne_twister_engine using values produced by seq.generate().
     */
    BOOST_RANDOM_DETAIL_SEED_SEQ_SEED(mersenne_twister_engine, SeeqSeq, seq)
    {
        detail::seed_array_int<w>(seq, x);
        i = n;

        normalize_state();
    }

    /** Sets the state of the generator using values from an iterator range. */
    template<class It>
    void seed(It& first, It last)
    {
        detail::fill_array_int<w>(first, last, x);
        i = n;

        normalize_state();
    }

    /** Returns the smallest value that the generator can produce. */
    static result_type min BOOST_PREVENT_MACRO_SUBSTITUTION ()
    { return 0; }
    /** Returns the largest value that the generator can produce. */
    static result_type max BOOST_PREVENT_MACRO_SUBSTITUTION ()
    { return boost::low_bits_mask_t<w>::sig_bits; }

    /** Produces the next value of the generator. */
    result_type operator()();

    /** Fills a range with random values */
    template<class Iter>
    void generate(Iter first, Iter last)
    { detail::generate_from_int(*this, first, last); }

    /**
     * Advances the state of the generator by @c z steps.  Equivalent to
     *
     * @code
     * for(unsigned long long i = 0; i < z; ++i) {
     *     gen();
     * }
     * @endcode
     */
    void discard(boost::uintmax_t z)
    {
#ifndef BOOST_RANDOM_MERSENNE_TWISTER_DISCARD_THRESHOLD
#define BOOST_RANDOM_MERSENNE_TWISTER_DISCARD_THRESHOLD 10000000
#endif
        if(z > BOOST_RANDOM_MERSENNE_TWISTER_DISCARD_THRESHOLD) {
            discard_many(z);
        } else {
            for(boost::uintmax_t j = 0; j < z; ++j) {
                (*this)();
            }
        }
    }

#ifndef BOOST_RANDOM_NO_STREAM_OPERATORS
    /** Writes a mersenne_twister_engine to a @c std::ostream */
    template<class CharT, class Traits>
    friend std::basic_ostream<CharT,Traits>&
    operator<<(std::basic_ostream<CharT,Traits>& os,
               const mersenne_twister_engine& mt)
    {
        mt.print(os);
        return os;
    }

    /** Reads a mersenne_twister_engine from a @c std::istream */
    template<class CharT, class Traits>
    friend std::basic_istream<CharT,Traits>&
    operator>>(std::basic_istream<CharT,Traits>& is,
               mersenne_twister_engine& mt)
    {
        for(std::size_t j = 0; j < mt.state_size; ++j)
            is >> mt.x[j] >> std::ws;
        // MSVC (up to 7.1) and Borland (up to 5.64) don't handle the template
        // value parameter "n" available from the class template scope, so use
        // the static constant with the same value
        mt.i = mt.state_size;
        return is;
    }
#endif

    /**
     * Returns true if the two generators are in the same state,
     * and will thus produce identical sequences.
     */
    friend bool operator==(const mersenne_twister_engine& x_,
                           const mersenne_twister_engine& y_)
    {
        if(x_.i < y_.i) return x_.equal_imp(y_);
        else return y_.equal_imp(x_);
    }

    /**
     * Returns true if the two generators are in different states.
     */
    friend bool operator!=(const mersenne_twister_engine& x_,
                           const mersenne_twister_engine& y_)
    { return !(x_ == y_); }

private:
    /// \cond show_private

    void twist();

    /**
     * Does the work of operator==.  This is in a member function
     * for portability.  Some compilers, such as msvc 7.1 and
     * Sun CC 5.10 can't access template parameters or static
     * members of the class from inline friend functions.
     *
     * requires i <= other.i
     */
    bool equal_imp(const mersenne_twister_engine& other) const
    {
        UIntType back[n];
        std::size_t offset = other.i - i;
        for(std::size_t j = 0; j + offset < n; ++j)
            if(x[j] != other.x[j+offset])
                return false;
        rewind(&back[n-1], offset);
        for(std::size_t j = 0; j < offset; ++j)
            if(back[j + n - offset] != other.x[j])
                return false;
        return true;
    }

    /**
     * Does the work of operator<<.  This is in a member function
     * for portability.
     */
    template<class CharT, class Traits>
    void print(std::basic_ostream<CharT, Traits>& os) const
    {
        UIntType data[n];
        for(std::size_t j = 0; j < i; ++j) {
            data[j + n - i] = x[j];
        }
        if(i != n) {
            rewind(&data[n - i - 1], n - i);
        }
        os << data[0];
        for(std::size_t j = 1; j < n; ++j) {
            os << ' ' << data[j];
        }
    }

    /**
     * Copies z elements of the state preceding x[0] into
     * the array whose last element is last.
     */
    void rewind(UIntType* last, std::size_t z) const
    {
        const UIntType upper_mask = (~static_cast<UIntType>(0)) << r;
        const UIntType lower_mask = ~upper_mask;
        UIntType y0 = x[m-1] ^ x[n-1];
        if(y0 & (static_cast<UIntType>(1) << (w-1))) {
            y0 = ((y0 ^ a) << 1) | 1;
        } else {
            y0 = y0 << 1;
        }
        for(std::size_t sz = 0; sz < z; ++sz) {
            UIntType y1 =
                rewind_find(last, sz, m-1) ^ rewind_find(last, sz, n-1);
            if(y1 & (static_cast<UIntType>(1) << (w-1))) {
                y1 = ((y1 ^ a) << 1) | 1;
            } else {
                y1 = y1 << 1;
            }
            *(last - sz) = (y0 & upper_mask) | (y1 & lower_mask);
            y0 = y1;
        }
    }

    /**
     * Converts an arbitrary array into a valid generator state.
     * First we normalize x[0], so that it contains the same
     * value we would get by running the generator forwards
     * and then in reverse.  (The low order r bits are redundant).
     * Then, if the state consists of all zeros, we set the
     * high order bit of x[0] to 1.  This function only needs to
     * be called by seed, since the state transform preserves
     * this relationship.
     */
    void normalize_state()
    {
        const UIntType upper_mask = (~static_cast<UIntType>(0)) << r;
        const UIntType lower_mask = ~upper_mask;
        UIntType y0 = x[m-1] ^ x[n-1];
        if(y0 & (static_cast<UIntType>(1) << (w-1))) {
            y0 = ((y0 ^ a) << 1) | 1;
        } else {
            y0 = y0 << 1;
        }
        x[0] = (x[0] & upper_mask) | (y0 & lower_mask);

        // fix up the state if it's all zeroes.
        for(std::size_t j = 0; j < n; ++j) {
            if(x[j] != 0) return;
        }
        x[0] = static_cast<UIntType>(1) << (w-1);
    }

    /**
     * Given a pointer to the last element of the rewind array,
     * and the current size of the rewind array, finds an element
     * relative to the next available slot in the rewind array.
     */
    UIntType
    rewind_find(UIntType* last, std::size_t size, std::size_t j) const
    {
        std::size_t index = (j + n - size + n - 1) % n;
        if(index < n - size) {
            return x[index];
        } else {
            return *(last - (n - 1 - index));
        }
    }

    /**
     * Optimized algorithm for large jumps.
     *
     * Hiroshi Haramoto, Makoto Matsumoto, and Pierre L'Ecuyer. 2008.
     * A Fast Jump Ahead Algorithm for Linear Recurrences in a Polynomial
     * Space. In Proceedings of the 5th international conference on
     * Sequences and Their Applications (SETA '08).
     * DOI=10.1007/978-3-540-85912-3_26
     */
    void discard_many(boost::uintmax_t z)
    {
        // Compute the minimal polynomial, phi(t)
        // This depends only on the transition function,
        // which is constant.  The characteristic
        // polynomial is the same as the minimal
        // polynomial for a maximum period generator
        // (which should be all specializations of
        // mersenne_twister.)  Even if it weren't,
        // the characteristic polynomial is guaranteed
        // to be a multiple of the minimal polynomial,
        // which is good enough.
        detail::polynomial phi = get_characteristic_polynomial();

        // calculate g(t) = t^z % phi(t)
        detail::polynomial g = mod_pow_x(z, phi);

        // h(s_0, t) = \sum_{i=0}^{2k-1}o(s_i)t^{2k-i-1}
        detail::polynomial h;
        const std::size_t num_bits = w*n - r;
        for(std::size_t j = 0; j < num_bits * 2; ++j) {
            // Yes, we're advancing the generator state
            // here, but it doesn't matter because
            // we're going to overwrite it completely
            // in reconstruct_state.
            if(i >= n) twist();
            h[2*num_bits - j - 1] = x[i++] & UIntType(1);
        }
        // g(t)h(s_0, t)
        detail::polynomial gh = g * h;
        detail::polynomial result;
        for(std::size_t j = 0; j <= num_bits; ++j) {
            result[j] = gh[2*num_bits - j - 1];
        }
        reconstruct_state(result);
    }
    static detail::polynomial get_characteristic_polynomial()
    {
        const std::size_t num_bits = w*n - r;
        detail::polynomial helper;
        helper[num_bits - 1] = 1;
        mersenne_twister_engine tmp;
        tmp.reconstruct_state(helper);
        // Skip the first num_bits elements, since we
        // already know what they are.
        for(std::size_t j = 0; j < num_bits; ++j) {
            if(tmp.i >= n) tmp.twist();
            if(j == num_bits - 1)
                assert((tmp.x[tmp.i] & 1) == 1);
            else
                assert((tmp.x[tmp.i] & 1) == 0);
            ++tmp.i;
        }
        detail::polynomial phi;
        phi[num_bits] = 1;
        detail::polynomial next_bits = tmp.as_polynomial(num_bits);
        for(std::size_t j = 0; j < num_bits; ++j) {
            int val = next_bits[j] ^ phi[num_bits-j-1];
            phi[num_bits-j-1] = val;
            if(val) {
                for(std::size_t k = j + 1; k < num_bits; ++k) {
                    phi[num_bits-k-1] ^= next_bits[k-j-1];
                }
            }
        }
        return phi;
    }
    detail::polynomial as_polynomial(std::size_t size) {
        detail::polynomial result;
        for(std::size_t j = 0; j < size; ++j) {
            if(i >= n) twist();
            result[j] = x[i++] & UIntType(1);
        }
        return result;
    }
    void reconstruct_state(const detail::polynomial& p)
    {
        const UIntType upper_mask = (~static_cast<UIntType>(0)) << r;
        const UIntType lower_mask = ~upper_mask;
        const std::size_t num_bits = w*n - r;
        for(std::size_t j = num_bits - n + 1; j <= num_bits; ++j)
            x[j % n] = p[j];

        UIntType y0 = 0;
        for(std::size_t j = num_bits + 1; j >= n - 1; --j) {
            UIntType y1 = x[j % n] ^ x[(j + m) % n];
            if(p[j - n + 1])
                y1 = (y1 ^ a) << UIntType(1) | UIntType(1);
            else
                y1 = y1 << UIntType(1);
            x[(j + 1) % n] = (y0 & upper_mask) | (y1 & lower_mask);
            y0 = y1;
        }
        i = 0;
    }

    /// \endcond

    // state representation: next output is o(x(i))
    //   x[0]  ... x[k] x[k+1] ... x[n-1]   represents
    //  x(i-k) ... x(i) x(i+1) ... x(i-k+n-1)

    UIntType x[n];
    std::size_t i;
};

/// \cond show_private

#ifndef BOOST_NO_INCLASS_MEMBER_INITIALIZATION
//  A definition is required even for integral static constants
#define BOOST_RANDOM_MT_DEFINE_CONSTANT(type, name)                         \
template<class UIntType, std::size_t w, std::size_t n, std::size_t m,       \
    std::size_t r, UIntType a, std::size_t u, UIntType d, std::size_t s,    \
    UIntType b, std::size_t t, UIntType c, std::size_t l, UIntType f>       \
const type mersenne_twister_engine<UIntType,w,n,m,r,a,u,d,s,b,t,c,l,f>::name
BOOST_RANDOM_MT_DEFINE_CONSTANT(std::size_t, word_size);
BOOST_RANDOM_MT_DEFINE_CONSTANT(std::size_t, state_size);
BOOST_RANDOM_MT_DEFINE_CONSTANT(std::size_t, shift_size);
BOOST_RANDOM_MT_DEFINE_CONSTANT(std::size_t, mask_bits);
BOOST_RANDOM_MT_DEFINE_CONSTANT(UIntType, xor_mask);
BOOST_RANDOM_MT_DEFINE_CONSTANT(std::size_t, tempering_u);
BOOST_RANDOM_MT_DEFINE_CONSTANT(UIntType, tempering_d);
BOOST_RANDOM_MT_DEFINE_CONSTANT(std::size_t, tempering_s);
BOOST_RANDOM_MT_DEFINE_CONSTANT(UIntType, tempering_b);
BOOST_RANDOM_MT_DEFINE_CONSTANT(std::size_t, tempering_t);
BOOST_RANDOM_MT_DEFINE_CONSTANT(UIntType, tempering_c);
BOOST_RANDOM_MT_DEFINE_CONSTANT(std::size_t, tempering_l);
BOOST_RANDOM_MT_DEFINE_CONSTANT(UIntType, initialization_multiplier);
BOOST_RANDOM_MT_DEFINE_CONSTANT(UIntType, default_seed);
BOOST_RANDOM_MT_DEFINE_CONSTANT(UIntType, parameter_a);
BOOST_RANDOM_MT_DEFINE_CONSTANT(std::size_t, output_u );
BOOST_RANDOM_MT_DEFINE_CONSTANT(std::size_t, output_s);
BOOST_RANDOM_MT_DEFINE_CONSTANT(UIntType, output_b);
BOOST_RANDOM_MT_DEFINE_CONSTANT(std::size_t, output_t);
BOOST_RANDOM_MT_DEFINE_CONSTANT(UIntType, output_c);
BOOST_RANDOM_MT_DEFINE_CONSTANT(std::size_t, output_l);
BOOST_RANDOM_MT_DEFINE_CONSTANT(bool, has_fixed_range);
#undef BOOST_RANDOM_MT_DEFINE_CONSTANT
#endif

template<class UIntType,
         std::size_t w, std::size_t n, std::size_t m, std::size_t r,
         UIntType a, std::size_t u, UIntType d, std::size_t s,
         UIntType b, std::size_t t,
         UIntType c, std::size_t l, UIntType f>
void
mersenne_twister_engine<UIntType,w,n,m,r,a,u,d,s,b,t,c,l,f>::twist()
{
    const UIntType upper_mask = (~static_cast<UIntType>(0)) << r;
    const UIntType lower_mask = ~upper_mask;

    const std::size_t unroll_factor = 6;
    const std::size_t unroll_extra1 = (n-m) % unroll_factor;
    const std::size_t unroll_extra2 = (m-1) % unroll_factor;

    // split loop to avoid costly modulo operations
    {  // extra scope for MSVC brokenness w.r.t. for scope
        for(std::size_t j = 0; j < n-m-unroll_extra1; j++) {
            UIntType y = (x[j] & upper_mask) | (x[j+1] & lower_mask);
            x[j] = x[j+m] ^ (y >> 1) ^ ((x[j+1]&1) * a);
        }
    }
    {
        for(std::size_t j = n-m-unroll_extra1; j < n-m; j++) {
            UIntType y = (x[j] & upper_mask) | (x[j+1] & lower_mask);
            x[j] = x[j+m] ^ (y >> 1) ^ ((x[j+1]&1) * a);
        }
    }
    {
        for(std::size_t j = n-m; j < n-1-unroll_extra2; j++) {
            UIntType y = (x[j] & upper_mask) | (x[j+1] & lower_mask);
            x[j] = x[j-(n-m)] ^ (y >> 1) ^ ((x[j+1]&1) * a);
        }
    }
    {
        for(std::size_t j = n-1-unroll_extra2; j < n-1; j++) {
            UIntType y = (x[j] & upper_mask) | (x[j+1] & lower_mask);
            x[j] = x[j-(n-m)] ^ (y >> 1) ^ ((x[j+1]&1) * a);
        }
    }
    // last iteration
    UIntType y = (x[n-1] & upper_mask) | (x[0] & lower_mask);
    x[n-1] = x[m-1] ^ (y >> 1) ^ ((x[0]&1) * a);
    i = 0;
}
/// \endcond

template<class UIntType,
         std::size_t w, std::size_t n, std::size_t m, std::size_t r,
         UIntType a, std::size_t u, UIntType d, std::size_t s,
         UIntType b, std::size_t t,
         UIntType c, std::size_t l, UIntType f>
inline typename
mersenne_twister_engine<UIntType,w,n,m,r,a,u,d,s,b,t,c,l,f>::result_type
mersenne_twister_engine<UIntType,w,n,m,r,a,u,d,s,b,t,c,l,f>::operator()()
{
    if(i == n)
        twist();
    // Step 4
    UIntType z = x[i];
    ++i;
    z ^= ((z >> u) & d);
    z ^= ((z << s) & b);
    z ^= ((z << t) & c);
    z ^= (z >> l);
    return z;
}

/**
 * The specializations \mt11213b and \mt19937 are from
 *
 *  @blockquote
 *  "Mersenne Twister: A 623-dimensionally equidistributed
 *  uniform pseudo-random number generator", Makoto Matsumoto
 *  and Takuji Nishimura, ACM Transactions on Modeling and
 *  Computer Simulation: Special Issue on Uniform Random Number
 *  Generation, Vol. 8, No. 1, January 1998, pp. 3-30.
 *  @endblockquote
 */
typedef mersenne_twister_engine<uint32_t,32,351,175,19,0xccab8ee7,
    11,0xffffffff,7,0x31b6ab00,15,0xffe50000,17,1812433253> mt11213b;

/**
 * The specializations \mt11213b and \mt19937 are from
 *
 *  @blockquote
 *  "Mersenne Twister: A 623-dimensionally equidistributed
 *  uniform pseudo-random number generator", Makoto Matsumoto
 *  and Takuji Nishimura, ACM Transactions on Modeling and
 *  Computer Simulation: Special Issue on Uniform Random Number
 *  Generation, Vol. 8, No. 1, January 1998, pp. 3-30.
 *  @endblockquote
 */
typedef mersenne_twister_engine<uint32_t,32,624,397,31,0x9908b0df,
    11,0xffffffff,7,0x9d2c5680,15,0xefc60000,18,1812433253> mt19937;

#if !defined(BOOST_NO_INT64_T) && !defined(BOOST_NO_INTEGRAL_INT64_T)
typedef mersenne_twister_engine<uint64_t,64,312,156,31,
    UINT64_C(0xb5026f5aa96619e9),29,UINT64_C(0x5555555555555555),17,
    UINT64_C(0x71d67fffeda60000),37,UINT64_C(0xfff7eee000000000),43,
    UINT64_C(6364136223846793005)> mt19937_64;
#endif

/// \cond show_deprecated

template<class UIntType,
         int w, int n, int m, int r,
         UIntType a, int u, std::size_t s,
         UIntType b, int t,
         UIntType c, int l, UIntType v>
class mersenne_twister :
    public mersenne_twister_engine<UIntType,
        w, n, m, r, a, u, ~(UIntType)0, s, b, t, c, l, 1812433253>
{
    typedef mersenne_twister_engine<UIntType,
        w, n, m, r, a, u, ~(UIntType)0, s, b, t, c, l, 1812433253> base_type;
public:
    mersenne_twister() {}
    BOOST_RANDOM_DETAIL_GENERATOR_CONSTRUCTOR(mersenne_twister, Gen, gen)
    { seed(gen); }
    BOOST_RANDOM_DETAIL_ARITHMETIC_CONSTRUCTOR(mersenne_twister, UIntType, val)
    { seed(val); }
    template<class It>
    mersenne_twister(It& first, It last) : base_type(first, last) {}
    void seed() { base_type::seed(); }
    BOOST_RANDOM_DETAIL_GENERATOR_SEED(mersenne_twister, Gen, gen)
    {
        detail::generator_seed_seq<Gen> seq(gen);
        base_type::seed(seq);
    }
    BOOST_RANDOM_DETAIL_ARITHMETIC_SEED(mersenne_twister, UIntType, val)
    { base_type::seed(val); }
    template<class It>
    void seed(It& first, It last) { base_type::seed(first, last); }
};

/// \endcond

} // namespace random

using random::mt11213b;
using random::mt19937;
using random::mt19937_64;

} // namespace boost

BOOST_RANDOM_PTR_HELPER_SPEC(boost::mt11213b)
BOOST_RANDOM_PTR_HELPER_SPEC(boost::mt19937)
BOOST_RANDOM_PTR_HELPER_SPEC(boost::mt19937_64)

#include <boost/random/detail/enable_warnings.hpp>

#endif // BOOST_RANDOM_MERSENNE_TWISTER_HPP


#define BOOST_TEST_MODULE ForwardFlipperTest

#include "ForwardFlipper.hpp"
#include <vector>
#include <cstddef>

#define BOOST_TEST_DYN_LINK
#include <boost/test/unit_test.hpp>
#include <boost/algorithm/string.hpp>

BOOST_AUTO_TEST_SUITE(ForwardFlipperSuite)

using namespace com::github::nimelo;
using namespace std;

BOOST_AUTO_TEST_CASE(TestScenario_1)
{
    // Arrange
    size_t n = 6, p = 2;
    size_t expected_distance = 1;

    ForwardFlipper flipper;

    // Act
    auto actual_distance = flipper.distance(n, p);

    // Assert
    BOOST_CHECK_EQUAL(expected_distance, actual_distance);
}

BOOST_AUTO_TEST_CASE(TestScenario_2)
{
    // Arrange
    size_t n = 5, p = 4;
    size_t expected_distance = 2;

    ForwardFlipper flipper;

    // Act
    auto actual_distance = flipper.distance(n, p);

    // Assert
    BOOST_CHECK_EQUAL(expected_distance, actual_distance);
}

BOOST_AUTO_TEST_SUITE_END()


///////////////////////////////////////////////////////////////////////////////
// lost_df_accumulator.hpp                                                   //
//                                                                           //
//  Copyright 2010 Erwann Rogard. Distributed under the Boost                //
//  Software License, Version 1.0. (See accompanying file                    //
//  LICENSE_1_0.txt or copy at http://www.boost.org/LICENSE_1_0.txt)         //
///////////////////////////////////////////////////////////////////////////////
#ifndef BOOST_STATISTICS_DETAIL_NON_PARAMETRIC_CONTINGENCY_TABLE_PEARSON_CHISQ_INDEPENDENCE_AUX_LOST_DF_ACCUMULATOR_HPP_ER_2010
#define BOOST_STATISTICS_DETAIL_NON_PARAMETRIC_CONTINGENCY_TABLE_PEARSON_CHISQ_INDEPENDENCE_AUX_LOST_DF_ACCUMULATOR_HPP_ER_2010
#include <boost/mpl/detail/wrapper.hpp>
#include <boost/statistics/detail/non_parametric/contingency_table/factor/levels.hpp>

namespace boost {
namespace statistics{
namespace detail{
namespace contingency_table{
namespace pearson_chi_square_statistic{
namespace independence_between_aux{

    template<typename A,typename N>
    struct lost_df_accumulator{

        lost_df_accumulator(const A& acc, N& df) : acc_(acc), df_(df){}

        template<typename Key>
        void operator()(const boost::mpl::detail::wrapper<Key>& wrapper)const
        {
            namespace ct = contingency_table;
            this->df_ += (ct::extract::levels<Key>( this->acc_ ).size() - 1);
        }

        private:
        const A& acc_;
        mutable N& df_;
    };

    template<typename A,typename N>
    independence_between_aux::lost_df_accumulator<A,N>
    make_lost_df_accumulator(const A& acc, N& df)
    {
        namespace aux = independence_between_aux;
        return aux::lost_df_accumulator<A,N>(acc, df);
    }

}// independence_between_aux
}// pearson_chi_square_statistic
}// contingency_table
}// detail
}// statistics
}// boost

#endif


// Copyright (c) 2020 Graphcore Ltd. All rights reserved.
#include <poplibs_support/LogArithmetic.hpp>
#include <poplibs_support/TestDevice.hpp>
#include <poplibs_test/CTCLoss.hpp>
#include <poplibs_test/CTCUtil.hpp>
#include <poplibs_test/Embedding.hpp>
#include <poplibs_test/MatrixTransforms.hpp>
#include <poplibs_test/Util.hpp>
#include <popnn/CTCLoss.hpp>
#include <popnn/codelets.hpp>
#include <popops/codelets.hpp>
#include <poputil/VertexTemplates.hpp>
#include <poputil/exceptions.hpp>

#include <poplar/Engine.hpp>
#include <poplar/Graph.hpp>
#include <poplar/IPUModel.hpp>

#include <boost/multi_array.hpp>
#include <boost/optional/optional_io.hpp>
#include <boost/program_options.hpp>
#include <boost/test/tools/floating_point_comparison.hpp>

#include <fstream>
#include <iomanip>
#include <random>

namespace po = boost::program_options;

using namespace poplar;
using namespace poplar::program;
using namespace poplibs_test::ctc;
using namespace poplibs_test;
using namespace poplibs_test::matrix;
using namespace poplibs_test::util;
using namespace poplibs_support;
using namespace poputil;

// Default tolerances used in tests
#define FLOAT_REL_TOL 0.04
#define HALF_REL_TOL 0.1
#define FLOAT_ABS_TOL 1e-3
#define HALF_ABS_TOL 1e-2

// The result returned by the codelet / reduction stages can be checked more
// precisely for all tests - especially those with larger lengths
#define CODELET_TEST_FLOAT_REL_TOL 0.01
#define CODELET_TEST_HALF_REL_TOL 0.1
#define CODELET_TEST_FLOAT_ABS_TOL 1e-6
#define CODELET_TEST_HALF_ABS_TOL 1e-5

// Mask results that aren't valid due to the time of the batch entry
boost::multi_array<double, 2>
maskResults(const boost::multi_array<double, 2> &in, size_t validTimesteps) {
  auto out = in;
  const auto timeSteps = out.shape()[0];
  const auto symbols = out.shape()[1];
  for (unsigned sym = 0; sym < symbols; sym++) {
    for (unsigned t = validTimesteps; t < timeSteps; t++) {
      out[t][sym] = 0;
    }
  }
  return out;
}
// Print a sequence, inserting `-` for the blank symbol
void print(const std::string &prefix, const std::vector<unsigned> &symbols,
           unsigned blank, bool verbose = true) {
  if (!verbose) {
    return;
  }
  std::cout << prefix << " ";
  for (auto &i : symbols) {
    if (i == blank) {
      std::cout << "- ";
    } else {
      std::cout << i << " ";
    }
  }
  std::cout << "\n";
}

// Print the matrix `in`
template <typename FPType>
void print(const std::string &prefix, const boost::multi_array<FPType, 2> &in,
           unsigned blank, bool verbose = true) {
  if (!verbose) {
    return;
  }
  std::cout << prefix << "\n        ";
  for (unsigned i = 0; i < in.size(); i++) {
    std::cout << std::setw(11) << (std::string{"t"} + std::to_string(i));
  }

  for (unsigned i = 0; i < in[0].size(); i++) {
    if (i == blank) {
      std::cout << "\nIndex:-  ";
    } else {
      std::cout << "\nIndex:" << i << "  ";
    }
    for (unsigned j = 0; j < in.size(); j++) {
      std::cout << std::setw(10) << std::setprecision(4) << in[j][i] << ",";
    }
  }
  std::cout << "\n\n";
}

// Struct and function to return the test inputs
template <typename FPType> struct InputSequence {
  // Input, always of max size (time) but with only inputLength valid values,
  // the rest padded as blank
  boost::multi_array<FPType, 2> input;
  unsigned inputLength;
  // Labels, of the randomly chosen size for this batch
  std::vector<unsigned> labels;
  bool isLogits;
};

template <typename FPType>
std::pair<FPType, boost::multi_array<FPType, 2>>
gradReference(const InputSequence<FPType> &test_, unsigned blankClass,
              unsigned numClasses, bool testReducedCodeletGradient,
              bool verbose) {
  auto test = test_;
  if (test.isLogits) { // Convert to log probs
    test.input = log::log(transpose(log::softMax(transpose(test.input))));
  }
  auto paddedSequence = extendedLabels(test.labels, blankClass);
  auto in = transpose(test.input);
  boost::multi_array<FPType, 2> logSequence(
      boost::extents[paddedSequence.size()][in.shape()[1]]);
  poplibs_test::embedding::multiSlice(in, paddedSequence, logSequence);

  auto alphaLog =
      alpha(logSequence, paddedSequence, blankClass, test.inputLength);
  auto betaLog =
      beta(logSequence, paddedSequence, blankClass, test.inputLength);
  auto expandedGradient =
      expandedGrad(logSequence, alphaLog, betaLog, paddedSequence, blankClass,
                   test.inputLength);
  auto negLogLoss =
      loss(logSequence, paddedSequence, blankClass, test.inputLength);
  auto gradient =
      grad(logSequence, in, alphaLog, betaLog, paddedSequence, numClasses,
           blankClass, test.inputLength, testReducedCodeletGradient);

  return {negLogLoss, transpose(gradient)};
}

std::vector<std::pair<double, boost::multi_array<double, 2>>>
gradIPU(const std::vector<InputSequence<double>> &inputs, unsigned maxLabels,
        unsigned blankSymbol, std::size_t numClasses, Type inType, Type outType,
        OptionFlags planOpts, OptionFlags debugOpts,
        const DeviceType &deviceType, boost::optional<unsigned> tiles,
        bool ignoreData, bool profile,
        boost::optional<std::string> &profileDir) {

  auto device = createTestDevice(deviceType, 1, tiles);
  const auto &target = device.getTarget();
  Graph graph(target);
  popnn::addCodelets(graph);
  popops::addCodelets(graph);

  const auto maxT = inputs[0].input.shape()[0];
  const auto batchSize = inputs.size();

  // Create the inputs to the gradient function
  const auto plan = popnn::ctc::plan(graph, inType, outType, batchSize, maxT,
                                     maxLabels, numClasses, planOpts);

  auto data = popnn::ctc::createDataInput(graph, inType, batchSize, maxT,
                                          numClasses, plan, "DataInput");
  auto labels = popnn::ctc::createLabelsInput(graph, UNSIGNED_INT, batchSize,
                                              maxLabels, plan, "LabelsInput");

  auto dataLengths = graph.addVariable(UNSIGNED_INT, {batchSize});
  auto labelLengths = graph.addVariable(UNSIGNED_INT, {batchSize});
  graph.setTileMapping(dataLengths, 0);
  graph.setTileMapping(labelLengths, 0);

  // Write the inputs
  Sequence uploadProg, downloadProg;
  std::vector<std::pair<std::string, char *>> tmap;
  std::unique_ptr<char[]> rawDataLengths, rawLabelLengths;
  std::vector<std::unique_ptr<char[]>> rawData(batchSize), rawLabels(batchSize);
  rawDataLengths = allocateHostMemoryForTensor(
      dataLengths, "dataLengths", graph, uploadProg, downloadProg, tmap);
  rawLabelLengths = allocateHostMemoryForTensor(
      labelLengths, "labelLengths", graph, uploadProg, downloadProg, tmap);

  for (unsigned i = 0; i < batchSize; i++) {
    rawData[i] = allocateHostMemoryForTensor(data.slice(i, i + 1, 1),
                                             "data_" + std::to_string(i), graph,
                                             uploadProg, downloadProg, tmap);
    rawLabels[i] = allocateHostMemoryForTensor(
        labels.slice(i, i + 1, 0), "labels_" + std::to_string(i), graph,
        uploadProg, downloadProg, tmap);
    if (!ignoreData) {
      copy(target, inputs[i].input, inType, rawData[i].get());
      copy(target, inputs[i].labels, labels.elementType(), rawLabels[i].get());
    }
  }
  std::vector<unsigned> initLabelLengths(batchSize), initDataLengths(batchSize);
  for (unsigned i = 0; i < batchSize; i++) {
    initLabelLengths[i] = inputs[i].labels.size();
    initDataLengths[i] = inputs[i].inputLength;
  }
  copy(target, initLabelLengths, labelLengths.elementType(),
       rawLabelLengths.get());
  copy(target, initDataLengths, dataLengths.elementType(),
       rawDataLengths.get());

  // Create gradient
  Sequence prog;
  const auto result = [&]() {
    const auto layer = "ctc_loss";
    if (inputs[0].isLogits) {
      return popnn::ctc::calcLossAndGradientLogits(
          graph, outType, data, labels, dataLengths, labelLengths, prog,
          blankSymbol, plan, layer, debugOpts);
    } else {
      return popnn::ctc::calcLossAndGradientLogProbabilities(
          graph, outType, data, labels, dataLengths, labelLengths, prog,
          blankSymbol, plan, layer, debugOpts);
    }
  }();

  auto lossResult = result.first;
  auto gradResult = result.second;
  // Create handles for reading the result
  std::vector<std::unique_ptr<char[]>> rawLossResult(batchSize);
  std::vector<std::unique_ptr<char[]>> rawGradResult(batchSize);
  if (!ignoreData) {
    for (unsigned i = 0; i < batchSize; i++) {
      rawLossResult[i] = allocateHostMemoryForTensor(
          lossResult.slice(i, i + 1, 0), "result_loss_" + std::to_string(i),
          graph, uploadProg, downloadProg, tmap);
      rawGradResult[i] = allocateHostMemoryForTensor(
          gradResult.slice(i, i + 1, 1), "result_grad_" + std::to_string(i),
          graph, uploadProg, downloadProg, tmap);
    }
  }

  // Run input, gradient, output
  OptionFlags engineOptions;
  if (profile || profileDir) {
    engineOptions.set("debug.instrumentCompute", "true");
    if (profileDir) {
      engineOptions.set("autoReport.all", "true");
      engineOptions.set("autoReport.directory", *profileDir);
    }
  }

  Sequence s = [&]() {
    if (ignoreData) {
      // Because the input data has constraints of what is valid, we can't
      // ignore the uploadProg without reasonable likelihood of encountering an
      // exception or unexpected behaviour.
      return Sequence{uploadProg, prog};
    } else {
      return Sequence{uploadProg, prog, downloadProg};
    }
  }();
  Engine engine(graph, s, engineOptions);
  attachStreams(engine, tmap);
  device.bind([&](const Device &d) {
    engine.load(d);
    engine.run();
  });

  // Fetch the result
  std::vector<std::pair<double, boost::multi_array<double, 2>>> output(
      batchSize);
  if (!ignoreData) {
    for (unsigned i = 0; i < batchSize; i++) {
      copy(target, outType, rawLossResult[i].get(), &output[i].first, 1);
      output[i].second.resize(boost::extents[maxT][numClasses]);
      copy(target, outType, rawGradResult[i].get(), output[i].second);
    }
  }

  if (profile && deviceType != DeviceType::Cpu) {
    engine.printProfileSummary(std::cout,
                               OptionFlags{{"showExecutionSteps", "true"}});
  }
  return output;
}

int main(int argc, char **argv) {
  // Default input parameters.
  DeviceType deviceType = DeviceType::IpuModel2;
  boost::optional<std::string> profileDir;
  boost::optional<std::string> planConstraints;
  boost::optional<unsigned> minRandomTime = boost::none;
  boost::optional<unsigned> fixedTime = boost::none;
  unsigned maxTime = 15;
  boost::optional<unsigned> minRandomLabelLength = boost::none;
  boost::optional<unsigned> fixedLabelLength = boost::none;
  unsigned maxLabelLength = 5;
  unsigned blankClass = 0;
  unsigned numClasses = 4;
  unsigned batchSize = 1;
  boost::optional<unsigned> tiles = boost::none;
  Type inType = FLOAT;
  Type partialsType = FLOAT;
  Type outType = FLOAT;
  bool isLogits = true;
  bool testReducedCodeletGradient = false;

  po::options_description desc("Options");
  // clang-format off
  desc.add_options()
    ("help", "Produce help message")
    ("device-type",
     po::value<DeviceType>(&deviceType)->default_value(deviceType),
     deviceTypeHelp)
    ("tiles-per-ipu",po::value<boost::optional<unsigned>>(&tiles),
      "Number of tiles per IPU")
    ("profile", "Show profile report")
    ("profile-dir",
     po::value<decltype(profileDir)>(&profileDir)
      ->default_value(boost::none),
     "Write profile files to the specified directory.")
    ("plan-constraints", po::value(&planConstraints),
     "JSON constraints for planner, e.g. {\"parallel\": {\"batch\": 1}}")
    ("in-type", po::value(&inType)->default_value(inType),
     "Input data type")
    ("partials-type", po::value(&partialsType)->default_value(partialsType),
     "Input data type")
    ("out-type", po::value(&outType)->default_value(outType),
     "Output data type")
    ("batch", po::value(&batchSize)->default_value(batchSize),
     "Batch size")
    ("label-length", po::value(&fixedLabelLength),
     "If set, forces every label to be of length `label-length`")
    ("min-label-length",
      po::value(&minRandomLabelLength),
     "If set, minimum randomly generated label length")
    ("max-label-length",
      po::value(&maxLabelLength)->default_value(maxLabelLength),
     "Max test length (labels)")
    ("time", po::value(&fixedTime),
     "If set, forces every sequence to be of length `time`")
    ("min-time", po::value(&minRandomTime),
     "If set, minimum randomly generated time length")
    ("max-time", po::value(&maxTime)->default_value(maxTime),
     "Max test length (time)")
    ("blank-class", po::value(&blankClass)->default_value(blankClass),
     "Index of the blank symbol. Range 0 to (num-classes-1)")
    ("num-classes", po::value(&numClasses)->default_value(numClasses),
     "Classes in the alphabet including blank")
    ("ignore-data", "Ignore data, to check execution time")
    ("logit-inputs", po::value(&isLogits)->default_value(isLogits),
     "pass logit inputs to ctc loss api, otherwise convert to logProbs prior")
    ("test-reduced-codelet-result",
        po::value(&testReducedCodeletGradient)->
            default_value(testReducedCodeletGradient),
     "Test the reduced result: alpha * beta / probability, omitting any further"
     " processing")
    ("disable-always-satisfiable-error", "Disable the check when validating time"
    " and labelLength before generating random labels. This check ensures that the"
    " label is always representable for given t and labelLength."
    "\nThe length of t required to represent a given label depends on the number"
    " of duplicate classes in the label, this check assumes the worst case where"
    " every class is a duplicate."
    "\nSpecifically:"
    "\n  2 * t - 1 >= labelLength")
    ("plan-only", "Only plan the requested passes, don't build or run a graph")
    ("verbose", "Provide debug printout");
  // clang-format on

  po::variables_map vm;
  try {
    po::store(po::parse_command_line(argc, argv, desc), vm);
    if (vm.count("help")) {
      std::cout << desc << "\n\n";
      return 1;
    }

  } catch (std::exception &e) {
    std::cerr << "error parsing command line: " << e.what() << "\n";
    return 1;
  }

  const bool profile = vm.count("profile");
  const bool verbose = vm.count("verbose");
  const bool ignoreData = vm.count("ignore-data");
  const bool planOnly = vm.count("plan-only");
  const bool disableAlwaysSatisfiableError =
      vm.count("disable-always-satisfiable-error");

  // Needed to set default arguments.
  po::notify(vm);
  // Pick up on some parameters that are easy to get wrong
  if (blankClass >= numClasses) {
    throw poputil::poplibs_error("The blank class must be in the range 0 to "
                                 "(number of classes - 1)");
  }
  if (!minRandomTime && !fixedTime) {
    fixedTime = maxTime;
  }
  if (!minRandomLabelLength && !fixedLabelLength) {
    fixedLabelLength = maxLabelLength;
  }
  validateTimeAndLabelBounds(minRandomTime, fixedTime, maxTime,
                             minRandomLabelLength, fixedLabelLength,
                             maxLabelLength);

  poplar::OptionFlags planOpts;
  if (planConstraints) {
    planOpts.set("planConstraints", *planConstraints);
  }
  planOpts.set("partialsType", partialsType.toString());

  poplar::OptionFlags debugOpts;
  if (testReducedCodeletGradient) {
    debugOpts.set("returnReducedCodeletGradient", "true");
  }

  if (planOnly) {
    auto device = createTestDevice(deviceType, 1, tiles);
    const auto &target = device.getTarget();
    Graph graph(target);

    const auto plan =
        popnn::ctc::plan(graph, inType, outType, batchSize, maxTime,
                         maxLabelLength, numClasses, planOpts);

    std::cout << plan << std::endl;
    std::cout << "No test run - plan only" << std::endl;
    return 0;
  }

  RandomUtil rand{42};

  // For test call the reference function for each batch input
  std::vector<InputSequence<double>> tests;
  std::vector<std::pair<double, boost::multi_array<double, 2>>> references;
  for (unsigned i = 0; i < batchSize; i++) {
    const auto [t, labelLength] = getRandomSize(
        minRandomTime, fixedTime, maxTime, minRandomLabelLength,
        fixedLabelLength, maxLabelLength, disableAlwaysSatisfiableError, rand);
    auto [input, label] = getRandomTestInput<double>(
        t, maxTime, labelLength, numClasses, blankClass, isLogits, rand);
    tests.push_back(InputSequence<double>{input, t, label, isLogits});

    if (verbose) {
      std::cout << "\nBatch:" << i << " Time:" << tests[i].inputLength
                << " Label length:" << tests[i].labels.size();
    }
    print(" Test sequence[" + std::to_string(tests[i].labels.size()) + "] ",
          tests[i].labels, blankClass, verbose);
    if (tests[i].isLogits) {
      print("Logits in", tests[i].input, blankClass, verbose);
      print("Log Softmax in",
            log::log(transpose(log::softMax(transpose(tests[i].input)))),
            blankClass, verbose);
    } else {
      print("Log Softmax in", tests[i].input, blankClass, verbose);
    }
    if (!ignoreData) {
      references.push_back(
          gradReference<double>(tests[i], blankClass, numClasses,
                                testReducedCodeletGradient, verbose));
      references.back().second =
          maskResults(references.back().second, tests[i].inputLength);
      if (verbose) {
        std::cout << "Reference loss = " << references.back().first << "\n";
      }
    }
  }
  auto outputs = gradIPU(tests, maxLabelLength, blankClass, numClasses, inType,
                         outType, planOpts, debugOpts, deviceType, tiles,
                         ignoreData, profile, profileDir);

  for (unsigned i = 0; i < batchSize; i++) {
    outputs[i].second = maskResults(outputs[i].second, tests[i].inputLength);
    if (verbose) {
      std::cout << "Result loss = " << outputs[i].first << "\n";
    }
    print("Result gradient, batch:" + std::to_string(i), outputs[i].second,
          blankClass, verbose);
  }
  double relativeTolerance, absoluteTolerance;
  if (testReducedCodeletGradient) {
    relativeTolerance = outType == FLOAT ? CODELET_TEST_FLOAT_REL_TOL
                                         : CODELET_TEST_HALF_REL_TOL;
    absoluteTolerance = outType == FLOAT ? CODELET_TEST_FLOAT_ABS_TOL
                                         : CODELET_TEST_HALF_ABS_TOL;
  } else {
    relativeTolerance = outType == FLOAT ? FLOAT_REL_TOL : HALF_REL_TOL;
    absoluteTolerance = outType == FLOAT ? FLOAT_ABS_TOL : HALF_ABS_TOL;
  }
  bool success = true;
  if (!ignoreData) {
    for (unsigned i = 0; i < batchSize; i++) {
      bool batchSuccess = checkIsClose(outputs[i].first, references[i].first,
                                       relativeTolerance) &&
                          checkIsClose("Batch:" + std::to_string(i) + " result",
                                       outputs[i].second, references[i].second,
                                       relativeTolerance, absoluteTolerance);
      if (!batchSuccess) {
        success = false;
      }
    }
  }
  if (!success) {
    std::cerr << "Data mismatch\n";
  }
  return !success;
}