#pragma once

#include <map>
#include <random>
#include <vector>

#include "Libs/Optimize/Container/GenericContainer.h"
#include "Libs/Optimize/Domain/ParticleDomain.h"
#include "Libs/Optimize/Neighborhood/ParticleNeighborhood.h"
#include "Observer.h"
#include "ParticleEvents.h"
#include "itkCommand.h"
#include "itkDataObject.h"
#include "itkEventObject.h"
#include "itkObjectFactory.h"
#include "itkPoint.h"
#include "itkWeakPointer.h"
#include "vnl/vnl_inverse.h"
#include "vnl/vnl_matrix_fixed.h"
#include "vnl/vnl_vector_fixed.h"

namespace shapeworks {
/** \class ParticleSystem
 *  \brief  A facade class managing interactions with a particle system.
 *
 * A particle system consists of a set of particle locations and domains in
 * which those locations are defined.  For example, a particle system may
 * simply be a set of points in 3D Euclidean space that are constrained to a
 * specified bounding box.  A more complex example is a set of 3D points
 * constrained to a given surface.  The itkParticleSystem class is also
 * designed to manage multiple sets of points across multiple domains.  For
 * example, one may create a system of dozens of point sets that are each
 * defined in different spaces.  The itkParticleSystem class can also maintain
 * a set of transformations associated with each domain that map that
 * particular domain into another coordinate frame, for example, a common
 * coordinate frame.
 */
class ParticleSystem : public itk::DataObject {
 public:
  static constexpr int VDimension = 3;

  /** Standard class typedefs */
  typedef ParticleSystem Self;
  typedef DataObject Superclass;
  typedef itk::SmartPointer<Self> Pointer;
  typedef itk::SmartPointer<const Self> ConstPointer;
  typedef itk::WeakPointer<const Self> ConstWeakPointer;

  /** Method for creation through the object factory. */
  itkNewMacro(Self);

  /** Run-time type information (and related methods). */
  itkTypeMacro(ParticleSystem, itk::DataObject);

  /** Dimensionality of the domain of the particle system. */
  itkStaticConstMacro(Dimension, unsigned int, VDimension);

  /** Define the base domain type. */
  using DomainType = shapeworks::ParticleDomain;

  /** Point type used to store particle locations. */
  typedef itk::Point<double, VDimension> PointType;

  /** Point container type.  One is associated with each domain.*/
  typedef GenericContainer<PointType> PointContainerType;

  /** Defines a transform class type.  One is associated with each
      domain. Transforms are affine for simplicity.  This could be extended by
      using the itk::Transform base type so that a variety of transform types
      may be used. */
  //  typedef Transform<double, VDimension, VDimension> TransformType;
  typedef vnl_matrix_fixed<double, VDimension + 1, VDimension + 1> TransformType;
  typedef vnl_vector_fixed<double, VDimension> VectorType;
  typedef vnl_matrix<double> VnlMatrixType;

  /** Register an attribute object with this particle system.  This action adds
      the attribute object as an observer for the particle system so that it
      will receive published updates.  For example, a gradient attribute object
      may need to know when particle positions have changed so that it can
      update its internal state appropriately (e.g. precompute a gradient that
      corresponds to the new particle position).  Internally, this method
      manages registering the appropriate itk::Command and itk::Object as an
      observer of this itk::Object (AddObserver).  Relevant state information
      is passed to objects in particle events.
  */
  void RegisterObserver(Observer *attr);

  /** Invokes the set event on all particle positions, resetting them to their
      current value.  This method may be called to synchronize positional
      information among various observers which may have gone out of sync. */
  void SynchronizePositions() {
    for (unsigned int d = 0; d < this->GetNumberOfDomains(); d++) {
      for (unsigned int p = 0; p < this->GetNumberOfParticles(d); p++) {
        this->SetPosition(this->GetPosition(p, d), p, d);
      }
    }
  }

  /** Returns the number of particles in domain k. */
  unsigned long int GetNumberOfParticles(unsigned int d = 0) const { return m_Positions[d]->GetSize(); }

