// Copyright 2018 The Simons Foundation, Inc. - All Rights Reserved.
//
// 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 NETKET_PY_METROPOLIS_HASTINGS_HPP
#define NETKET_PY_METROPOLIS_HASTINGS_HPP

#include <pybind11/pybind11.h>
#include "exchange_kernel.hpp"
#include "metropolis_hastings.hpp"

namespace py = pybind11;

namespace netket {

void AddMetropolisHastings(py::module& subm) {
  auto mh =
      py::class_<MetropolisHastings, AbstractSampler>(subm,
                                                      "MetropolisHastings",
                                                      R"EOF(

            ``MetropolisHastings`` is a generic Metropolis-Hastings sampler using
            a local transition kernel to perform moves in the Markov Chain.
            The transition kernel is used to generate
            a proposed state $$ s^\prime $$, starting from the current state $$ s $$.
            The move is accepted with probability

            $$
            A(s\rightarrow s^\prime) = \mathrm{min}\left (1,\frac{P(s^\prime)}{P(s)} F(e^{L(s,s^\prime)})\right),
            $$

            where the probability being sampled is $$ F(\Psi(s)) $$ (by default $$ F(x)=|x|^2 $$)
            and $L(s,s^\prime)$ is a correcting factor computed by the transition kernel.

                 )EOF")
          .def(py::init([](AbstractMachine& machine,
                           MetropolisHastings::TransitionKernel tk,
                           Index n_chains, nonstd::optional<Index> sweep_size,
                           nonstd::optional<Index> batch_size) {
                 return MetropolisHastings(
                     machine, tk, n_chains,
                     sweep_size.value_or(machine.Nvisible()),
                     batch_size.value_or(n_chains));
               }),
               py::keep_alive<1, 2>(), py::arg("machine"),
               py::arg("transition_kernel"), py::arg("n_chains") = 16,
               py::arg("sweep_size") = py::none(),
               py::arg("batch_size") = py::none(),
               R"EOF(
             Constructs a new ``MetropolisHastings`` sampler given a machine and
             a transition kernel.

             Args:
                 machine: A machine $$\Psi(s)$$ used for the sampling.
                          The probability distribution being sampled
                          from is $$F(\Psi(s))$$, where the function
                          $$F(X)$$, is arbitrary, by default $$F(X)=|X|^2$$.
                 transition_kernel: A function to generate a transition.
                          This should take as an input the current state (in batches)
                          and return a modified state (also in batches).
                          This function must also return an array containing the
                          `log_prob_corrections` $$L(s,s^\prime)$$.
                 n_chains: The number of Markov Chain to be run in parallel on a single process.
                 sweep_size: The number of exchanges that compose a single sweep.
                             If None, sweep_size is equal to the number of degrees of freedom (n_visible).
                 batch_size: The batch size to be used when calling log_val on the given Machine.
                             If None, batch_size is equal to the number Markov chains (n_chains).

             Examples:
                 Sampling from a RBM machine in a 1D lattice of spin 1/2, using
                 nearest-neighbours exchanges with a custom kernel.

                 ```python
                 import netket as nk
                 import numpy as np

                 # 1D Lattice
                 g = nk.graph.Hypercube(length=20, n_dim=1, pbc=True)

                 # Hilbert space of spins on the graph
                 # with total Sz equal to 0
                 hi = nk.hilbert.Spin(s=0.5, graph=g, total_sz=0)

                 # Heisenberg hamiltonian
                 ha = nk.operator.Heisenberg(hilbert=hi)

                 # Symmetric RBM Spin Machine
                 ma = nk.machine.RbmSpin(alpha=1, hilbert=hi)
                 ma.init_random_parameters(seed=1234, sigma=0.01)

                 # Defining a custom kernel for MetropolisHastings
                 # Notice that this sampler exchanges two random sites
                 # thus preserving the total magnetization
                 # Also notice that it is not recommended to define custom kernels in python
                 # For speed reasons it is better to define exchange kernels using CustomSampler
                 def exchange_kernel(v, vnew, logprobcorr):

                     vnew[:, :] = v[:, :]
                     logprobcorr[:] = 0.0

                     rands = np.random.randint(v.shape[1], size=(v.shape[0], 2))

                     for i in range(v.shape[0]):
                         iss = rands[i, 0]
                         jss = rands[i, 1]

                         vnew[i, iss], vnew[i, jss] = vnew[i, jss], vnew[i, iss]


                 sa = nk.sampler.MetropolisHastings(ma, exchange_kernel, n_chains=16, sweep_size=20)


                 ```
             )EOF");

  AddAcceptance(mh);

  subm.def(
      "mh_acceptance_kernel",
      [](Eigen::Ref<RowMatrix<double>> state,
         Eigen::Ref<RowMatrix<double>> state1, Eigen::Ref<VectorXcd> log_val,
         Eigen::Ref<VectorXcd> log_val1, Eigen::Ref<VectorXd> log_prob_corr,
         Eigen::Ref<VectorXd> rand_for_acceptance,
         double machine_pow) -> Index {
        Index accepted = 0;

        for (Index i = 0; i < state.rows(); i++) {
          double prob =
              std::exp(machine_pow *
                       (log_val1(i) - log_val(i) + log_prob_corr(i)).real());

          bool accept = prob > rand_for_acceptance(i);

          if (accept) {
            log_val(i) = log_val1(i);
            state.row(i) = state1.row(i);
            accepted += 1;
          }
        }
        return accepted;
      });

  auto mh_pt =
      py::class_<MetropolisHastingsPt, AbstractSampler>(subm,
                                                        "MetropolisHastingsPt",
                                                        R"EOF(
            This sampler performs parallel-tempering
            moves in addition to the moves implemented in `MetropolisHastings`.
            The number of replicas can be $$ N_{\mathrm{rep}} $$ chosen by the user.
            )EOF")
          .def(py::init([](AbstractMachine& machine,
                           MetropolisHastings::TransitionKernel tk,
                           Index n_replicas,
                           nonstd::optional<Index> sweep_size) {
                 return MetropolisHastingsPt(
                     machine, tk, n_replicas,
                     sweep_size.value_or(machine.Nvisible()));
               }),
               py::keep_alive<1, 2>(), py::arg("machine"),
               py::arg("transition_kernel"), py::arg("n_replicas") = 16,
               py::arg("sweep_size") = py::none(),
               R"EOF(
            Constructs a new ``MetropolisHastingsPt`` sampler.

            Args:
                machine: A machine $$\Psi(s)$$ used for the sampling.
                         The probability distribution being sampled
                         from is $$F(\Psi(s))$$, where the function
                         $$F(X)$$, is arbitrary, by default $$F(X)=|X|^2$$.
                transition_kernel: A function to generate a transition.
                         This should take as an input the current state (in batches)
                         and return a modified state (also in batches).
                         This function must also return an array containing the
                         `log_prob_corrections` $$L(s,s^\prime)$$.
                n_replicas: The number of replicas used for parallel tempering.
                sweep_size: The number of exchanges that compose a single sweep.
                            If None, sweep_size is equal to the number of degrees
                            of freedom (n_visible).

            )EOF");

  AddAcceptance(mh_pt);
  AddSamplerStats(mh_pt);
}
}  // namespace netket
#endif
