// Copyright (c) 2018, Tom Westerhout
// All rights reserved.
//
// Redistribution and use in source and binary forms, with or without
// modification, are permitted provided that the following conditions are met:
//
// * Redistributions of source code must retain the above copyright notice, this
//   list of conditions and the following disclaimer.
//
// * Redistributions in binary form must reproduce the above copyright notice,
//   this list of conditions and the following disclaimer in the documentation
//   and/or other materials provided with the distribution.
//
// * Neither the name of the copyright holder nor the names of its
//   contributors may be used to endorse or promote products derived from
//   this software without specific prior written permission.
//
// THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
// AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
// IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
// DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
// FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
// DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
// SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
// CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
// OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
// OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.

#pragma once

#include "detail/random.hpp"
#include "config.h"
#include <boost/align/aligned_allocator.hpp>
#include <gsl/gsl-lite.hpp>
#include <mkl_cblas.h>
#include <mkl_vsl.h>
#include <cmath>
#include <functional>
#include <iostream>
#include <type_traits>

TCM_NAMESPACE_BEGIN

namespace detail {
/// Checks that the last `vslRng*` operation succeeded. Throws an exception
/// otherwise.
inline auto check_vsl_status_after_generation(int const status) -> void
{
    if (status == VSL_STATUS_OK) { return; }
    char const* const msg = [](auto code) {
        switch (code) {
        case VSL_ERROR_NULL_PTR: return "`stream` is a NULL pointer";
        case VSL_RNG_ERROR_BAD_STREAM:
            return "`stream` is not a valid random stream";
        case VSL_RNG_ERROR_BAD_UPDATE:
            return "callback function for an abstract BRNG returns an "
                   "invalid number of updated entries in a buffer, "
                   "that is, < 0 or > nmax";
        case VSL_RNG_ERROR_NO_NUMBERS:
            return "callback function for an abstract BRNG returns 0 "
                   "as the number of updated entries in a buffer.";
        case VSL_RNG_ERROR_QRNG_PERIOD_ELAPSED:
            return "period of the generator has been exceeded";
        case VSL_RNG_ERROR_NONDETERM_NRETRIES_EXCEEDED:
            return "number of retries to generate a random number by "
                   "using non-deterministic random number generator "
                   "exceeds threshold";
        case VSL_RNG_ERROR_ARS5_NOT_SUPPORTED:
            return "ARS-5 random number generator is not supported on "
                   "the CPU running the application";
        default: return "unknown error code: a bug in Intel MKL?";
        } // end switch
    }(status);
    throw std::runtime_error{msg};
}

/// Uses Intel MKL to generate numbers distributed according to
/// `Gamma(shape, scale)`.
///
/// https://en.wikipedia.org/wiki/Gamma_distribution.
template <class RealT> class buffered_gamma_distribution_t { // {{{
    static_assert(std::is_same_v<RealT, float> || std::is_same_v<RealT, double>,
                  "tcm::detail::buffered_gamma_distribution_t currently only "
                  "supports `float`s and `double`s.");

  public:
    /// Generation method.
    ///
    /// Possible values are:
    ///
    ///   * `VSL_RNG_METHOD_GAMMA_GNORM`
    ///   * `VSL_RNG_METHOD_GAMMA_GNORM_ACCURATE`
    ///
    /// See https://software.intel.com/en-us/mkl-developer-reference-c-vrnggamma
    /// for more info.
    static constexpr MKL_INT method = VSL_RNG_METHOD_GAMMA_GNORM_ACCURATE;

    static constexpr size_t default_buffer_size = 1024;

    /// Buffer type.
    ///
    /// Performance of Intel MKL's random number generation increases
    /// considerably when multiple variates a generated at once (see
    /// https://software.intel.com/en-us/mkl-vsperfdata-uniform).
    using buffer_type =
        std::vector<RealT, boost::alignment::aligned_allocator<RealT, 64u>>;

    using stream_type = VSLStreamStatePtr;

    /// Distribution parameters.
    class param_type {
        RealT _shape; ///< Shape parameter
        RealT _scale; ///< Scale parameter

      public:
        constexpr explicit param_type(RealT const shape,
                                      RealT const scale) TCM_NOEXCEPT
            : _shape{shape}
            , _scale{scale}
        {
            TCM_ASSERT(shape > 0, "`shape` must be positive");
            TCM_ASSERT(scale > 0, "`scale` must be positive");
        }

        constexpr param_type(param_type const&) noexcept = default;
        constexpr param_type(param_type&&) noexcept      = default;
        constexpr param_type& operator=(param_type const&) noexcept = default;
        constexpr param_type& operator=(param_type&&) noexcept = default;

        constexpr auto shape() const noexcept { return _shape; }
        constexpr auto scale() const noexcept { return _scale; }
    };

    static_assert(std::is_trivially_copyable_v<param_type>);

  private:
    stream_type _stream; ///< Intel MKL random number generator stream
                         ///<
    buffer_type _buffer; ///< Output buffer where the generated
                         ///< numbers are written to
    size_t _i; ///< Current position in `_buffer`. `_i == _buffer.size()`
               ///< indicates that the buffer needs to be refilled.
    param_type _params; ///< Shape and scale parameters
                        ///<

  public:
    /// Constructs a new Gamma distribution and allocates memory for internal
    /// buffers.
    buffered_gamma_distribution_t(
        stream_type const stream, RealT const alpha, RealT const beta,
        size_t const buffer_size = default_buffer_size)
        : _stream{stream}
        , _buffer(buffer_size)
        , _i{buffer_size}
        , _params{alpha, beta}
    {
        TCM_ASSERT(stream != nullptr, "`stream` should not be null");
        TCM_ASSERT(buffer_size > 0, "buffer should not be empty");
    }

    /// Deleted copy constructor to avoid unintentionally copying potentially
    /// big buffers.
    buffered_gamma_distribution_t(buffered_gamma_distribution_t const&) =
        delete;

    /// Move constructor.
    buffered_gamma_distribution_t(buffered_gamma_distribution_t&&) noexcept =
        default;

    /// Deleted copy assignment operator to avoid unintentionally copying
    /// potentially big buffers.
    buffered_gamma_distribution_t&
    operator=(buffered_gamma_distribution_t const&) = delete;

    /// Move assignment operator.
    buffered_gamma_distribution_t&
    operator=(buffered_gamma_distribution_t&&) noexcept = default;

    constexpr auto param() const noexcept -> param_type { return _params; }

    constexpr auto param(param_type const& params) noexcept -> void
    {
        _params = params;
        _i      = _buffer.size();
    }

#if 0 // This is probably too low-level and should not be exposed.
    constexpr auto buffer() noexcept -> gsl::span<RealT> { return {_buffer}; }

    constexpr auto buffer() const noexcept -> gsl::span<RealT const>
    {
        return {_buffer};
    }
#endif

    constexpr auto stream() const noexcept -> VSLStreamStatePtr
    {
        return _stream;
    }

    auto resize(size_t const buffer_size) -> void
    {
        TCM_ASSERT(buffer_size > 0, "buffer should not be empty");
        // TODO(twesterhout): We could implement an optimisation here which
        // would copy the unused part of the buffer. This could potentially
        // avoid generating some random numbers.
        _buffer.resize(buffer_size);
        _i = buffer_size;
    }

  private:
    static auto vsl_generate(VSLStreamStatePtr const stream, MKL_INT const n,
                             float* const r, float const alpha, float const a,
                             float const beta) -> int
    {
        return vsRngGamma(method, stream, n, r, alpha, a, beta);
    }

    static auto vsl_generate(VSLStreamStatePtr const stream, MKL_INT const n,
                             double* const r, double const alpha,
                             double const a, double const beta) -> int
    {
        return vdRngGamma(method, stream, n, r, alpha, a, beta);
    }

    static auto fill(gsl::span<RealT> out, VSLStreamStatePtr const stream,
                     param_type const& params) -> void
    {
        // NOTE: Intel MKL calls the scale parameter `beta`. Do not confuse
        // `beta` with the _inverse_ scale parameter.
        auto const status =
            vsl_generate(stream, static_cast<MKL_INT>(out.size()), out.data(),
                         /*alpha = */ params.shape(),
                         /*displacement = */ RealT{0},
                         /*beta = */ params.scale());
        check_vsl_status_after_generation(status);
    }

  public:
    auto operator()(gsl::span<RealT> out, param_type const& params) const
        -> void
    {
        fill(out, _stream, params);
    }

    auto operator()(gsl::span<RealT> out) const -> void
    {
        (*this)(out, _params);
    }

    auto operator()() -> RealT
    {
        TCM_ASSERT(_i <= _buffer.size(), "Pre-condition violated");
        if (_i == _buffer.size()) {
            fill(_buffer, _stream, _params);
            _i = 0;
        }
        return _buffer[_i++];
    }
};
// }}}

/// Uses Intel MKL to generate numbers distributed according to N(mu, sigma^2).
///
/// See https://en.wikipedia.org/wiki/Normal_distribution.
template <class RealT> class buffered_gauss_distribution_t { // {{{
    static_assert(std::is_same_v<RealT, float> || std::is_same_v<RealT, double>,
                  "tcm::detail::buffered_gauss_distribution_t currently only "
                  "supports `float`s and `double`s.");

  public:
    /// Generation method.
    ///
    /// Possible values are:
    ///
    ///   * `VSL_RNG_METHOD_GAUSSIAN_BOXMULLER`
    ///   * `VSL_RNG_METHOD_GAUSSIAN_BOXMULLER2`
    ///   * `VSL_RNG_METHOD_GAUSSIAN_ICDF`
    ///
    /// See https://software.intel.com/en-us/mkl-developer-reference-c-vrnggaussian
    /// for more info.
    static constexpr MKL_INT method = VSL_RNG_METHOD_GAUSSIAN_BOXMULLER;

    static constexpr size_t default_buffer_size = 1024;

    /// Buffer type.
    ///
    /// Performance of Intel MKL's random number generation increases
    /// considerably when multiple variates a generated at once (see
    /// https://software.intel.com/en-us/mkl-vsperfdata-uniform).
    using buffer_type =
        std::vector<RealT, boost::alignment::aligned_allocator<RealT, 64u>>;

    using stream_type = VSLStreamStatePtr;

    class param_type {
        RealT _mu;    ///< Mean
        RealT _sigma; ///< Standard deviation

      public:
        constexpr explicit param_type(RealT const mu,
                                      RealT const sigma) TCM_NOEXCEPT
            : _mu{mu}
            , _sigma{sigma}
        {
            TCM_ASSERT(sigma > 0, "`sigma` must be positive");
        }

        constexpr param_type(param_type const&) noexcept = default;
        constexpr param_type(param_type&&) noexcept      = default;
        constexpr param_type& operator=(param_type const&) noexcept = default;
        constexpr param_type& operator=(param_type&&) noexcept = default;

        constexpr auto mu() const noexcept { return _mu; }
        constexpr auto sigma() const noexcept { return _sigma; }
    };

  private:
    stream_type _stream; ///< Intel MKL random number generator stream.
                         ///<
    buffer_type _buffer; ///< Output buffer where the generated numbers
                         ///< are written.
    size_t _i; ///< Current position in `_buffer`. `_i == _buffer.size()`
               ///< means that the buffer needs to be refilled.
    param_type _params; ///< Distribution parameters.
                        ///<

  public:
    buffered_gauss_distribution_t(
        stream_type const stream, RealT const mu, RealT const sigma,
        size_t const buffer_size = default_buffer_size) TCM_NOEXCEPT
        : _stream{stream}
        , _buffer(buffer_size)
        , _i{buffer_size}
        , _params{mu, sigma}
    {
        TCM_ASSERT(stream != nullptr, "`stream` should not be null");
        TCM_ASSERT(buffer_size > 0, "buffer should not be empty");
    }

    buffered_gauss_distribution_t(buffered_gauss_distribution_t const&) =
        delete;

    buffered_gauss_distribution_t(buffered_gauss_distribution_t&&) noexcept =
        default;

    buffered_gauss_distribution_t&
    operator=(buffered_gauss_distribution_t const&) = delete;

    buffered_gauss_distribution_t&
    operator=(buffered_gauss_distribution_t&&) noexcept = default;

    constexpr auto param() const noexcept -> param_type { return _params; }
    constexpr auto param(param_type const params) noexcept -> void
    {
        _params = params;
        _i      = _buffer.size();
    }

#if 0
    constexpr auto buffer() noexcept -> gsl::span<RealT> { return {_buffer}; }
    constexpr auto buffer() const noexcept -> gsl::span<RealT const>
    {
        return {_buffer};
    }
#endif

    auto resize(size_t const buffer_size) -> void
    {
        TCM_ASSERT(buffer_size > 0, "buffer should not be empty");
        // TODO(twesterhout): We could implement an optimisation here which
        // would copy the unused part of the buffer. This could potentially
        // avoid generating some random numbers.
        _buffer.resize(buffer_size);
        _i = buffer_size;
    }

    constexpr auto stream() const noexcept -> VSLStreamStatePtr
    {
        return _stream;
    }

  private:
    static auto vsl_generate(VSLStreamStatePtr const stream, MKL_INT const n,
                             float* const r, float const a, float const sigma)
        -> int
    {
        return vsRngGaussian(method, stream, n, r, a, sigma);
    }

    static auto vsl_generate(VSLStreamStatePtr const stream, MKL_INT const n,
                             double* const r, double const a,
                             double const sigma) -> int
    {
        return vdRngGaussian(method, stream, n, r, a, sigma);
    }

    static auto fill(gsl::span<RealT> out, VSLStreamStatePtr stream,
                     param_type const& params) -> void
    {
        auto const status =
            vsl_generate(stream, static_cast<MKL_INT>(out.size()), out.data(),
                         /*a = */ params.mu(),
                         /*sigma = */ params.sigma());
        check_vsl_status_after_generation(status);
    }

  public:
    auto operator()(gsl::span<RealT> out, param_type const& params) const
        -> void
    {
        fill(out, _stream, params);
    }

    auto operator()(gsl::span<RealT> out) const -> void
    {
        (*this)(out, _params);
    }

    auto operator()() -> RealT
    {
        TCM_ASSERT(_i <= _buffer.size(), "Pre-condition violated");
        if (_i == _buffer.size()) {
            fill(_buffer, _stream, _params);
            _i = 0;
        }
        return _buffer[_i++];
    }
};
// }}}

template <class RealT> class buffered_tsallis_distribution_t { // {{{
    static_assert(std::is_same_v<RealT, float> || std::is_same_v<RealT, double>,
                  "Currently, only `float`s and `double`s are supported.");

  private:
    /// Calculation of \f$p\f$ given \f$q_V\f$ (see `Tsallis_RNG` function in
    /// [@Schanze2006] for an explanation).
    ///
    /// [@Schanze2006]: Thomas Schanze, "An exact D-dimensional Tsallis random
    ///                 number generator for generalized simulated annealing", 2006.
    static constexpr auto get_p(RealT const q_V) noexcept -> RealT
    {
        using R = RealT;
        return (R{3} - q_V) / (R{2} * (q_V - R{1}));
    }

    /// Calculation of \f$s\f$ given \f$q_V\f$ and \f$t_V\f$ (see `Tsallis_RNG`
    /// function in [@Schanze2006] for an explanation).
    ///
    /// [@Schanze2006]: Thomas Schanze, "An exact D-dimensional Tsallis random
    ///                 number generator for generalized simulated annealing", 2006.
    static /*constexpr*/ auto get_s(RealT const q_V, RealT const t_V) noexcept
        -> RealT
    {
        using R = RealT;
        return std::sqrt(R{2} * (q_V - R{1}))
               / std::pow(t_V, R{1} / (R{3} - q_V));
    }

  public:
    static constexpr size_t default_buffer_size = 1024;

    /// Buffer type.
    ///
    /// Performance of Intel MKL's random number generation increases
    /// considerably when multiple variates a generated at once (see
    /// https://software.intel.com/en-us/mkl-vsperfdata-uniform).
    using buffer_type =
        std::vector<RealT, boost::alignment::aligned_allocator<RealT, 64u>>;

    class param_type {
        RealT _q_V; ///< Visiting distribution shape parameter
        RealT _t_V; ///< Visiting temperature
        RealT _s;   ///< Variable `s` from `Tsallis_RNG` in [@Schanze2006].

      public:
        explicit param_type(RealT const q_V, RealT const t_V) TCM_NOEXCEPT
            : _q_V{q_V}
            , _t_V{t_V}
        {
            TCM_ASSERT(1 < q_V && q_V < 3, "`q_V` must be in (1, 3)");
            TCM_ASSERT(t_V > 0, "`t_V` must be positive");
            _s = get_s(q_V, t_V);
        }

        constexpr param_type(param_type const&) noexcept = default;
        constexpr param_type(param_type&&) noexcept      = default;
        constexpr param_type& operator=(param_type const&) noexcept = default;
        constexpr param_type& operator=(param_type&&) noexcept = default;

        constexpr auto q_V() const noexcept { return _q_V; }
        constexpr auto t_V() const noexcept { return _t_V; }
        constexpr auto s() const noexcept { return _s; }
    };

  private:
    buffered_gamma_distribution_t<RealT> _gamma_dist;
    buffered_gauss_distribution_t<RealT> _gauss_dist;
    param_type                           _params;

  public:
    buffered_tsallis_distribution_t(
        VSLStreamStatePtr const stream, RealT const q_V, RealT const t_V,
        size_t const buffer_size = default_buffer_size)
        : _gamma_dist{stream, get_p(q_V), RealT{1}, buffer_size}
        , _gauss_dist{stream, RealT{0}, RealT{1}, buffer_size}
        , _params{q_V, t_V}
    {}

    constexpr auto param() const noexcept -> param_type { return _params; }
    constexpr auto param(param_type const params) noexcept -> void
    {
        _gamma_dist.param(
            typename buffered_gamma_distribution_t<RealT>::param_type{
                get_p(params.q_V()), RealT{1}});
        // _gauss_dist doesn't depend on params
        _params = params;
    }

    constexpr auto stream() const noexcept -> VSLStreamStatePtr
    {
        return _gamma_dist.stream();
    }

  private:
    auto fill(gsl::span<RealT> out) -> void
    {
        using R      = RealT;
        using Params = typename buffered_gauss_distribution_t<R>::param_type;

        // Updates internal buffer and is the reason `fill` is not const.
        auto const u = _gamma_dist();
        auto const y = _params.s() * std::sqrt(u);

        // TODO(twesterhout): Profile what is more efficient: using different
        // sigma or calling cblas_?scal.
        _gauss_dist(out, Params{R{0}, R{1} / y});

        // NOLINTNEXTLINE(readability-braces-around-statements)
        // if constexpr (std::is_same_v<R, float>) {
        //     cblas_sscal(static_cast<MKL_INT>(out.size()), R{1} / y, out.data(),
        //                 1);
        // }
        // else {
        //     cblas_dscal(static_cast<MKL_INT>(out.size()), R{1} / y, out.data(),
        //                 1);
        // }
    }

  public:
    // Generates an N-dimensional sample
    auto operator()(gsl::span<RealT> out) { fill(out); }

    // Generates a 1-dimensional sample
    auto operator()() -> RealT
    {
        auto const u = _gamma_dist();
        auto const y = _params.s() * std::sqrt(u);
        auto const x = _gauss_dist();
        return x / y;
    }
};
// }}}

template class buffered_gamma_distribution_t<float>;
template class buffered_gamma_distribution_t<double>;

template class buffered_gauss_distribution_t<float>;
template class buffered_gauss_distribution_t<double>;

template class buffered_tsallis_distribution_t<float>;
template class buffered_tsallis_distribution_t<double>;

} // namespace detail

