#ifndef PROJECTION_HPP
#define PROJECTION_HPP

/**
 * \file projection.hpp
 * \brief Utilities for conversion from coordinates in an image to
 * geographical coordinates.
 * \author Le Bars, Yoann
 * \version 1.0
 * \date 2013/06/27
 * \date 2013/06/28
 * \date 2014/03/06
 */

#include <boost/concept_check.hpp>
#include <cassert>
#include <vector>
#include <eigen3/Eigen/Dense>

/// \brief Namespace for projection computations.
namespace Projection {
  /// \brief Class for 2D points.
  class Point2D {
    public:
      /**
       * \brief Construct a point with two double.
       * \param _x Abscissa.
       * \param _y Ordinate.
       */
      Point2D (double _x = 0., double _y = 0.): x_ (_x), y_ (_y) {}

      /**
       * \brief Copy constructor.
       * \param p Point2D to be copied.
       */
      Point2D (const Point2D &p): x_ (p.x_), y_ (p.y_) {}

      /// \brief Destructor.
      ~Point2D () {}

      /// \brief Access to abscissa, modification is not possible.
      double x () const {return x_;}

      /// \brief Access to abscissa, modification is possible.
      double &x () {return x_;}

      /// \brief Access to ordinate, modification is not possible.
      double y () const {return y_;}

      /// \brief Access to ordinate, modification is possible.
      double &y () {return y_;}

      /**
       * \brief Copy operator.
       * \param p Point2D to be copied.
       * \return A reference to current point.
       */
      Point2D &operator = (const Point2D &p) {
        x_ = p.x_;
        y_ = p.y_;
        return *this;
      }

    private:
      /// \brief Abscissa.
      double x_;

      /// \brief Ordinate;
      double y_;
  };

  /// \brief Type for projection coefficients.
  typedef Eigen::Matrix<double, 3, 2> Coefficients;

  /**
   * \brief Compute the projection coefficients.
   * \param r1 Vector containing reference points in image coordinates.
   * \param r2 Vector containing reference points in geographical coordinates.
   * \return Vector containing projection coefficients.
   *
   * Information on coefficients can be found at:
   * <http://en.wikipedia.org/wiki/World_file>
   */
  inline Coefficients computeCoefficients (const std::vector<Point2D> &r1,
                                           const std::vector<Point2D> &r2) {
    assert(r1.size() == 3);
    assert(r2.size() == 3);
    /* Matrix to compute coefficients. */
    Eigen::Matrix3d a;
    a << r1[0].x(), r1[0].y(), 1.,
         r1[1].x(), r1[1].y(), 1.,
         r1[2].x(), r1[2].y(), 1.;
    /* Matrix of points in geographical coordinates. */
    Coefficients b;
    b << r2[0].x(), r2[0].y(),
         r2[1].x(), r2[1].y(),
         r2[2].x(), r2[2].y();
//     return a.fullPivLu().solve(b);
    return a.jacobiSvd(Eigen::ComputeThinU | Eigen::ComputeThinV).solve(b);
  }
}

#endif  // #ifndef PROJECTION_HPP


#ifndef NBT_SIMULATOR_HPP
#define NBT_SIMULATOR_HPP

#include <vector>
#include <queue>
#include <cstdint>

#include <Eigen>
#include "integrator.hpp"
#include "dynamics_engine.hpp"
#include "rigidbody.hpp"
#include "octree.hpp"

/**
 * Simulator object to control simulations.
 * Add objects using addObject().
 * Step simulation using step().
 */
class Simulator {
    private:
        Integrator* const integrator;         //!< Integrator used in this simulation.
        DynamicsEngine* const dynamicsEngine; //!< dynamicsEngine use in this simulation.

        double timeStep; //!< dt value used in integrators.

        // Structure of arrays for object properties
        Eigen::Matrix<double, 1, Eigen::Dynamic> m;     //!< Mass of each object packed into a 1 x N vector.
        Eigen::Matrix<double, 1, Eigen::Dynamic> r;     //!< Radius of each object packed into a 1 x N vector.
        Eigen::Matrix<double, 3, Eigen::Dynamic> pos;   //!< 3D position of each object packed into a 3 x N matrix.
        Eigen::Matrix<double, 3, Eigen::Dynamic> v;     //!< 3D velocity of each object packed into a 3 x N matrix.
        Eigen::Matrix<double, 3, Eigen::Dynamic> a;     //!< 3D acceleration of each object packed into a 3 x N matrix.

