#include <boost/python.hpp>
#include <boost/python/def.hpp>
#include <scitbx/array_family/shared.h>
#include <scitbx/array_family/flex_types.h>
#include <scitbx/constants.h>
#include <scitbx/vec2.h>
#include <boost_adaptbx/python_streambuf.h>
#include <cstdint>
#include <cctype>
#include <fstream>
#include <vector>
#include <limits>
#include <dxtbx/error.h>
#include "compression.h"

namespace dxtbx { namespace boost_python {

  using boost::uint32_t;

  bool is_big_endian() {
    uint32_t val = 1;
    char *buff = (char *)&val;

    return (buff[0] == 0);
  }

  scitbx::af::shared<int> read_uint8(boost_adaptbx::python::streambuf &input,
                                     size_t count) {
    scitbx::af::shared<int> result;
    boost_adaptbx::python::streambuf::istream is(input);
    std::vector<unsigned char> data;
    data.resize(count);

    is.read((char *)&data[0], count * sizeof(unsigned char));

    for (size_t j = 0; j < count; j++) {
      result.push_back((int)data[j]);
    }

    return result;
  }

  scitbx::af::shared<int> read_uint16(boost_adaptbx::python::streambuf &input,
                                      size_t count) {
    scitbx::af::shared<int> result;
    boost_adaptbx::python::streambuf::istream is(input);
    std::vector<unsigned short> data;
    data.resize(count);

    is.read((char *)&data[0], count * sizeof(unsigned short));

    for (size_t j = 0; j < count; j++) {
      result.push_back((int)data[j]);
    }

    return result;
  }

  scitbx::af::shared<int> read_uint32(boost_adaptbx::python::streambuf &input,
                                      size_t count) {
    scitbx::af::shared<int> result;
    boost_adaptbx::python::streambuf::istream is(input);
    std::vector<unsigned int> data;
    data.resize(count);

    is.read((char *)&data[0], count * sizeof(unsigned int));

    for (size_t j = 0; j < count; j++) {
      DXTBX_ASSERT(data[j] <= std::numeric_limits<int>::max());
      result.push_back((int)data[j]);
    }

    return result;
  }

  scitbx::af::shared<int> read_uint16_bs(boost_adaptbx::python::streambuf &input,
                                         size_t count) {
    scitbx::af::shared<int> result;
    boost_adaptbx::python::streambuf::istream is(input);
    std::vector<unsigned short> data;
    data.resize(count);

    is.read((char *)&data[0], count * sizeof(unsigned short));

    /* swap bytes */

    for (size_t j = 0; j < count; j++) {
      unsigned short x = data[j];
      data[j] = x >> 8 | x << 8;
    }

    for (size_t j = 0; j < count; j++) {
      result.push_back((int)data[j]);
    }

    return result;
  }

  scitbx::af::shared<int> read_uint32_bs(boost_adaptbx::python::streambuf &input,
                                         size_t count) {
    scitbx::af::shared<int> result;
    boost_adaptbx::python::streambuf::istream is(input);
    std::vector<unsigned int> data;
    data.resize(count);

    is.read((char *)&data[0], count * sizeof(unsigned int));

    /* swap bytes */

    for (size_t j = 0; j < count; j++) {
      unsigned int x = data[j];
      data[j] = (x << 24) | (x << 8 & 0xff0000) | (x >> 8 & 0xff00) | (x >> 24);
    }

    for (size_t j = 0; j < count; j++) {
      DXTBX_ASSERT(data[j] <= std::numeric_limits<int>::max());
      result.push_back((int)data[j]);
    }

    return result;
  }

  scitbx::af::shared<int> read_int16(boost_adaptbx::python::streambuf &input,
                                     size_t count) {
    scitbx::af::shared<int> result;
    boost_adaptbx::python::streambuf::istream is(input);
    std::vector<short> data;
    data.resize(count);

    is.read((char *)&data[0], count * sizeof(short));

    for (size_t j = 0; j < count; j++) {
      result.push_back((int)data[j]);
    }

    return result;
  }

  scitbx::af::shared<int> read_int32(boost_adaptbx::python::streambuf &input,
                                     size_t count) {
    scitbx::af::shared<int> result;
    boost_adaptbx::python::streambuf::istream is(input);
    std::vector<int> data;
    data.resize(count);

    is.read((char *)&data[0], count * sizeof(int));

    for (size_t j = 0; j < count; j++) {
      result.push_back((int)data[j]);
    }

    return result;
  }

  scitbx::af::shared<double> read_float32(boost_adaptbx::python::streambuf &input,
                                          size_t count) {
    scitbx::af::shared<double> result;
    boost_adaptbx::python::streambuf::istream is(input);
    std::vector<float> data;
    data.resize(count);

    is.read((char *)&data[0], count * sizeof(float));

    for (size_t j = 0; j < count; j++) {
      result.push_back((double)data[j]);
    }

    return result;
  }