  /**  Add/Set/Remove a single particle position.  The actual position added or
      set will be returned.  If, for example, the domain imposes any
      constraints on this location (e.g. out-of-bounds, projection onto a
      surface), then the method returns the location after constraints are
      applied.  Both the Add and Set methods require a domain index.  If
      unspecified, the domain index is zero.  The Set method requires a
      specific particle index.  If the given index does not exist, an exception
      is thrown.  Set is intended to be used for moving particles. The
      SetTransformedPosition sets the position using a Point transformed by the
      m_Transform associated with the given domain.*/
  const PointType &AddPosition(const PointType &, unsigned int d = 0);
  const PointType &SetPosition(const PointType &, unsigned long int k, unsigned int d = 0);

  void RemovePosition(unsigned long int k, unsigned int d = 0);

  /** Return a position with index k from domain d.  Note the order in which the 2
      integers must be specified!   The domain number is specified second and
      the position index within the domain is specified first.  Note that the
      domain index may be omitted if the particle system only contains a single
      domain.  GetTransformedPosition returns the position premultiplied by the
      transform matrix for the given domain.*/
  PointType &GetPosition(unsigned long int k, unsigned int d = 0) { return m_Positions[d]->operator[](k); }
  const PointType &GetPosition(unsigned long int k, unsigned int d = 0) const { return m_Positions[d]->operator[](k); }
  PointType GetTransformedPosition(unsigned long int k, unsigned int d = 0) const {
    return this->TransformPoint(m_Positions[d]->operator[](k), m_Transforms[d] * m_PrefixTransforms[d]);
  }
  PointType GetPrefixTransformedPosition(unsigned long int k, unsigned int d = 0) const {
    return this->TransformPoint(m_Positions[d]->operator[](k), m_PrefixTransforms[d]);
  }

  /** Doubles the number of particles of the system by
      splitting each particle into 2 particles.  Each new particle position is
      added to the system at a random epsilon distance on the surface from the
      original particle.  The new particles are added using the AddPosition
      method, so all appropriate callbacks are invoked. */

  void SplitAllParticles(double epsilon);
  void SplitParticle(double epsilon, unsigned int idx, unsigned int d = 0);
  void AdvancedAllParticleSplitting(double epsilon, unsigned int domains_per_shape, unsigned int dom_to_process);
  // Debug function
  void PrintParticleSystem();

  //! Get the neighborhood object associated with domain k
  std::shared_ptr<ParticleNeighborhood> GetNeighborhood(unsigned int k) const { return m_Neighborhoods[k]; }

  using PointVectorType = std::vector<ParticlePointIndexPair>;

  /** Return the neighborhood of points with radius r around point p in domain
      k.  This is just a convenience method to avoid exposing the underlying
      Neighborhood objects. FindTransformedNeighborhoodPoints returns the list
      with all points transformed by the transform associated with the given
      domain.*/
  PointVectorType FindNeighborhoodPoints(const PointType &p, int idx, double r, unsigned int d = 0) const {
    return m_Neighborhoods[d]->find_neighborhood_points(p, idx, r);
  }
  inline PointVectorType FindNeighborhoodPoints(const PointType &p, int idx, std::vector<double> &w,
                                                std::vector<double> &distances, double r, unsigned int d = 0) const {
    return m_Neighborhoods[d]->find_neighborhood_points(p, idx, w, distances, r);
  }
  inline PointVectorType FindNeighborhoodPoints(const PointType &p, int idx, std::vector<double> &w, double r,
                                                unsigned int d = 0) const {
    return m_Neighborhoods[d]->find_neighborhood_points(p, idx, w, r);
  }
  inline PointVectorType FindNeighborhoodPoints(unsigned int idx, double r, unsigned int d = 0) const {
    return m_Neighborhoods[d]->find_neighborhood_points(GetPosition(idx, d), idx, r);
  }
  inline PointVectorType FindNeighborhoodPoints(unsigned int idx, std::vector<double> &w,
                                                std::vector<double> &distances, double r, unsigned int d = 0) const {
    return m_Neighborhoods[d]->find_neighborhood_points(GetPosition(idx, d), idx, w, distances, r);
  }
  inline PointVectorType FindNeighborhoodPoints(unsigned int idx, std::vector<double> &w, double r,
                                                unsigned int d = 0) const {
    return m_Neighborhoods[d]->find_neighborhood_points(GetPosition(idx, d), idx, w, r);
  }

