/*
 * detector.cc
 *
 *  Copyright (C) 2013 Diamond Light Source
 *
 *  Author: James Parkhurst
 *
 *  This code is distributed under the BSD license, a copy of which is
 *  included in the root directory of this package.
 */
#include <boost/python.hpp>
#include <boost/python/def.hpp>
#include <string>
#include <iostream>
#include <memory>
#include <sstream>
#include <boost_adaptbx/std_pair_conversion.h>
#include <scitbx/array_family/boost_python/flex_wrapper.h>
#include <scitbx/array_family/simple_io.h>
#include <scitbx/array_family/simple_tiny_io.h>
#include <scitbx/constants.h>
#include <dxtbx/model/detector.h>
#include <dxtbx/model/boost_python/to_from_dict.h>
#include <dxtbx/model/boost_python/pickle_suite.h>
#include <dxtbx/format/image.h>

namespace dxtbx { namespace model { namespace boost_python {
  template <>
  boost::python::dict to_dict<Detector::Node>(const Detector::Node &obj) {
    boost::python::dict result = to_dict<Panel>(obj);
    boost::python::list children;
    for (std::size_t i = 0; i < obj.size(); ++i) {
      Detector::const_node_pointer node = obj[i];
      if (node->is_panel()) {
        std::size_t idx = node->index();
        boost::python::dict panel;
        panel["panel"] = idx;
        children.append(panel);
      } else {
        children.append(to_dict(*node));
      }
    }
    result["children"] = children;
    return result;
  }

  namespace detector_detail {

    using format::Image;
    using scitbx::deg_as_rad;

    std::string detector_to_string(const Detector &detector) {
      std::stringstream ss;
      ss << detector;
      return ss.str();
    }

    static void detector_set_item(Detector &d, std::size_t i, const Panel &v) {
      d[i] = v;
    }

    static Panel *detector_get_item(Detector &d, std::size_t i) {
      return d.at(i);
    }

    static scitbx::af::shared<std::string> get_names(const Detector &d) {
      scitbx::af::shared<std::string> result(d.size());
      for (std::size_t i = 0; i < result.size(); ++i) {
        result[i] = d[i].get_name();
      }
      return result;
    }

    static void rotate_around_origin(Detector &detector,
                                     vec3<double> axis,
                                     double angle,
                                     bool deg) {
      double angle_rad = deg ? deg_as_rad(angle) : angle;
      detector.rotate_around_origin(axis, angle_rad);
    }

    void node_from_dict(boost::python::dict obj,
                        Detector::Node *node,
                        boost::python::list panels,
                        scitbx::af::ref<bool> used) {
      if (obj.contains("panel")) {
        std::size_t idx = boost::python::extract<std::size_t>(obj["panel"]);
        DXTBX_ASSERT(used.size() == boost::python::len(panels));
        DXTBX_ASSERT(idx < boost::python::len(panels));
        DXTBX_ASSERT(used[idx] == false);
        used[idx] = true;
        boost::python::dict panel_dict =
          boost::python::extract<boost::python::dict>(panels[idx]);
        Panel *panel = from_dict<Panel>(panel_dict);
        node->add_panel(*panel, idx);
        delete panel;
      } else {
        Panel *panel = from_dict<Panel>(obj);
        node = node->add_group(*panel);
        for (std::size_t i = 0; i < boost::python::len(obj["children"]); ++i) {
          boost::python::dict child =
            boost::python::extract<boost::python::dict>(obj["children"][i]);
          node_from_dict(child, node, panels, used);
        }
        delete panel;
      }
    }

    static Detector detector_deepcopy(const Detector &detector,
                                      boost::python::object dict) {
      return Detector(detector);
    }

