//==============================================================================
//         Copyright 2003 - 2011 LASMEA UMR 6602 CNRS/Univ. Clermont II
//         Copyright 2009 - 2011 LRI    UMR 8623 CNRS/Univ Paris Sud XI
//
//          Distributed under the Boost Software License, Version 1.0.
//                 See accompanying file LICENSE.txt or copy at
//                     http://www.boost.org/LICENSE_1_0.txt
//==============================================================================
#ifndef BOOST_SIMD_ARITHMETIC_FUNCTIONS_SIMD_SSE_AVX_ADDS_HPP_INCLUDED
#define BOOST_SIMD_ARITHMETIC_FUNCTIONS_SIMD_SSE_AVX_ADDS_HPP_INCLUDED
#ifdef BOOST_SIMD_HAS_AVX_SUPPORT
#include <boost/simd/arithmetic/functions/adds.hpp>
#include <boost/dispatch/meta/scalar_of.hpp>
#include <boost/dispatch/meta/downgrade.hpp>
#include <boost/simd/include/functions/simd/bitwise_and.hpp>
#include <boost/simd/include/functions/simd/minus.hpp>
#include <boost/simd/include/constants/digits.hpp>

namespace boost { namespace simd { namespace ext
{
  BOOST_SIMD_FUNCTOR_IMPLEMENTATION( boost::simd::tag::adds_, boost::simd::tag::avx_
                            , (A0)
                            , ((simd_<int_<A0>,boost::simd::tag::avx_>))
                              ((simd_<int_<A0>,boost::simd::tag::avx_>))
                            )
  {
    typedef A0 result_type;
    BOOST_SIMD_FUNCTOR_CALL_REPEAT(2)
    {
      typedef typename meta::scalar_of<A0>::type             sctype;
      typedef simd::native<sctype, boost::simd::tag::sse_ >  svtype;
      svtype a00 = _mm256_extractf128_si256(a0, 0);
      svtype a10 = _mm256_extractf128_si256(a1, 0);
      svtype  r0 = adds(a00,a10);
      result_type that  = _mm256_castsi128_si256(r0);
      svtype a01 = _mm256_extractf128_si256(a0, 1);
      svtype a11 = _mm256_extractf128_si256(a1, 1);
      svtype r1 = adds(a01,a11);
      that = _mm256_insertf128_si256(that, r1, 1);
      return that;
    }
  };
 BOOST_SIMD_FUNCTOR_IMPLEMENTATION( boost::simd::tag::adds_, boost::simd::tag::avx_
                            , (A0)
                            , ((simd_<uint_<A0>,boost::simd::tag::avx_>))
                              ((simd_<uint_<A0>,boost::simd::tag::avx_>))
                            )
  {
    typedef A0 result_type;
    BOOST_SIMD_FUNCTOR_CALL_REPEAT(2)
    {
      typedef typename meta::scalar_of<A0>::type             sctype;
      typedef simd::native<sctype, boost::simd::tag::sse_ >  svtype;
      svtype a00 = _mm256_extractf128_si256(a0, 0);
      svtype a10 = _mm256_extractf128_si256(a1, 0);
      svtype  r0 = adds(a00,a10);
      result_type that  = _mm256_castsi128_si256(r0);
      svtype a01 = _mm256_extractf128_si256(a0, 1);
      svtype a11 = _mm256_extractf128_si256(a1, 1);
      svtype r1 = adds(a01,a11);
      that = _mm256_insertf128_si256(that, r1, 1);
      return that;
    }
  };
} } }

#endif
#endif


#ifndef UI_SDL_GL_HPP_
#define UI_SDL_GL_HPP_

#include <cstdint>
#include <fstream>
#include <list>
#include <math.h>
#include <sstream>
#include <stdexcept>
#include <string>
#include <utility>
#include <vector>
#include "glm/gtc/matrix_transform.hpp"
#include <GL/glew.h>
#include <SDL2/SDL.h>
#include <boost/algorithm/string.hpp>
#include <boost/optional.hpp>
#include <glm/gtc/type_ptr.hpp>
#include "keyboard.hpp"
#include "misc.hpp"
#include "world.hpp"

class ui_sdl_gl : public ui {
public:
  static const int WINDOW_WIDTH  = 1280;
  static const int WINDOW_HEIGHT = 800;

  ui_sdl_gl(std::string title)
    : geometry_mesh_size_(0),
      angles_last_read_ms_(misc::get_time_ms())
  {
    //
    // init SDL
    //

    if (SDL_Init(SDL_INIT_VIDEO) != 0)
      throw std::runtime_error(std::string("error, SDL_Init(): ") + SDL_GetError());

    window_ = SDL_CreateWindow(
        title.c_str(),
        SDL_WINDOWPOS_CENTERED,
        SDL_WINDOWPOS_CENTERED,
        WINDOW_WIDTH,
        WINDOW_HEIGHT,
        SDL_WINDOW_OPENGL | SDL_WINDOW_SHOWN);

    if (!window_)
      throw std::runtime_error(std::string("error, SDL_CreateWindow(): ") + SDL_GetError());

    SDL_GL_SetAttribute(SDL_GL_CONTEXT_MAJOR_VERSION, 3);
    SDL_GL_SetAttribute(SDL_GL_CONTEXT_MINOR_VERSION, 3);
    SDL_GL_SetAttribute(SDL_GL_DOUBLEBUFFER, 1);

    context_ = SDL_GL_CreateContext(window_);

    if (!context_)
      throw std::runtime_error(std::string("error, SDL_GL_CreateContext(): ") + SDL_GetError());

    //
    // init glew
    //

    glewExperimental = true;

    if (glewInit() != GLEW_OK)
      throw std::runtime_error(std::string("error, glewInit() failed"));

    //
    // init opengl and shader program
    //

    program_id_ = create_shader_program("vertex_shader.glsl", "fragment_shader.glsl");

    if (!program_id_) {
	//printf("error when building shader program\n");
 	throw std::runtime_error(std::string("error when building shader program"));
    }

    glUseProgram(program_id_);
    glEnable(GL_DEPTH_TEST);
    glEnable(GL_BLEND);
    glBlendFunc(GL_SRC_ALPHA,GL_ONE_MINUS_SRC_ALPHA);
    glCullFace(GL_BACK);
    glClearColor(0.1, 0.1, 0.1, 0);

    uni_model_id_ = glGetUniformLocation(program_id_, "model");
    uni_view_id_ = glGetUniformLocation(program_id_, "view");
    uni_projection_id_ = glGetUniformLocation(program_id_, "projection");
    uni_color_id_ = glGetUniformLocation(program_id_, "color");

    projection_matrix_ =
        glm::perspective(70.0f, (float)WINDOW_WIDTH / (float)WINDOW_HEIGHT, 0.1f, 100.0f);

    glUniformMatrix4fv(uni_projection_id_, 1, GL_FALSE, &projection_matrix_[0][0]);

    //
    // get geometry data
    //

    std::vector<GLfloat> model_vertices, model_normals;

    if(!get_geometry_data_from_file("mask.obj", model_vertices, model_normals))
      throw std::runtime_error(std::string("error when loading model geometry"));

    geometry_mesh_size_ = model_vertices.size();

    //
    // upload geometry
    //

    // create vertex array object
    GLuint vao;
    glGenVertexArrays(1, &vao);
    glBindVertexArray(vao);

    // create vertex buffer objects
    GLuint vbo[2];
    glGenBuffers(2, vbo);

    // put vertices in first buffer
    glBindBuffer(GL_ARRAY_BUFFER, vbo[0]);
    glBufferData(GL_ARRAY_BUFFER, model_vertices.size() * sizeof(GLfloat), &model_vertices[0],
        GL_STATIC_DRAW);
    glVertexAttribPointer(0, 3, GL_FLOAT, GL_FALSE, 0, 0);
    glEnableVertexAttribArray(0);

    // put normals in second buffer
    glBindBuffer(GL_ARRAY_BUFFER, vbo[1]);
    glBufferData(GL_ARRAY_BUFFER, model_normals.size() * sizeof(GLfloat), &model_normals[0],
        GL_STATIC_DRAW);
    glVertexAttribPointer(1, 3, GL_FLOAT, GL_FALSE, 0, 0);
    glEnableVertexAttribArray(1);
  }

  ~ui_sdl_gl() {
    SDL_DestroyWindow(window_);
    SDL_GL_DeleteContext(context_);
    SDL_Quit();
  }

  virtual void poll_events() {
    events_.resize(64);

    int events_returned = SDL_PeepEvents(
        &events_.front(),
        events_.capacity(),
        SDL_GETEVENT,
        SDL_FIRSTEVENT,
        SDL_LASTEVENT);

    if (events_returned == -1) {
      INFO("error, SDL_PeepEvents(): " << SDL_GetError());
      events_returned = 0;
    }

    events_.resize(events_returned);

    // show/hide cursor
    for (auto& e : events_) {
      if (e.type != SDL_WINDOWEVENT)
        continue;

      if (e.window.event == SDL_WINDOWEVENT_FOCUS_LOST)
        SDL_ShowCursor(SDL_TRUE);

      if (e.window.event == SDL_WINDOWEVENT_FOCUS_GAINED)
        SDL_ShowCursor(SDL_FALSE);
    }
  }

  virtual bool check_event_quit() {
    for (auto& e : events_)
      if (e.type == SDL_QUIT)
        return true;

    return false;
  }

  virtual bool check_event_button_released(keyboard::button button) {
    for (auto& e : events_) {
      if (e.type != SDL_KEYUP)
        continue;

      int code = e.key.keysym.scancode;

      for (int i = 1; i != keyboard::button::end; i <<= 1) {
        std::vector<int> codes = get_sdl_scancodes(i);

        for (auto &c : codes)
          if (code == c && i == button)
            return true;
      }
    }

    return false;
  }

  virtual int get_pressed_buttons() {
    SDL_PumpEvents();
    const Uint8* state = SDL_GetKeyboardState(nullptr);

    int pressed = 0;

    for (int i = 1; i != keyboard::button::end; i <<= 1) {
      std::vector<int> codes = get_sdl_scancodes(i);

      for (auto& c : codes)
        if (state[c])
          pressed |= i;
    }

    return pressed;
  }

  virtual void get_delta_angles(float& horizontal, float& vertical) {
    // if window is not in focus
    if (!(SDL_GetWindowFlags(window_) & SDL_WINDOW_MOUSE_FOCUS) ||
        !(SDL_GetWindowFlags(window_) & SDL_WINDOW_INPUT_FOCUS)) {
      horizontal = vertical = 0.0;
      angles_last_read_ms_ = misc::get_time_ms();

      return;
    }

    int x, y;
    SDL_GetMouseState(&x, &y);

    float delta_time = (float) (misc::get_time_ms() - angles_last_read_ms_) / 1000;
    horizontal = 0.05 * delta_time * float(WINDOW_WIDTH / 2 - x);
    vertical = 0.05 * delta_time * float(WINDOW_HEIGHT / 2 - y);

    SDL_WarpMouseInWindow(window_, WINDOW_WIDTH / 2, WINDOW_HEIGHT / 2);
    angles_last_read_ms_ = misc::get_time_ms();
  }

  virtual void draw_clear() {
    glClear(GL_COLOR_BUFFER_BIT | GL_DEPTH_BUFFER_BIT);
  }