  //  inline int FindNeighborhoodPoints(const PointType &p,  double r, PointVectorType &vec, unsigned int d = 0) const
  //  {  return m_Neighborhoods[d]->FindNeighborhoodPoints(p, r, vec); }

  //   PointVectorType FindTransformedNeighborhoodPoints(const PointType &p, double r, unsigned int d = 0) const
  //   {
  //     PointVectorType ans = m_Neighborhoods[d]
  //       ->FindNeighborhoodPoints(this->TransformPoint(p, InverseTransform[d]), r);
  //     for (unsigned int i = 0; i < ans.size(); i++)
  //       {
  //       ans.Point[i] = this->TransformPoint(ans.Point[i], m_Transform[d]);
  //       }
  //     return ans;
  //   }

  /** Add a domain to the particle system.  This method also allocates space
      for a list of positions that are contained within the domain, and a
      default neighborhood calculator.  The final, optional argument indicates
      the calling thread id.*/
  void AddDomain(DomainType::Pointer input);

  /** Return an iterator that points to the first element of the list of the
      domains. */
  std::vector<DomainType::Pointer>::const_iterator GetDomainsBegin() const { return m_Domains.begin(); }

  /** Return an iterator that points one position past the last element of the list of the
      domains. */
  std::vector<DomainType::Pointer>::const_iterator GetDomainsEnd() const { return m_Domains.end(); }

  /** Return the i'th domain object. */
  DomainType *GetDomain(unsigned int i) { return m_Domains[i].get(); }

  /** API for the single domain case. */
  DomainType *GetDomain() { return m_Domains[0].get(); }

  /** Return the i'th domain object. */
  const DomainType *GetDomain(unsigned int i) const { return m_Domains[i].get(); }

  /** API for the single domain case. */
  const DomainType *GetDomain() const { return m_Domains[0].get(); }

  /** Returns the number of domains contained in the particle system. */
  unsigned int GetNumberOfDomains() const { return m_Domains.size(); }

  /** Set the transform associated with domain i. This method will also compute
      and set the corresponding inverse transform if possible.  If the inverse
      transform cannot be computed, the exception is quietly handled by this
      method and an identity transform is set in its place.  The calling
      program is responsible for knowing when an inverse transform cannot be
      computed. */
  void SetTransform(unsigned int i, const TransformType &);
  void SetTransform(const TransformType &p) { this->SetTransform(0, p); }
  void SetPrefixTransform(unsigned int i, const TransformType &);
  void SetPrefixTransform(const TransformType &p) { this->SetPrefixTransform(0, p); }

  /** Return an iterator that points to the first element of the list of the
      transforms. */
  std::vector<TransformType>::const_iterator GetTransformsBegin() const { return m_Transforms.begin(); }

  /** Return an iterator that points one position past the last element of the list of the
      transforms. */
  std::vector<TransformType>::const_iterator GetTransformsEnd() const { return m_Transforms.end(); }

  /** Return the i'th transform object. */
  const TransformType &GetTransform(unsigned int i) const { return m_Transforms[i]; }

  /** API for the single transform case. */
  const TransformType &GetTransform() const { return m_Transforms[0]; }

  /** Return the i'th transform object. */
  TransformType GetTransform(unsigned int i) { return m_Transforms[i]; }

  /** API for the single transform case. */
  TransformType GetTransform() { return m_Transforms[0]; }

  /** Return the i'th transform object. */
  const TransformType &GetPrefixTransform(unsigned int i) const { return m_PrefixTransforms[i]; }

  /** API for the single transform case. */
  const TransformType &GetPrefixTransform() const { return m_PrefixTransforms[0]; }

  /** Return the i'th transform object. */
  TransformType GetPrefixTransform(unsigned int i) { return m_PrefixTransforms[i]; }

  /** API for the single transform case. */
  TransformType GetPrefixTransform() { return m_PrefixTransforms[0]; }