        uint64_t iteration = 0;                 //!< Current iteration of the simulation.

        RigidbodyIdx nextIdx = 0;               //!< Index of next available column in the structure of arrays. Also serves as a counter of active objects.
        Rigidbody    nextID = 0;                //!< Next available ID to be assigned to a newly created Rigidbody.
        std::vector<Rigidbody>    idx2id;       //!< Maps index to associated ID
        std::vector<RigidbodyIdx> id2idx;       //!< Maps ID to associated index
        std::queue<Rigidbody> availableUsedIDs; //!< Stores IDs of destroyed objects for reallocation

        /*! Returns slice of array structure component with only active objects. */
        Eigen::Ref<Eigen::MatrixXd> active(Eigen::Ref<Eigen::MatrixXd> mat);

        /*! Computes force between each object using #forceComputer. #a is updated. */
        void updateAccelerations();
    public:
        const Rigidbody maxObjects;         //!< Maximum number of objects in the simulation. Sets the dimensions of the sstructure of arrays.

        /*! Constructs a Simulator object. */
        Simulator(double timeStep, uint64_t maxObjects, Integrator* integrator, DynamicsEngine* dynamicsEngine);

        /*! Destroys a Simulator object and deallocates all used memory. */
        ~Simulator();

        /*! Returns sum of kinetic energies of the particles in the system. */
        double totalKineticEnergy();

        /*! Returns sum of potential energies of each particle in the sytem */
        double totalPotentialEnergy();

        /*! Returns the sum of totalKineticEnergy() and totalPotentialEnergy() */
        double totalEnergy();

        /*! Returns sum of moments of inertia of each particle in the system */
        double totalMomentOfInertia();

        /*! Returns number of active objects */
        Rigidbody nObjects();

        /*! Adds an object to the simulation. Can be done during the simulation if the #maxObjects is not met. */
        Rigidbody addObject(double m, double r, const Eigen::Vector3d& p0, const Eigen::Vector3d& v0);

        /*! Deletes an object from the simulation. */
        void delObject(Rigidbody id);

        /*! Objects combine into id1 and momentum is conserved. */
        void collideObject(Rigidbody id1, Rigidbody id2); // TODO collideObject remember to conserve momentum

        /*! Returns if rigidbody with this id exists. */
        bool rb_exists(Rigidbody id);

        /*! Returns the position vector of a rigidbody. */
        Eigen::Ref<const Eigen::Vector3d> rb_pos(Rigidbody id);

        /*! Returns the velocity vector of a rigidbody. */
        Eigen::Ref<const Eigen::Vector3d> rb_v(Rigidbody id);

        /*! Returns the acceleration vector of a rigidbody. */
        Eigen::Ref<const Eigen::Vector3d> rb_a(Rigidbody id);

        /*! Returns the mass of a rigidbody. */
        double rb_m(Rigidbody id);

        /*! Returns the radius of a rigidbody. */
        double rb_r(Rigidbody id);

        /*! Returns Eigen::Matrix ref of object positions */
        Eigen::Ref<const Eigen::Matrix3Xd> activePos();

        /*! Returns Eigen::Matrix ref of object velocities */
        Eigen::Ref<const Eigen::Matrix3Xd> activeV();

        /*! Returns Eigen::Matrix ref of object accelerations */
        Eigen::Ref<const Eigen::Matrix3Xd> activeA();

        /*! Returns Eigen::Matrix ref of object masses */
        Eigen::Ref<const Eigen::RowVectorXd> activeM();

        /*! Returns Eigen::Matrix ref of object radii */
        Eigen::Ref<const Eigen::RowVectorXd> activeR();

        /*! Steps simulation */
        void step();
};

#endif


/* Copyright (C) 2019-2020 IBM Corp.
 * This program is Licensed under the Apache License, Version 2.0
 * (the "License"); you may not use this file except in compliance
 * with the License. You may obtain a copy of the License at
 *   http://www.apache.org/licenses/LICENSE-2.0
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License. See accompanying LICENSE file.
 */

// This is a sample program for education purposes only.
// It attempts to show the various basic mathematical
// operations that can be performed on both ciphertexts
// and plaintexts.