  virtual void draw_world(world& world, uint8_t perspective_player_id, bool draw_phantoms) {
    // get player to render scene for
    boost::optional<player&> cam_player = world.get_player(perspective_player_id);

    glm::vec3 cam_position = glm::vec3(
        cam_player.get().get_x(),
        cam_player.get().get_y(),
        cam_player.get().get_z());

    glm::vec3 cam_direction(
        cos(cam_player.get().get_vert_angel()) * cos(cam_player.get().get_horz_angel()),
        sin(cam_player.get().get_vert_angel()),
        -(cos(cam_player.get().get_vert_angel()) * sin(cam_player.get().get_horz_angel())));

    // set view matrix
    glm::mat4 view = glm::lookAt(cam_position, cam_position + cam_direction, glm::vec3(0, 1, 0));

    glUniformMatrix4fv(uni_view_id_, 1, GL_FALSE, &view[0][0]);

    // draw rest of players
    for (player& p : world.get_players()) {
      if (p.get_id() == perspective_player_id)
        continue;

      // set player color
      glm::vec4 color = glm::vec4(
          (p.get_color_AABBGGRR() & 0xff) / 255.f,
          ((p.get_color_AABBGGRR() >> 8) & 0xff) / 255.f,
          ((p.get_color_AABBGGRR() >> 16) & 0xff) / 255.f,
          (draw_phantoms) ? 0.2 : 1.0);

      glUniform4fv(uni_color_id_, 1, &color[0]);

      // set model matrix
      glm::mat4 model =
        glm::translate(glm::mat4(1.0), glm::vec3(p.get_x(), p.get_y(), p.get_z())) *
        glm::rotate(glm::mat4(1.0), p.get_horz_angel(), glm::vec3(0, 1, 0)) *
        glm::rotate(glm::mat4(1.0), p.get_vert_angel(), glm::vec3(0, 0, 1));

      glUniformMatrix4fv(uni_model_id_, 1, GL_FALSE, &model[0][0]);

      // draw player
      glDrawArrays(GL_TRIANGLES, 0, geometry_mesh_size_ / 2);
    }
  }

  virtual void draw_update() {
    SDL_GL_SwapWindow(window_);
  }

private:
  std::vector<int> get_sdl_scancodes(int button) {
    switch (button) {
      case keyboard::button::up:         return std::vector<int> { SDL_SCANCODE_SPACE };
      case keyboard::button::down:       return std::vector<int> { SDL_SCANCODE_LCTRL };
      case keyboard::button::forward:    return std::vector<int> { SDL_SCANCODE_UP,
                                                                   SDL_SCANCODE_W };
      case keyboard::button::backward:   return std::vector<int> { SDL_SCANCODE_DOWN,
                                                                   SDL_SCANCODE_S };
      case keyboard::button::left:       return std::vector<int> { SDL_SCANCODE_LEFT };
      case keyboard::button::right:      return std::vector<int> { SDL_SCANCODE_RIGHT };
      case keyboard::button::step_left:  return std::vector<int> { SDL_SCANCODE_A };
      case keyboard::button::step_right: return std::vector<int> { SDL_SCANCODE_D };
      case keyboard::button::quit:       return std::vector<int> { SDL_SCANCODE_ESCAPE };
      case keyboard::button::f1:         return std::vector<int> { SDL_SCANCODE_F1 };
      case keyboard::button::f2:         return std::vector<int> { SDL_SCANCODE_F2 };
    }

    return std::vector<int>();
  }

  bool get_geometry_data_from_file(const char* file, std::vector<GLfloat>& vertices,
      std::vector<GLfloat>& normals) {
    std::ifstream file_stream(file);

    if (!file_stream.is_open()) {
      INFO("error, could not open file: " << file);
      return false;
    }

    std::vector<std::string> vertex_lines, normal_lines;

    while (file_stream.good()) {
      std::string line;
      getline(file_stream, line);

      // read all strings holding vertices and normals
      if (line[0] == 'v') {
        if (line[1] == 'n')
          normal_lines.push_back(line);
        else
          vertex_lines.push_back(line);

        continue;
      }

      // read strings holding faces and put there values into argument vectors
      if (line[0] == 'f') {
        std::vector<std::string> edges;
        boost::split(edges, line, boost::is_any_of(" ")); // split "f 1//1 2//2 4//3"

        for (size_t i = 1; i < edges.size(); i++) {
          std::vector<std::string> indices;
          boost::split(indices, edges[i], boost::is_any_of("/")); // split "1//1"

          GLfloat temp;

          std::istringstream vertex_stream(vertex_lines[stoi(indices[0]) - 1]);
          vertex_stream.ignore(32, ' '); // skip "v"
          while (vertex_stream >> temp)
            vertices.push_back(temp);

          std::istringstream normal_stream(normal_lines[stoi(indices[2]) - 1]);
          normal_stream.ignore(32, ' '); // skip "vn"
          while (normal_stream >> temp)
            normals.push_back(temp);
        }
      }
    }

    file_stream.close();

    return true;
  }

  GLuint create_shader_program(const char* vertex_file_path, const char* fragment_file_path) {
    std::string vertex_shader_code = misc::get_file_content(vertex_file_path);
    if (vertex_shader_code.empty()) {
      INFO("error, got no data from: " << vertex_file_path);
      return 0;
    }

    std::string fragment_shader_code = misc::get_file_content(fragment_file_path);
    if (fragment_shader_code.empty()) {
      INFO("error, got no data from: " << fragment_file_path);
      return 0;
    }

    GLuint vertex_shader_id = glCreateShader(GL_VERTEX_SHADER);
    GLuint fragment_shader_id = glCreateShader(GL_FRAGMENT_SHADER);

    GLint result = GL_FALSE;
    int info_log_length;

    // compile vertex shader
    char const* vertex_source_ptr = vertex_shader_code.c_str();
    glShaderSource(vertex_shader_id, 1, &vertex_source_ptr , NULL);
    glCompileShader(vertex_shader_id);

    // check vertex shader
    glGetShaderiv(vertex_shader_id, GL_COMPILE_STATUS, &result);
    if (result != GL_TRUE) {
      glGetShaderiv(vertex_shader_id, GL_INFO_LOG_LENGTH, &info_log_length);
      std::vector<char> msg(info_log_length);
      glGetShaderInfoLog(vertex_shader_id, info_log_length, NULL, &msg[0]);
      INFO("error, compiling vertex shader: " << &msg[0]);
      return 0;
    }

    // compile fragment shader
    char const* fragment_source_ptr = fragment_shader_code.c_str();
    glShaderSource(fragment_shader_id, 1, &fragment_source_ptr , NULL);
    glCompileShader(fragment_shader_id);

    // check fragment shader
    glGetShaderiv(fragment_shader_id, GL_COMPILE_STATUS, &result);
    if (result != GL_TRUE) {
      glGetShaderiv(fragment_shader_id, GL_INFO_LOG_LENGTH, &info_log_length);
      std::vector<char> msg(info_log_length);
      glGetShaderInfoLog(fragment_shader_id, info_log_length, NULL, &msg[0]);
      INFO("error, compiling fragment shader: " << &msg[0]);
      glDeleteShader(vertex_shader_id);
      return 0;
    }

    // link program
    GLuint program_id = glCreateProgram();
    glAttachShader(program_id, vertex_shader_id);
    glAttachShader(program_id, fragment_shader_id);

    glLinkProgram(program_id);

    glDeleteShader(vertex_shader_id);
    glDeleteShader(fragment_shader_id);

    // check program
    glGetProgramiv(program_id, GL_LINK_STATUS, &result);
    if (result != GL_TRUE) {
      glGetProgramiv(program_id, GL_INFO_LOG_LENGTH, &info_log_length);
      std::vector<char> msg(std::max(info_log_length, int(1)));
      glGetProgramInfoLog(program_id, info_log_length, NULL, &msg[0]);
      INFO("error, linking shaders: " << &msg[0]);
      return 0;
    }

    return program_id;
  }

  // sdl
  SDL_Window* window_;
  SDL_GLContext context_;
  std::vector<SDL_Event> events_;

  // gl
  GLuint program_id_;
  GLuint uni_model_id_;
  GLuint uni_view_id_;
  GLuint uni_projection_id_;
  GLuint uni_color_id_;

  // other
  size_t geometry_mesh_size_;
  uint64_t angles_last_read_ms_;
  glm::mat4 projection_matrix_;
};

#endif // UI_SDL_GL_HPP_


#include <cctbx/boost_python/flex_fwd.h>
#include <boost/python.hpp>
#include <simtbx/diffBragg/src/diffBragg.h>
#include <simtbx/nanoBragg/nanoBragg.h>
#include <iostream>

using namespace boost::python;
namespace simtbx{
namespace diffBragg{
namespace boost_python { namespace {

  void (simtbx::nanoBragg::diffBragg::*add_diffBragg_spots_A)() = &simtbx::nanoBragg::diffBragg::add_diffBragg_spots;
  void (simtbx::nanoBragg::diffBragg::*add_diffBragg_spots_B)(const nanoBragg::af::shared<size_t>&) = &simtbx::nanoBragg::diffBragg::add_diffBragg_spots;
  void (simtbx::nanoBragg::diffBragg::*add_diffBragg_spots_C)(const nanoBragg::af::shared<size_t>&, boost::python::list per_pix_nominal_hkl)
        = &simtbx::nanoBragg::diffBragg::add_diffBragg_spots;

  static void  set_diffuse_gamma(simtbx::nanoBragg::diffBragg& diffBragg, boost::python::tuple const& values) {
      double g0 = boost::python::extract<double>(values[0]);
      double g1 = boost::python::extract<double>(values[1]);
      double g2 = boost::python::extract<double>(values[2]);
      diffBragg.db_cryst.anisoG <<  g0,0,0,
                                    0,g1,0,
                                    0,0,g2;
  }

  static boost::python::tuple get_diffuse_gamma(simtbx::nanoBragg::diffBragg const& diffBragg) {
      return boost::python::make_tuple(diffBragg.db_cryst.anisoG(0,0),
                                      diffBragg.db_cryst.anisoG(1,1),
                                      diffBragg.db_cryst.anisoG(2,2));
  }

  static void  set_diffuse_sigma(simtbx::nanoBragg::diffBragg& diffBragg, boost::python::tuple const& values) {
      double sig0 = boost::python::extract<double>(values[0]);
      double sig1 = boost::python::extract<double>(values[1]);
      double sig2 = boost::python::extract<double>(values[2]);
      double sig_array[3] = {sig0,sig1,sig2};

      diffBragg.db_cryst.anisoU <<  sig0*sig0,0,0,
                                    0,sig1*sig1,0,
                                    0,0,sig2*sig2;
      diffBragg.db_cryst.dU_dsigma.clear();
      diffBragg.db_cryst.dU_dsigma.resize(3);

      for (int i=0; i<3; i++){
        diffBragg.db_cryst.dU_dsigma[i] <<  0,0,0,
                                    0,0,0,
                                    0,0,0;
        diffBragg.db_cryst.dU_dsigma[i](i,i) = 2.*sig_array[i];
      }
  }