template <class EnergyFn, class WrapFn> class sa_chain_t;

struct sa_pars_t {
    float    q_V;
    float    q_A;
    float    t_0;
    unsigned dimension;
    unsigned num_iterations;
};

class sa_buffers_t { // {{{
  public:
    template <class EnergyFn, class WrapFn> friend class sa_chain_t;

    using buffer_type =
        std::vector<float, boost::alignment::aligned_allocator<float, 64u>>;

    struct state_type {
        float       value;
        buffer_type buffer;

      private:
        /// Rounds x up to the closest multiple of 64.
        static constexpr auto round_up_to_64(size_t const x) noexcept -> size_t
        {
            constexpr size_t N = 64;
            return (x + N - 1) & N;
        }

      public:
        /// Allocates a buffer of the specified size.
        explicit state_type(size_t const size)
            : value{std::numeric_limits<float>::infinity()}, buffer{}
        {
            resize(size);
        }

        state_type(state_type const&)     = default;
        state_type(state_type&&) noexcept = default;
        state_type& operator=(state_type const&) = default;
        state_type& operator=(state_type&&) noexcept = default;

        /// Ensures buffer is big enough to hold `size` elements.
        ///
        /// We use a small trick here. We round `size` up to a multiple of 64 so
        /// that we can safely use SIMD instructions without worrying about
        /// writing past the end of the buffer.
        auto resize(size_t const size) -> void
        {
            buffer.reserve(round_up_to_64(size));
            buffer.resize(size);
        }
    };