#include <iostream>
#include <boost/lexical_cast.hpp>
#include <helib/helib.h>
#include <time.h>

int main(int argc, char* argv[])
{
  /*  Example of BGV scheme  */

  if( argc != 3 && argc != 4 )
  {std::cout << "not enough args: need p and m" << std::endl;
		exit(-1);
  }
  // Plaintext prime modulus
  unsigned long p = boost::lexical_cast<long>(argv[1]);
  // Cyclotomic polynomial - defines phi(m)
  unsigned long m = boost::lexical_cast<long>(argv[2]);
  // Hensel lifting (default = 1)
  unsigned long r = 1;
  // Number of bits of the modulus chain
  unsigned long bits = argc == 3 ? 300 : boost::lexical_cast<long>(argv[3]);
  // Number of columns of Key-Switching matrix (default = 2 or 3)
  unsigned long c = 2;

	clock_t start = clock();
  // Initialize context
  // This object will hold information about the algebra created from the
  // previously set parameters
  helib::Context context(m, p, r);
  std::cout << "Initialising context object..." << (double)(clock()-start)/CLOCKS_PER_SEC << std::endl;
  // Modify the context, adding primes to the modulus chain
  // This defines the ciphertext space
  buildModChain(context, bits, c);
	std::cout << "Building modulus chain..." << (double)(clock()-start)/CLOCKS_PER_SEC << std::endl;

  // Print the context
  context.zMStar.printout();
  std::cout << std::endl;

  // Print the security level
  std::cout << "Security: " << context.securityLevel() << std::endl;

  //FIXME: really should make this a separate progrm but too lazy for that
  if( argc == 3)
    exit(0);

  // Secret key management
  std::cout << "Creating secret key..." << std::endl;
  // Create a secret key associated with the context
  helib::SecKey secret_key(context);
  // Generate the secret key
  secret_key.GenSecKey();
  // Compute key-switching matrices that we need
  helib::addSome1DMatrices(secret_key);
  std::cout << "Generating key-switching matrices..." << (double)(clock()-start)/CLOCKS_PER_SEC << std::endl;

  // Public key management
  // Set the secret key (upcast: SecKey is a subclass of PubKey)
  const helib::PubKey& public_key = secret_key;

  // Get the EncryptedArray of the context
  const helib::EncryptedArray& ea = *(context.ea);

  // Get the number of slot (phi(m))
  long nslots = ea.size();
  std::cout << "Number of slots: " << nslots << std::endl;

  // Create a vector of long with nslots elements
  helib::Ptxt<helib::BGV> ptxt(context);
  // Set it with numbers 0..nslots - 1
  // ptxt = [0] [1] [2] ... [nslots-2] [nslots-1]
  for (int i = 0; i < ptxt.size(); ++i) {
    ptxt[i] = i;
  }

  // Create a ciphertext object
	start = clock();
  helib::Ctxt ctxt(public_key);
  // Encrypt the plaintext using the public_key
  public_key.Encrypt(ctxt, ptxt);
	std::cout << "Done encrypting: " << (double)(clock()-start)/CLOCKS_PER_SEC << std::endl;

  // ********** Operations **********
  // Ciphertext and plaintext operations are performed
  // "entry-wise".

  // Square the ciphertext
  // [0] [1] [2] [3] [4] ... [nslots-1]
  // -> [0] [1] [4] [9] [16] ... [(nslots-1)*(nslots-1)]
	start = clock();
	ctxt.multiplyBy(ctxt);
  std::cout << "Power: " << (double)(clock()-start)/CLOCKS_PER_SEC << std::endl;

  // Raise the copy to the exponent 2
  // Note: 0 is a special case because 0^n = 0 for any power n
  //ctxt.power(2);

  // Subtract it from itself (result should be 0)
	start = clock();
  ctxt += ctxt;
	std::cout << "Subtract: " << (double)(clock()-start)/CLOCKS_PER_SEC << std::endl;

  // Create a plaintext for decryption
  start = clock();
  helib::Ptxt<helib::BGV> plaintext_result(context);
  // Decrypt the modified ciphertext
  secret_key.Decrypt(plaintext_result, ctxt);
	std::cout << "Decryt: " << (double)(clock()-start)/CLOCKS_PER_SEC << std::endl;

  // Print the decrypted plaintext
  std::cout << "Decrypted Result: " << plaintext_result << std::endl;

  return 0;
}