  scitbx::af::flex_int uncompress(const boost::python::object &packed,
                                  const int &slow,
                                  const int &fast) {
    std::string strpacked = boost::python::extract<std::string>(packed);
    std::size_t sz_buffer = strpacked.size();

    scitbx::af::flex_int z((scitbx::af::flex_grid<>(slow, fast)),
                           scitbx::af::init_functor_null<int>());
    int *begin = z.begin();

    unsigned int nn = dxtbx::boost_python::cbf_decompress(
      strpacked.c_str(), sz_buffer, begin, slow * fast);

    DXTBX_ASSERT(nn == (slow * fast));

    return z;
  }

  PyObject *compress(const scitbx::af::flex_int z) {
    const int *begin = z.begin();
    std::size_t sz = z.size();

    std::vector<char> packed = dxtbx::boost_python::cbf_compress(begin, sz);

    return PyBytes_FromStringAndSize(&*packed.begin(), packed.size());
  }

  double distance_between_points(scitbx::vec2<int> const &a,
                                 scitbx::vec2<int> const &b) {
    return std::sqrt(
      (std::pow(double(b[0] - a[0]), 2) + std::pow(double(b[1] - a[1]), 2)));
  }

  void radial_average(scitbx::af::versa<double, scitbx::af::flex_grid<> > &data,
                      scitbx::af::versa<bool, scitbx::af::flex_grid<> > &mask,
                      scitbx::vec2<int> const &beam_center,
                      scitbx::af::shared<double> sums,
                      scitbx::af::shared<double> sums_sq,
                      scitbx::af::shared<int> counts,
                      double pixel_size,
                      double distance,
                      scitbx::vec2<int> const &upper_left,
                      scitbx::vec2<int> const &lower_right) {
    std::size_t extent = sums.size();
    double extent_in_mm = extent * pixel_size;
    double extent_two_theta =
      std::atan(extent_in_mm / distance) * 180 / scitbx::constants::pi;

    for (std::size_t y = upper_left[1]; y < lower_right[1]; y++) {
      for (std::size_t x = upper_left[0]; x < lower_right[0]; x++) {
        double val = data(x, y);
        if (val > 0 && mask(x, y)) {
          scitbx::vec2<int> point((int)x, (int)y);
          double d_in_mm = distance_between_points(point, beam_center) * pixel_size;
          double twotheta = std::atan(d_in_mm / distance) * 180 / scitbx::constants::pi;
          std::size_t bin =
            (std::size_t)std::floor(twotheta * extent / extent_two_theta);
          if (bin >= extent) continue;
          sums[bin] += val;
          sums_sq[bin] += val * val;
          counts[bin]++;
        }
      }
    }
  }

  // Python entry point to decompress Rigaku Oxford Diffractometer TY6 compression
  scitbx::af::flex_int uncompress_rod_TY6(const boost::python::object &data,
                                          const boost::python::object &offsets,
                                          const int &slow,
                                          const int &fast) {
    // Cannot I extract const char* directly?
    std::string str_data = boost::python::extract<std::string>(data);
    std::string str_offsets = boost::python::extract<std::string>(offsets);

    scitbx::af::flex_int z((scitbx::af::flex_grid<>(slow, fast)),
                           scitbx::af::init_functor_null<int>());

    dxtbx::boost_python::rod_TY6_decompress(
      z.begin(), str_data.c_str(), str_offsets.c_str(), slow, fast);

    return z;
  }

  void init_module() {
    using namespace boost::python;
    def("read_uint8", read_uint8, (arg("file"), arg("count")));
    def("read_uint16", read_uint16, (arg("file"), arg("count")));
    def("read_uint32", read_uint32, (arg("file"), arg("count")));
    def("read_uint16_bs", read_uint16_bs, (arg("file"), arg("count")));
    def("read_uint32_bs", read_uint32_bs, (arg("file"), arg("count")));
    def("read_int16", read_int16, (arg("file"), arg("count")));
    def("read_int32", read_int32, (arg("file"), arg("count")));
    def("read_float32", read_float32, (arg("file"), arg("count")));
    def("is_big_endian", is_big_endian);
    def("uncompress", &uncompress, (arg_("packed"), arg_("slow"), arg_("fast")));
    def("compress", &compress);
    def("radial_average",
        &radial_average,
        (arg("data"),
         arg("beam_center"),
         arg("sums"),
         arg("sums_sq"),
         arg("counts"),
         arg("pixel_size"),
         arg("distance"),
         arg("upper_left"),
         arg("lower_right")));
    def("uncompress_rod_TY6",
        &uncompress_rod_TY6,
        (arg_("data"), arg_("offsets"), arg_("slow"), arg_("fast")));
  }

  BOOST_PYTHON_MODULE(dxtbx_ext) {
    init_module();
  }

}}  // namespace dxtbx::boost_python
