#pragma once
#include <vector>
#include "DisentangledCorrespondenceFunction.h"
#include "Libs/Optimize/Matrix/LegacyShapeMatrix.h"
#include "VectorFunction.h"

namespace shapeworks {

/**
 * \class DisentangledCorrespondenceFunction
 * This class implements the Correspondence Term where the entropy computation is disentangled 
 * across time and shape domain, and builds a Sapatiotemporal SSM.
 */
class DisentangledCorrespondenceFunction : public VectorFunction {
 public:
  constexpr static unsigned int VDimension = 3;
  constexpr static unsigned int Dimension = VDimension;

  typedef LegacyShapeMatrix ShapeMatrixType;

  typedef typename ShapeMatrixType::DataType DataType;

  /** Vector & Point types. */
  typedef VectorFunction::VectorType VectorType;
  typedef ParticleSystem::PointType PointType;
  typedef vnl_vector<DataType> vnl_vector_type;
  typedef vnl_matrix<DataType> vnl_matrix_type;
  typedef std::shared_ptr<std::vector<vnl_matrix_type>> shared_vnl_matrix_array_type;

  /// Factory method for creating instances
  static std::shared_ptr<DisentangledCorrespondenceFunction> New() {
    return std::make_shared<DisentangledCorrespondenceFunction>();
  }

  /** The first argument is a pointer to the particle system.  The second
      argument is the index of the domain within that particle system.  The
      third argument is the index of the particle location within the given
      domain. */
  virtual VectorType evaluate(unsigned int, unsigned int, const ParticleSystem*, double&, double&) const;
  virtual VectorType evaluate(unsigned int a, unsigned int b, const ParticleSystem* c, double& d) const {
    double e;
    return this->evaluate(a, b, c, d, e);
  }

  virtual double energy(unsigned int a, unsigned int b, const ParticleSystem* c) const {
    double e, d;
    this->evaluate(a, b, c, d, e);
    return e;
  }

  /** Write the first n modes to +- 3 std dev and the mean of the model
      described by the covariance matrix.  The string argument is a prefix to
      the file names. */
  void WriteModes(const std::string&, int) const;

  /**Access the shape matrix. */
  void SetShapeMatrix(ShapeMatrixType* s) { m_ShapeMatrix = s; }
  ShapeMatrixType* GetShapeMatrix() { return m_ShapeMatrix.GetPointer(); }
  const ShapeMatrixType* GetShapeMatrix() const { return m_ShapeMatrix.GetPointer(); }

  /** Called before each iteration of a solver. */
  virtual void before_iteration() {
    m_ShapeMatrix->before_iteration();

    if (m_Counter == 0) {
      this->ComputeCovarianceMatrices();
    }
  }

  /** Called after each iteration of the solver. */
  virtual void after_iteration() override {
    m_ShapeMatrix->after_iteration();
    // Update the annealing parameter.
    if (m_HoldMinimumVariance != true && !m_UseMeanEnergy) {
      m_Counter++;
      if (m_Counter >= m_RecomputeCovarianceInterval) {
        m_Counter = 0;
        m_MinimumVariance *= m_MinimumVarianceDecayConstant;
      }
    }
  }

  /** */
  void SetMinimumVariance(double d) { m_MinimumVariance = d; }
  double GetMinimumVariance() const { return m_MinimumVariance; }

  void SetMinimumVarianceDecay(double initial_value, double final_value, double time_period) {
    m_MinimumVarianceDecayConstant = exp(log(final_value / initial_value) / time_period);
    m_MinimumVariance = initial_value;
    m_HoldMinimumVariance = false;
  }
  bool GetMinimumVarianceDecayConstant() const { return m_MinimumVarianceDecayConstant; }

  void PrintShapeMatrix() { m_ShapeMatrix->PrintMatrix(); }

  void UseMeanenergy() { m_UseMeanEnergy = true; }
  void UseEntropy() { m_UseMeanEnergy = false; }

  /** */
  bool GetHoldMinimumVariance() const { return m_HoldMinimumVariance; }
  void SetHoldMinimumVariance(bool b) { m_HoldMinimumVariance = b; }

  void SetRecomputeCovarianceInterval(int i) { m_RecomputeCovarianceInterval = i; }
  int GetRecomputeCovarianceInterval() const { return m_RecomputeCovarianceInterval; }

  std::shared_ptr<VectorFunction> clone() override {
    auto copy = DisentangledCorrespondenceFunction::New();

    copy->m_Shape_PointsUpdate = this->m_Shape_PointsUpdate;
    copy->m_Time_PointsUpdate = this->m_Time_PointsUpdate;
    copy->m_MinimumVariance = this->m_MinimumVariance;
    copy->m_MinimumEigenValue_shape_cohort = this->m_MinimumEigenValue_shape_cohort;
    copy->m_MinimumEigenValue_time_cohort = this->m_MinimumEigenValue_time_cohort;

    copy->m_CurrentEnergy = this->m_CurrentEnergy;
    copy->m_HoldMinimumVariance = this->m_HoldMinimumVariance;
    copy->m_MinimumVarianceDecayConstant = this->m_MinimumVarianceDecayConstant;
    copy->m_RecomputeCovarianceInterval = this->m_RecomputeCovarianceInterval;
    copy->m_Counter = m_Counter;

    copy->domain_number_ = this->domain_number_;
    copy->particle_system_ = this->particle_system_;
    copy->m_ShapeMatrix = this->m_ShapeMatrix;


    copy->m_InverseCovMatrices_time_cohort = this->m_InverseCovMatrices_time_cohort;
    copy->m_InverseCovMatrices_shape_cohort = this->m_InverseCovMatrices_shape_cohort;

    copy->m_points_mean_time_cohort = this->m_points_mean_time_cohort;
    copy->m_points_mean_shape_cohort = this->m_points_mean_shape_cohort;

    return copy;
  }