  static boost::python::tuple get_diffuse_sigma(simtbx::nanoBragg::diffBragg const& diffBragg) {
      double sig0 = sqrt(diffBragg.db_cryst.anisoU(0,0));
      double sig1 = sqrt(diffBragg.db_cryst.anisoU(1,1));
      double sig2 = sqrt(diffBragg.db_cryst.anisoU(2,2));
      return boost::python::make_tuple(sig0, sig1, sig2);
  }

  static void  set_Ndef(simtbx::nanoBragg::diffBragg& diffBragg, boost::python::tuple const& values) {
      diffBragg.set_ncells_def_values(values);
  }

  static boost::python::tuple get_Ndef(simtbx::nanoBragg::diffBragg const& diffBragg) {
      return boost::python::make_tuple(diffBragg.Nd, diffBragg.Ne, diffBragg.Nf);
  }

  static void  set_Nabc_aniso(simtbx::nanoBragg::diffBragg& diffBragg, boost::python::tuple const& values) {
      diffBragg.isotropic_ncells=false;
      diffBragg.set_ncells_values(values);
  }

  static boost::python::tuple get_Nabc_aniso(simtbx::nanoBragg::diffBragg const& diffBragg) {
      return boost::python::make_tuple(diffBragg.Na, diffBragg.Nb, diffBragg.Nc);
  }

  static void  set_Nabc(simtbx::nanoBragg::diffBragg& diffBragg, double const& value) {
      // NOTE this is the old way, assumming isotropic Nabc, for anisotropic use set_Nabc_aniso
      diffBragg.isotropic_ncells=true;
      diffBragg.set_value(9, value); //  9 means Ncells parameter
  }

  static double get_Nabc(simtbx::nanoBragg::diffBragg const& diffBragg) {
      return diffBragg.Na;
  }

  boost::python::tuple get_atom_data(simtbx::nanoBragg::diffBragg &diffBragg){
    boost::python::tuple values;
    values = boost::python::make_tuple(0,0);
    return values;
  }
  //TODO override the set_sources function (or xray_beams) property in nanoBragg in order
  // to set the fpfdp accordingly (if Fhkl2 is set)
  static void set_atom_data(simtbx::nanoBragg::diffBragg & diffBragg,
            boost::python::tuple const& atom_XYZBO){
            /// atom XYZBO is tuple of x coord, ycoords, z coords, B-factors, occupancy
            // e.g. 5 parameters per atom
   boost::python::list atom_x = boost::python::extract<boost::python::list>(atom_XYZBO[0]);
   boost::python::list atom_y = boost::python::extract<boost::python::list>(atom_XYZBO[1]);
   boost::python::list atom_z = boost::python::extract<boost::python::list>(atom_XYZBO[2]);
   boost::python::list atom_B = boost::python::extract<boost::python::list>(atom_XYZBO[3]);
   boost::python::list atom_O = boost::python::extract<boost::python::list>(atom_XYZBO[4]);
    int natoms = boost::python::len(atom_x);
    SCITBX_ASSERT( boost::python::len(atom_y)== natoms);
    SCITBX_ASSERT( boost::python::len(atom_z)== natoms);
    SCITBX_ASSERT( boost::python::len(atom_B)== natoms);
    SCITBX_ASSERT( boost::python::len(atom_O)== natoms);
    diffBragg.db_cryst.atom_data.clear();
    for (int i=0; i < natoms; i++){
        diffBragg.db_cryst.atom_data.push_back( boost::python::extract<double>(atom_x[i]));
        diffBragg.db_cryst.atom_data.push_back(  boost::python::extract<double>(atom_y[i]));
        diffBragg.db_cryst.atom_data.push_back(  boost::python::extract<double>(atom_z[i]));
        diffBragg.db_cryst.atom_data.push_back( boost::python::extract<double>(atom_B[i]));
        diffBragg.db_cryst.atom_data.push_back(  boost::python::extract<double>(atom_O[i]));
    }
  }

  boost::python::tuple get_fpfdp_derivs(simtbx::nanoBragg::diffBragg & diffBragg){
        boost::python::tuple vals;
        vals = boost::python::make_tuple(0,0);
        //TODO implement
        return vals;
  }

  static void set_fpfdp_derivs(simtbx::nanoBragg::diffBragg & diffBragg,
                boost::python::tuple const& derivs) {
    boost::python::list fprime_derivs = boost::python::extract<boost::python::list>(derivs[0]);
    boost::python::list fdblprime_derivs = boost::python::extract<boost::python::list>(derivs[1]);
    int num_fprime = boost::python::len(fprime_derivs);
    int num_fdblprime =   boost::python::len(fdblprime_derivs);
    SCITBX_ASSERT(num_fprime % diffBragg.sources==0);
    SCITBX_ASSERT(num_fdblprime % diffBragg.sources == 0 );
    diffBragg.db_cryst.fpfdp_derivs.clear();
    //diffBragg.fdp_derivs.clear();
    int n_deriv = num_fprime / diffBragg.sources;
    for (int i_deriv=0; i_deriv < n_deriv; i_deriv++ ){
        for (int i=0; i<diffBragg.sources; i++){
            int idx=diffBragg.sources*i_deriv + i;
            double val = boost::python::extract<double>(fprime_derivs[idx]);
            diffBragg.db_cryst.fpfdp_derivs.push_back(val);
            val = boost::python::extract<double>(fdblprime_derivs[idx]);
            diffBragg.db_cryst.fpfdp_derivs.push_back(val);
        }
    }
  }

  static void set_fpfdp(simtbx::nanoBragg::diffBragg & diffBragg,
                boost::python::tuple const& fprime_fdblprime) {
    boost::python::list fprime = boost::python::extract<boost::python::list>(fprime_fdblprime[0]);
    boost::python::list fdblprime = boost::python::extract<boost::python::list>(fprime_fdblprime[1]);
    int num_fprime = boost::python::len(fprime);
    int num_fdblprime =   boost::python::len(fdblprime);
    SCITBX_ASSERT(num_fprime==diffBragg.sources);
    SCITBX_ASSERT(num_fdblprime==diffBragg.sources);
    diffBragg.db_cryst.fpfdp.clear();
    for (int i=0; i<diffBragg.sources; i++){
        double val = boost::python::extract<double>(fprime[i]);
        diffBragg.db_cryst.fpfdp.push_back(val);
        val = boost::python::extract<double>(fdblprime[i]);
        diffBragg.db_cryst.fpfdp.push_back(val);
    }
  }

  boost::python::tuple get_fpfdp(simtbx::nanoBragg::diffBragg const& diffBragg){
    // TODO implement
    boost::python::list fp;
    boost::python::list fdp;
    boost::python::tuple values;
    int n = diffBragg.db_cryst.fpfdp.size()/2;
    for(int i=0; i < n; i++ ){
        double fp_val = diffBragg.db_cryst.fpfdp[2*i];
        double fdp_val = diffBragg.db_cryst.fpfdp[2*i+1];
        fp.append(fp_val);
        fdp.append(fdp_val);
    }
    values = boost::python::make_tuple(fp,fdp);
    return values;
  }

  // change of basis operator see dxtbx.model.crystal.h
  static nanoBragg::mat3 get_O(simtbx::nanoBragg::diffBragg& diffBragg){
    return diffBragg.Omatrix;
  }

  void set_O(simtbx::nanoBragg::diffBragg& diffBragg, nanoBragg::mat3 const& value){
    diffBragg.Omatrix = value;
    diffBragg.db_cryst.eig_O << value[0], value[1], value[2],
            value[3], value[4], value[5],
            value[6], value[7], value[8];
    //std::cout << "eigO: " << diffBragg.eig_O << std::endl;
  }

  static nanoBragg::mat3 get_U(simtbx::nanoBragg::diffBragg& diffBragg){
    return diffBragg.Umatrix;
  }

  void set_U(simtbx::nanoBragg::diffBragg& diffBragg, nanoBragg::mat3 const& value){
    diffBragg.Umatrix = value;
    diffBragg.db_cryst.eig_U << value[0], value[1], value[2],
            value[3], value[4], value[5],
            value[6], value[7], value[8];
  }

  static nanoBragg::mat3 get_B(simtbx::nanoBragg::diffBragg& diffBragg){
    return diffBragg.Bmatrix;
  }

  void set_B(simtbx::nanoBragg::diffBragg& diffBragg, nanoBragg::mat3 const& value){
      diffBragg.Bmatrix = value;
      diffBragg.a_star[1] = value(0,0);  // initialize Angstrom version, update at the end
      diffBragg.a_star[2] = value(1,0);
      diffBragg.a_star[3] = value(2,0);
      diffBragg.b_star[1] = value(0,1);
      diffBragg.b_star[2] = value(1,1);
      diffBragg.b_star[3] = value(2,1);
      diffBragg.c_star[1] = value(0,2);
      diffBragg.c_star[2] = value(1,2);
      diffBragg.c_star[3] = value(2,2);

      diffBragg.user_cell = 0;
      diffBragg.user_matrix = 1;
      diffBragg.init_cell();

       diffBragg.db_cryst.eig_B << diffBragg.a0[1], diffBragg.b0[1], diffBragg.c0[1],
            diffBragg.a0[2], diffBragg.b0[2], diffBragg.c0[2],
            diffBragg.a0[3], diffBragg.b0[3], diffBragg.c0[3];
  }

  //void set_roi(simtbx::nanoBragg::diffBragg& diffBragg, boost::python::tuple const& tupleoftuples){
  //    boost::python::tuple frange;
  //    boost::python::tuple srange;
  //    frange = extract<boost::python::tuple>(tupleoftuples[0]);
  //    srange = extract<boost::python::tuple>(tupleoftuples[1]);
  //    diffBragg.roi_xmin=extract<int>(frange[0]);
  //    diffBragg.roi_xmax=extract<int>(frange[1]);
  //    diffBragg.roi_ymin=extract<int>(srange[0]);
  //    diffBragg.roi_ymax=extract<int>(srange[1]);
  //     /* fix this up so its not out-of-range */
  //    diffBragg.init_beamcenter();
  //    diffBragg.init_raw_pixels_roi();
  //    diffBragg.initialize_managers();
  //}

  /* region of interest in pixels */
  //static boost::python::tuple get_roi(simtbx::nanoBragg::diffBragg const& diffBragg) {
  //    boost::python::tuple frange;
  //    boost::python::tuple srange;
  //    frange = boost::python::make_tuple(diffBragg.roi_xmin,diffBragg.roi_xmax);
  //    srange = boost::python::make_tuple(diffBragg.roi_ymin,diffBragg.roi_ymax);
  //    return boost::python::make_tuple(frange,srange);
  //}

  static boost::python::tuple get_reference_origin(simtbx::nanoBragg::diffBragg & diffBragg){
    return boost::python::make_tuple(diffBragg.O_reference[0]*1000, diffBragg.O_reference[1]*1000, diffBragg.O_reference[2]*1000);
  }

  //simtbx::nanoBragg::af::shared<double> get_deriv_pix(simtbx::nanoBragg::diffBragg& diffBragg, int refine_id){
  //  simtbx::nanoBragg::af::shared<double> pixels = diffBragg.get_derivative_pixels(refine_id);
  //  // slice them
  //  //return pixels[simtbx::nanoBragg::af::slice(0,diffBragg.Npix_to_model)];
  //  return pixels[0:diffBragg.Npix_to_model];
  //}

