// 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.

#include <pybind11/pybind11.h>
#include <pybind11/numpy.h>
#include <pybind11/complex.h>
#include <pybind11/stl.h>
#include <pybind11/pytypes.h>
#include <vector>
#include <complex>
#include <iostream>
#if defined(_OPENMP)
#include <omp.h>
#endif
#include "_cppkernels/simulator.hpp"

namespace py = pybind11;

using c_type = std::complex<double>;
using ArrayType = std::vector<c_type, aligned_allocator<c_type,64>>;
using MatrixType = std::vector<ArrayType>;
using QuRegs = std::vector<std::vector<unsigned>>;

template <class QR>
void emulate_math_wrapper(Simulator &sim,
			  py::function const& pyfunc,
			  QR const& qr,
			  std::vector<unsigned> const& ctrls){
     auto f = [&](std::vector<int>& x) {
		   pybind11::gil_scoped_acquire acquire;
		   x = pyfunc(x).cast<std::vector<int>>();
	      };
     pybind11::gil_scoped_release release;
     sim.emulate_math(f, qr, ctrls);
}

PYBIND11_MODULE(_cppsim, m) {
     m.doc() = "C++ simulator backend for ProjectQ";
     
     py::class_<Simulator>(m, "Simulator")
	  .def(py::init<unsigned>())
	  .def("allocate_qubit", &Simulator::allocate_qubit)
	  .def("deallocate_qubit", &Simulator::deallocate_qubit)
	  .def("get_classical_value", &Simulator::get_classical_value)
	  .def("is_classical", &Simulator::is_classical)
	  .def("measure_qubits", &Simulator::measure_qubits_return)
	  .def("apply_controlled_gate", &Simulator::apply_controlled_gate<MatrixType>)
	  .def("emulate_math", &emulate_math_wrapper<QuRegs>)
	  .def("emulate_math_addConstant", &Simulator::emulate_math_addConstant<QuRegs>)
	  .def("emulate_math_addConstantModN", &Simulator::emulate_math_addConstantModN<QuRegs>)
	  .def("emulate_math_multiplyByConstantModN", &Simulator::emulate_math_multiplyByConstantModN<QuRegs>)
	  .def("get_expectation_value", &Simulator::get_expectation_value)
	  .def("apply_qubit_operator", &Simulator::apply_qubit_operator)
	  .def("emulate_time_evolution", &Simulator::emulate_time_evolution)
	  .def("get_probability", &Simulator::get_probability)
	  .def("get_amplitude", &Simulator::get_amplitude)
	  .def("set_wavefunction", &Simulator::set_wavefunction)
	  .def("collapse_wavefunction", &Simulator::collapse_wavefunction)
	  .def("run", &Simulator::run)
	  .def("cheat", &Simulator::cheat)
	  ;
}