  private:
    using distribution_type = detail::buffered_tsallis_distribution_t<float>;

    state_type        _current;
    state_type        _proposed;
    state_type        _best;
    distribution_type _tsallis;

  public:
    sa_buffers_t(sa_pars_t const&  params,
                 VSLStreamStatePtr stream = random_stream());

    sa_buffers_t(sa_buffers_t const&)     = delete;
    sa_buffers_t(sa_buffers_t&&) noexcept = default;
    sa_buffers_t& operator=(sa_buffers_t const&) = delete;
    sa_buffers_t& operator=(sa_buffers_t&&) = default;

    auto reset(sa_pars_t const& params) -> void;
    auto reset(distribution_type::param_type const& params) -> void;

    template <class Initialise> auto guess(Initialise&& init) -> void;
};

inline sa_buffers_t::sa_buffers_t(sa_pars_t const&  params,
                                  VSLStreamStatePtr stream)
    : _current{params.dimension}
    , _proposed{params.dimension}
    , _best{params.dimension}
    , _tsallis{stream, params.q_V, params.t_0}
{}

inline auto sa_buffers_t::reset(sa_pars_t const& params) -> void
{
    _current.resize(params.dimension);
    _proposed.resize(params.dimension);
    _best.resize(params.dimension);
    _tsallis.param(distribution_type::param_type{params.q_V, params.t_0});
}

inline auto sa_buffers_t::reset(distribution_type::param_type const& params)
    -> void
{
    _tsallis.param(params);
}

template <class Initialise> auto sa_buffers_t::guess(Initialise&& init) -> void
{
    std::forward<Initialise>(init)(gsl::span<float>{_current.buffer});
}
// }}}