  DisentangledCorrespondenceFunction() {
    // m_MinimumVarianceBase = 1.0;//exp(log(1.0e-5)/10000.0);
    m_HoldMinimumVariance = true;
    m_MinimumVariance = 1.0e-5;
    m_MinimumEigenValue = 0.0;
    m_MinimumVarianceDecayConstant = 1.0;  // log(2.0) / 50000.0;
    m_RecomputeCovarianceInterval = 1;
    m_Counter = 0;
    m_UseMeanEnergy = true;
    m_InverseCovMatrices_time_cohort = std::make_shared<std::vector<Eigen::MatrixXd>>();
    m_InverseCovMatrices_shape_cohort = std::make_shared<std::vector<Eigen::MatrixXd>>();
    m_points_mean_time_cohort = std::make_shared<std::vector<vnl_matrix_type>>();
    m_points_mean_shape_cohort = std::make_shared<std::vector<vnl_matrix_type>>();
    m_Time_PointsUpdate = std::make_shared<std::vector<vnl_matrix_type>>();
    m_Shape_PointsUpdate = std::make_shared<std::vector<vnl_matrix_type>>();
  }
  ~DisentangledCorrespondenceFunction() override = default;

 protected:
  DisentangledCorrespondenceFunction(const DisentangledCorrespondenceFunction&) = delete;
  DisentangledCorrespondenceFunction& operator=(const DisentangledCorrespondenceFunction&) = delete;
  typename ShapeMatrixType::Pointer m_ShapeMatrix;

  // Computes Covariance Matrices across time and shape domain and then generate gradient updates for them.
  virtual void ComputeCovarianceMatrices();

  // Initialize size and clear relevant variables from previous iteration, before gradient updates computation.
  void Initialize(){
    const unsigned int total_time_points = m_ShapeMatrix->GetDomainsPerShape();
    int total_subjects = m_ShapeMatrix->cols();
    m_points_mean_time_cohort->clear();
    m_points_mean_shape_cohort->clear();
    m_InverseCovMatrices_time_cohort->clear();
    m_InverseCovMatrices_shape_cohort->clear();
    m_Shape_PointsUpdate->clear();
    m_Time_PointsUpdate->clear();
    m_MinimumEigenValue_shape_cohort.resize(total_subjects, 0.0);
    m_MinimumEigenValue_time_cohort.resize(total_time_points, 0.0);
    for(int i = 0; i < total_time_points; ++i){
      Eigen::MatrixXd temp_cov_matrix;
      vnl_matrix_type temp_points_matrix;
      vnl_matrix_type temp_mean_matrix;
      m_InverseCovMatrices_time_cohort->push_back(temp_cov_matrix);
      m_points_mean_time_cohort->push_back(temp_mean_matrix);
      m_Time_PointsUpdate->push_back(temp_points_matrix);
    }
    for(int i = 0; i < total_subjects; ++i){
      Eigen::MatrixXd temp_cov_matrix;
      vnl_matrix_type temp_points_matrix;
      vnl_matrix_type temp_mean_matrix;
      m_InverseCovMatrices_shape_cohort->push_back(temp_cov_matrix);
      m_points_mean_shape_cohort->push_back(temp_mean_matrix);
      m_Shape_PointsUpdate->push_back(temp_points_matrix);
    }
  }

  double m_MinimumVariance;
  double m_MinimumEigenValue;
  std::vector<double> m_MinimumEigenValue_time_cohort;
  std::vector<double> m_MinimumEigenValue_shape_cohort;
  double m_CurrentEnergy;
  bool m_HoldMinimumVariance;
  double m_MinimumVarianceDecayConstant;
  int m_RecomputeCovarianceInterval;
  int m_Counter;
  bool m_UseMeanEnergy;

  // Inverse Covariance matrices across time and shape cohort
  std::shared_ptr<std::vector<Eigen::MatrixXd>> m_InverseCovMatrices_time_cohort; // T obj matrices each of dimensionality dM X N
  std::shared_ptr<std::vector<Eigen::MatrixXd>> m_InverseCovMatrices_shape_cohort; // N obj matrices each of dimensionality dM X T
  // mean vectors across time and shape cohort
  shared_vnl_matrix_array_type m_points_mean_time_cohort; // T mean vectors each of dimension dM
  shared_vnl_matrix_array_type m_points_mean_shape_cohort; // N mean vectors each of dimension dM
  // Matrices for Gradient Updates across time and shape cohort
  shared_vnl_matrix_array_type m_Time_PointsUpdate; // T update matrices each of dimensionality dM X N
  shared_vnl_matrix_array_type m_Shape_PointsUpdate; // N update matrices each of dimensionality dM X T
};

}  // namespace shapeworks