  /** Return an iterator that points to the first element of the list of the
      inverse transforms. */
  std::vector<TransformType>::const_iterator GetInverseTransformsBegin() const { return m_InverseTransforms.begin(); }

  /** Return an iterator that points one position past the last element of the list of the
      transforms. */
  std::vector<TransformType>::const_iterator GetInverseTransformsEnd() const { return m_InverseTransforms.end(); }

  /** Return the i'th transform object. */
  const TransformType &GetInverseTransform(unsigned int i) const { return m_InverseTransforms[i]; }

  /** API for the single transform case. */
  const TransformType &GetInverseTransform() const { return m_InverseTransforms[0]; }

  /** Return the i'th transform object. */
  const TransformType &GetInversePrefixTransform(unsigned int i) const { return m_InversePrefixTransforms[i]; }

  /** API for the single transform case. */
  const TransformType &GetInversePrefixTransform() const { return m_InversePrefixTransforms[0]; }

  /** Return the array of particle positions. */
  const std::vector<PointContainerType::Pointer> &GetPositions() const { return m_Positions; }
  const PointContainerType::Pointer &GetPositions(unsigned int d) const { return m_Positions[d]; }

  /** Adds a list of points to the specified domain.  The arguments are the
     std::vector of points and the domain number. */
  void AddPositionList(const std::vector<PointType> &, unsigned int d = 0);

  /** Transforms a point using the given transform. NOTE: Scaling is not
      currently implemented. (This method may be converted to virtual and
      overridden if tranform type is generalized.)*/
  PointType TransformPoint(const PointType &, const TransformType &) const;

  /** Transforms a vector using the given transform.   Only the rotational part
      of the transform is applied. NOTE: Scaling is not currently supported.*/
  VectorType TransformVector(const VectorType &, const TransformType &) const;

  /** Transforms the derivative of normals to new space. */
  VnlMatrixType TransformNormalDerivative(const VnlMatrixType &, const TransformType &) const;

  /** Returns the inverse of a transformation matrix.*/
  inline TransformType InvertTransform(const TransformType &T) const {
    // Note, vnl_inverse is optimized for small matrices 1x1 - 4x4
    return vnl_inverse(T);
  }

  /** Flag/Unflag a domain.  Flagging a domain has different meanings according
      to the application using this particle system. */
  void FlagDomain(unsigned int i) {
    // ensure large enough
    while (i >= this->m_DomainFlags.size()) {
      m_DomainFlags.push_back(false);
    }

    // set the flag
    m_DomainFlags[i] = true;
  }
  void UnflagDomain(unsigned int i) { m_DomainFlags[i] = false; }
  bool GetDomainFlag(unsigned int i) const {
    if (i >= m_DomainFlags.size()) {
      // not set
      return false;
    }
    return m_DomainFlags[i];
  }
  const std::vector<bool> &GetDomainFlags() const { return m_DomainFlags; }
  void SetDomainFlags() {
    for (unsigned int i = 0; i < m_DomainFlags.size(); i++) {
      m_DomainFlags[i] = true;
    }
  }
  void ResetDomainFlags() {
    for (unsigned int i = 0; i < m_DomainFlags.size(); i++) {
      m_DomainFlags[i] = false;
    }
  }

  /** Determine which shapes are fixed. A shape is fixed if all its domains are flagged.
      Returns a per-shape vector of booleans and sets has_fixed to true if any shape is fixed. */
  std::vector<bool> GetFixedShapes(bool &has_fixed) const {
    unsigned int num_shapes = GetNumberOfDomains() / m_DomainsPerShape;
    std::vector<bool> is_fixed(num_shapes, false);
    has_fixed = false;
    for (unsigned int shape = 0; shape < num_shapes; shape++) {
      bool shape_is_fixed = true;
      for (unsigned int d = 0; d < m_DomainsPerShape; d++) {
        if (!GetDomainFlag(shape * m_DomainsPerShape + d)) {
          shape_is_fixed = false;
          break;
        }
      }
      if (shape_is_fixed && m_DomainsPerShape > 0) {
        is_fixed[shape] = true;
        has_fixed = true;
      }
    }
    return is_fixed;
  }