    Detector *detector_from_dict(Detector *result, boost::python::dict obj) {
      boost::python::list panels =
        boost::python::extract<boost::python::list>(obj["panels"]);
      if (obj.contains("hierarchy")) {
        boost::python::dict hierarchy =
          boost::python::extract<boost::python::dict>(obj["hierarchy"]);
        scitbx::af::shared<bool> used(boost::python::len(panels), false);
        DXTBX_ASSERT(!hierarchy.contains("panel"));
        Panel *panel = from_dict<Panel>(hierarchy);
        std::swap(*((Panel *)result->root()), *panel);
        for (std::size_t i = 0; i < boost::python::len(hierarchy["children"]); ++i) {
          boost::python::dict child =
            boost::python::extract<boost::python::dict>(hierarchy["children"][i]);
          node_from_dict(child, result->root(), panels, used.ref());
        }
        delete panel;
        for (std::size_t i = 0; i < used.size(); ++i) {
          DXTBX_ASSERT(used[i] == true);
        }
      } else {
        for (std::size_t i = 0; i < boost::python::len(panels); ++i) {
          boost::python::dict panel_dict =
            boost::python::extract<boost::python::dict>(panels[i]);
          Panel *panel = from_dict<Panel>(panel_dict);
          result->add_panel(*panel);
          delete panel;
        }
      }
      return result;
    }

    Detector *detector_from_dict(Detector *result,
                                 boost::python::dict obj,
                                 const Image<double> &dx,
                                 const Image<double> &dy) {
      boost::python::list panels =
        boost::python::extract<boost::python::list>(obj["panels"]);
      if (obj.contains("hierarchy")) {
        boost::python::dict hierarchy =
          boost::python::extract<boost::python::dict>(obj["hierarchy"]);
        scitbx::af::shared<bool> used(boost::python::len(panels), false);
        DXTBX_ASSERT(!hierarchy.contains("panel"));
        Panel *panel = from_dict<Panel>(hierarchy);
        std::swap(*((Panel *)result->root()), *panel);
        for (std::size_t i = 0; i < boost::python::len(hierarchy["children"]); ++i) {
          boost::python::dict child =
            boost::python::extract<boost::python::dict>(hierarchy["children"][i]);
          node_from_dict(child, result->root(), panels, used.ref());
        }
        delete panel;
        for (std::size_t i = 0; i < used.size(); ++i) {
          DXTBX_ASSERT(used[i] == true);
        }
      } else {
        for (std::size_t i = 0; i < boost::python::len(panels); ++i) {
          boost::python::dict panel_dict =
            boost::python::extract<boost::python::dict>(panels[i]);
          Panel *panel = panel_from_dict_with_offset(
            panel_dict, dx.tile(i).data(), dy.tile(i).data());
          result->add_panel(*panel);
          delete panel;
        }
      }
      return result;
    }

    struct DetectorPickleSuite : boost::python::pickle_suite {
      static boost::python::tuple getstate(boost::python::object obj) {
        unsigned int version = 3;
        const Detector &detector = boost::python::extract<const Detector &>(obj);
        boost::python::dict data;

        // Convert panel array into python list of panels
        boost::python::list panels;
        for (std::size_t i = 0; i < detector.size(); ++i) {
          panels.append(detector[i]);
        }
        data["panels"] = panels;

        // Convert hierarchy into dict
        data["hierarchy"] = to_dict(*detector.root());

        // Return the tuple
        return boost::python::make_tuple(version, data);
      }

      static void setstate(boost::python::object obj, boost::python::tuple state) {
        Detector *detector = boost::python::extract<Detector *>(obj);
        DXTBX_ASSERT(len(state) == 2);
        unsigned int version = boost::python::extract<unsigned int>(state[0]);
        DXTBX_ASSERT(version == 3);

        // Extract data from state object
        boost::python::dict data =
          boost::python::extract<boost::python::dict>(state[1]);
        boost::python::list panels =
          boost::python::extract<boost::python::list>(data["panels"]);
        boost::python::dict hierarchy =
          boost::python::extract<boost::python::dict>(data["hierarchy"]);

        DXTBX_ASSERT(!hierarchy.contains("panel"));
        Panel *panel = from_dict<Panel>(hierarchy);
        std::swap(*((Panel *)detector->root()), *panel);
        copy_node(detector->root(), hierarchy, panels);
        delete panel;

        for (std::size_t i = 0; i < detector->size(); ++i) {
          DXTBX_ASSERT(detector->at(i) != NULL);
        }
      }