template <class EnergyFn, class WrapFn> class sa_chain_t { // {{{
  public:
    using energy_fn_type    = EnergyFn;
    using wrap_fn_type      = WrapFn;
    using tsallis_dist_type = detail::buffered_tsallis_distribution_t<float>;

  private:
    energy_fn_type _energy_fn; ///< Energy function `f`.
                               ///<
    wrap_fn_type _wrap_fn;     ///< Wrapping function `wrap` which ensures that
                               ///< `wrap(x + dx)` is in the domain of `f`.
    sa_buffers_t&    _buffers;
    sa_pars_t const& _params;
    size_t           _i; ///< Current iteration.
                         ///<

  public:
    sa_chain_t(sa_buffers_t& buffers, sa_pars_t const& params,
               energy_fn_type energy_fn, wrap_fn_type wrap_fn);

    sa_chain_t(sa_chain_t const&) = delete;
    sa_chain_t(sa_chain_t&&)      = delete;
    sa_chain_t& operator=(sa_chain_t const&) = delete;
    sa_chain_t& operator=(sa_chain_t&&) = delete;

    inline auto    operator()() -> void;
    constexpr auto restart() noexcept -> void;
    constexpr auto current() const noexcept -> sa_buffers_t::state_type const&;
    constexpr auto best() const noexcept -> sa_buffers_t::state_type const&;

  private:
    constexpr auto t_0() const noexcept { return _params.t_0; }
    constexpr auto q_V() const noexcept { return _params.q_V; }
    constexpr auto q_A() const noexcept { return _params.q_A; }

    /// Returns the dimension of the parameter space.
    auto dim() const noexcept { return _buffers._current.buffer.size(); }

    /// Calculates the visiting temperature `t_V` for the given iteration.
    inline auto temperature(size_t i) const noexcept -> float;

    /// Sets the best state equal to current.
    inline auto update_best() noexcept -> void;

    /// Accepts the move, i.e. `proposed` becomes `current`
    inline auto accept_full() noexcept -> void;
    inline auto accept_one(size_t i, float const x) noexcept -> void;

    /// Rejects the move.
    inline auto reject_full() noexcept -> void;
    inline auto reject_one(size_t i, float const x) noexcept -> void;

    /// Returns a uniform random number
    ///
    /// TODO(twesterhout): This function is terribly inefficient.
    inline auto uniform() -> float;

    template <class Accept, class Reject>
    inline auto accept_or_reject(float const dE, float const t_A,
                                 Accept&& accept, Reject&& reject) -> void;

    inline auto generate_full() -> void;
    inline auto generate_one(size_t const i) -> std::tuple<float, float>;
};
// }}}

// {{{ sa_chain_t IMPLEMENTATION
template <class E, class W>
inline sa_chain_t<E, W>::sa_chain_t(sa_buffers_t&    buffers,
                                    sa_pars_t const& params,
                                    energy_fn_type   energy_fn,
                                    wrap_fn_type     wrap_fn)
    : _energy_fn{std::move(energy_fn)}
    , _wrap_fn{std::move(wrap_fn)}
    , _buffers{buffers}
    , _params{params}
    , _i{0}
{
    _buffers._current.value = _energy_fn(_buffers._current.buffer);
    update_best();
}

template <class E, class W>
inline auto sa_chain_t<E, W>::temperature(std::size_t const i) const noexcept
    -> float
{
    auto const num = t_0() * (std::pow(2.0f, q_V() - 1.0f) - 1.0f);
    auto const den = std::pow(static_cast<float>(2 + i), q_V() - 1.0f) - 1.0f;
    return num / den;
}

template <class E, class W>
inline auto sa_chain_t<E, W>::update_best() noexcept -> void
{
    using std::begin, std::end;
    std::copy(begin(_buffers._current.buffer), end(_buffers._current.buffer),
              begin(_buffers._best.buffer));
    _buffers._best.value = _buffers._current.value;
}

template <class E, class W>
TCM_NOINLINE auto sa_chain_t<E, W>::uniform() -> float
{
    float      r;
    auto const status =
        vsRngUniform(VSL_RNG_METHOD_UNIFORM_STD, _buffers._tsallis.stream(), 1,
                     &r, 0.0f, 1.0f);
    detail::check_vsl_status_after_generation(status);
    return r;
}

template <class E, class W>
template <class Accept, class Reject>
inline auto sa_chain_t<E, W>::accept_or_reject(float const dE, float const t_A,
                                               Accept&& accept, Reject&& reject)
    -> void
{
#if 0
    auto const accept = [this]() {
        using std::swap;
        swap(_buffers._current, _buffers._proposed);
        if (_buffers._current.value < _buffers._best.value) { update_best(); }
    };
    auto const reject = []() {};
#endif

    // Always accept moves that reduce the energy
    if (dE < 0.0f) {
        std::forward<Accept>(accept)();
        return;
    }

    // Eq. (5)
    //
    // pqv_temp = (q_A - 1.0) * (e - self.energy_state.current_energy) / ( self.temperature_step + 1.)
    //
    // auto const factor = 1.0f + (q_A() - 1.0f) * dE / t_A;
    auto const factor = 1.0f + (q_A() - 1.0f) * dE / t_A;
    auto const P_qA =
        factor <= 0.0f ? 0.0f : std::pow(factor, 1.0f / (1.0f - q_A()));
    if (uniform() <= P_qA) {
        // std::fprintf(stderr, "Ping!\n");
        std::forward<Accept>(accept)();
    }
    else {
        // std::fprintf(stderr, "T_A = %f, factor = %f, r = %f, P_qA = %f\n",
        //              static_cast<double>(t_A), static_cast<double>(factor),
        //              static_cast<double>(r), static_cast<double>(P_qA));
        std::forward<Reject>(reject)();
    }
}

template <class E, class W>
inline auto sa_chain_t<E, W>::generate_full() -> void
{
    using std::begin, std::end;
    _buffers._tsallis(_buffers._proposed.buffer);
    std::transform(
        begin(_buffers._current.buffer), end(_buffers._current.buffer),
        begin(_buffers._proposed.buffer), begin(_buffers._proposed.buffer),
        [this](auto const x, auto const dx) -> float {
            return _wrap_fn(x + dx);
        });
    _buffers._proposed.value = _energy_fn(_buffers._proposed.buffer);
}

template <class E, class W>
inline auto sa_chain_t<E, W>::generate_one(std::size_t const i)
    -> std::tuple<float, float>
{
    auto const x = _wrap_fn(_buffers._current.buffer[i] + _buffers._tsallis());
    auto       temp_buffer = _buffers._current.buffer;
    temp_buffer[i]         = x;
    auto const value       = _energy_fn(temp_buffer);
    // auto const value =
    //     _energy_fn(i, x, _buffers._current.value, _buffers._current.buffer);
    return std::make_tuple(x, value);
}

template <class E, class W>
constexpr auto sa_chain_t<E, W>::current() const noexcept
    -> sa_buffers_t::state_type const&
{
    return _buffers._current;
}

template <class E, class W>
constexpr auto sa_chain_t<E, W>::best() const noexcept
    -> sa_buffers_t::state_type const&
{
    return _buffers._best;
}

template <class E, class W>
constexpr auto sa_chain_t<E, W>::restart() noexcept -> void
{
    _i = 0;
}

template <class E, class W>
TCM_NOINLINE auto sa_chain_t<E, W>::operator()() -> void
{
    // (iv) Calculate new temperature...
    auto const t_V = temperature(_i);
    auto const t_A = t_V / static_cast<float>(_i + 1);
    _buffers._tsallis.param(typename tsallis_dist_type::param_type{q_V(), t_V});

    // Markov chain at constant temperature
    for (auto j = size_t{0}; j < dim(); ++j) {
        auto const accept = [this]() {
            using std::swap;
            swap(_buffers._current, _buffers._proposed);
            if (_buffers._current.value < _buffers._best.value) {
                update_best();
            }
            std::cerr << _buffers._current.value << '\n';
        };
        auto const reject = []() {};
        generate_full();
        accept_or_reject(_buffers._proposed.value - _buffers._current.value,
                         t_A, accept, reject);
    }
    for (auto j = size_t{0}; j < dim(); ++j) {
        auto const [x, value] = generate_one(j);
        auto const accept     = [this, j, x = x, value = value]() {
            _buffers._current.buffer[j] = x;
            _buffers._current.value     = value;
            if (_buffers._current.value < _buffers._best.value) {
                update_best();
            }
            std::cerr << _buffers._current.value << '\n';
        };
        auto const reject = []() {};
        accept_or_reject(value - _buffers._current.value, t_A, accept, reject);
    }

    // NOTE: Don't forget this!
    ++_i;

#if 0
        for (auto j = std::size_t{0}; j < 2 * dim(); ++j) {
            // (ii) Then randomly generate x...
            if (j < dim()) { generate_full(); }
            else {
                generate_one(j - dim());
            }
            // (iii)
            auto const proposed_energy = std::get<0>(_proposed);
            auto const current_energy  = std::get<0>(_current);
            accept_or_reject(proposed_energy - current_energy, t_A);
            auto const proposed_energy = std::get<0>(_proposed);
            auto const current_energy  = std::get<0>(_current);
            if (proposed_energy < current_energy) { accept(); }
            else {
                real_type r;
                int       status = vsRngUniform(
                    VSL_RNG_METHOD_UNIFORM_STD_ACCURATE, _stream, 1, &r,
                    /*a = */ real_type{0}, /*b = */ real_type{1});
                if (status != VSL_STATUS_OK)
                    throw std::runtime_error{"vsRngUniform failed."};

                // Eq. (5)
                auto const factor = real_type{1}
                                    + (_params.q_A - 1)
                                          * (proposed_energy - current_energy)
                                          / t_A;
                auto const P_qA =
                    factor <= 0
                        ? real_type{0}
                        : std::pow(factor,
                                   real_type{1} / (real_type{1} - _params.q_A));
                if (r <= P_qA) {
                    std::fprintf(stderr, "Ping!\n");
                    accept();
                }
                else {
                    std::fprintf(
                        stderr, "T_V = %f, factor = %f, r = %f, P_qA = %f\n",
                        static_cast<double>(t_V),
                        static_cast<double>(factor), static_cast<double>(r),
                        static_cast<double>(P_qA));
                    reject();
                }
            }
        }
#endif
}
// }}}