  void set_reference_origin(simtbx::nanoBragg::diffBragg & diffBragg, boost::python::tuple const& values){
    diffBragg.O_reference[0] = extract<double>(values[0])/1000.;
    diffBragg.O_reference[1] = extract<double>(values[1])/1000.;
    diffBragg.O_reference[2] = extract<double>(values[2])/1000.;
  }

  void set_lambda_coef(simtbx::nanoBragg::diffBragg& diffBragg, boost::python::tuple const& values){
      double coef0 = extract<double>(values[0]);
      double coef1 = extract<double>(values[1]);
      diffBragg.lambda_managers[0]->value = coef0;
      diffBragg.lambda_managers[1]->value = coef1;
      diffBragg.use_lambda_coefficients = true;
  }

  static boost::python::tuple get_lambda_coef(simtbx::nanoBragg::diffBragg const& diffBragg) {
      double coef0=diffBragg.lambda_managers[0]->value;
      double coef1=diffBragg.lambda_managers[1]->value;
      return boost::python::make_tuple(coef0, coef1);
  }

  bool get_use_diffuse(simtbx::nanoBragg::diffBragg& diffBragg){
      return diffBragg.db_flags.use_diffuse;
  }

  void set_use_diffuse(simtbx::nanoBragg::diffBragg& diffBragg, bool val){
      diffBragg.db_flags.use_diffuse = val ;
  }

  bool get_gamma_miller_units(simtbx::nanoBragg::diffBragg& diffBragg){
      return diffBragg.db_flags.gamma_miller_units;
  }

  void set_gamma_miller_units(simtbx::nanoBragg::diffBragg& diffBragg, bool val){
      diffBragg.db_flags.gamma_miller_units = val ;
  }

  bool get_only_diffuse(simtbx::nanoBragg::diffBragg& diffBragg){
      return diffBragg.db_flags.only_diffuse;
  }

  void set_only_diffuse(simtbx::nanoBragg::diffBragg& diffBragg, bool val){
      diffBragg.db_flags.only_diffuse = val ;
  }

  bool get_wavelen_img_flag(simtbx::nanoBragg::diffBragg& diffBragg){
    return diffBragg.db_flags.wavelength_img;
  }

  void set_wavelen_img_flag(simtbx::nanoBragg::diffBragg& diffBragg, bool val){
    diffBragg.db_flags.wavelength_img = val;
  }

  bool get_refine_Icell(simtbx::nanoBragg::diffBragg& diffBragg){
      return diffBragg.db_flags.refine_Icell;
  }

  void set_refine_Icell(simtbx::nanoBragg::diffBragg& diffBragg, bool val){
      diffBragg.db_flags.refine_Icell = val ;
  }

  static boost::python::tuple get_origin(simtbx::nanoBragg::diffBragg& diffBragg){
    return boost::python::make_tuple(diffBragg.pix0_vector[1]*1000,
                                    diffBragg.pix0_vector[2]*1000, diffBragg.pix0_vector[3]*1000);
  }

  static boost::python::list get_close_distances(simtbx::nanoBragg::diffBragg& diffBragg){
      boost::python::list detList;
      for (int pid=0; pid < diffBragg.db_det.close_distances.size(); pid++)
          detList.append( diffBragg.db_det.close_distances[pid]);
      return detList;
  }

  static void  set_Fhkl_tuple(simtbx::nanoBragg::diffBragg& diffBragg, boost::python::tuple const& value) {
      //TODO nanoBragg set as well ?
      diffBragg.pythony_indices = extract<nanoBragg::indices >(value[0]);
      diffBragg.pythony_amplitudes = extract<nanoBragg::af::shared<double> >(value[1]);
      diffBragg.init_Fhkl();
      diffBragg.complex_miller = false;
      if (boost::python::len(value)==3){
          if (value[2]!=boost::python::object()){ // check if it is not None
              diffBragg.pythony_amplitudes2 = extract<nanoBragg::af::shared<double> >(value[2]);
              diffBragg.init_Fhkl2();
              diffBragg.complex_miller = true;
          }
      }
      diffBragg.linearize_Fhkl();
  }

  static boost::python::tuple get_Fhkl_tuple(simtbx::nanoBragg::diffBragg diffBragg) {
      int h,k,l;
      double temp;
      int hkls = diffBragg.h_range*diffBragg.k_range*diffBragg.l_range;
      diffBragg.pythony_indices = nanoBragg::indices(hkls,nanoBragg::af::init_functor_null<scitbx::vec3<int> >());
      diffBragg.pythony_amplitudes = nanoBragg::af::shared<double>(hkls,nanoBragg::af::init_functor_null<double>());
      diffBragg.pythony_amplitudes2 = nanoBragg::af::shared<double>(hkls,nanoBragg::af::init_functor_null<double>());
      int i=0;
      for(h=diffBragg.h_min;h<=diffBragg.h_max;++h){
          for(k=diffBragg.k_min;k<=diffBragg.k_max;++k){
              for(l=diffBragg.l_min;l<=diffBragg.l_max;++l){
                  if ( (h<=diffBragg.h_max) && (h>=diffBragg.h_min) && (k<=diffBragg.k_max)
                    && (k>=diffBragg.k_min) && (l<=diffBragg.l_max) && (l>=diffBragg.l_min)  ) {
                      /* populate all stored values */
                      nanoBragg::miller_t hkl (h,k,l);
                      diffBragg.pythony_indices[i] = hkl;
                      temp = diffBragg.Fhkl[h-diffBragg.h_min][k-diffBragg.k_min][l-diffBragg.l_min];
                      diffBragg.pythony_amplitudes[i] = temp;
                      if (diffBragg.complex_miller){
                          temp = diffBragg.Fhkl2[h-diffBragg.h_min][k-diffBragg.k_min][l-diffBragg.l_min];
                          diffBragg.pythony_amplitudes2[i] = temp;
                      }
                      ++i;
                  }
              }
          }
      }
      if (diffBragg.complex_miller)
        return boost::python::make_tuple(diffBragg.pythony_indices,diffBragg.pythony_amplitudes, diffBragg.pythony_amplitudes2);
      else
        return boost::python::make_tuple(diffBragg.pythony_indices,diffBragg.pythony_amplitudes);
  }