      static void copy_node(Detector::Node::pointer self,
                            boost::python::dict node,
                            boost::python::list panels) {
        for (std::size_t i = 0; i < boost::python::len(node["children"]); ++i) {
          boost::python::dict child =
            boost::python::extract<boost::python::dict>(node["children"][i]);
          if (child.contains("panel")) {
            std::size_t index = boost::python::extract<std::size_t>(child["panel"]);
            Panel panel = boost::python::extract<Panel>(panels[index]);
            self->add_panel(panel, index);
          } else {
            Panel *panel = from_dict<Panel>(child);
            copy_node(self->add_group(*panel), child, panels);
            delete panel;
          }
        }
      }

      static bool getstate_manages_dict() {
        return true;
      }
    };
  }  // namespace detector_detail

  Detector *detector_from_dict_with_offset(boost::python::dict obj,
                                           const Image<double> &dx,
                                           const Image<double> &dy) {
    Detector *result = new Detector();
    return detector_detail::detector_from_dict(result, obj, dx, dy);
  }

  template <>
  boost::python::dict to_dict<Detector>(const Detector &obj) {
    boost::python::dict result;
    boost::python::list panels;
    for (std::size_t i = 0; i < obj.size(); ++i) {
      panels.append(to_dict(obj[i]));
    }
    result["panels"] = panels;
    result["hierarchy"] = to_dict(*obj.root());
    return result;
  }

  template <>
  Detector *from_dict<Detector>(boost::python::dict obj) {
    Detector *result = new Detector();
    return detector_detail::detector_from_dict(result, obj);
  }

  boost::python::tuple get_ray_intersection_py(const Detector &self,
                                               scitbx::af::const_ref<vec3<double>> s1) {
    /**
     * Wrapper to convert coord_type to Python Tuple
     */

    scitbx::af::shared<Detector::coord_type> result = self.get_ray_intersection(s1);

    scitbx::af::shared<vec2<double>> xy =
      scitbx::af::shared<vec2<double>>(result.size());
    scitbx::af::shared<std::size_t> panel =
      scitbx::af::shared<std::size_t>(result.size());

    for (std::size_t i = 0; i < result.size(); ++i) {
      panel[i] = result[i].first;
      xy[i] = result[i].second;
    }
    return boost::python::make_tuple(panel, xy);
  }

