/*
 * partiality_calculator.h
 *
 *  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.
 */
#ifndef DIALS_ALGORITHMS_PROFILE_MODEL_GAUSSIAN_RS_PARTIALITY_CALCULATOR_H
#define DIALS_ALGORITHMS_PROFILE_MODEL_GAUSSIAN_RS_PARTIALITY_CALCULATOR_H

#include <cmath>
#include <boost/shared_ptr.hpp>
#include <boost/math/special_functions/erf.hpp>
#include <scitbx/constants.h>
#include <scitbx/vec3.h>
#include <scitbx/array_family/tiny_types.h>
#include <scitbx/array_family/ref_reductions.h>
#include <dxtbx/model/beam.h>
#include <dxtbx/model/detector.h>
#include <dxtbx/model/goniometer.h>
#include <dxtbx/model/scan.h>
#include <dials/array_family/scitbx_shared_and_versa.h>
#include <dials/algorithms/profile_model/gaussian_rs/coordinate_system.h>

namespace dials {
  namespace algorithms {
    namespace profile_model {
      namespace gaussian_rs {

  // Use a load of stuff from other namespaces
  using boost::math::erf;
  using dxtbx::model::BeamBase;
  using dxtbx::model::Detector;
  using dxtbx::model::Goniometer;
  using dxtbx::model::Scan;
  using scitbx::vec2;
  using scitbx::vec3;
  using scitbx::af::double2;
  using scitbx::af::double3;
  using scitbx::af::double4;
  using scitbx::af::int6;
  using scitbx::af::max;
  using scitbx::af::min;
  using std::ceil;
  using std::floor;

  /**
   * Interface for bounding box calculator.
   */
  class PartialityCalculatorIface {
  public:
    virtual ~PartialityCalculatorIface() {}

    virtual double single(vec3<double> s1, double frame, int6 bbox) const = 0;

    virtual af::shared<double> array(const af::const_ref<vec3<double> > &s1,
                                     const af::const_ref<double> &frame,
                                     const af::const_ref<int6> &bbox) const = 0;
  };

  /** Calculate the partiality for each reflection */
  class PartialityCalculator3D : public PartialityCalculatorIface {
  public:
    /**
     * Initialise the partiality calculation.
     * @param beam The beam parameters
     * @param goniometer The goniometer parameters
     * @param scan The scan parameters
     * @param sigma_mosaicity The xds sigma_mosaicity parameter
     */
    PartialityCalculator3D(const BeamBase &beam,
                           const Goniometer &gonio,
                           const Scan &scan,
                           double sigma_m)
        : s0_(beam.get_s0()),
          m2_(gonio.get_rotation_axis()),
          scan_(scan),
          sigma_m_(1, sigma_m) {
      DIALS_ASSERT(sigma_m > 0.0);
    }

    /**
     * Initialise the partiality calculation.
     * @param beam The beam parameters
     * @param goniometer The goniometer parameters
     * @param scan The scan parameters
     * @param sigma_mosaicity The xds sigma_mosaicity parameter
     */
    PartialityCalculator3D(const BeamBase &beam,
                           const Goniometer &gonio,
                           const Scan &scan,
                           const af::const_ref<double> &sigma_m)
        : s0_(beam.get_s0()),
          m2_(gonio.get_rotation_axis()),
          scan_(scan),
          sigma_m_(sigma_m.begin(), sigma_m.end()) {
      DIALS_ASSERT(sigma_m.all_gt(0.0));
      DIALS_ASSERT(sigma_m.size() == scan.get_num_images());
      DIALS_ASSERT(sigma_m.size() > 0);
    }

    /**
     * Calculate the Partiality of the reflection.
     *
     * @param s1 The diffracted beam vector
     * @param frame The frame number
     * @param bbox The bounding box
     * @returns The partiality as a fraction of the total phi extent
     */
    virtual double single(vec3<double> s1, double frame, int6 bbox) const {
      // Ensure our values are ok
      DIALS_ASSERT(s1.length_sq() > 0);
      DIALS_ASSERT(bbox[4] < bbox[5]);

      // Get the mosaicity at this point
      double sigma_m = 0.0;
      if (sigma_m_.size() == 1) {
        sigma_m = sigma_m_[0];
      } else {
        int frame0 = scan_.get_array_range()[0];
        int index = (int)std::floor(frame) - frame0;
        if (index < 0) {
          sigma_m = sigma_m_.front();
        } else if (index >= sigma_m_.size()) {
          sigma_m = sigma_m_.back();
        } else {
          sigma_m = sigma_m_[index];
        }
      }

      // Get the rotation angle
      double phi = scan_.get_angle_from_array_index(frame);
      double phia = scan_.get_angle_from_array_index(bbox[4]);
      double phib = scan_.get_angle_from_array_index(bbox[5]);

      // Compute the partiality
      double zeta = profile_model::gaussian_rs::zeta_factor(m2_.normalize(), s0_, s1);
      double c = std::abs(zeta) / (sqrt(2.0) * sigma_m);
      double p = 0.5 * (erf(c * (phib - phi)) - erf(c * (phia - phi)));
      DIALS_ASSERT(p >= 0.0 && p <= 1.0);
      return p;
    }