  void diffBragg_init_module() {
    using namespace boost::python;
    typedef return_value_policy<return_by_value> rbv;
    typedef default_call_policies dcp;
    typedef return_internal_reference<> rir;
    class_<simtbx::nanoBragg::diffBragg, bases<simtbx::nanoBragg::nanoBragg> >
            ("diffBragg", no_init)
      /* constructor that takes a dxtbx detector and beam model */
      .def(init<const dxtbx::model::Detector&,
                const dxtbx::model::Beam&,
                int >(
        (arg_("detector"),
         arg_("beam"),
         arg_("verbose")=0),
        "nanoBragg simulation initialized from dxtbx detector and beam objects"))

      /* function for doing some differentiation */
      //.def("add_diffBragg_spots",&simtbx::nanoBragg::diffBragg::add_diffBragg_spots,
      //  "gives derivative of average photon count w.r.t. parameter of choice")

      .def("add_diffBragg_spots", add_diffBragg_spots_A,
        "gives derivative of average photon count w.r.t. parameter of choice")

      .def("add_diffBragg_spots", add_diffBragg_spots_B,
        "gives derivative of average photon count w.r.t. parameter of choice")

      .def("add_diffBragg_spots", add_diffBragg_spots_C,
        "gives derivative of average photon count w.r.t. parameter of choice")

      /* Run this before running add_diffBragg_spots */
      .def("vectorize_umats",&simtbx::nanoBragg::diffBragg::vectorize_umats,
        "caches the UMATS")

      .def("initialize_managers",&simtbx::nanoBragg::diffBragg::initialize_managers,
        "sets up the managers")

      .def("refine", &simtbx::nanoBragg::diffBragg::refine,
        "sets the refinement flag to true")

      .def("fix", &simtbx::nanoBragg::diffBragg::fix,
        "sets the refinement flag to false")

      .def("let_loose", &simtbx::nanoBragg::diffBragg::let_loose,
        "sets the refinement flag to True(but doesnt initialize)")

      .def("get_value", &simtbx::nanoBragg::diffBragg::get_value, "get value of the refinement parameter")

      .def("set_value", &simtbx::nanoBragg::diffBragg::set_value, "set value of the refinement parameter")

      .def("set_ncells_values", &simtbx::nanoBragg::diffBragg::set_ncells_values, "set Ncells values as a 3-tuple (Na, Nb, Nc)")

      .def("get_ncells_values", &simtbx::nanoBragg::diffBragg::get_ncells_values, "get Ncells values as a 3-tuple (Na, Nb, Nc)")

      .def("add_diffBragg_spots_full", &simtbx::nanoBragg::diffBragg::add_diffBragg_spots_full, "forward model and gradients at every pixel")

      .def("get_ncells_derivative_pixels", &simtbx::nanoBragg::diffBragg::get_ncells_derivative_pixels, "get derivatives of intensity w.r.t (Na, Nb, Nc)")

      .def("get_diffuse_gamma_derivative_pixels", &simtbx::nanoBragg::diffBragg::get_diffuse_gamma_derivative_pixels, "get derivatives of intensity w.r.t. diffuse (gamma_a, gamma_b, gamma_c)")

      .def("get_diffuse_sigma_derivative_pixels", &simtbx::nanoBragg::diffBragg::get_diffuse_sigma_derivative_pixels, "get derivatives of intensity w.r.t. diffuse (sigma_a, sigma_b, sigma_c)")

      .def("get_fp_fdp_derivative_pixels", &simtbx::nanoBragg::diffBragg::get_fp_fdp_derivative_pixels, "get derivatives of intensity w.r.t c,d that describe fprime and fdblprime (see diffBragg.utils)")

      .def("get_ncells_def_derivative_pixels", &simtbx::nanoBragg::diffBragg::get_ncells_def_derivative_pixels, "get derivatives of intensity w.r.t (Nd, Ne, Nf)")

      //.def("get_Na_derivative_pixels", &simtbx::nanoBragg::diffBragg::get_Na_derivative_pixels, "get derivatives of intensity w.r.t (Nd, Ne, Nf)")

      .def("get_aniso_eta_deriv_pixels", &simtbx::nanoBragg::diffBragg::get_aniso_eta_deriv_pixels, "get derivatives of intensity w.r.t anisotropic mosaicity model")

      .def("get_aniso_eta_second_deriv_pixels", &simtbx::nanoBragg::diffBragg::get_aniso_eta_second_deriv_pixels, "get second derivatives of intensity w.r.t anisotropic mosaicity model")

      .def("get_ncells_def_second_derivative_pixels", &simtbx::nanoBragg::diffBragg::get_ncells_def_second_derivative_pixels, "get 2nd derivatives of intensity w.r.t (Nd, Ne, Nf)")

      .def("get_ncells_second_derivative_pixels", &simtbx::nanoBragg::diffBragg::get_ncells_second_derivative_pixels, "get second derivatives of intensity w.r.t (Na, Nb, Nc)")

      .def("get_lambda_derivative_pixels", &simtbx::nanoBragg::diffBragg::get_lambda_derivative_pixels, "get derivatives of intensity w.r.t (lambda0, lambda1) as a 2-tuple")

      .def("print_if_refining", &simtbx::nanoBragg::diffBragg::print_if_refining, "print a statement if the refiner flag is set")

      .def("quick_Fhkl_update", &simtbx::nanoBragg::diffBragg::quick_Fcell_update, "update Fcell amplitudes without reallocating memory")
      .def("update_xray_beams", &simtbx::nanoBragg::diffBragg::update_xray_beams, "update xray beams without reallocating memory")

      .def("set_ucell_derivative_matrix",  &simtbx::nanoBragg::diffBragg::set_ucell_derivative_matrix, "Boo-ya")

      .def("set_ucell_second_derivative_matrix",  &simtbx::nanoBragg::diffBragg::set_ucell_second_derivative_matrix,
            "Two-ya")

      //.def("get_ucell_derivative_matrix",  &simtbx::nanoBragg::diffBragg::get_ucell_derivative_matrix, "scooby snacks")

      //.def("get_derivative_pixels", get_deriv_pix,
      //      "gets the manager raw image containing first derivatives")

      .def("__get_derivative_pixels", &simtbx::nanoBragg::diffBragg::get_derivative_pixels,
            "gets the manager raw image containing first derivatives")

      .def("get_second_derivative_pixels", &simtbx::nanoBragg::diffBragg::get_second_derivative_pixels,
            "gets the manager raw image containing second derivatives")

      .def("zero_raw_pixel_rois", &simtbx::nanoBragg::diffBragg::zero_raw_pixel_rois,
           "reallocate the raw image ROIs (they are usually tiny)")

      .def("get_origin", get_origin, "get panel origin")

      .def("shift_origin_z",
            &simtbx::nanoBragg::diffBragg::shift_originZ,
            "shift the Z-component of the pixel0 vector in every detector panel")

      .def("update_dxtbx_geoms",
            &simtbx::nanoBragg::diffBragg::update_dxtbx_geoms,
             (arg_("detector"),
              arg_("beam"),
              arg_("panel_id")=0,
              arg_("panel_rot_angO")=0,
              arg_("panel_rot_angF")=0,
              arg_("panel_rot_angS")=0,
              arg_("panel_offsetX")=0,
              arg_("panel_offsetY")=0,
              arg_("panel_offsetZ")=0, arg_("force")=true),
           "update the geometries with new dxtbx models, number of pixels should remain constant")

      .def("free_Fhkl2",&simtbx::nanoBragg::diffBragg::free_Fhkl2)

      .def("show_fp_fdp", &simtbx::nanoBragg::diffBragg::show_fp_fdp)

      .def("show_heavy_atom_data", &simtbx::nanoBragg::diffBragg::show_heavy_atom_data)

#ifdef NANOBRAGG_HAVE_CUDA
      .def("gpu_free",&simtbx::nanoBragg::diffBragg::gpu_free)
#endif

      .def("set_mosaic_blocks_prime",
             &simtbx::nanoBragg::diffBragg::set_mosaic_blocks_prime,
             "enter a list of unitary matrix derivatives Uprime")

      .def("set_mosaic_blocks_dbl_prime",
             &simtbx::nanoBragg::diffBragg::set_mosaic_blocks_dbl_prime,
             "enter a list of unitary matrix 2nd derivatives Uprime")

      .def("show_timings",
           &simtbx::nanoBragg::diffBragg::show_timing_stats,
           (arg_("MPI_RANK")=0),
           "Shows a few key accumulated timings (in milliseconds) of specific code portions")

      //.add_property("region_of_interest",
      //       make_function(&get_roi,rbv()),
      //       make_function(&set_roi,dcp()),
      //       "region of interest on detector: fast_min fast_max slow_min slow_max")

      .add_property("close_distances",
             make_function(&get_close_distances,rbv()))
             //"get the effective detector distance for each panel")

      .add_property("raw_pixels_roi",
                     make_getter(&simtbx::nanoBragg::diffBragg::raw_pixels_roi,rbv()),
                     make_setter(&simtbx::nanoBragg::diffBragg::raw_pixels_roi,dcp()),
                    "raw pixels from region of interest only")

      .add_property("oversample_omega",
                     make_getter(&simtbx::nanoBragg::diffBragg::oversample_omega,rbv()),
                     make_setter(&simtbx::nanoBragg::diffBragg::oversample_omega,dcp()),
                    "whether to use an average solid angle correction per pixel, or one at the sub pixel level")

      .add_property("track_Fhkl",
                     make_getter(&simtbx::nanoBragg::diffBragg::track_Fhkl,rbv()),
                     make_setter(&simtbx::nanoBragg::diffBragg::track_Fhkl,dcp()),
                    "whether to track Fhkl per pixel")

      .add_property("fprime_fdblprime",
             make_function(&get_fpfdp,rbv()),
             make_function(&set_fpfdp,dcp()),
             "fprime and fdoubleprime, passed as two lists whose length should be the same")

      .add_property("fprime_fdblprime_derivs",
             make_function(&get_fpfdp_derivs,rbv()), // TODO implement me
             make_function(&set_fpfdp_derivs,dcp()),
             "fprime and fdoubleprime derivatives, as a 2-tuple. e.g. 1st tuple element is fprime derivs, length should be multiple of number of sources")

      .add_property("heavy_atom_data",
             make_function(&get_atom_data,rbv()),
             make_function(&set_atom_data,dcp()),
             "5 tuple of fracx_list,fracy_list,fracz_list,Bfactor_list,Occupancy_list, each list is same length")

      .add_property("no_Nabc_scale",
                     make_getter(&simtbx::nanoBragg::diffBragg::no_Nabc_scale,rbv()),
                     make_setter(&simtbx::nanoBragg::diffBragg::no_Nabc_scale,dcp()),
                    "toggle off the Nabc scale factor in the forward model (such that it is replaced entirely by spot_scale)")

      .add_property("__number_of_pixels_modeled_using_diffBragg", // protect this by making it a long name
                     make_getter(&simtbx::nanoBragg::diffBragg::Npix_to_model,rbv()),
                     make_setter(&simtbx::nanoBragg::diffBragg::Npix_to_model,dcp()),
                    "num pix modeled most recently")

      .add_property("compute_curvatures",
                     make_getter(&simtbx::nanoBragg::diffBragg::compute_curvatures,rbv()),
                     make_setter(&simtbx::nanoBragg::diffBragg::compute_curvatures,dcp()),
                    "Whether to compute the curvatures")

      .add_property("only_save_omega_kahn",
                     make_getter(&simtbx::nanoBragg::diffBragg::only_save_omega_kahn,rbv()),
                     make_setter(&simtbx::nanoBragg::diffBragg::only_save_omega_kahn,dcp()),
                    "ONLY simulate the kahn polarization correction and solid angle components")

      .add_property("reference_origin",
            make_function(&get_reference_origin, rbv()),
            make_function(&set_reference_origin, dcp()),
            "reference origin for doing panel rotations for a group of panels")

      .add_property("Umatrix",
             make_function(&get_U,rbv()),
             make_function(&set_U,dcp()),
             "rotation portion of Amatrix (dxtbx format, C.get_U())")

      .add_property("Bmatrix",
             make_function(&get_B,rbv()),
             make_function(&set_B,dcp()),
             "unit cell portion of Amatrix (dxtbx format; C.get_B())")

      .add_property("Omatrix",
             make_function(&get_O,rbv()),
             make_function(&set_O,dcp()),
             "change of basis Omatrix")

      .add_property("Ncells_abc",
             make_function(&get_Nabc,rbv()),
             make_function(&set_Nabc,dcp()),
             "number of mosaic domains along an axis")

      .add_property("Ncells_def",
             make_function(&get_Ndef,rbv()),
             make_function(&set_Ndef,dcp()),
             "off-axis components in mosaic domain size matrix")

      .add_property("Ncells_abc_aniso",
             make_function(&get_Nabc_aniso,rbv()),
             make_function(&set_Nabc_aniso,dcp()),
             "number of mosaic domains along an axis, can be anisotropic")

      .add_property("update_oversample_during_refinement",
                     make_getter(&simtbx::nanoBragg::diffBragg::update_oversample_during_refinement,rbv()),
                     make_setter(&simtbx::nanoBragg::diffBragg::update_oversample_during_refinement,dcp()),
                     "Allows oversample to change as Ncells abc changes")

      .add_property("isotropic_ncells",
                     make_getter(&simtbx::nanoBragg::diffBragg::isotropic_ncells,rbv()),
                     make_setter(&simtbx::nanoBragg::diffBragg::isotropic_ncells,dcp()),
                     "refine (Na, Nb, Nc) as three independent parameters")

      .add_property("has_anisotropic_mosaic_spread",
                     make_getter(&simtbx::nanoBragg::diffBragg::modeling_anisotropic_mosaic_spread,rbv()),
                     make_setter(&simtbx::nanoBragg::diffBragg::modeling_anisotropic_mosaic_spread,dcp()),
                     "compute derivatives for a single mosaic spread parameter")

      .add_property("use_lambda_coefficients",
                     make_getter(&simtbx::nanoBragg::diffBragg::use_lambda_coefficients,rbv()),
                     make_setter(&simtbx::nanoBragg::diffBragg::use_lambda_coefficients,dcp()),
                     "represent wavelength using coefficients a + b*source where a,b are coefficients and source is the source index")

      .add_property("Fhkl_tuple",
            make_function(&get_Fhkl_tuple,rbv()),
            make_function(&set_Fhkl_tuple,dcp()),
            "hkl and Freal and Fcomplex as a 3-tuple of (indices ,flex-double, flex-double). experimental, also accepts (indices, flex-double, None) for mod only")

      .add_property("lambda_coefficients",
             make_function(&get_lambda_coef,rbv()),
             make_function(&set_lambda_coef,dcp()),
             "coefficients for source_lambda refinement: `lambda = coef0 + coef1*source`  where `source` is the source index")

     // CUDA PROPERTIES
      .add_property("use_cuda",
             make_getter(&simtbx::nanoBragg::diffBragg::use_cuda,rbv()),
             make_setter(&simtbx::nanoBragg::diffBragg::use_cuda,dcp()),
             "use GPU acceleration")

      .add_property("record_time",
             make_getter(&simtbx::nanoBragg::diffBragg::record_time,rbv()),
             make_setter(&simtbx::nanoBragg::diffBragg::record_time,dcp()),
             "store timing data for key aspects of the code")

      .add_property("Npix_to_allocate",
            make_getter(&simtbx::nanoBragg::diffBragg::Npix_to_allocate,rbv()),
            make_setter(&simtbx::nanoBragg::diffBragg::Npix_to_allocate,dcp()),
            "number of pixels per parameter to allocate in the GPU (default value of -1 means auto-determine)")

      .add_property("use_diffuse",
            make_function(&get_use_diffuse,rbv()),
            make_function(&set_use_diffuse,dcp()),
            "sim with diffuse")

      .add_property("gamma_miller_units",
            make_function(&get_gamma_miller_units,rbv()),
            make_function(&set_gamma_miller_units,dcp()),
            "use Miller index units for gamma parameters")

      .add_property("only_diffuse",
            make_function(&get_only_diffuse,rbv()),
            make_function(&set_only_diffuse,dcp()),
            "sim with ONLY diffuse")

      .add_property("refine_Icell",
            make_function(&get_refine_Icell,rbv()),
            make_function(&set_refine_Icell,dcp()),
            "Assume whatever is stored in the structure factor array is an intensity (allowed to be negative)")

      .add_property("diffuse_gamma",
            make_function(&get_diffuse_gamma,rbv()),
            make_function(&set_diffuse_gamma,dcp()),
            "the 3 diffuse gamma factors")

      .add_property("diffuse_sigma",
            make_function(&get_diffuse_sigma,rbv()),
            make_function(&set_diffuse_sigma,dcp()),
            "set the 3 diffuse sigma factors")

      .add_property("most_recent_kernel_used_GPU",
            make_getter(&simtbx::nanoBragg::diffBragg::last_kernel_on_GPU,rbv()),
            "if True, the most recent kernel call utilized a GPU")
      .add_property("store_ave_wavelength_image",
            make_function(&get_wavelen_img_flag,rbv()),
            make_function(&set_wavelen_img_flag,dcp()),
            "if True, then record the average wavelength per pixel, weighted by Bragg intensity")

      //.add_property("ave_wavelength",
      //      make_function(&simtbx::nanoBragg::diffBragg::ave_wavelength_img , rbv()),
      //      "return flex array containing average wavelen per pixel")
      .def("ave_wavelength_image",
            &simtbx::nanoBragg::diffBragg::ave_wavelength_img,
            "return flex array containing average wavelen per pixel")

    ; // end of diffBragg extention

  } // end of diffBragg_init_module
} // end of namespace
} // end of boost python namespace
} // end of refine
} // end of simtbx