inline auto simulated_annealing_buffers(sa_pars_t const& params)
    -> sa_buffers_t&
{
    constexpr auto default_params = sa_pars_t{/*q_V=*/2.62f,
                                              /*q_A=*/-5.0f,
                                              /*t_0=*/5230.0f,
                                              /*dimension=*/1000000,
                                              /*num_iterations=*/1000};

    thread_local sa_buffers_t buffers{default_params, random_stream()};
    buffers.reset(params);
    return buffers;
}

template <class EnergyFn, class WrapFn, class Initialise, class Finalise>
inline auto minimise_using_simulated_annealing(EnergyFn energy, WrapFn wrap,
                                               sa_pars_t const& params,
                                               Initialise&&     input,
                                               Finalise&&       output) -> void
{
    auto& buffers = simulated_annealing_buffers(params);
    buffers.guess(std::forward<Initialise>(input));

    sa_chain_t chain{buffers, params, std::move(energy), std::move(wrap)};
    for (auto i = 0u; i < params.num_iterations; ++i) {
        chain();
    }
    auto const& best = chain.best();
    std::forward<Finalise>(output)(gsl::span<float const>{best.buffer});
}

TCM_NAMESPACE_END


/****************************************************************************
 * Copyright (c) 2021 by the Picasso authors                                *
 * All rights reserved.                                                     *
 *                                                                          *
 * This file is part of the Picasso library. Picasso is distributed under a *
 * BSD 3-clause license. For the licensing terms see the LICENSE file in    *
 * the top-level directory.                                                 *
 *                                                                          *
 * SPDX-License-Identifier: BSD-3-Clause                                    *
 ****************************************************************************/

#ifndef PICASSO_PARTICLELEVELSET_HPP
#define PICASSO_PARTICLELEVELSET_HPP

#include <Picasso_FieldManager.hpp>
#include <Picasso_LevelSet.hpp>
#include <Picasso_Types.hpp>

#include <Cajita.hpp>

#include <ArborX.hpp>

#include <boost/property_tree/ptree.hpp>

#include <cfloat>
#include <cmath>

//---------------------------------------------------------------------------//
// ArborX Data
//---------------------------------------------------------------------------//
// Search primitives. We build the tree from particles of the color we want
// the level set for.
namespace Picasso
{
template <class CoordinateSlice, class ViewType>
struct ParticleLevelSetPrimitiveData
{
    using memory_space = typename CoordinateSlice::memory_space;
    using size_type = typename CoordinateSlice::size_type;
    CoordinateSlice x;
    ViewType c;
    int num_color;
};

// Predicate storage - store the the 3d index of the mesh entity we are going to
// search the tree with along with its coordinates.
template <class IndexType>
struct ParticleLevelSetPredicateStorage
{
    using size_type = IndexType;
    size_type i;
    size_type j;
    size_type k;
    float x;
    float y;
    float z;
};

// Search predicates. We search the tree with the mesh entities on which we
// want to build the level set.
template <class LocalMesh, class EntityType>
struct ParticleLevelSetPredicateData
{
    using memory_space = typename LocalMesh::memory_space;
    using size_type = int;
    using entity_type = EntityType;
    LocalMesh local_mesh;
    size_type size;
    size_type i_size;
    size_type ij_size;

    template <class LocalGrid>
    ParticleLevelSetPredicateData( const LocalMesh& lm,
                                   const LocalGrid& local_grid )
        : local_mesh( lm )
    {
        auto ghost_entities = local_grid.indexSpace(
            Cajita::Ghost(), entity_type(), Cajita::Local() );
        size = ghost_entities.size();
        i_size = ghost_entities.extent( Dim::I );
        ij_size = i_size * ghost_entities.extent( Dim::J );
    }

    KOKKOS_FUNCTION void
    convertIndexTo3d( size_type i,
                      ParticleLevelSetPredicateStorage<size_type>& ijk ) const
    {
        ijk.k = i / ij_size;
        size_type k_size = ijk.k * ij_size;
        ijk.j = ( i - k_size ) / i_size;
        ijk.i = i - ijk.j * i_size - k_size;
    }
};

//---------------------------------------------------------------------------//
// Query callback. When the query occurs we store the distance to the closest
// point. We don't use the tree graph so we don't insert anything into the
// graph.
template <class CoordinateSlice, class DistanceEstimateView>
struct ParticleLevelSetCallback
{
    ParticleLevelSetPrimitiveData<
        CoordinateSlice,
        Kokkos::View<int*, typename CoordinateSlice::memory_space>>
        primitive_data;
    DistanceEstimateView distance_estimate;
    float radius;

    template <typename Predicate>
    KOKKOS_FUNCTION void operator()( Predicate const& predicate,
                                     int primitive_index ) const
    {
        // Get the actual index of the particle.
        auto p = primitive_data.c( primitive_index );

        // Get the predicate storage.
        auto storage = getData( predicate );

        // Compute the distance from the grid entity to the particle sphere.
        float dx = static_cast<float>( primitive_data.x( p, 0 ) ) - storage.x;
        float dy = static_cast<float>( primitive_data.x( p, 1 ) ) - storage.y;
        float dz = static_cast<float>( primitive_data.x( p, 2 ) ) - storage.z;
        distance_estimate( storage.i, storage.j, storage.k, 0 ) =
            sqrt( dx * dx + dy * dy + dz * dz ) - radius;
    }
};

} // end namespace Picasso

//---------------------------------------------------------------------------//
// ArborX traits.
namespace ArborX
{

// Create the primitives we build the tree from. These are the particle
// coordinates of the color we build the level set for.
template <class CoordinateSlice, class ViewType>
struct AccessTraits<
    Picasso::ParticleLevelSetPrimitiveData<CoordinateSlice, ViewType>,
    PrimitivesTag>
{
    using primitive_data =
        Picasso::ParticleLevelSetPrimitiveData<CoordinateSlice, ViewType>;
    using memory_space = typename primitive_data::memory_space;
    using size_type = typename primitive_data::size_type;
    static size_type size( const primitive_data& data )
    {
        return data.num_color;
    }
    static KOKKOS_FUNCTION Point get( const primitive_data& data, size_type i )
    {
        // Get the actual index of the particle.
        auto p = data.c( i );

        // Return a point made from the particles.
        return { static_cast<float>( data.x( p, 0 ) ),
                 static_cast<float>( data.x( p, 1 ) ),
                 static_cast<float>( data.x( p, 2 ) ) };
    }
};

// Create the predicates we search the tree with. These are the mesh entities
// on which we build the level set.
template <class LocalMesh, class EntityType>
struct AccessTraits<
    Picasso::ParticleLevelSetPredicateData<LocalMesh, EntityType>,
    PredicatesTag>
{
    using predicate_data =
        Picasso::ParticleLevelSetPredicateData<LocalMesh, EntityType>;
    using entity_type = typename predicate_data::entity_type;
    using memory_space = typename predicate_data::memory_space;
    using size_type = typename predicate_data::size_type;
    static size_type size( const predicate_data& data ) { return data.size; }
    static KOKKOS_FUNCTION auto get( const predicate_data& data, size_type i )
    {
        // Get the entity index.
        Picasso::ParticleLevelSetPredicateStorage<size_type> storage;
        data.convertIndexTo3d( i, storage );
        int index[3] = { storage.i, storage.j, storage.k };

        // Get the coordinates of the entity.
        double x[3];
        data.local_mesh.coordinates( entity_type(), index, x );
        storage.x = x[0];
        storage.y = x[1];
        storage.z = x[2];

        // Find the nearest particle to the entity. Attach the entity index to
        // use in the callback.
        return attach( nearest( Point{ storage.x, storage.y, storage.z }, 1 ),
                       storage );
    }
};

} // end namespace ArborX

