// Copyright (c) 2013-2020, SIB - Swiss Institute of Bioinformatics and
//                          Biozentrum - University of Basel
// 
// 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.


#include <promod3/core/superpose.hh>
#include <Eigen/SVD>
#include <Eigen/Geometry>
#include <map>

namespace{

inline geom::Vec3 EigenVec3ToVec3(const promod3::core::EVec3& vec){
  return geom::Vec3(vec.data());
}

inline geom::Mat3 EigenMat3ToMat3(const promod3::core::EMat3& mat){
  geom::Mat3 return_mat;
  for(int i=0;i<3;++i){
    for(int j=0;j<3;++j){
      return_mat(i,j) = mat(i,j);
    }
  }
  return return_mat;
}

// get RMSD and rotation matrix using the Theobald method
// Both position matrices are expected to have the same size
// and to have their average position at the origin 
void TheobaldRMSD(const promod3::core::EMatX3& pos_one,
                  const promod3::core::EMatX3& pos_two,
                  Real& rmsd,
                  promod3::core::EMat3& rot){

  if(pos_one.rows() < 3){
    throw promod3::Error("Observed superposition with < 3 positions to "
                         "superpose!");
  }

  // using floats when calculating M might lead to numerical instabilities
  // later on, so let's cast to double precision
  Eigen::Matrix<double,3,3> M =
  pos_one.cast<double>().transpose() * pos_two.cast<double>();

  // using floats for the squared norm is fine
  double GA = pos_one.squaredNorm();
  double GB = pos_two.squaredNorm();

  Eigen::Matrix<double,4,4> K;
  K(0,0)  =  M(0,0) + M(1,1) + M(2,2);
  K(0,1)  =  M(1,2) - M(2,1);
  K(0,2)  =  M(2,0) - M(0,2);
  K(0,3)  =  M(0,1) - M(1,0);

  K(1,0)  =  K(0,1);
  K(1,1)  =  M(0,0) - M(1,1) - M(2,2);
  K(1,2)  =  M(0,1) + M(1,0);
  K(1,3)  =  M(0,2) + M(2,0);

  K(2,0)  =  K(0,2);
  K(2,1)  =  K(1,2);
  K(2,2)  = -M(0,0) + M(1,1) - M(2,2);
  K(2,3)  =  M(1,2) + M(2,1);

  K(3,0)  =  K(0,3);
  K(3,1)  =  K(1,3);
  K(3,2)  =  K(2,3);
  K(3,3)  = -M(0,0) - M(1,1) + M(2,2);

  double C0 = K.determinant();
  double C1 = -8.0*M.determinant();
  double C2 = -2.0*M.squaredNorm();
  double lambda = 0.5 * (GA + GB);
  double a, b, d, lambda_2;
  for(int i = 0; i < 50; ++i){
    lambda_2 = lambda * lambda;
    b = (lambda_2 + C2) * lambda;
    a = b + C1;
    d = (a*lambda + C0) / (2.0*lambda_2*lambda + b + a);
    lambda -= d;
    if(std::abs(d) < 1e-6){
      break;
    }
  }

  double msd = (GA + GB - 2.0 * lambda) / pos_one.rows();
  if(msd < 1e-4){
    // The algorithm never really goes to zero... if msd is super small we just
    // assign zero. 1e-4 corresponds to an rmsd of 0.01
    rmsd = 0.0;
  }
  else{
    rmsd = std::sqrt(msd);
  }

  K -= lambda*Eigen::Matrix<double,4,4>::Identity();

  double helper[6];
  helper[0] = K(2,2)*K(3,3) - K(3,2)*K(2,3);
  helper[1] = K(2,1)*K(3,3) - K(3,1)*K(2,3);
  helper[2] = K(2,1)*K(3,2) - K(3,1)*K(2,2);
  helper[3] = K(2,0)*K(3,3) - K(3,0)*K(2,3);
  helper[4] = K(2,0)*K(3,2) - K(3,0)*K(2,2);
  helper[5] = K(2,0)*K(3,1) - K(3,0)*K(2,1);

  double q1 =  K(1,1)*helper[0] - K(1,2)*helper[1] + K(1,3)*helper[2];
  double q2 = -K(1,0)*helper[0] + K(1,2)*helper[3] - K(1,3)*helper[4];
  double q3 =  K(1,0)*helper[1] - K(1,1)*helper[3] + K(1,3)*helper[5];
  double q4 = -K(1,0)*helper[2] + K(1,1)*helper[4] - K(1,2)*helper[5]; 
  double norm = q1*q1 + q2*q2 + q3*q3 + q4*q4;

  if(norm < 1e-6){
    q1 =  K(0,1)*helper[0] - K(0,2)*helper[1] + K(0,3)*helper[2];
    q2 = -K(0,0)*helper[0] + K(0,2)*helper[3] - K(0,3)*helper[4];
    q3 =  K(0,0)*helper[1] - K(0,1)*helper[3] + K(0,3)*helper[5];
    q4 = -K(0,0)*helper[2] + K(0,1)*helper[4] - K(0,2)*helper[5];
    norm = q1*q1 + q2*q2 +q3*q3 + q4*q4;

    if (norm < 1e-6){
      helper[0] = K(0,2)*K(1,3) - K(0,3)*K(1,2);
      helper[1] = K(0,1)*K(1,3) - K(0,3)*K(1,1);
      helper[2] = K(0,1)*K(1,2) - K(0,2)*K(1,1);
      helper[3] = K(0,0)*K(1,3) - K(0,3)*K(1,0);
      helper[4] = K(0,0)*K(1,2) - K(0,2)*K(1,0);
      helper[5] = K(0,0)*K(1,1) - K(0,1)*K(1,0);

      q1 =  K(3,1)*helper[0] - K(3,2)*helper[1] + K(3,3)*helper[2];
      q2 = -K(3,0)*helper[0] + K(3,2)*helper[3] - K(3,3)*helper[4];
      q3 =  K(3,0)*helper[1] - K(3,1)*helper[3] + K(3,3)*helper[5];
      q4 = -K(3,0)*helper[2] + K(3,1)*helper[4] - K(3,2)*helper[5];
      norm = q1*q1 + q2*q2 + q3*q3 + q4*q4;

      if (norm < 1e-6){
        q1 =  K(2,1)*helper[0] - K(2,2)*helper[1] + K(2,3)*helper[2];
        q2 = -K(2,0)*helper[0] + K(2,2)*helper[3] - K(2,3)*helper[4];
        q3 =  K(2,0)*helper[1] - K(2,1)*helper[3] + K(2,3)*helper[5];
        q4 = -K(2,0)*helper[2] + K(2,1)*helper[4] - K(2,2)*helper[5];
        norm = q1*q1 + q2*q2 + q3*q3 + q4*q4;
        if (norm < 1e-6){
          // this should not happen
          rot = promod3::core::EMat3::Identity();
          return;
        }
      }
    }
  }

  norm = 1.0 / std::sqrt(norm);
  q1 *= norm; q2 *= norm; q3 *= norm; q4 *= norm;
  rot(0,0) = 1.0 - 2.0*(q3*q3 + q4*q4);
  rot(0,1) =       2.0*(q2*q3 - q1*q4);
  rot(0,2) =       2.0*(q2*q4 + q3*q1);
  rot(1,0) =       2.0*(q2*q3 + q1*q4);
  rot(1,1) = 1.0 - 2.0*(q2*q2 + q4*q4);
  rot(1,2) =       2.0*(q3*q4 - q2*q1);
  rot(2,0) =       2.0*(q2*q4 - q3*q1);
  rot(2,1) =       2.0*(q3*q4 + q2*q1);
  rot(2,2) = 1.0 - 2.0*(q2*q2 + q3*q3);
}

// calculate average row of the point set
inline promod3::core::ERVec3 AverageRow(promod3::core::EMatX3& pos) {
  // return pos.colwise().sum()/pos.rows();
  // NOTE: the colwise-operator is opt-flag-dependent and we hence do it manually
  promod3::core::ERVec3 avg_row = promod3::core::ERVec3::Zero();
  for (uint i = 0; i < pos.rows(); ++i) {
    avg_row += pos.row(i);
  }
  return avg_row / pos.rows();
}

inline void CenterPositions(promod3::core::EMatX3& pos){
  promod3::core::ERVec3 avg = AverageRow(pos);
  for(uint i = 0; i < pos.rows(); ++i){
    pos.row(i) -= avg;
  }
}

inline void ApplySuperposition(promod3::core::EMatX3& pos,
                               promod3::core::ERVec3& avg_one,
                               promod3::core::ERVec3& avg_two,
                               promod3::core::EMat3& rotation){

  uint num_rows = pos.rows();
  //must be implemented in one operation...
  for (uint i = 0; i < num_rows; ++i){
    pos.row(i) -= avg_one;
  }
  for(uint i = 0; i < num_rows; ++i){
    pos.row(i) = (rotation * pos.row(i).transpose()).transpose();
  }
  for (uint i = 0; i < num_rows; ++i){
    pos.row(i) += avg_two;
  } 
}

// The Superpose function performs all tasks required for a superposition.
// After calling the function, you can construct the transformation matrix 
// using avg_one, avg_two and rotation.
// 
// During function execution, pos_one and pos_two are shifted, such that their
// avg position lies on the origin
//
// If apply_rotation is true, the calculated rotation matrix gets applied on
// pos one, the two position matrices are therefore superposed but still have
// their avg on the origin.
//
// - mean position of pos_one gets filled into avg_one
// - mean position of pos_two gets filled into avg_two
// - the mean gets subtracted from both, pos_one and pos_two
// - the rotation matrix to rotate pos_one onto pos_two gets filled
//   and is optionally applied 
// - the rmsd gets set into the rmsd variable

void Superpose(promod3::core::EMatX3& pos_one,
               promod3::core::EMatX3& pos_two,
               promod3::core::ERVec3& avg_one,
               promod3::core::ERVec3& avg_two,
               Real& rmsd,
               promod3::core::EMat3& rotation,
               bool apply_rotation = false){

  if(pos_one.rows() != pos_two.rows()){
    throw promod3::Error("Cannot superpose positions of different size!");
  }

  avg_one = AverageRow(pos_one);
  avg_two = AverageRow(pos_two);

  // TheobaldRMSD only determines the rotational component of the superposition
  // we need to shift the centers of the two point sets onto origin
  for (uint i = 0; i < pos_one.rows(); ++i){
    pos_one.row(i) -= avg_one;
    pos_two.row(i) -= avg_two;
  }

  TheobaldRMSD(pos_one,pos_two,rmsd,rotation);

  if(apply_rotation){
    for(uint i = 0; i < pos_one.rows(); ++i){
      pos_one.row(i) = (rotation * pos_one.row(i).transpose()).transpose();
    }
  }
}


// Does the same as the Superpose function but without touching pos_one and
// pos_two

// - the squared distance from every element in pos_one and pos_two after
//   the final superposition gets filled into squared_distances
// - the selectiton of the superposed position gets filled into indices
// - mean position of pos_one selection gets filled into avg_one
// - mean position of pos_two selection gets filled into avg_two
// - the rotation matrix to rotate the pos_one selection onto the 
//   pos_two selection gets filled
void SuperposeIterative(const promod3::core::EMatX3& pos_one,
                        const promod3::core::EMatX3& pos_two,
                        uint max_iterations, 
                        Real distance_thresh,
                        std::vector<Real>& squared_distances,
                        std::vector<uint>& indices,
                        promod3::core::ERVec3& avg_one,
                        promod3::core::ERVec3& avg_two,
                        promod3::core::EMat3& rotation){

  if(pos_one.rows() != pos_two.rows()){
    throw promod3::Error("Position data must be of consistent size!");
  }

  if(max_iterations == 0){
    throw promod3::Error("max_iterations must be at least 1!");
  }

  uint num_rows = pos_one.rows();

  if(indices.empty()){
    //there are no idx, so we use all positions for the initial superposition
    indices.resize(num_rows);
    for(uint i = 0; i < num_rows; ++i){
      indices[i] = i;
    }
  }
  else{
    //its not empty! let's quickly check whether there are at least 3 positions
    //and whether the indices are all valid
    if(indices.size() < 3){
      throw promod3::Error("Must have at least 3 start indices for iterative " 
                           "Superposition!");
    }
    for(std::vector<uint>::iterator i = indices.begin();
        i != indices.end(); ++i){
      if((*i) >= num_rows){
        throw promod3::Error("Invalid index in iterative Superposition!");
      }
    }
  }

  Real squared_dist_thresh = distance_thresh * distance_thresh;
  std::vector<uint> new_indices;
  new_indices.reserve(num_rows);
  squared_distances.resize(num_rows);
  promod3::core::ERVec3 temp_vec;
  Real rmsd;

  for(uint iteration = 0; iteration < max_iterations; ++iteration){

    if(indices.size() < 3) break; //the thing is not really superposable...

    promod3::core::EMatX3 temp_pos_one = 
    promod3::core::EMatX3::Zero(indices.size(), 3);
    promod3::core::EMatX3 temp_pos_two = 
    promod3::core::EMatX3::Zero(indices.size(), 3);

    for(uint i = 0; i < indices.size(); ++i){
      temp_pos_one.row(i) = pos_one.row(indices[i]);
      temp_pos_two.row(i) = pos_two.row(indices[i]);
    }

    Superpose(temp_pos_one, temp_pos_two, avg_one, avg_two, rmsd, rotation, false);

    for(uint i = 0; i < num_rows; ++i){
      temp_vec = pos_one.row(i) - avg_one;
      temp_vec = (rotation * temp_vec.transpose()).transpose() + avg_two;
      squared_distances[i] = (temp_vec - pos_two.row(i)).squaredNorm();
    }

    new_indices.clear();
    for(uint i = 0; i < num_rows; ++i){
      if(squared_distances[i] < squared_dist_thresh){
        new_indices.push_back(i);
      }
    }

    if(new_indices == indices) break; //nothing changes anymore

    indices = new_indices;
  }
}

// Get an ost transform from the parameters you get when calling the
// Superpose function.

inline geom::Mat4 GenerateOstTransform(const promod3::core::ERVec3& avg_one,
                                       const promod3::core::ERVec3& avg_two,
                                       promod3::core::EMat3& rotation){

  //there are three transformation to be applied to reach pos_two from pos_one.
  //shift pos one to center, apply estimated rotation and shift onto average of pos_two
  promod3::core::EVec3 translation = rotation * (-avg_one.transpose()) + 
                                     avg_two.transpose();

  geom::Mat3 rot = EigenMat3ToMat3(rotation);
  geom::Vec3 trans = EigenVec3ToVec3(translation.transpose());

  geom::Mat4 result;
  result.PasteRotation(rot);
  result.PasteTranslation(trans);

  return result;
}

}