BOOST_PYTHON_MODULE(simtbx_diffBragg_ext)
{
  simtbx::diffBragg::boost_python::diffBragg_init_module();
}


#pragma once

/*
Game Tree Search Algorithms
Copyright (C) Adam Stelmaszczyk <stelmaszczyk.adam@gmail.com>

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 3 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, see <http://www.gnu.org/licenses/>.
*/

#include <boost/math/distributions/binomial.hpp>
#include <unordered_map>
#include <unordered_set>
#include <sys/time.h>
#include <algorithm>
#include <iostream>
#include <assert.h>
#include <fstream>
#include <sstream>
#include <iomanip>
#include <memory>
#include <random>
#include <vector>

using namespace std;

static const int MAX_SIMULATIONS = 10000000;
static const double UCT_C = sqrt(2);
static const double WIN_SCORE = 1;
static const double DRAW_SCORE = 0.5;
static const double LOSE_SCORE = 0;

static const int MAX_DEPTH = 20;
static const int INF = 2147483647;

struct Random {

    virtual ~Random() {}

    int uniform(int min, int max) const {
        mt19937 engine;
        uniform_int_distribution<int> distribution(min, max);
        return distribution(engine);
    }
};

struct Timer {
    double start_time;

    virtual ~Timer() {}

    void start() {
        start_time = get_time();
    }

    double get_time() const {
        timeval tv;
        gettimeofday(&tv, 0);
        return tv.tv_sec + tv.tv_usec * 1e-6;
    }

    double seconds_elapsed() const {
        return get_time() - start_time;
    }

    bool exceeded(double seconds) const {
        return seconds_elapsed() > seconds;
    }

    friend ostream &operator<<(ostream &os, const Timer &timer) {
        return os << setprecision(2) << fixed << timer.seconds_elapsed() << "s";
    }
};

template<class M>
struct Move {
    virtual ~Move() {}

    virtual void read(istream &stream) = 0;

    virtual ostream &to_stream(ostream &os) const = 0;

    friend ostream &operator<<(ostream &os, const Move &move) {
        return move.to_stream(os);
    }

    virtual bool operator==(const M &rhs) const = 0;

    virtual size_t hash() const = 0;
};

enum TTEntryType { EXACT_VALUE, LOWER_BOUND, UPPER_BOUND };

template<class M>
struct TTEntry {
    M move;
    int depth;
    int value;
    TTEntryType value_type;

    TTEntry() {}

    virtual ~TTEntry() {}

    TTEntry(const M &move, int depth, int value, TTEntryType value_type) :
            move(move), depth(depth), value(value), value_type(value_type) {}

    ostream &to_stream(ostream &os) {
        return os << "move: " << move << " depth: " << depth << " value: " << value << " value_type: " << value_type;
    }

    friend ostream &operator<<(ostream &os, const TTEntry &entry) {
        return entry.to_stream(os);
    }
};

template<class S, class M>
struct State {
    unsigned visits = 0;
    double score = 0;
    char player_to_move = 0;
    S *parent = nullptr;
    unordered_map<size_t, shared_ptr<S>> children = unordered_map<size_t, shared_ptr<S>>();

    State(char player_to_move) : player_to_move(player_to_move) {}

    virtual ~State() {}

    void update_stats(double result) {
        score += result;
        ++visits;
    }

    double get_uct(double c) const {
        assert(visits > 0);
        double parent_visits = 0.0;
        if (parent != nullptr) {
            parent_visits = parent->visits;
        }
        return score / visits + c * sqrt(log(parent_visits) / visits);
    }

    shared_ptr<S> create_child(const M &move) {
        S child = clone();
        child.make_move(move);
        child.parent = (S*) this;
        return make_shared<S>(child);
    }

    S* add_child(const M &move) {
        const auto child = create_child(move);
        const auto key = move.hash();
        const auto pair = children.insert({key, child});
        const auto it = pair.first;
        return it->second.get();
    }

    S* get_child(const M &move) const {
        const auto key = move.hash();
        const auto it = children.find(key);
        if (it == children.end()) {
            return nullptr;
        }
        return it->second.get();
    }

    virtual string to_executable_format() const {
        stringstream ss;
        ss << *this;
        return ss.str();
    }

    virtual void swap_players() {}

    virtual S clone() const = 0;

    virtual int get_goodness() const = 0;

    virtual vector<M> get_legal_moves(int max_moves) const = 0;

    virtual char get_enemy(char player) const = 0;

    virtual bool is_terminal() const = 0;

    virtual bool is_winner(char player) const = 0;

    virtual void make_move(const M &move) = 0;

    virtual void undo_move(const M &move) = 0;

    virtual ostream &to_stream(ostream &os) const = 0;

    friend ostream &operator<<(ostream &os, const State &state) {
        return state.to_stream(os);
    }

    virtual bool operator==(const S &other) const = 0;

    virtual size_t hash() const = 0;
};

template<class S, class M>
struct Algorithm {

    stringstream log;

    Algorithm() { }

    Algorithm(const Algorithm& algorithm) {}

    virtual ~Algorithm() {}

    string read_log() {
        string result = log.str();
        log.str("");
        return result;
    }

    virtual void reset() {}

    virtual M get_move(const S *state) = 0;

    virtual string get_name() const = 0;

    friend ostream &operator<<(ostream &os, const Algorithm &algorithm) {
        os << algorithm.get_name();
        return os;
    }
};

template<class S, class M>
struct Human : public Algorithm<S, M> {

    Human() : Algorithm<S, M>() {}

    M get_move(const S *state) override {
        const vector<M> &legal_moves = state->get_legal_moves();
        if (legal_moves.empty()) {
            stringstream stream;
            state->to_stream(stream);
            throw invalid_argument("Given state is terminal:\n" + stream.str());
        }
        while (true) {
            M move = M();
            move.read();
            if (find(legal_moves.begin(), legal_moves.end(), move) != legal_moves.end()) {
                return move;
            } else {
                cout << "Move " << move << " is not legal" << endl;
            }
        }
    }

    string get_name() const {
        return "Human";
    }
};

template<class S, class M>
struct Executable : public Algorithm<S, M> {

    const string executable;

    Executable(string executable) : Algorithm<S, M>(), executable(executable) {}

    M get_move(const S *state) override {
        const vector<M> &legal_moves = state->get_legal_moves();
        if (legal_moves.empty()) {
            stringstream stream;
            state->to_stream(stream);
            throw invalid_argument("Given state is terminal:\n" + stream.str());
        }

        stringstream cmd;
        cmd << "echo \"" << state->to_executable_format() << "\" | " << executable;

        const string output = run_cmd(cmd.str());
        stringstream stream_to_read(output);

        M move = M();
        move.read(stream_to_read);
        if (find(legal_moves.begin(), legal_moves.end(), move) != legal_moves.end()) {
            return move;
        } else {
            stringstream message;
            message << "Legal moves: ";
            for (auto move : legal_moves) {
                message << move << ", ";
            }
            message << endl;
            message << "Move " << move << " is not legal, state:" << endl;
            message << *state;
            throw runtime_error(message.str());
        }
    }

    string run_cmd(string cmd) {
        stringstream result;
        const int BUFFER_SIZE = 128;
        char buffer[BUFFER_SIZE];
        shared_ptr<FILE> pipe(popen(cmd.c_str(), "r"), pclose);
        if (!pipe) {
            throw runtime_error("popen() failed");
        }
        while (!feof(pipe.get())) {
            if (fgets(buffer, BUFFER_SIZE, pipe.get()) != NULL) {
                result << buffer;
            }
        }
        return result.str();
    }

    string get_name() const {
        return "Executable " + executable;
    }
};

template<class M>
struct MinimaxResult {
    int goodness;
    M best_move;
    bool completed;
};

template<class S, class M>
struct Minimax : public Algorithm<S, M> {
    unordered_map<size_t, TTEntry<M>> transposition_table;
    const double MAX_SECONDS;
    const int MAX_MOVES;
    function<vector<M>(const S*, int)> get_legal_moves;
    function<int(const S*)> get_goodness;
    Timer timer;
    int scout_cuts;
    int beta_cuts, cut_bf_sum;
    int tt_hits, tt_exacts, tt_cuts;
    int nodes, leafs;

    Minimax(double max_seconds = 1, int max_moves = INF, function<vector<M>(const S*, int)> get_legal_moves = nullptr, function<int(const S*)> get_goodness = nullptr) :
            Algorithm<S, M>(),
            transposition_table(unordered_map<size_t, TTEntry<M>>(1000000)),
            MAX_SECONDS(max_seconds),
            MAX_MOVES(max_moves),
            get_legal_moves(get_legal_moves),
            get_goodness(get_goodness),
            timer(Timer()) {}

    void reset() {
        transposition_table.clear();
    }

    M get_move(const S *state) override {
        if (state->is_terminal()) {
            stringstream stream;
            state->to_stream(stream);
            throw invalid_argument("Given state is terminal:\n" + stream.str());
        }
        if (get_legal_moves == nullptr) {
            get_legal_moves = &State<S,M>::get_legal_moves;
        }
        if (get_goodness == nullptr) {
            get_goodness = &State<S,M>::get_goodness;
        }
        timer.start();

        const auto moves = get_legal_moves(state, MAX_MOVES);
        this->log << "moves: " << moves.size() << endl;
        for (const auto move : moves) {
            this->log << move << ", ";
        }
        this->log << endl;

        M best_move;
        for (int max_depth = 1; max_depth <= MAX_DEPTH; ++max_depth) {
            scout_cuts = 0;
            beta_cuts = 0;
            cut_bf_sum = 0;
            tt_hits = 0;
            tt_exacts = 0;
            tt_cuts = 0;
            nodes = 0;
            leafs = 0;
            S clone = state->clone();
            auto result = minimax(&clone, max_depth, -INF, INF);
            if (result.completed) {
                best_move = result.best_move;
                this->log << "goodness: " << result.goodness
                << " time: " << timer
                << " move: " << best_move
                << " nodes: " << nodes
                << " leafs: " << leafs
                << " scout_cuts: " << scout_cuts
                << " beta_cuts: " << beta_cuts
                << " cutBF: " << (double) cut_bf_sum / beta_cuts
                << " tt_hits: " << tt_hits
                << " tt_exacts: " << tt_exacts
                << " tt_cuts: " << tt_cuts
                << " tt_size: " << transposition_table.size()
                << " max_depth: " << max_depth << endl;
            }
            if (timer.exceeded(MAX_SECONDS)) {
                break;
            }
        }
        return best_move;
    }