//---------------------------------------------------------------------------//
namespace Picasso
{
//---------------------------------------------------------------------------//
// Particle level set. Composes a signed distance function for particles of a
// given color.
template <class MeshType, class SignedDistanceLocation>
class ParticleLevelSet
{
  public:
    using mesh_type = MeshType;
    using memory_space = typename mesh_type::memory_space;
    using location_type = SignedDistanceLocation;
    using entity_type = typename location_type::entity_type;
    using halo_type = Cajita::Halo<memory_space>;
    using level_set = LevelSet<MeshType, SignedDistanceLocation>;

    /*!
      \brief Construct the level set with particles of a given color. If the
      input color is negative then all particles will be used in the level set
      regardless of color.
      \param ptree Level set settings.
      \param mesh The mesh over which to build the signed distnace function.
      \param color The particle color over which to build the level set. Use
      -1 if the level set is to be built over all particles.
    */
    ParticleLevelSet( const boost::property_tree::ptree& ptree,
                      const std::shared_ptr<MeshType>& mesh, const int color )
        : _color( color )
        , _ls( createLevelSet<SignedDistanceLocation>( ptree, mesh ) )
    {
        // Extract parameters.
        const auto& params = ptree.get_child( "particle_level_set" );

        // Particles have an analytic spherical level set. Get the radius as a
        // fraction of the cell size.
        _dx = mesh->localGrid()->globalGrid().globalMesh().cellSize( 0 );
        _radius = _dx * params.get<double>( "particle_radius", 0.5 );
    }

    Kokkos::View<const int*, memory_space> colorIndex() const
    {
        return _color_indices;
    }

    int colorCount() const { return _color_count; }

    /*!
      \brief Update the set of particle indices for the color we are building
      the set for. This operation is needed any time the particle population
      is updated (e.g. after a redistribution) and therefore is separated from
      the distance estimate because it may only be called periodically.
      \param exec_space The execution space to use for parallel kernels.
      \param c_p A view or slice containing the particle colors. The number
      and order of particles with respect to these colors must remain
      consistent in between calls to this function (e.g. in subsequent calls
      to estimateSignedDistance()).
    */
    template <class ExecutionSpace, class ParticleColors>
    void updateParticleColors( const ExecutionSpace& exec_space,
                               const ParticleColors& c_p )
    {
        // Initialize color indices.
        _color_indices = Kokkos::View<int*, memory_space>(
            Kokkos::ViewAllocateWithoutInitializing( "color_indices" ),
            c_p.size() );

        // If we dont have any particles on this rank we have nothing to do.
        if ( 0 == c_p.size() )
        {
            _color_count = 0;
            return;
        }

        // If the color is not negative then count how many particles have the
        // color we want and compute the indices.
        else if ( _color >= 0 )
        {
            Kokkos::View<int, memory_space> color_count( "color_count" );
            int color = _color;
            auto color_ind = _color_indices;
            Kokkos::parallel_for(
                "Picasso::ParticleLevelSet::CountColor",
                Kokkos::RangePolicy<ExecutionSpace>( exec_space, 0,
                                                     c_p.size() ),
                KOKKOS_LAMBDA( const int p ) {
                    if ( color == c_p( p ) )
                    {
                        color_ind(
                            Kokkos::atomic_fetch_add( &color_count(), 1 ) ) = p;
                    }
                } );
            Kokkos::deep_copy( _color_count, color_count );
            _color_indices = color_ind;
        }

        // Otherwise a negative color means all particles are included in the
        // level set so we can more efficiently generate the array.
        else
        {
            int color_count = c_p.size();
            auto color_ind = _color_indices;
            Kokkos::parallel_for(
                "Picasso::ParticleLevelSet::FillColorIndices",
                Kokkos::RangePolicy<ExecutionSpace>( exec_space, 0,
                                                     c_p.size() ),
                KOKKOS_LAMBDA( const int p ) { color_ind( p ) = p; } );

            _color_count = color_count;
            _color_indices = color_ind;
        }
    }

    /*!
      \brief Compute the signed distance function estimate from the current
      particle locations. All the positive values will be correct but the
      negative values will only have the correct sign until redistanced.
      \param exec_space The execution space to use for parallel kernels.
      \param x_p A view or slice of particle positions. If the grid is
      adaptive these positions must be in the logical frame. The number and
      order of particles with respect to these positions must be consistent
      with the colors provided to the last call to updateParticleColors().
    */
    template <class ExecutionSpace, class ParticlePositions>
    void estimateSignedDistance( const ExecutionSpace& exec_space,
                                 const ParticlePositions& x_p )
    {
        // Distance estimate.
        auto distance_estimate = _ls->getDistanceEstimate();

        // View of the distance estimate.
        auto estimate_view = distance_estimate->view();

        // Local mesh.
        auto local_grid = distance_estimate->layout()->localGrid();
        auto local_mesh = Cajita::createLocalMesh<memory_space>( *local_grid );

        // If we have no particles of the given color on this rank then we are
        // in a region of positive distance. Estimate the signed distance
        // everywhere to be the minimum halo width as this is the minimum
        // distance to a particle surface we could have resolved that wouldn't
        // show up in the min-reduce operation.
        if ( 0 == _color_count )
        {
            double min_dist = _dx * local_grid->haloCellWidth();
            Cajita::ArrayOp::assign( *distance_estimate, min_dist,
                                     Cajita::Ghost() );
        }

        // Otherwise we have particles so build a tree from the particles of
        // the given color and estimate the distance.
        else
        {
            // Build the tree
            ParticleLevelSetPrimitiveData<ParticlePositions,
                                          Kokkos::View<int*, memory_space>>
                primitive_data;
            primitive_data.x = x_p;
            primitive_data.c = _color_indices;
            primitive_data.num_color = _color_count;
            ArborX::BVH<memory_space> bvh( exec_space, primitive_data );

            // Make the search predicates.
            ParticleLevelSetPredicateData<decltype( local_mesh ), entity_type>
                predicate_data( local_mesh, *local_grid );

            // Make the distance callback.
            ParticleLevelSetCallback<ParticlePositions,
                                     decltype( estimate_view )>
                distance_callback;
            distance_callback.primitive_data = primitive_data;
            distance_callback.distance_estimate = estimate_view;
            distance_callback.radius = static_cast<float>( _radius );

            // Query the particle tree with the mesh entities to find the
            // closest particle and compute the initial signed distance
            // estimate. Dummy arguments are needed even though we don't care
            // about the output.
            bvh.query( exec_space, predicate_data, distance_callback );
        }

        // Do a reduction to get the minimum distance within the minimum halo
        // width.
        _ls->getHalo()->scatter( exec_space, Cajita::ScatterReduce::Min(),
                                 *distance_estimate );
    }

    // Get the particle radius.
    double particleRadius() const { return _radius; }

    // Get the level set.
    std::shared_ptr<level_set> levelSet() const { return _ls; }

