// 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 <cmath>
#include <cassert>
#include <promod3/core/geom_base.hh>
#include <promod3/core/message.hh>

namespace promod3 { namespace core {

void EvaluateGromacsPosRule(uint rule, uint number, const geom::Vec3* anchors,
                            geom::Vec3* positions) {
  // adapted from ost::mol::mm::GromacsPositionRuleEvaluator::EvaluatePosRule
  switch(rule) {

    case 1: {
      assert(number >= 1);
      geom::Vec3 i_j = geom::Normalize(anchors[0]-anchors[1]);
      geom::Vec3 i_k = geom::Normalize(anchors[0]-anchors[2]);
      geom::Vec3 temp = geom::Normalize(i_j+i_k);
      positions[0] = anchors[0] + temp;
      break;
    }

    case 2: {
      assert(number >= 1);
      Real x = 0.333806859234;
      Real y = 0.942641491092;

      geom::Vec3 i = anchors[0];
      geom::Vec3 j = anchors[1];
      geom::Vec3 k = anchors[2];

      geom::Vec3 j_i = geom::Normalize(i-j);
      geom::Vec3 q = geom::Normalize(geom::Cross(geom::Cross(k-j, j_i), j_i));

      positions[0] = i + x*j_i + y*q;
      break;
    }

    case 3: {
      assert(number >= 2);
      Real x = 0.5;
      Real y = 0.866025403784;

      geom::Vec3 i = anchors[0];
      geom::Vec3 j = anchors[1];
      geom::Vec3 k = anchors[2];
      
      geom::Vec3 j_i = geom::Normalize(i-j);
      geom::Vec3 q = geom::Normalize(geom::Cross(geom::Cross(k-j, j_i), j_i));

      positions[0] = i + x*j_i - y*q;
      positions[1] = i + x*j_i + y*q;
      break;
    }

    case 4: {
      assert(number >= 2);
      Real x = 0.333313247568;
      Real y = 0.942816142732;
      Real z = 0.866025403784;


      geom::Vec3 i = anchors[0];
      geom::Vec3 j = anchors[1];
      geom::Vec3 k = anchors[2];
      
      geom::Vec3 j_i = geom::Normalize(i-j);
      geom::Vec3 q = geom::Normalize(geom::Cross(geom::Cross(k-j, j_i), j_i));

      positions[0] = i + x*j_i + y*q;

      geom::Vec3 temp = geom::Normalize(geom::Cross(j_i, positions[0]-i));

      positions[1] = i + x*j_i - 0.5*q + z*temp;
      if (number == 3) positions[2] = i + x*j_i - 0.5*q - z*temp;
      break;
    }

    case 5: {
      assert(number >= 1);
      geom::Vec3 i = anchors[0];
      geom::Vec3 i_j = geom::Normalize(i-anchors[1]);
      geom::Vec3 i_k = geom::Normalize(i-anchors[2]);
      geom::Vec3 i_l = geom::Normalize(i-anchors[3]);
      geom::Vec3 temp = geom::Normalize(i_j+i_k+i_l);
      positions[0] = i + temp;
      break;
    }

    case 6: {
      assert(number >= 2);
      Real x = 0.81649043092;
      Real y = 0.577358966516;

      geom::Vec3 i = anchors[0];
      geom::Vec3 i_j = geom::Normalize(i-anchors[1]);
      geom::Vec3 i_k = geom::Normalize(i-anchors[2]);
      geom::Vec3 temp_one = geom::Normalize(i_j+i_k);
      geom::Vec3 temp_two = geom::Normalize(geom::Cross(i_j, i_k));
      positions[0] = i + y*temp_one + x*temp_two;
      positions[1] = i + y*temp_one - x*temp_two;
      break;
    }

    case 8: {
      assert(number >= 2);
      Real x = 1.36*0.522498564716;
      Real y = 1.36*0.852640164354;

      geom::Vec3 i = anchors[0];
      geom::Vec3 j = anchors[1];
      geom::Vec3 k = anchors[2];
      
      geom::Vec3 j_i = geom::Normalize(i-j);
      geom::Vec3 q = geom::Normalize(geom::Cross(geom::Cross(k-j, j_i), j_i));

      positions[0] = i + x*j_i - y*q;
      positions[1] = i + x*j_i + y*q;
      break;
    }

    case 9: {
      assert(number >= 3);
      Real x1 = 1.23*0.51503807491;
      Real x2 = 1.25*0.422618261741;
      Real y1 = 1.23*0.857167300702;
      Real y2 = 1.25*0.906307787037;

      geom::Vec3 i = anchors[0];
      geom::Vec3 j = anchors[1];
      geom::Vec3 k = anchors[2];
      
      geom::Vec3 j_i = geom::Normalize(i-j);
      geom::Vec3 q = geom::Normalize(geom::Cross(geom::Cross(k-j, j_i), j_i));

      positions[0] = i + x1*j_i - y1*q;
      positions[1] = i + x2*j_i + y2*q;

      // we have to construct a third position hanging at the second one...

      Real x = 0.333806859234;
      Real y = 0.942641491092;

      i = positions[1];
      j = anchors[0];
      k = anchors[1];

      j_i = geom::Normalize(i-j);
      q = geom::Normalize(geom::Cross(geom::Cross(k-j,j_i),j_i));

      positions[2] = i + x*j_i + y*q;
      break;
    }

    default: {
      std::stringstream ss;
      ss << "Gromacs position rule " << rule << " is not implemented!";
      throw promod3::Error(ss.str());
    }
  }
}

void ConstructCTerminalOxygens(const geom::Vec3& c_pos, const geom::Vec3& ca_pos,
                               const geom::Vec3& n_pos, geom::Vec3& o_pos,
                               geom::Vec3& oxt_pos) {
  // use rule 8 above
  geom::Vec3 anchors[] = {c_pos, ca_pos, n_pos};
  geom::Vec3 positions[2];
  EvaluateGromacsPosRule(8, 2, anchors, positions);
  o_pos = positions[0];
  oxt_pos = positions[1];
}

void ConstructAtomPos(const geom::Vec3& A, const geom::Vec3& B,
                      const geom::Vec3& C, Real bond_length, Real angle,
                      Real dihedral, geom::Vec3& pos) {

  Real x,xx,x1;
  x = std::tan((Real(M_PI)-angle)*Real(0.5));
  xx = x*x;
  x1 = Real(1.0)+xx;
  Real bond_length_x_sin_angle = bond_length*Real(2.0)*x/x1;
  Real bond_length_x_cos_angle = bond_length*(Real(1.0)-xx)/x1;

  x = std::tan(Real(0.5)*dihedral);
  xx = x*x;
  x1 = Real(1.0)+xx;
  Real sin_dihedral = Real(2.0)*x/x1;
  Real cos_dihedral = (Real(1.0)-xx)/x1;

  Real ab[] = {B[0]-A[0],B[1]-A[1],B[2]-A[2]};

  Real norm_bc[] = {C[0]-B[0],C[1]-B[1],C[2]-B[2]};
  Real a = norm_bc[0] * norm_bc[0];
  a += norm_bc[1] * norm_bc[1];
  a += norm_bc[2] * norm_bc[2];
  a = Real(1.0) / std::sqrt(a);

  norm_bc[0]*=a;
  norm_bc[1]*=a;
  norm_bc[2]*=a;

  Real norm_n[] = {ab[1]*norm_bc[2]-norm_bc[1]*ab[2],
                   ab[2]*norm_bc[0]-norm_bc[2]*ab[0],
                   ab[0]*norm_bc[1]-norm_bc[0]*ab[1]};

  a = norm_n[0]*norm_n[0];  
  a += norm_n[1]*norm_n[1];  
  a += norm_n[2]*norm_n[2];  
  a = Real(1.0) / std::sqrt(a);

  norm_n[0]*=a;
  norm_n[1]*=a;
  norm_n[2]*=a;

  Real n_x_bc[] = {norm_n[1]*norm_bc[2]-norm_bc[1]*norm_n[2],
                   norm_n[2]*norm_bc[0]-norm_bc[2]*norm_n[0],
                   norm_n[0]*norm_bc[1]-norm_bc[0]*norm_n[1]};

  pos[0] = norm_bc[0]*bond_length_x_cos_angle + 
           n_x_bc[0]*cos_dihedral*bond_length_x_sin_angle + 
           norm_n[0]*sin_dihedral*bond_length_x_sin_angle + C[0];

  pos[1] = norm_bc[1]*bond_length_x_cos_angle + 
           n_x_bc[1]*cos_dihedral*bond_length_x_sin_angle + 
           norm_n[1]*sin_dihedral*bond_length_x_sin_angle + C[1];

  pos[2] = norm_bc[2]*bond_length_x_cos_angle + 
           n_x_bc[2]*cos_dihedral*bond_length_x_sin_angle + 
           norm_n[2]*sin_dihedral*bond_length_x_sin_angle + C[2];

}

void ConstructAtomPos(const geom::Vec3& A, const geom::Vec3& B,
                      const geom::Vec3& C, Real bond_length,
                      Real B_C_bond_length, Real angle, Real dihedral,
                      geom::Vec3& pos) {

  Real x,xx,x1;
  x = std::tan((Real(M_PI)-angle)*Real(0.5));
  xx = x*x;
  x1 = Real(1.0)+xx;
  Real bond_length_x_sin_angle = bond_length*Real(2.0)*x/x1;
  Real bond_length_x_cos_angle = bond_length*(Real(1.0)-xx)/x1;

  x = std::tan(dihedral*Real(0.5));
  xx = x*x;
  x1 = Real(1.0)+xx;
  Real sin_dihedral = Real(2.0)*x/x1;
  Real cos_dihedral = (Real(1.0)-xx)/x1;

  Real ab[] = {B[0]-A[0],B[1]-A[1],B[2]-A[2]};

  Real one_over_B_C_length = 1.0/B_C_bond_length;

  Real norm_bc[] = {(C[0]-B[0]) * one_over_B_C_length,
                    (C[1]-B[1]) * one_over_B_C_length,
                    (C[2]-B[2]) * one_over_B_C_length};

  Real norm_n[] = {ab[1]*norm_bc[2]-norm_bc[1]*ab[2],
                   ab[2]*norm_bc[0]-norm_bc[2]*ab[0],
                   ab[0]*norm_bc[1]-norm_bc[0]*ab[1]};

  Real a = norm_n[0]*norm_n[0];  
  a += norm_n[1]*norm_n[1];  
  a += norm_n[2]*norm_n[2];  
  a = Real(1.0) / std::sqrt(a);


  norm_n[0]*=a;
  norm_n[1]*=a;
  norm_n[2]*=a;

  Real n_x_bc[] = {norm_n[1]*norm_bc[2]-norm_bc[1]*norm_n[2],
                   norm_n[2]*norm_bc[0]-norm_bc[2]*norm_n[0],
                   norm_n[0]*norm_bc[1]-norm_bc[0]*norm_n[1]};

  pos[0] = norm_bc[0]*bond_length_x_cos_angle + 
           n_x_bc[0]*cos_dihedral*bond_length_x_sin_angle + 
           norm_n[0]*sin_dihedral*bond_length_x_sin_angle + C[0];

  pos[1] = norm_bc[1]*bond_length_x_cos_angle + 
           n_x_bc[1]*cos_dihedral*bond_length_x_sin_angle + 
           norm_n[1]*sin_dihedral*bond_length_x_sin_angle + C[1];

  pos[2] = norm_bc[2]*bond_length_x_cos_angle + 
           n_x_bc[2]*cos_dihedral*bond_length_x_sin_angle + 
           norm_n[2]*sin_dihedral*bond_length_x_sin_angle + C[2];
}

void ConstructCBetaPos(const geom::Vec3& n_pos, const geom::Vec3& ca_pos, 
                       const geom::Vec3& c_pos, geom::Vec3& cb_pos) {


  Real n_ca[] = {ca_pos[0]-n_pos[0],ca_pos[1]-n_pos[1],ca_pos[2]-n_pos[2]};
  Real c_ca[] = {ca_pos[0]-c_pos[0],ca_pos[1]-c_pos[1],ca_pos[2]-c_pos[2]};

  //normalize the vector n_ca
  Real temp = n_ca[0] * n_ca[0];
  temp += n_ca[1] * n_ca[1];
  temp += n_ca[2] * n_ca[2];
  temp = 1.0/std::sqrt(temp);

  n_ca[0] *= temp;
  n_ca[1] *= temp;
  n_ca[2] *= temp;

  //normalize the vector c_ca
  temp = c_ca[0] * c_ca[0];
  temp += c_ca[1] * c_ca[1];
  temp += c_ca[2] * c_ca[2];
  temp = 1.0/std::sqrt(temp);

  c_ca[0] *= temp;
  c_ca[1] *= temp;
  c_ca[2] *= temp;

  //let's get a vector perpendicular to the previous two
  Real cross[] = {n_ca[1]*c_ca[2]-c_ca[1]*n_ca[2],
                  n_ca[2]*c_ca[0]-c_ca[2]*n_ca[0],
                  n_ca[0]*c_ca[1]-c_ca[0]*n_ca[1]};

  //and normalize it
  temp = cross[0] * cross[0];
  temp += cross[1] * cross[1];
  temp += cross[2] * cross[2];
  temp = 1.0/std::sqrt(temp);

  cross[0] *= temp;
  cross[1] *= temp;
  cross[2] *= temp;

  //let's assign the values of a vector lying in the plance defined by the three atom
  //positions to n_ca
  n_ca[0] += c_ca[0];
  n_ca[1] += c_ca[1];
  n_ca[2] += c_ca[2];

  //and normalize it
  temp = n_ca[0] * n_ca[0];
  temp += n_ca[1] * n_ca[1];
  temp += n_ca[2] * n_ca[2];
  temp = 1.0/std::sqrt(temp);

  n_ca[0] *= temp;
  n_ca[1] *= temp;
  n_ca[2] *= temp;

  //The vector connecting ca_pos with cb_pos can now be described as a linear combination
  //between n_ca and cross   => a*n_ca + b*cross
  //with a and b chosen, such that the new vector is tilted by an angle x from n_ca
  //and the length is l.
  //this gives us following equations:
  // tan(x) = b/a
  // a**2 + b**2 = l**2
  // 
  // => a = sqrt(l**2/(1+tan(x)**2)) and b = sqrt(l**2/(1+1/tan(x)**2))
  //
  // with x = -54 degree and l=1.5 this gives:
  //
  // a = 0.88168
  // b = 1.2135

  //the cbeta position can now be calculate as the linear combination of the
  //values in cross and n_ca
  cb_pos[0] = ca_pos[0] + 0.88168 * n_ca[0] + 1.2135 * cross[0];
  cb_pos[1] = ca_pos[1] + 0.88168 * n_ca[1] + 1.2135 * cross[1];
  cb_pos[2] = ca_pos[2] + 0.88168 * n_ca[2] + 1.2135 * cross[2];
}

geom::Mat4 RotationAroundLine(const geom::Vec3& axis, const geom::Vec3& anchor,
                              Real angle) {

  Real x,xx,x1;
  x = std::tan(angle*Real(0.5));
  xx = x*x;
  x1 = Real(1.0)+xx;
  Real sin_angle = Real(2.0)*x/x1;
  Real cos_angle = (Real(1.0)-xx)/x1;

  geom::Vec3 n_axis = geom::Normalize(axis);
  Real one_m_cos = Real(1.0)-cos_angle;
  Real u2 = n_axis[0]*n_axis[0];
  Real v2 = n_axis[1]*n_axis[1];
  Real w2 = n_axis[2]*n_axis[2];
  Real u = n_axis[0];
  Real v = n_axis[1];
  Real w = n_axis[2];
  Real a = anchor[0];
  Real b = anchor[1];
  Real c = anchor[2];

  return geom::Mat4(
    u2+(v2+w2)*cos_angle,
    u*v*one_m_cos-w*sin_angle,
    u*w*one_m_cos+v*sin_angle,
    (a*(v2+w2)-u*(b*v+c*w))*one_m_cos+(b*w-c*v)*sin_angle,

    u*v*one_m_cos+w*sin_angle,
    v2+(u2+w2)*cos_angle,
    v*w*one_m_cos-u*sin_angle,
    (b*(u2+w2)-v*(a*u+c*w))*one_m_cos+(c*u-a*w)*sin_angle,

    u*w*one_m_cos-v*sin_angle,
    v*w*one_m_cos+u*sin_angle,
    w2+(u2+v2)*cos_angle,
    (c*(u2+v2)-w*(a*u+b*v))*one_m_cos+(a*v-b*u)*sin_angle,
    0,0,0,1);
}

geom::Mat3 RotationAroundAxis(const geom::Vec3& axis, Real angle) {
  Real x,xx,x1;
  x = std::tan(angle*Real(0.5));
  xx = x*x;
  x1 = Real(1.0)+xx;
  Real sin_ang = Real(2.0)*x/x1;
  Real cos_ang = (Real(1.0)-xx)/x1;
  Real one_m_cos = Real(1.0)-cos_ang;
  geom::Vec3 normalized_axis = geom::Normalize(axis);
  return geom::Mat3(cos_ang+normalized_axis[0]*normalized_axis[0]*one_m_cos,
                    normalized_axis[0]*normalized_axis[1]*one_m_cos-normalized_axis[2]*sin_ang,
                    normalized_axis[0]*normalized_axis[2]*one_m_cos+normalized_axis[1]*sin_ang,

                    normalized_axis[1]*normalized_axis[0]*one_m_cos+normalized_axis[2]*sin_ang,
                    cos_ang+normalized_axis[1]*normalized_axis[1]*one_m_cos,
                    normalized_axis[1]*normalized_axis[2]*one_m_cos-normalized_axis[0]*sin_ang,

                    normalized_axis[2]*normalized_axis[0]*one_m_cos-normalized_axis[1]*sin_ang,
                    normalized_axis[2]*normalized_axis[1]*one_m_cos+normalized_axis[0]*sin_ang,
                    cos_ang+normalized_axis[2]*normalized_axis[2]*one_m_cos);
}

}}