    // Find Minimax value of the given tree,
    // Minimax value lies within a range of [alpha; beta] window.
    // Whenever alpha >= beta, further checks of children in a node can be pruned.
    MinimaxResult<M> minimax(S *state, int depth, int alpha, int beta) {
        ++nodes;
        const int alpha_original = alpha;

        M best_move;
        if (depth == 0 || state->is_terminal()) {
            ++leafs;
            return {get_goodness(state), best_move, false};
        }

        TTEntry<M> entry;
        const bool entry_found = get_tt_entry(state, entry);
        if (entry_found && entry.depth >= depth) {
            ++tt_hits;
            if (entry.value_type == TTEntryType::EXACT_VALUE) {
                ++tt_exacts;
                return {entry.value, entry.move, true};
            }
            if (entry.value_type == TTEntryType::LOWER_BOUND && alpha < entry.value) {
                alpha = entry.value;
            }
            if (entry.value_type == TTEntryType::UPPER_BOUND && beta > entry.value) {
                beta = entry.value;
            }
            if (alpha >= beta) {
                ++tt_cuts;
                return {entry.value, entry.move, true};
            }
        }

        int max_goodness = -INF;

        bool completed = true;
        const auto legal_moves = get_legal_moves(state, MAX_MOVES);
        assert(legal_moves.size() > 0);
        for (int i = 0; i < legal_moves.size(); i++) {
            const auto move = legal_moves[i];
            state->make_move(move);
            int goodness;
            if (i > 0) {
                // null window search
                goodness = -minimax(
                    state,
                    depth - 1,
                    -alpha - 1,
                    -alpha
                ).goodness;
                if (alpha < goodness && goodness < beta) {
                    // failed high, do a full re-search
                    goodness = -minimax(
                        state,
                        depth - 1,
                        -beta,
                        -goodness
                    ).goodness;
                } else {
                    scout_cuts++;
                }
            }
            else {
                goodness = -minimax(
                    state,
                    depth - 1,
                    -beta,
                    -alpha
                ).goodness;
            }
            state->undo_move(move);
            if (timer.exceeded(MAX_SECONDS)) {
                completed = false;
                break;
            }
            if (max_goodness < goodness) {
                max_goodness = goodness;
                best_move = move;
                if (max_goodness >= beta) {
                    ++beta_cuts;
                    cut_bf_sum += i + 1;
                    break;
                }
            }
            if (alpha < max_goodness) {
                alpha = max_goodness;
            }
        }

        if (completed) {
            update_tt(state, alpha_original, beta, max_goodness, best_move, depth);
        }

        return {max_goodness, best_move, completed};
    }

    bool get_tt_entry(const S *state, TTEntry<M> &entry) const {
        const auto key = state->hash();
        const auto it = transposition_table.find(key);
        if (it == transposition_table.end()) {
            return false;
        }
        entry = it->second;
        return true;
    }

    void add_tt_entry(const S *state, const TTEntry<M> &entry) {
        const auto key = state->hash();
        transposition_table.insert({key, entry});
    }

    void update_tt(const S *state, int alpha, int beta, int max_goodness, const M &best_move, int depth) {
        TTEntryType value_type;
        if (max_goodness <= alpha) {
            value_type = TTEntryType::UPPER_BOUND;
        }
        else if (max_goodness >= beta) {
            value_type = TTEntryType::LOWER_BOUND;
        }
        else {
            value_type = TTEntryType::EXACT_VALUE;
        }
        const TTEntry<M> entry = {best_move, depth, max_goodness, value_type};
        add_tt_entry(state, entry);
    }

    string get_name() const {
        return "Minimax";
    }
};

template<class S, class M>
struct MonteCarloTreeSearch : public Algorithm<S, M> {
    const double max_seconds;
    const int max_simulations;
    const bool block;
    const Random random;

    MonteCarloTreeSearch(double max_seconds = 1,
                         int max_simulations = MAX_SIMULATIONS,
                         bool block = false) :
        Algorithm<S, M>(),
        max_seconds(max_seconds),
        block(block),
        max_simulations(max_simulations) {}

    M get_move(const S *root) override {
        if (root->is_terminal()) {
            stringstream stream;
            root->to_stream(stream);
            throw invalid_argument("Given state is terminal:\n" + stream.str());
        }
        Timer timer;
        timer.start();
        int simulation = 0;
        S clone = root->clone();
        while (simulation < max_simulations && !timer.exceeded(max_seconds)) {
            monte_carlo_tree_search(&clone);
            ++simulation;
        }
        this->log << "ratio: " << root->score / root->visits << endl;
        this->log << "simulations: " << simulation << endl;
        const auto legal_moves = root->get_legal_moves();
        this->log << "moves: " << legal_moves.size() << endl;
        for (const auto move : legal_moves) {
            this->log << "move: " << move;
            const auto child = root->get_child(move);
            if (child != nullptr) {
                this->log << " score: " << child->score
                << " visits: " << child->visits
                << " UCT: " << child->get_uct(UCT_C);
            }
            this->log << endl;
        }
        return get_most_visited_move(&clone);
    }

    void monte_carlo_tree_search(S *root) const {
        S *current = tree_policy(root, root);
        const auto result = rollout(current, root);
        propagate_up(current, result);
    }

    void propagate_up(S *current, double result) const {
        current->update_stats(result);
        if (current->parent) {
            propagate_up(current->parent, result);
        }
    }

    S* tree_policy(S *state, const S *root) const {
        if (state->is_terminal()) {
            return state;
        }
        const M move = get_tree_policy_move(state, root);
        const auto child = state->get_child(move);
        if (child == nullptr) {
            return state->add_child(move);
        }
        return tree_policy(child, root);
    }

    M get_most_visited_move(const S *state) const {
        const auto legal_moves = state->get_legal_moves();
        assert(legal_moves.size() > 0);
        M best_move;
        double max_visits = -INF;
        for (const auto move : legal_moves) {
            const auto child = state->get_child(move);
            if (child != nullptr) {
                const auto visits = child->visits;
                if (max_visits < visits) {
                    max_visits = visits;
                    best_move = move;
                }
            }
        }
        assert(max_visits != -INF);
        return best_move;
    }

    M get_best_move(S *state, const S *root) const {
        const auto legal_moves = state->get_legal_moves();
        assert(legal_moves.size() > 0);
        M best_move;
        if (state->player_to_move == root->player_to_move) {
            // maximize
            double best_uct = -INF;
            for (const auto move : legal_moves) {
                const auto child = state->get_child(move);
                if (child != nullptr) {
                    const auto uct = child->get_uct(UCT_C);
                    if (best_uct < uct) {
                        best_uct = uct;
                        best_move = move;
                    }
                } else {
                    return move;
                }
            }
        }
        else {
            // minimize
            double best_uct = INF;
            for (const auto move : legal_moves) {
                const auto child = state->get_child(move);
                if (child != nullptr) {
                    const auto uct = child->get_uct(-UCT_C);
                    if (best_uct > uct) {
                        best_uct = uct;
                        best_move = move;
                    }
                } else {
                    return move;
                }
            }
        }
        return best_move;
    }

    M get_random_move(const S *state) const {
        const auto legal_moves = state->get_legal_moves();
        assert(legal_moves.size() > 0);
        const int index = random.uniform(0, legal_moves.size() - 1);
        return legal_moves[index];
    }

    shared_ptr<M> get_winning_move(const S *state) const {
        const auto current_player = state->player_to_move;
        const auto legal_moves = state->get_legal_moves();
        S clone = state->clone();
        assert(legal_moves.size() > 0);
        for (const M &move : legal_moves) {
            clone.make_move(move);
            if (clone.is_winner(current_player)) {
                return make_shared<M>(move);
            }
            clone.undo_move(move);
        }
        return nullptr;
    }

    shared_ptr<M> get_blocking_move(const S *state) const {
        const auto current_player = state->player_to_move;
        const auto enemy = state->get_enemy(current_player);
        S enemy_state = state->clone();
        enemy_state.player_to_move = enemy;
        return get_winning_move(&enemy_state);
    }

    M get_tree_policy_move(S *state, const S *root) const {
        // If player has a winning move he makes it.
        auto move_ptr = get_winning_move(state);
        if (move_ptr != nullptr) {
            return *move_ptr;
        }
        if (block) {
            // If player has a blocking move he makes it.
            move_ptr = get_blocking_move(state);
            if (move_ptr != nullptr) {
                return *move_ptr;
            }
        }
        return get_best_move(state, root);
    }

    M get_default_policy_move(const S *state) const {
        // If player has a winning move he makes it.
        auto move_ptr = get_winning_move(state);
        if (move_ptr != nullptr) {
            return *move_ptr;
        }
        // If player has a blocking move he makes it.
        move_ptr = get_blocking_move(state);
        if (move_ptr != nullptr) {
            return *move_ptr;
        }
        return get_random_move(state);
    }

    double rollout(S *current, const S *root) const {
        if (current->is_terminal()) {
            if (current->is_winner(root->player_to_move)) {
                return WIN_SCORE;
            }
            if (current->is_winner(root->get_enemy(root->player_to_move))) {
                return LOSE_SCORE;
            }
            return DRAW_SCORE;
        }
        M move = get_default_policy_move(current);
        current->make_move(move);
        auto result = rollout(current, root);
        current->undo_move(move);
        return result;
    }

    string get_name() const {
        return "MonteCarloTreeSearch";
    }

};

struct OutcomeCounts {
    int wins = 0;
    int draws = 0;
    int loses = 0;
};

template<class S, class M>
struct Tester {
    S *root = nullptr;
    Algorithm<S, M> &algorithm_1;
    Algorithm<S, M> &algorithm_2;
    const int MATCHES;
    const bool VERBOSE;
    const bool SAVE;
    const double SIGNIFICANCE_LEVEL = 0.005; // two sided 99% confidence interval

    Tester(S *state,
           Algorithm<S, M> &algorithm_1,
           Algorithm<S, M> &algorithm_2,
           int matches = INF,
           bool verbose = false,
           bool save = false
    ) : root(state), algorithm_1(algorithm_1), algorithm_2(algorithm_2), MATCHES(matches), VERBOSE(verbose), SAVE(save) {}

    virtual ~Tester() {}