  private:
    int _color;
    double _radius;
    double _dx;
    Kokkos::View<int*, memory_space> _color_indices;
    int _color_count;
    std::shared_ptr<level_set> _ls;
};

//---------------------------------------------------------------------------//
/*!
  \brief Create a particle level set over particles of the given color. A
  color of -1 will create the level set over all particles.
  \param ptree Level set settings.
  \param mesh The mesh over which to build the signed distance function.
  \param color The particle color over which to build the level set. Use -1 if
  the level set is to be built over all particles.
*/
template <class SignedDistanceLocation, class MeshType>
std::shared_ptr<ParticleLevelSet<MeshType, SignedDistanceLocation>>
createParticleLevelSet( const boost::property_tree::ptree& ptree,
                        const std::shared_ptr<MeshType>& mesh, const int color )
{
    return std::make_shared<ParticleLevelSet<MeshType, SignedDistanceLocation>>(
        ptree, mesh, color );
}

//---------------------------------------------------------------------------//

} // end namespace Picasso

#endif // end PICASSO_PARTICLELEVELSET_HPP


/* FreeCT_ICD is MBIR CT reconstruction Software */
/* Copyright (C) 2018  John Hoffman, Stefano Young, Frederic Noo */

/* This program is free software; you can redistribute it and/or */
/* modify it under the terms of the GNU General Public License */
/* as published by the Free Software Foundation; either version 2 */
/* of the License, or (at your option) any later version. */

/* This program is distributed in the hope that it will be useful, */
/* but WITHOUT ANY WARRANTY; without even the implied warranty of */
/* MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the */
/* GNU General Public License for more details. */

/* You should have received a copy of the GNU General Public License */
/* along with this program; if not, write to the Free Software */
/* Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston, MA  02110-1301, USA. */

/* Questions and comments should be directed to */
/* jmhoffman@mednet.ucla.edu with "FreeCT_ICD" in the subject line*/

#include <iostream>
#include <fstream>
#include <string>
#include <vector>
#include <math.h>

#include <chrono>

#include <omp.h>
#include <boost/numeric/ublas/vector_sparse.hpp>
#include <boost/numeric/ublas/vector.hpp>
#include <boost/numeric/ublas/io.hpp>
namespace ublas = boost::numeric::ublas;

#include "spinner.h"
#include "recon_structs.h"
#include "icd_iteration.h"
#include "penalties.h"

#define OMP_N_THREADS 11

void icd_iteration(const struct recon_params * rp, struct ct_data * data){

    size_t data_size = rp->Readings*rp->n_channels*rp->Nrows_projection;

    // Allocate sinogram estimate (all zeros)
    double * sinogram_estimate = new double[rp->Readings*rp->n_channels*rp->Nrows_projection]();
    double * reconstructed_image= new double[rp->num_voxels_x*rp->num_voxels_y*rp->num_voxels_z];

    // Copy the float recon volume into the vector array (if uninitialized, will just copy zeros);
    for (int i=0; i<rp->num_voxels_x; i++){
        for (int j=0; j<rp->num_voxels_y; j++){
            for (int k=0; k<rp->num_voxels_z; k++){
                size_t idx=i+j*rp->num_voxels_x+k*rp->num_voxels_x*rp->num_voxels_y;
                reconstructed_image[idx]=(double)data->recon_volume[idx];
            }
        }
    }

    std::ifstream file(rp->matrix_path, std::ios_base::binary);

    // If WFBP was used to inialize the reconstructions, we need to initialize our sinogram estimate.
    if (rp->wfbp_initialize){
        std::cout << "Initializing sinogram estimate..." << std::endl;

        // run a forward projection to initialize
        init_spinner();
        for (int j=0; j<rp->num_voxels_y; j++){
            update_spinner(j,rp->num_voxels_x);
            for (int i=0; i<rp->num_voxels_x; i++){
                // Extract column of projection matrix
                size_t nnz;
                file.read((char*)&nnz, sizeof(nnz));
                int num_nonzeros = (int)nnz; // cast to int to avoid potential issues
                struct pair{
                    int index;
                    float value;
                };
                //std::vector<pair> nonzeros(num_nonzeros);

                struct pair * nonzeros = new struct pair[num_nonzeros];
                if (num_nonzeros > 0)
                    file.read((char*)&nonzeros[0], num_nonzeros*sizeof(pair));

                // Loop over all slices for current x,y
                for (int k=0; k<rp->num_voxels_z; k++){

                    size_t central_idx=data->slice_indices[k];
                    // If using flying focal spots, need to guarantee that offset corresponds to the first projection in a FFS "stack"
                    if (rp->Zffs || rp->Phiffs){
                        int n_focal_spots=pow(2.0,rp->Zffs)*pow(2.0,rp->Phiffs);
                        int mod=central_idx%n_focal_spots;
                        central_idx=central_idx-mod;
                    }

                    int offset = (central_idx - rp->num_views_for_system_matrix/2)*rp->n_channels*rp->Nrows_projection; // offset is a projection index

#pragma omp parallel num_threads(OMP_N_THREADS)
                    {
#pragma omp for
                    for (int m = 0; m<num_nonzeros; m++){
                        int index = nonzeros[m].index + offset; // Raw data index
                        if ((index > -1) && (index < data_size)){
                            size_t voxel_idx=i+j*rp->num_voxels_x+k*rp->num_voxels_x*rp->num_voxels_y;
                            sinogram_estimate[index] = sinogram_estimate[index] + reconstructed_image[voxel_idx]*nonzeros[m].value;
                        }
                    }
                    }
                }
                delete[] nonzeros;
            }
        }
        destroy_spinner();
        file.clear();
        file.seekg(0, std::ios_base::beg);
    }

    // Write pre-iteration reconstruction to disk
    std::ostringstream recon_path;
    recon_path << rp->output_dir << "/reconstructions/iteration0.rcn";
    std::ofstream recon_file(recon_path.str(), std::ios_base::binary);
    recon_file.write((char*)&reconstructed_image[0], rp->num_voxels_x*rp->num_voxels_y*rp->num_voxels_z*sizeof(reconstructed_image[0]));
    recon_file.close();
    std::cout << "Wrote initial image to disk." << std::endl;

    // Write reconstruction to disk
    std::ostringstream sino_est_path;
    sino_est_path << rp->output_dir << "/reconstructions/sino_estimation" << ".rcn";
    std::ofstream sino_file(sino_est_path.str(), std::ios_base::binary);
    sino_file.write((char*)sinogram_estimate, rp->Readings*rp->Nrows_projection*rp->n_channels*sizeof(double));
    sino_file.close();
    std::cout << "Wrote initial sinogram to disk." << std::endl;

    //tk end debugging

    ublas::compressed_vector<float> col(rp->Readings*rp->n_channels*rp->Nrows_projection);

    // Initialize iterative parameters
    // Current implementation limited to 2D (hard coded)
    struct iterative_params ip;
    initialize_2d_weights(&ip);
    ip.lambda = rp->lambda;
    ip.Nx=rp->num_voxels_x;
    ip.Ny=rp->num_voxels_y;
    ip.Nz=rp->num_voxels_z;
    ip.delta  = rp->delta;

    for (int n = 0; n < rp->num_iterations; n++){

        std::cout << "Iteration #" << n+1 << std::endl;
        std::chrono::high_resolution_clock::time_point start=std::chrono::high_resolution_clock::now();

        double fov_limit=(rp->acquisition_fov/2.0)*(rp->acquisition_fov/2.0);

        init_spinner();
        for (int j = 0; j < rp->num_voxels_y; j++){
            update_spinner(j,rp->num_voxels_y);
            double y = (j - rp->center_voxel_y)*rp->voxel_size_y;
            for (int i = 0; i < rp->num_voxels_x; i++){

                double x = (i - rp->center_voxel_x)*rp->voxel_size_x;

                size_t nnz;
                file.read((char*)&nnz, sizeof(nnz));

                int num_nonzeros = (int)nnz; // cast to int to avoid potential issues

                struct pair{
                    int index;
                    float value;
                };

                //std::vector<pair> nonzeros(num_nonzeros);

                struct pair * nonzeros = new struct pair[num_nonzeros];

                if (num_nonzeros > 0)
                    file.read((char*)&nonzeros[0], num_nonzeros*sizeof(pair));

                if ((x*x + y*y) < fov_limit){

                    int q0 = i + rp->num_voxels_x*j;

                    for (int k = 0; k < (rp->num_voxels_z); k++){

                        /// Grab the Z slice locations (spatial+idx)
                        double curr_slice_location=data->slice_locations[k];
                        size_t central_idx=data->slice_indices[k];

                        int q = q0 + rp->num_voxels_x*rp->num_voxels_y*k;

                        // This is the key spot to select slice location (done via the "central_idx" variable)
                        // If using flying focal spots, need to guarantee that offset corresponds to the first projection in a FFS "stack"
                        int start_idx=(central_idx - rp->num_views_for_system_matrix/2);
                        int offset = start_idx*rp->n_channels*rp->Nrows_projection;

                        double alpha = 0.0;
                        double beta  = 0.0;

#pragma omp parallel num_threads(OMP_N_THREADS)
                        {
#pragma omp for reduction(+:alpha,beta)
                            for (int m = 0; m<num_nonzeros; m++){
                                int index = nonzeros[m].index + offset;

                                if ((index > -1) && (index < data_size)){
                                    alpha += nonzeros[m].value * nonzeros[m].value;
                                    beta  += nonzeros[m].value * ((double)data->raw[index] - sinogram_estimate[index]);
                                }
                            }
                        }

                        ip.alpha = alpha;
                        ip.beta  = beta;

                        // Apply selected penalty functions
                        double pixel_update=0.0;
                        /* Quadratic */
                        if (rp->penalty.compare("quadratic")==0){
                            pixel_update=quadratic(q,&ip,reconstructed_image);
                        }
                        /* Edge Preserving*/
                        else if(rp->penalty.compare("edge-preserving")==0){
                            pixel_update=edge_preserving(q,&ip,reconstructed_image);
                        }
                        else{
                            std::cout << "Unrecognized penalty selected. Exiting." << std::endl;
                            exit(1);
                        }

                        //Enforce positivity
                        if (pixel_update+reconstructed_image[q]<0)
                            pixel_update = -reconstructed_image[q];

                        //Update image
                        reconstructed_image[q] += pixel_update;

                        //Update the forward-projection data
#pragma omp parallel num_threads(OMP_N_THREADS)
                        {
#pragma omp for
                            for (int m = 0; m<num_nonzeros; m++){
                                int index = nonzeros[m].index + offset;

                                if ((index > -1) && (index < data_size))
                                    sinogram_estimate[index] += pixel_update*nonzeros[m].value;
                            }
                        }

                    }
                }
                delete[] nonzeros;
            }
        }

        destroy_spinner();

        // Close up our timer
        std::chrono::high_resolution_clock::time_point end=std::chrono::high_resolution_clock::now();
        auto duration=std::chrono::duration_cast<std::chrono::seconds>(end-start).count();
        std::cout << duration << " s" << std::endl;

        // Write reconstruction to disk
        std::ostringstream recon_path;
        recon_path << rp->output_dir << "/reconstructions/iteration" << n+1 << ".rcn";
        std::ofstream recon_file(recon_path.str(), std::ios_base::binary);
        recon_file.write((char*)&reconstructed_image[0], rp->num_voxels_x*rp->num_voxels_y*rp->num_voxels_z*sizeof(reconstructed_image[0]));
        recon_file.close();

        // "Rewind" the matrix file for the next iteration
        file.clear();
        file.seekg(0, std::ios_base::beg);
    }

    // Copy the final reconstructed volume back into our data structure
    for (int i=0; i<rp->num_voxels_x; i++){
        for (int j=0; j<rp->num_voxels_y; j++){
            for (int k=0; k<rp->num_voxels_z; k++){
                size_t idx=i+j*rp->num_voxels_x+k*rp->num_voxels_x*rp->num_voxels_y;
                data->recon_volume[idx]=(float)reconstructed_image[idx];
            }
        }
    }

}