  /** The following methods provide functionality for specifying particle
      indices that are fixed landmarks.  SetPosition() calls to these particle
      indices will silently fail. For simplicity, only one list of indices is
      maintained for all dimensions.  If particle index n is flagged, for
      example, then particle index n in all domains is fixed.*/
  void SetFixedParticleFlag(unsigned int d, unsigned int i) { m_FixedParticleFlags[d][i] = true; }
  void ResetFixedParticleFlag(unsigned int d, unsigned int i) { m_FixedParticleFlags[d][i] = false; }
  bool GetFixedParticleFlag(unsigned int d, unsigned int i) const { return m_FixedParticleFlags[d][i]; }
  void ResetFixedParticleFlags() {
    for (unsigned d = 0; d < m_FixedParticleFlags.size(); d++) {
      for (unsigned int i = 0; i < m_FixedParticleFlags[d].size(); i++) m_FixedParticleFlags[d][i] = false;
    }
  }

  void SetDomainsPerShape(unsigned int num) {
    m_DomainsPerShape = num;
    m_FixedParticleFlags.resize(m_DomainsPerShape);
  }
  unsigned int GetDomainsPerShape() const { return m_DomainsPerShape; }

  /** Set the number of domains.  This method modifies the size of the
    m_Domains, m_Positions, and m_Transform lists. */
  void SetNumberOfDomains(unsigned int);

  // Returns the maximum distance between nearest neighbors in domain dom
  double ComputeMaxDistNearestNeighbors(size_t dom);

  void SetFieldAttributes(const std::vector<std::string> &field_attributes) { m_FieldAttributes = field_attributes; }

  const std::vector<std::string> &GetFieldAttributes() const { return m_FieldAttributes; }

 protected:
  ParticleSystem();
  void PrintSelf(std::ostream &os, itk::Indent indent) const;
  virtual ~ParticleSystem(){};

  /** Return the i'th transform object. */
  TransformType &GetInverseTransform(unsigned int i) { return m_InverseTransforms[i]; }

  /** API for the single transform case. */
  TransformType &GetInverseTransform() { return m_InverseTransforms[0]; }

  /** Return the i'th transform object. */
  TransformType &GetInversePrefixTransform(unsigned int i) { return m_InversePrefixTransforms[i]; }

  /** API for the single transform case. */
  TransformType &GetInversePrefixTransform() { return m_InversePrefixTransforms[0]; }

 private:
  ParticleSystem(const Self &);  // purposely not implemented
  void operator=(const Self &);  // purposely not implemented

  /** The 2D array of particle positions.  1st array axis is the domain number.
      These values may only be modified by the ParticleSystem class itself. */
  std::vector<PointContainerType::Pointer> m_Positions;

  /** The set of particle domain definitions. */
  std::vector<DomainType::Pointer> m_Domains;

  /** Set number of domains per shape -- Praful */
  unsigned int m_DomainsPerShape;

  /** The set of domain neighborhood objects. */
  std::vector<std::shared_ptr<ParticleNeighborhood>> m_Neighborhoods;

  /** The set of domain transform objects */
  std::vector<TransformType> m_Transforms;

  /** The set of inverse transform objects */
  std::vector<TransformType> m_InverseTransforms;

  /** The set of domain transform objects */
  std::vector<TransformType> m_PrefixTransforms;

  /** The set of domain transform objects */
  std::vector<TransformType> m_InversePrefixTransforms;

  /** A counter used to assign indicies for new particle locations. */
  std::vector<unsigned long int> m_IndexCounters;

  /** An array that indicates which domains are flagged.  A domain may be
      flagged for various purposes depending on the application.  This array
      was created to flag which domains are "fixed" during optimization, for
      example. */
  std::vector<bool> m_DomainFlags;

  /** An array that indicates which particle indices (in any domain) will
      not respond to SetPosition. */
  std::vector<std::vector<bool>> m_FixedParticleFlags;

  std::vector<std::string> m_FieldAttributes;

  std::mt19937 m_rand{42};
};

}  // end namespace shapeworks