  void export_detector() {
    using namespace boost::python;
    using namespace detector_detail;

    class_<Detector::Node, bases<Panel>>("DetectorNode", no_init)
      .def("add_group",
           (Detector::Node::pointer (Detector::Node::*)())&Detector::Node::add_group,
           return_internal_reference<>())
      .def("add_group",
           (Detector::Node::pointer (Detector::Node::*)(
             const Panel &))&Detector::Node::add_group,
           return_internal_reference<>())
      .def("add_panel",
           (Detector::Node::pointer (Detector::Node::*)())&Detector::Node::add_panel,
           return_internal_reference<>())
      .def("add_panel",
           (Detector::Node::pointer (Detector::Node::*)(
             const Panel &))&Detector::Node::add_panel,
           return_internal_reference<>())
      .def("parent",
           (Detector::Node::pointer (Detector::Node::*)())&Detector::Node::parent,
           return_internal_reference<>())
      .def("root",
           (Detector::Node::pointer (Detector::Node::*)())&Detector::Node::root,
           return_internal_reference<>())
      .def("__getitem__",
           (Detector::Node::pointer (Detector::Node::*)(std::size_t))&Detector::Node::
           operator[],
           return_internal_reference<>())
      .def("empty", &Detector::Node::empty)
      .def("__len__", &Detector::Node::size)
      .def("index", &Detector::Node::index)
      .def("is_panel", &Detector::Node::is_panel)
      .def("is_group", &Detector::Node::is_group)
      .def("set_frame",
           &Detector::Node::set_frame,
           (arg("fast_axis"), arg("slow_axis"), arg("origin")))
      .def("set_local_frame",
           &Detector::Node::set_local_frame,
           (arg("fast_axis"), arg("slow_axis"), arg("origin")))
      .def("set_parent_frame",
           &Detector::Node::set_parent_frame,
           (arg("fast_axis"), arg("slow_axis"), arg("origin")))
      .def("__eq__", &Detector::Node::operator==)
      .def("__ne__", &Detector::Node::operator!=)
      .def("is_similar_to",
           &Detector::Node::is_similar_to,
           (arg("rhs"),
            arg("fast_axis_tolerance") = 1e-6,
            arg("slow_axis_tolerance") = 1e-6,
            arg("origin_tolerance") = 1e-6,
            arg("static_only") = false,
            arg("ignore_trusted_range") = false))
      .def("__iter__", iterator<Detector::Node, return_internal_reference<>>())
      .def("children", iterator<Detector::Node, return_internal_reference<>>());

    // Export a Detector base class
    class_<Detector, std::shared_ptr<Detector>>("Detector")
      .def(init<const Panel &>())
      .def("hierarchy",
           (Detector::node_pointer (Detector::*)())&Detector::root,
           return_internal_reference<>())
      .def("add_group",
           (Detector::node_pointer (Detector::*)())&Detector::add_group,
           return_internal_reference<>())
      .def("add_group",
           (Detector::node_pointer (Detector::*)(const Panel &))&Detector::add_group,
           return_internal_reference<>())
      .def("add_panel",
           (Detector::node_pointer (Detector::*)())&Detector::add_panel,
           return_internal_reference<>())
      .def("add_panel",
           (Detector::node_pointer (Detector::*)(const Panel &))&Detector::add_panel,
           return_internal_reference<>())
      .def("__len__", &Detector::size)
      .def("__setitem__", &detector_set_item)
      .def("__getitem__", &detector_get_item, return_internal_reference<>())
      .def("__iter__", iterator<Detector, return_internal_reference<>>())
      .def("__eq__", &Detector::operator==)
      .def("__ne__", &Detector::operator!=)
      .def("is_similar_to",
           &Detector::is_similar_to,
           (arg("other"),
            arg("fast_axis_tolerance") = 1e-6,
            arg("slow_axis_tolerance") = 1e-6,
            arg("origin_tolerance") = 1e-6,
            arg("static_only") = false,
            arg("ignore_trusted_range") = false))
      .def("get_max_resolution", &Detector::get_max_resolution, (arg("s0")))
      .def("get_max_inscribed_resolution",
           &Detector::get_max_inscribed_resolution,
           (arg("s0")))
      .def("get_ray_intersection",
           (Detector::coord_type (Detector::*)(vec3<double>) const)
             & Detector::get_ray_intersection,
           (arg("s1")))
      .def("get_ray_intersection", &get_ray_intersection_py, (arg("s1")))
      .def("get_ray_intersection",
           (scitbx::af::shared<vec2<double>> (Detector::*)(
             scitbx::af::const_ref<vec3<double>>, scitbx::af::const_ref<std::size_t>)
              const)
             & Detector::get_ray_intersection,
           (arg("s1"), arg("panel")))
      .def("get_panel_intersection", &Detector::get_panel_intersection, (arg("s1")))
      //.def("do_panels_intersect",
      //  &Detector::do_panels_intersect)
      .def("get_names", &get_names)
      .def("rotate_around_origin",
           &rotate_around_origin,
           (arg("axis"), arg("angle"), arg("deg") = true))
      .def("has_projection_2d", &Detector::has_projection_2d)
      .def("__str__", &detector_to_string)
      .def("__deepcopy__", &detector_deepcopy)
      .def("__copy__", &detector_deepcopy)
      .def("to_dict", &to_dict<Detector>)
      .def("from_dict", &from_dict<Detector>, return_value_policy<manage_new_object>())
      .def("from_dict",
           &detector_from_dict_with_offset,
           return_value_policy<manage_new_object>())
      .staticmethod("from_dict")
      .def_pickle(DetectorPickleSuite());

    boost_adaptbx::std_pair_conversions::to_and_from_tuple<int, vec2<double>>();
  }

}}}  // namespace dxtbx::model::boost_python