namespace promod3{ namespace core{

Real SuperposedRMSD(EMatX3& pos_one, EMatX3& pos_two, 
                    bool apply_superposition){
  ERVec3 avg_one, avg_two;
  Real rmsd;
  EMat3 rotation;
  Superpose(pos_one, pos_two, avg_one, avg_two, rmsd, 
            rotation, apply_superposition);

  if(apply_superposition){
    for (uint i = 0; i < pos_one.rows(); ++i){
      pos_one.row(i) += avg_two;
      pos_two.row(i) += avg_two;
    }   
  }

  return rmsd;
}

geom::Mat4 MinRMSDSuperposition(EMatX3& pos_one, EMatX3& pos_two, 
                                bool apply_superposition){
  ERVec3 avg_one, avg_two;
  Real rmsd;
  EMat3 rotation;
  Superpose(pos_one, pos_two, avg_one, avg_two, rmsd, 
            rotation, apply_superposition);

  if(apply_superposition){
    for (uint i = 0; i < pos_one.rows(); ++i){
      pos_one.row(i) += avg_two;
      pos_two.row(i) += avg_two;
    }   
  }

  geom::Mat4 transform = GenerateOstTransform(avg_one, avg_two, rotation);
  return transform; 
}

std::pair<geom::Mat4,Real> Superposition(EMatX3& pos_one, EMatX3& pos_two,
                                         bool apply_superposition){
  ERVec3 avg_one, avg_two;
  Real rmsd;
  EMat3 rotation;
  Superpose(pos_one, pos_two, avg_one, avg_two, rmsd, 
            rotation, apply_superposition);

  if(apply_superposition){
    for (uint i = 0; i < pos_one.rows(); ++i){
      pos_one.row(i) += avg_two;
      pos_two.row(i) += avg_two;
    }   
  }

  geom::Mat4 transform = GenerateOstTransform(avg_one, avg_two, rotation);
  return std::make_pair(transform, rmsd); 
}

void FillRMSDMatrix(EMatX3List& position_list, ost::TriMatrix<Real>& data){

  if(position_list.empty()){
    throw promod3::Error("Empty position list in FillRMSDMatrix!");
  }

  // check whether all positions have the same size
  int num_pos = position_list.size();
  int pos_size = position_list[0].rows(); 

  if(num_pos != data.GetSize()) {
    throw promod3::Error("Number of positions and size of rmsd matrix is "
                         "inconsistent in FillRMSDMatrix function!");
  }

  if(pos_size < 3) {
    throw promod3::Error("Observed superposition with < 3 positions to "
                         "superpose!");
  }

  for(int i = 1; i < num_pos; ++i){
    if(position_list[i].rows() != pos_size){
      throw promod3::Error("Size inconsistency in FillRMSDMatrix!");
    }
  }

  // we first center all the positions
  for(int i = 0; i < num_pos; ++i){
    CenterPositions(position_list[i]);
  }

  // fill the data
  Real rmsd;
  EMat3 rot;
  for(int i = 0; i < num_pos; ++i){
    data.Set(i,i,Real(0.0));
    for(int j = i + 1; j < num_pos; ++j){
      TheobaldRMSD(position_list[i], position_list[j], rmsd, rot);
      data.Set(i, j, rmsd);
    }
  }
}

Real SuperposedRMSD(EMatX3& pos_one, EMatX3& pos_two, 
                    uint max_iterations,
                    Real distance_thresh, 
                    std::vector<uint>& indices,
                    bool apply_superposition){

  ERVec3 avg_one, avg_two;
  EMat3 rotation;
  std::vector<Real> squared_distances;
  SuperposeIterative(pos_one, pos_two, max_iterations, distance_thresh,
                     squared_distances, indices, avg_one, avg_two, rotation);

  if(apply_superposition){
    ApplySuperposition(pos_one, avg_one, avg_two, rotation);      
  }

  Real rmsd = 0.0;
  for(uint i = 0; i < indices.size(); ++i){
    rmsd += squared_distances[indices[i]];
  }
  rmsd = std::sqrt(rmsd / indices.size());

  return rmsd;
}

geom::Mat4 MinRMSDSuperposition(EMatX3& pos_one, EMatX3& pos_two, 
                                uint max_iterations,
                                Real distance_thresh, 
                                std::vector<uint>& indices,
                                bool apply_superposition){

  ERVec3 avg_one, avg_two;
  EMat3 rotation;
  std::vector<Real> squared_distances;
  SuperposeIterative(pos_one, pos_two, max_iterations, distance_thresh,
                     squared_distances, indices, avg_one, avg_two, rotation);

  if(apply_superposition){
    ApplySuperposition(pos_one, avg_one, avg_two, rotation);      
  }

  geom::Mat4 transform = GenerateOstTransform(avg_one, avg_two, rotation);
  
  return transform;
}

std::pair<geom::Mat4, Real> Superposition(EMatX3& pos_one, EMatX3& pos_two, 
                                          uint max_iterations,
                                          Real distance_thresh, 
                                          std::vector<uint>& indices,
                                          bool apply_superposition){

  ERVec3 avg_one, avg_two;
  EMat3 rotation;
  std::vector<Real> squared_distances;
  SuperposeIterative(pos_one, pos_two, max_iterations, distance_thresh,
                     squared_distances, indices, avg_one, avg_two, rotation);

  if(apply_superposition){
    ApplySuperposition(pos_one, avg_one, avg_two, rotation);      
  }

  Real rmsd = 0.0;
  for(uint i = 0; i < indices.size(); ++i){
    rmsd += squared_distances[indices[i]];
  }
  rmsd = std::sqrt(rmsd / indices.size());
  geom::Mat4 transform = GenerateOstTransform(avg_one, avg_two, rotation);
  
  return std::make_pair(transform, rmsd);
}

void RigidBlocks(EMatX3& pos_one, EMatX3& pos_two,
                 uint window_length,
                 uint max_iterations,
                 Real distance_thresh, 
                 std::vector<std::vector<uint> >& indices,
                 std::vector<geom::Mat4>& transformations){

  indices.clear();
  transformations.clear();

  std::map<std::vector<uint>, geom::Mat4> unique_blocks;
  geom::Mat4 t;
  uint num_rows = pos_one.rows();

  for(uint start_idx = 0; (start_idx + window_length - 1) < num_rows;
      ++start_idx){

    std::vector<uint> current_indices(window_length,0);
    for(uint i = 0; i < window_length; ++i) current_indices[i] = start_idx + i;

    t = MinRMSDSuperposition(pos_one, pos_two, max_iterations,
                             distance_thresh, current_indices, 
                             false);

    // only add if iterative superposition converged => number of indices
    // must be larger than 3
    if(current_indices.size() >= 3) unique_blocks[current_indices] = t;
  } 

  for(std::map<std::vector<uint>, geom::Mat4>::iterator i = unique_blocks.begin();
      i != unique_blocks.end(); ++i){
    indices.push_back(i->first);
    transformations.push_back(i->second);
  }
}

}} //ns
