Stream.hpp

/*
 * (c) Copyright 2021 CORSIKA Project, corsika-project@lists.kit.edu
 *
 * This software is distributed under the terms of the GNU General Public
 * Licence version 3 (GPL Version 3). See file LICENSE for a full version of
 * the license.
 */

/*
 * Stream.hpp
 *
 *  Created on: 23/02/2021
 *      Author: Antonio Augusto Alves Junior
 */

#pragma once

#include <stdint.h>

#include "detail/tbb/iterators.h"
#include "detail/Engine.hpp"
#include "detail/functors/EngineCaller.hpp"

namespace random_iterator {

  typedef detail::philox philox;
  typedef detail::threefry threefry;
  typedef detail::squares3_128 squares3_128;
  typedef detail::squares4_128 squares4_128;
  typedef detail::squares3_64 squares3_64;
  typedef detail::squares4_64 squares4_64;

#if RANDOM_ITERATOR_R123_USE_AES_NI
  typedef detail::ars ars;
#endif

  template <typename DistibutionType, typename EngineType>
  class Stream {
    typedef detail::EngineCaller<DistibutionType, EngineType> caller_type; 
    typedef random_iterator_tbb::counting_iterator<typename EngineType::advance_type>
        counting_iterator; 

  public:
    typedef EngineType engine_type;                          
    typedef typename engine_type::state_type state_type;     
    typedef typename engine_type::seed_type seed_type;       
    typedef typename engine_type::advance_type advance_type; 
    typedef typename engine_type::init_type init_type;       

    typedef DistibutionType distribution_type; 
    typedef typename distribution_type::result_type result_type; 

    typedef random_iterator_tbb::transform_iterator<caller_type, counting_iterator>
        iterator; 

    Stream() = delete;

    Stream(distribution_type const& distribution, seed_type seed, uint32_t stream = 0)
        : distribution_(distribution)
        , engine_(seed, stream)
        , seed_(seed)
        , stream_(stream) {}

    Stream(Stream<DistibutionType, EngineType> const& other)
        : distribution_(other.getDistribution())
        , engine_(other.getEngine())
        , seed_(other.getSeed())
        , stream_(other.getStream()) {}

    inline iterator begin() const {

      counting_iterator first(0);
      caller_type caller(distribution_, seed_, stream_);

      return iterator(first, caller);
    }

    inline iterator end() const {

      counting_iterator last(std::numeric_limits<advance_type>::max());
      caller_type caller(distribution_, seed_, stream_);

      return iterator(last, caller);
    }

    inline result_type operator[](advance_type n) const {

      caller_type caller(distribution_, seed_, stream_);

      return caller(n);
    }

    inline result_type operator()(void) { return distribution_(engine_); }

    inline distribution_type getDistribution() const { return distribution_; }

    inline void setDistribution(distribution_type dist) { distribution_ = dist; }

    inline engine_type getEngine() const { return engine_; }

    inline void setEngine(engine_type engine) { engine_ = engine; }

    inline const seed_type& getSeed() const { return seed_; }

    inline void setSeed(const seed_type& seed) { seed_ = seed; }

    inline uint32_t getStream() const { return stream_; }

    inline void setStream(uint32_t stream) { stream_ = stream; }

    friend inline std::ostream& operator<<(
        std::ostream& os, const Stream<DistibutionType, EngineType>& be) {
      return os << "engine: " << be.getEngine() << " stream: " << be.getStream()
                << " distribution: " << be.getDistribution();
    }

  private:
    distribution_type distribution_;
    engine_type engine_;
    seed_type seed_;
    uint32_t stream_;
  };

  template <typename DistibutionType, typename EngineType>
  Stream<DistibutionType, EngineType> make_stream(DistibutionType const& dist,
                                                  EngineType eng, uint32_t stream) {

    return Stream<DistibutionType, EngineType>(dist, eng.getSeed(), stream);
  }

} // namespace random_iterator