    OutcomeCounts start() {
        Timer all_timer;
        all_timer.start();
        OutcomeCounts outcome_counts;
        unordered_set<int> unique_game_hashes;
        const char enemy = root->get_enemy(root->player_to_move);
        for (int i = 1; i <= MATCHES; ++i) {
            int move_number = 1;
            auto current = root->clone();
            if (i % 4 == 0 || i % 4 == 2) {
                current.player_to_move = current.get_enemy(current.player_to_move);
            }
            if (i % 4 == 0 || i % 4 == 3) {
                current.swap_players();
            }
            if (VERBOSE) {
                cout << current << endl;
            }
            if (SAVE) {
                save_file(move_number, current);
            }
            int game_hash = current.hash();
            while (!current.is_terminal()) {
                auto &algorithm = (current.player_to_move == root->player_to_move) ? algorithm_1 : algorithm_2;
                if (VERBOSE) {
                    cout << current.player_to_move << " " << algorithm << endl;
                }
                algorithm.reset();
                Timer timer;
                timer.start();
                auto move = algorithm.get_move(&current);
                if (VERBOSE) {
                    cout << algorithm.read_log();
                    cout << timer << endl;
                }
                current.make_move(move);
                ++move_number;
                if (VERBOSE) {
                    cout << current << endl;
                }
                if (SAVE) {
                    save_file(move_number, current);
                }
                game_hash ^= current.hash();
            }
            cout << "Game " << i << ": ";
            auto insert = unique_game_hashes.insert(game_hash);
            if (!insert.second) {
                cout << "Not unique, not counting" << endl << endl;
                continue;
            }
            if (current.is_winner(root->player_to_move)) {
                ++outcome_counts.wins;
                cout << root->player_to_move << " " << algorithm_1 << " won" << endl;
            } else if (current.is_winner(enemy)) {
                ++outcome_counts.loses;
                cout << enemy << " " << algorithm_2 << " won" << endl;
            } else {
                ++outcome_counts.draws;
                cout << "draw" << endl;
            }
            const int unique_games_count = unique_game_hashes.size();
            cout << "Unique games: " << unique_games_count << endl;
            cout << root->player_to_move << " " << algorithm_1 << " wins: " << outcome_counts.wins << endl;
            cout << enemy << " " << algorithm_2 << " wins: " << outcome_counts.loses << endl;
            cout << "Draws: " << outcome_counts.draws << endl;
            const double successes = outcome_counts.wins + 0.5 * outcome_counts.draws;
            const double ratio = successes / unique_games_count;
            cout << "Ratio: " << ratio << endl;
            const double lower = boost::math::binomial_distribution<>::find_lower_bound_on_p(unique_games_count, successes, SIGNIFICANCE_LEVEL);
            const double upper = boost::math::binomial_distribution<>::find_upper_bound_on_p(unique_games_count, successes, SIGNIFICANCE_LEVEL);
            cout << "Lower confidence bound: " << lower << endl;
            cout << "Upper confidence bound: " << upper << endl;
            cout << endl;
            if (upper < 0.5 || lower > 0.5) {
                cout << "Total time: " << all_timer << endl;
                break;
            }
        }
        if (SAVE) {
            shell("convert -delay 100 -loop 0 $(ls -v *.gif) game.gif");
            shell("rm [0-9]*.gif");
        }
        return outcome_counts;
    }

    void save_file(int move_number, const S &state) const {
        const int FONT_SIZE = 16;
        stringstream ss;
        ss << state;
        vector<string> lines;
        string line;
        while (getline(ss, line)) {
            lines.push_back(line);
        }
        assert(lines.size() > 0);
        const int width = lines[0].length() * FONT_SIZE;
        const int height = lines.size() * FONT_SIZE;
        stringstream command;
        command << "convert -size " << width << "x" << height;
        command << " xc:black -font square.ttf -pointsize " << FONT_SIZE << " -fill white -draw \"";
        for (int y = 0; y < lines.size(); y++) {
            command << "text 0," << (y + 1) * FONT_SIZE << " '" << lines[y] << "' ";
        }
        command << "\" " << move_number << ".gif";
        shell(command.str());
    }

    void shell(const string &command) const {
        const int return_code = system(command.c_str());
        if (return_code != 0) {
            cout << "Command " << command << " returned " << return_code;
        }
    }
};


/**
*   ReadWKT.hpp
*
*   Copyright 2015 The Climate Corporation
*   Copyright 2015 Tony Stone
*
*   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.
*
*   Created by Tony Stone on 6/13/15.
*
*   MODIFIED 2015 BY Tony Stone. Modifications licensed under Apache License, Version 2.0.
*
*/

#ifndef __ReadWKT_HPP_
#define __ReadWKT_HPP_

#include "GeometryCollection.hpp"
#include <boost/geometry/io/wkt/read.hpp>
#include <boost/algorithm/string/trim.hpp>
#include <boost/algorithm/string/replace.hpp>
#include <boost/algorithm/string/erase.hpp>
#include <boost/algorithm/string/find_iterator.hpp>

#include <regex>

namespace geofeatures {
    namespace io {

        template <typename Geometry, typename = typename std::enable_if<std::is_same<Geometry,GeometryCollection<>>::value>::type>
        inline void readWKT(std::string const &wkt, Geometry &geometryCollection)
        {
            if (!boost::istarts_with(wkt, "GEOMETRYCOLLECTION")) {
                throw std::invalid_argument("Should start with 'GEOMETRYCOLLECTION'' in (" + wkt + ")");
            }

            GeometryCollection<>::value_type (*geometry)(std::string & wkt) =
                    [](std::string & wkt) -> GeometryCollection<>::value_type {

                        if (boost::istarts_with(wkt, "POINT")) {
                            geofeatures::Point geometry;

                            boost::geometry::read_wkt(wkt, geometry);

                            return geometry;
                        } else if (boost::istarts_with(wkt, "MULTIPOINT")) {
                            geofeatures::MultiPoint geometry;

                            boost::geometry::read_wkt(wkt, geometry);

                            return geometry;
                        } else if (boost::istarts_with(wkt, "LINESTRING")) {
                            geofeatures::LineString geometry;

                            boost::geometry::read_wkt(wkt, geometry);

                            return geometry;
                        } else if (boost::istarts_with(wkt, "MULTILINESTRING")) {
                            geofeatures::MultiLineString geometry;

                            boost::geometry::read_wkt(wkt, geometry);

                            return geometry;
                        } else if (boost::istarts_with(wkt, "POLYGON")) {
                            geofeatures::Polygon geometry;

                            boost::geometry::read_wkt(wkt, geometry);

                            return geometry;
                        } else if (boost::istarts_with(wkt, "MULTIPOLYGON")) {
                            geofeatures::MultiPolygon geometry;

                            boost::geometry::read_wkt(wkt, geometry);

                            return geometry;
                        }
                        return GeometryCollection<>::value_type();
                    };

            std::regex words_regex("\\b(EMPTY|POINT|MULTIPOINT|LINESTRING|MULTILINESTRING|POLYGON|MULTIPOLYGON)", std::regex_constants::icase);

            std::string::const_iterator collectionBegin = wkt.begin();
            std::string::const_iterator collectionEnd   = wkt.end();

            long beginOffset = wkt.find_first_of("(");
            long endPosition = wkt.find_last_of(")");

            std::advance(collectionBegin, beginOffset);

            auto words_begin = std::sregex_iterator(collectionBegin, collectionEnd, words_regex);
            auto words_end = std::sregex_iterator();

            long beginPosition = beginOffset;;

            for (std::sregex_iterator i = words_begin; i != words_end; ++i) {

                std::smatch token = *i;

                if (beginPosition == beginOffset) {
                    if (boost::iequals(token.str(), "EMPTY")) {
                        continue;   // Nothing more to do, this
                    }
                } else {
                    long nextPosition = token.position() + beginOffset;

                    std::string str(wkt.substr(beginPosition, nextPosition - beginPosition));
                    boost::algorithm::erase_last(str, ",");

                    geometryCollection.push_back(geometry(str));
                }
                beginPosition = token.position() + beginOffset;
            }
            if (beginPosition != beginOffset) {  // We had at least 1 so get the last
                std::string str(wkt.substr(beginPosition, endPosition - beginPosition));

                geometryCollection.push_back(geometry(str));
            }
        }
    } // namespace io

}   // namespace geofeatures

#endif //__ReadWKT_HPP_


/*
 * test_dstar.cpp
 *
 * Test the Graph Library for Autonomous and Dynamic Systems
 *
 * author:  Arnaud Degroote <arnaud.degroote@laas.fr>
 * created: 2013-09-10
 * license: BSD
 */
#define BOOST_TEST_MODULE const_string test
#include <boost/test/included/unit_test.hpp>

#include <string>
#include <sstream>

#include "gdalwrap/gdal.hpp"
#include "gladys/gladys.hpp"
#include "gladys/nav_graph.hpp"
#include "gladys/dstar.hpp"

using namespace gladys;

BOOST_AUTO_TEST_SUITE( dstar )

BOOST_AUTO_TEST_CASE( test_dstar )
{
    std::string region_path = "/tmp/test_dstar_raster_to_graph.tif";
    std::string weight_path = "/tmp/test_dstar_raster_to_graph_nav.tif";
    std::string robotm_path = "/tmp/robot.json";
    std::string graphv_path = "/tmp/test_dstar_raster_to_graph_nav.dot";

    // create a robot model (JSON configuration file)
    std::ofstream robot_cfg(robotm_path);
    robot_cfg<<"{\"robot\":{\"mass\":1.0,\"radius\":1.0,\"velocity\":1.0}}";
    robot_cfg.close();

    // create a region map (GeoTiff image)
    gdalwrap::gdal region;
    region.set_size(4, 9, 9);
    // name bands
    region.names = {"NO_3D_CLASS", "FLAT", "OBSTACLE", "ROUGH"};
    // add an obstacle at the center of the map
    region.bands[1].assign(9*9, 1);
    for ( int i=1 ; i < 9 ; i++ ) {
        region.bands[1][i+5*9] = 0.2 ;
        region.bands[2][i+5*9] = 0.8 ;
    }
    region.save(region_path);

    // create a navigation graph from the map
    weight_map wm(region_path, robotm_path);
    nav_graph ng(wm);
    std::ostringstream oss_graphviz;
    ng.write_graphviz(oss_graphviz);
    ng.write_graphviz(graphv_path); // for debug
    ng.save(weight_path);

    BOOST_TEST_MESSAGE( "oss_graphviz.size() = " << oss_graphviz.str().size() );
    BOOST_CHECK_EQUAL( oss_graphviz.str().size() , 8335 );

    point_xy_t p1 = {1, 1};
    point_xy_t p2 = {5, 9};
    dstar_search dstar(ng.get_graph(), ng.get_closest_vertex(p1), ng.get_closest_vertex(p2));
    path_t path = dstar.get_path();
    BOOST_TEST_MESSAGE( "path: " + to_string(path) );
    BOOST_CHECK_EQUAL( path.size() , 14 );
    point_xy_t check_point = {1.0, 6.5} ;
    BOOST_CHECK_EQUAL( path[7][0], check_point[0] );
    BOOST_CHECK_EQUAL( path[7][1], check_point[1] );

    // replan must be instant, and do not change the plan
    dstar.replan(ng.get_closest_vertex(p1));
    path = dstar.get_path();
    BOOST_TEST_MESSAGE( "path: " + to_string(path) );
    BOOST_CHECK_EQUAL( path.size() , 14 );

    dstar.replan(ng.get_closest_vertex(path[2]));
    path = dstar.get_path();
    BOOST_TEST_MESSAGE( "path: " + to_string(path) );
    BOOST_CHECK_EQUAL( path.size() , 12 );

}

BOOST_AUTO_TEST_SUITE_END();