    /**
     * Calculate the partiality for an array of reflections
     * @param s1 The array of diffracted beam vectors
     * @param frame The array of frame numbers.
     * @param bbox The array of bboxes
     */
    virtual af::shared<double> array(const af::const_ref<vec3<double> > &s1,
                                     const af::const_ref<double> &frame,
                                     const af::const_ref<int6> &bbox) const {
      DIALS_ASSERT(s1.size() == frame.size());
      DIALS_ASSERT(s1.size() == bbox.size());
      af::shared<double> result(s1.size(), af::init_functor_null<double>());
      for (std::size_t i = 0; i < s1.size(); ++i) {
        result[i] = single(s1[i], frame[i], bbox[i]);
      }
      return result;
    }

  private:
    vec3<double> s0_;
    vec3<double> m2_;
    Scan scan_;
    af::shared<double> sigma_m_;
  };

  /** Calculate the partiality for each reflection */
  class PartialityCalculator2D : public PartialityCalculatorIface {
  public:
    /**
     * Initialise the partiality calculation.
     * @param beam The beam parameters
     * @param sigma_mosaicity The xds sigma_mosaicity parameter
     */
    PartialityCalculator2D(const BeamBase &beam, double sigma_m) : s0_(beam.get_s0()) {}

    /**
     * Calculate the Partiality of the reflection.
     *
     * @param s1 The diffracted beam vector
     * @param frame The frame number
     * @param bbox The bounding box
     * @returns The partiality as a fraction of the total phi extent
     */
    virtual double single(vec3<double> s1, double frame, int6 bbox) const {
      // Ensure our values are ok
      DIALS_ASSERT(s1.length_sq() > 0);
      DIALS_ASSERT(bbox[4] < bbox[5]);

      // FIXME This is a placeholder

      // Return the fraction recorded
      return 1.0;
    }

    /**
     * Calculate the partiality for an array of reflections
     * @param s1 The array of diffracted beam vectors
     * @param frame The array of frame numbers.
     * @param bbox The array of bboxes
     */
    virtual af::shared<double> array(const af::const_ref<vec3<double> > &s1,
                                     const af::const_ref<double> &frame,
                                     const af::const_ref<int6> &bbox) const {
      DIALS_ASSERT(s1.size() == frame.size());
      DIALS_ASSERT(s1.size() == bbox.size());
      af::shared<double> result(s1.size(), af::init_functor_null<double>());
      for (std::size_t i = 0; i < s1.size(); ++i) {
        result[i] = single(s1[i], frame[i], bbox[i]);
      }
      return result;
    }

  private:
    vec3<double> s0_;
    vec3<double> m2_;
    Scan scan_;
  };

  /**
   * Class to help compute partiality for multiple experiments.
   */
  class PartialityMultiCalculator {
  public:
    /**
     * Add a bbox calculator to the list.
     */
    void push_back(boost::shared_ptr<PartialityCalculatorIface> obj) {
      compute_.push_back(obj);
    }

    /**
     * Get the number of calculators.
     */
    std::size_t size() const {
      return compute_.size();
    }

    /**
     * Calculate the rois for an array of reflections given by the array of
     * diffracted beam vectors and rotation angles.
     * @param id The experiment id
     * @param s1 The array of diffracted beam vectors
     * @param frame The array of frame numbers.
     * @param bbox The array of bounding boxes
     */
    af::shared<double> operator()(const af::const_ref<std::size_t> &id,
                                  const af::const_ref<vec3<double> > &s1,
                                  const af::const_ref<double> &frame,
                                  const af::const_ref<int6> &bbox) const {
      DIALS_ASSERT(s1.size() == id.size());
      DIALS_ASSERT(s1.size() == frame.size());
      DIALS_ASSERT(s1.size() == bbox.size());
      af::shared<double> result(s1.size(), af::init_functor_null<double>());
      for (std::size_t i = 0; i < s1.size(); ++i) {
        DIALS_ASSERT(id[i] < size());
        result[i] = compute_[id[i]]->single(s1[i], frame[i], bbox[i]);
      }
      return result;
    }

  private:
    std::vector<boost::shared_ptr<PartialityCalculatorIface> > compute_;
  };

}}}}  // namespace dials::algorithms::profile_model::gaussian_rs

#endif  // DIALS_ALGORITHMS_PROFILE_MODEL_GAUSSIAN_RS_PARTIALITY_CALCULATOR_H