// Boost.Units - A C++ library for zero-overhead dimensional analysis and
// unit/quantity manipulation and conversion
//
// Copyright (C) 2003-2008 Matthias Christian Schabel
// Copyright (C) 2008 Steven Watanabe
//
// 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_UNITS_CODATA_NEUTRON_CONSTANTS_HPP
#define BOOST_UNITS_CODATA_NEUTRON_CONSTANTS_HPP

#include <boost/units/static_constant.hpp>

#include <boost/units/systems/detail/constants.hpp>
#include <boost/units/systems/si/amount.hpp>
#include <boost/units/systems/si/area.hpp>
#include <boost/units/systems/si/electric_charge.hpp>
#include <boost/units/systems/si/energy.hpp>
#include <boost/units/systems/si/frequency.hpp>
#include <boost/units/systems/si/length.hpp>
#include <boost/units/systems/si/mass.hpp>
#include <boost/units/systems/si/magnetic_flux_density.hpp>
#include <boost/units/systems/si/time.hpp>
#include <boost/units/systems/si/wavenumber.hpp>

#include <boost/units/systems/si/codata/typedefs.hpp>

/// \file
/// CODATA recommended values of fundamental atomic and nuclear constants
/// CODATA 2006 values as of 2007/03/30

namespace boost {

namespace units {

namespace si {

namespace constants {

namespace codata {

/// CODATA recommended values of the fundamental physical constants: NIST SP 961

/// neutron mass
BOOST_UNITS_PHYSICAL_CONSTANT(m_n,quantity<mass>,1.674927211e-27*kilograms,8.4e-35*kilograms);
/// neutron-electron mass ratio
BOOST_UNITS_PHYSICAL_CONSTANT(m_n_over_m_e,quantity<dimensionless>,1838.6836605*dimensionless(),1.1e-6*dimensionless());
/// neutron-muon mass ratio
BOOST_UNITS_PHYSICAL_CONSTANT(m_n_over_m_mu,quantity<dimensionless>,8.89248409*dimensionless(),2.3e-7*dimensionless());
/// neutron-tau mass ratio
BOOST_UNITS_PHYSICAL_CONSTANT(m_n_over_m_tau,quantity<dimensionless>,0.528740*dimensionless(),8.6e-5*dimensionless());
/// neutron-proton mass ratio
BOOST_UNITS_PHYSICAL_CONSTANT(m_n_over_m_p,quantity<dimensionless>,1.00137841918*dimensionless(),4.6e-10*dimensionless());
/// neutron molar mass
BOOST_UNITS_PHYSICAL_CONSTANT(M_n,quantity<mass_over_amount>,1.00866491597e-3*kilograms/mole,4.3e-13*kilograms/mole);
/// neutron Compton wavelength
BOOST_UNITS_PHYSICAL_CONSTANT(lambda_C_n,quantity<length>,1.3195908951e-15*meters,2.0e-24*meters);
/// neutron magnetic moment
BOOST_UNITS_PHYSICAL_CONSTANT(mu_n,quantity<energy_over_magnetic_flux_density>,-0.96623641e-26*joules/tesla,2.3e-33*joules/tesla);
/// neutron g-factor
BOOST_UNITS_PHYSICAL_CONSTANT(g_n,quantity<dimensionless>,-3.82608545*dimensionless(),9.0e-7*dimensionless());
/// neutron-electron magnetic moment ratio
BOOST_UNITS_PHYSICAL_CONSTANT(mu_n_over_mu_e,quantity<dimensionless>,1.04066882e-3*dimensionless(),2.5e-10*dimensionless());
/// neutron-proton magnetic moment ratio
BOOST_UNITS_PHYSICAL_CONSTANT(mu_n_over_mu_p,quantity<dimensionless>,-0.68497934*dimensionless(),1.6e-7*dimensionless());
/// neutron-shielded proton magnetic moment ratio
BOOST_UNITS_PHYSICAL_CONSTANT(mu_n_over_mu_p_prime,quantity<dimensionless>,-0.68499694*dimensionless(),1.6e-7*dimensionless());
/// neutron gyromagnetic ratio
BOOST_UNITS_PHYSICAL_CONSTANT(gamma_n,quantity<frequency_over_magnetic_flux_density>,1.83247185e8/second/tesla,4.3e1/second/tesla);

} // namespace codata

} // namespace constants

} // namespace si

} // namespace units

} // namespace boost

#endif // BOOST_UNITS_CODATA_NEUTRON_CONSTANTS_HPP