// 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/loop/backbone.hh>
#include <promod3/core/message.hh>
#define BOOST_TEST_DYN_LINK
#include <boost/test/unit_test.hpp>

BOOST_AUTO_TEST_SUITE( loop );

using namespace promod3::loop;


namespace {

bool CheckBackbonePlausibility(const BackboneList& bb_list, bool skip_last){

  int correction;
  if(skip_last) correction = 1;
  else correction = 0;

  // the parameters from the backbone trace (N, CA, C atoms) are taken as 
  // averages from the BBTraceParam function. They might differ from amino 
  // acid to amino acid, the tolerances have therefore been adapted.
  // the parameters involving CB and O are chosen, so they match the 
  // corresponding construction functions

  // Python code to produce the parameters:

  // from promod3 import loop
  // 
  // olcs = "ACDEFGHIKLMNPQRSTVWY"
  // 
  // n_ca_bond = 0.0
  // ca_c_bond = 0.0
  // c_n_bond = 0.0
  // c_n_ca_angle = 0.0
  // n_ca_c_angle = 0.0
  // ca_c_n_angle = 0.0
  // 
  // for olc in olcs:
  //     param_tuple = loop.BBTraceParam(olc)
  //     n_ca_bond += param_tuple[0]
  //     ca_c_bond += param_tuple[1]
  //     c_n_bond += param_tuple[2]
  //     c_n_ca_angle += param_tuple[3]
  //     n_ca_c_angle += param_tuple[4]
  //     ca_c_n_angle += param_tuple[5]
  // 
  // print("Real d_n_ca = " + str(n_ca_bond/20) + ';')
  // print("Real d_ca_c = " + str(ca_c_bond/20) + ';')
  // print("Real d_c_n = " + str(c_n_bond/20) + ';')
  // print("Real d_c_o = 1.230;")
  // print("Real d_ca_cb = 1.5; // this value comes from the ost cb reconstruction alg")
  // print("Real a_c_n_ca = " + str(c_n_ca_angle/20) + ';')
  // print("Real a_n_ca_c = " + str(n_ca_c_angle/20) + ';')
  // print("Real a_ca_c_n = " + str(ca_c_n_angle/20) + ';')

  Real d_n_ca = 1.4612949967384339;
  Real d_ca_c = 1.5247449994087219;
  Real d_c_n = 1.3309999704360962;
  Real d_c_o = 1.230;
  Real d_ca_cb = 1.5; // this value comes from the ost cb reconstruction alg
  Real a_c_n_ca = 2.119999885559082;
  Real a_n_ca_c = 1.9374399960041047;
  Real a_ca_c_n = 2.0380001068115234;

  //check internal bond_lengths 
  for (uint i = 0; i < bb_list.size()-correction; ++i) {
    Real diff_n_ca = std::abs(geom::Distance(bb_list.GetN(i), bb_list.GetCA(i))
                              - d_n_ca);
    Real diff_ca_c = std::abs(geom::Distance(bb_list.GetCA(i), bb_list.GetC(i))
                              - d_ca_c);
    Real diff_c_o = std::abs(geom::Distance(bb_list.GetC(i), bb_list.GetO(i))
                             - d_c_o);
    Real diff_ca_cb = std::abs(geom::Distance(bb_list.GetCA(i), bb_list.GetCB(i))
                               - d_ca_cb);

    // ca_c values differ a bit more from amino acid to amino acid...
    if (diff_n_ca > 0.01 || diff_ca_c > 0.018 || diff_c_o > 0.01
        || diff_ca_cb > 0.01) return false;
  }

  //check peptide bonds
  for (uint i = 0; i < bb_list.size()-1-correction; ++i) {
    Real diff_c_n = std::abs(geom::Distance(bb_list.GetC(i),bb_list.GetN(i+1)) - d_c_n);
    if (diff_c_n > 0.01) return false;
  }

  //check internal angle
  for (uint i = 0; i < bb_list.size()-correction; ++i) {
    Real diff_n_ca_c = std::abs(geom::Angle(bb_list.GetN(i)-bb_list.GetCA(i),
                                            bb_list.GetC(i)-bb_list.GetCA(i))
                                - a_n_ca_c);
    // there is quit some variation. e.g. ALA has: 114.44 degrees, 
    // and valine 105.54 degrees
    if (diff_n_ca_c > 0.13) return false;
  }

  //check c_n_ca angle
  for(uint i = 1; i < bb_list.size()-correction; ++i){
    Real diff_c_n_ca = std::abs(geom::Angle(bb_list.GetC(i-1)-bb_list.GetN(i),
                                            bb_list.GetCA(i)-bb_list.GetN(i))
                                - a_c_n_ca);
    if (diff_c_n_ca > 0.05) return false;
  }

  //check ca_c_n angle
  for(uint i = 0; i < bb_list.size()-1-correction; ++i){
    Real diff_ca_c_n = std::abs(geom::Angle(bb_list.GetCA(i)-bb_list.GetC(i),
                                            bb_list.GetN(i+1)-bb_list.GetC(i))
                                - a_ca_c_n);
    if (diff_ca_c_n > 0.05) return false;
  }  

  return true;
}

void CheckEntity(const ost::mol::EntityHandle& ent) {
  // look at all residues
  ost::mol::ResidueHandleList res_list = ent.GetResidueList();
  for (uint res_idx = 0; res_idx < res_list.size(); ++res_idx) {
    // get residues
    const ost::mol::ResidueHandle& r = res_list[res_idx];
    const ost::mol::ResidueHandle r_next = r.GetNext();
    const ost::mol::ResidueHandle r_prev = r.GetPrev();
    // check types
    BOOST_CHECK(r.IsPeptideLinking());
    BOOST_CHECK(r.IsProtein());
    // check links
    if (r_next.IsValid() && r_next.GetNumber() - 1 == r.GetNumber()) {
      BOOST_CHECK(ost::mol::InSequence(r, r_next));
    }
    // check torsions
    if (r_prev.IsValid() && ost::mol::InSequence(r_prev, r)) {
      BOOST_CHECK(r.GetPhiTorsion().IsValid());
      BOOST_CHECK(r.GetOmegaTorsion().IsValid());
    }
    if (r_next.IsValid() && ost::mol::InSequence(r, r_next)) {
      BOOST_CHECK(r.GetPsiTorsion().IsValid());
    }
  }
}

void CheckEntityBB(const ost::mol::EntityHandle& ent,
                   const BackboneList& bb_list) {
  // look at all backbone atoms
  ost::mol::ResidueHandleList res_list = ent.GetResidueList();
  for (uint res_idx = 0; res_idx < res_list.size(); ++res_idx) {
    // get atoms
    const ost::mol::ResidueHandle& r = res_list[res_idx];
    ost::mol::AtomHandle n = r.FindAtom("N");
    ost::mol::AtomHandle ca = r.FindAtom("CA");
    ost::mol::AtomHandle c = r.FindAtom("C");
    ost::mol::AtomHandle o = r.FindAtom("O");
    ost::mol::AtomHandle cb = r.FindAtom("CB");
    // check validity
    BOOST_CHECK(n.IsValid());
    BOOST_CHECK(ca.IsValid());
    BOOST_CHECK(c.IsValid());
    BOOST_CHECK(o.IsValid());
    // compare pos
    BOOST_CHECK_EQUAL(n.GetPos(), bb_list.GetN(res_idx));
    BOOST_CHECK_EQUAL(ca.GetPos(), bb_list.GetCA(res_idx));
    BOOST_CHECK_EQUAL(c.GetPos(), bb_list.GetC(res_idx));
    BOOST_CHECK_EQUAL(o.GetPos(), bb_list.GetO(res_idx));
    // check CB
    if (bb_list.GetAA(res_idx) == ost::conop::GLY) {
      BOOST_CHECK(!cb.IsValid());
    } else {
      BOOST_CHECK(cb.IsValid());
      BOOST_CHECK_EQUAL(cb.GetPos(), bb_list.GetCB(res_idx));
    }
  }
}

}

BOOST_AUTO_TEST_CASE(test_initialization){

  String seq = "HELLYEAH";

  //if no dihedral angles are given,
  //the constructor should construct an alpha helix
  Real helical_phi = -60.0 / 360 * 2 * M_PI;
  Real helical_psi = -45.0 / 360 * 2 * M_PI;
  BackboneList bb_list(seq);

  for(uint i = 0; i < 8; ++i){
    if(i>0) BOOST_CHECK_CLOSE(bb_list.GetPhiTorsion(i),helical_phi,1e-3);
    if(i<7) BOOST_CHECK_CLOSE(bb_list.GetPsiTorsion(i),helical_psi,1e-3);
  }

  //let's create custom dihedrals...
  std::vector<Real> phi;
  std::vector<Real> psi;

  phi.push_back(1.0);
  phi.push_back(-1.0);
  phi.push_back(1.5);
  phi.push_back(1.0);
  phi.push_back(1.0);
  phi.push_back(0.2);
  phi.push_back(0.5);

  psi.push_back(1.0);
  psi.push_back(-1.0);
  psi.push_back(0.5);
  psi.push_back(-1.0);
  psi.push_back(-1.0);
  psi.push_back(0.5);
  psi.push_back(0.3);

  //an error should be thrown when the sizes of the sequence and the dihedrals 
  //doesn't match 
  BOOST_CHECK_THROW(BackboneList b(seq,phi,psi),promod3::Error);

  phi.push_back(-0.1);
  psi.push_back(-0.5);

  BackboneList bb_list_two(seq,phi,psi);

  for(uint i = 0; i < 8; ++i){
    if(i>0) BOOST_CHECK_CLOSE(bb_list_two.GetPhiTorsion(i),phi[i],1e-3);
    
    if(i<7) BOOST_CHECK_CLOSE(bb_list_two.GetPsiTorsion(i),psi[i],1e-3);    
  }

  //last oxygen should have the dihedral angle last_psi+pi  (n-ca-c-o)
  const uint last_idx = bb_list_two.size() - 1;
  BOOST_CHECK_CLOSE(geom::DihedralAngle(bb_list_two.GetN(last_idx),
                                        bb_list_two.GetCA(last_idx),
                                        bb_list_two.GetC(last_idx),
                                        bb_list_two.GetO(last_idx)),
                                        psi.back()+M_PI, 1e-3);
}

BOOST_AUTO_TEST_CASE(test_rotations){
  
  String seq  = "EYECATCHER";

  Real helical_phi = -60.0 / 360 * 2 * M_PI;
  Real helical_psi = -45.0 / 360 * 2 * M_PI;
  BackboneList bb_list(seq);

  BOOST_CHECK_CLOSE(bb_list.GetPhiTorsion(2),helical_phi,1e-3);
  BOOST_CHECK_CLOSE(bb_list.GetPsiTorsion(4),helical_psi,1e-3);
  BOOST_CHECK_CLOSE(bb_list.GetPhiTorsion(8),helical_phi,1e-3);
  BOOST_CHECK_CLOSE(bb_list.GetPsiTorsion(7),helical_psi,1e-3);
 
  bb_list.RotateAroundPhiTorsion(2,0.1);
  BOOST_CHECK_CLOSE(bb_list.GetPhiTorsion(2),helical_phi+0.1,1e-3);
  bb_list.RotateAroundPhiTorsion(2,-0.2);
  BOOST_CHECK_CLOSE(bb_list.GetPhiTorsion(2),helical_phi-0.1,1e-3);

  bb_list.RotateAroundPhiTorsion(8,0.1);
  BOOST_CHECK_CLOSE(bb_list.GetPhiTorsion(8),helical_phi+0.1,1e-3);
  bb_list.RotateAroundPhiTorsion(8,-0.2);
  BOOST_CHECK_CLOSE(bb_list.GetPhiTorsion(8),helical_phi-0.1,1e-3);

  bb_list.RotateAroundPsiTorsion(4,0.1);
  BOOST_CHECK_CLOSE(bb_list.GetPsiTorsion(4),helical_psi+0.1,1e-3);
  bb_list.RotateAroundPsiTorsion(4,-0.2);
  BOOST_CHECK_CLOSE(bb_list.GetPsiTorsion(4),helical_psi-0.1,1e-3);

  bb_list.RotateAroundPsiTorsion(7,0.1);
  BOOST_CHECK_CLOSE(bb_list.GetPsiTorsion(7),helical_psi+0.1,1e-3);
  bb_list.RotateAroundPsiTorsion(7,-0.2);
  BOOST_CHECK_CLOSE(bb_list.GetPsiTorsion(7),helical_psi-0.1,1e-3);

  bb_list.RotateAroundPhiPsiTorsion(2,0.1,0.1);
  BOOST_CHECK_CLOSE(bb_list.GetPhiTorsion(2),helical_phi,1e-3);
  BOOST_CHECK_CLOSE(bb_list.GetPsiTorsion(2),helical_psi+0.1,1e-3);

  bb_list.RotateAroundPhiPsiTorsion(2,-0.2,-0.2);
  BOOST_CHECK_CLOSE(bb_list.GetPhiTorsion(2),helical_phi-0.2,1e-3);
  BOOST_CHECK_CLOSE(bb_list.GetPsiTorsion(2),helical_psi-0.1,1e-3);

  bb_list.RotateAroundPhiPsiTorsion(8,0.1,0.1);
  BOOST_CHECK_CLOSE(bb_list.GetPhiTorsion(8),helical_phi,1e-3);
  BOOST_CHECK_CLOSE(bb_list.GetPsiTorsion(8),helical_psi+0.1,1e-3);

  bb_list.RotateAroundPhiPsiTorsion(8,-0.2,-0.2);
  BOOST_CHECK_CLOSE(bb_list.GetPhiTorsion(8),helical_phi-0.2,1e-3);
  BOOST_CHECK_CLOSE(bb_list.GetPsiTorsion(8),helical_psi-0.1,1e-3);

  bb_list.SetPhiTorsion(3,0.5);
  bb_list.SetPsiTorsion(3,1.0);
  BOOST_CHECK_CLOSE(bb_list.GetPhiTorsion(3),0.5,1e-3);
  BOOST_CHECK_CLOSE(bb_list.GetPsiTorsion(3),1.0,1e-3);

  bb_list.SetPhiTorsion(6,-0.5);
  bb_list.SetPsiTorsion(6,-1.0);
  BOOST_CHECK_CLOSE(bb_list.GetPhiTorsion(6),-0.5,1e-3);
  BOOST_CHECK_CLOSE(bb_list.GetPsiTorsion(6),-1.0,1e-3);

  bb_list.SetPhiPsiTorsion(3, -1.0, -0.5);
  BOOST_CHECK_CLOSE(bb_list.GetPhiTorsion(3), -1.0, 1e-3);
  BOOST_CHECK_CLOSE(bb_list.GetPsiTorsion(3), -0.5, 1e-3);

  bb_list.SetPhiPsiTorsion(6, -1.0, -0.5);
  BOOST_CHECK_CLOSE(bb_list.GetPhiTorsion(6), -1.0, 1e-3);
  BOOST_CHECK_CLOSE(bb_list.GetPsiTorsion(6), -0.5, 1e-3);

  //check whether errors get thrown when invalid indices are passed
  BOOST_CHECK_THROW(bb_list.SetPhiTorsion(-1,0.5),promod3::Error);
  BOOST_CHECK_THROW(bb_list.SetPsiTorsion(-1,0.5),promod3::Error);
  BOOST_CHECK_THROW(bb_list.SetPhiTorsion(20,0.5),promod3::Error);
  BOOST_CHECK_THROW(bb_list.SetPsiTorsion(20,0.5),promod3::Error);

  BOOST_CHECK_THROW(bb_list.RotateAroundPhiTorsion(-1,0.5),promod3::Error);
  BOOST_CHECK_THROW(bb_list.RotateAroundPsiTorsion(-1,0.5),promod3::Error);
  BOOST_CHECK_THROW(bb_list.RotateAroundPhiTorsion(20,0.5),promod3::Error);
  BOOST_CHECK_THROW(bb_list.RotateAroundPsiTorsion(20,0.5),promod3::Error);

  BOOST_CHECK_THROW(bb_list.SetPhiTorsion(0,0.5),promod3::Error);
  BOOST_CHECK_THROW(bb_list.RotateAroundPhiTorsion(0,0.5),promod3::Error);

  BOOST_CHECK_THROW(bb_list.SetPsiTorsion(seq.size()-1,0.5),promod3::Error);
  BOOST_CHECK_THROW(bb_list.RotateAroundPsiTorsion(seq.size()-1,0.5),promod3::Error);

  //check, whether some Backbone distortions have been introduced
  BOOST_CHECK(CheckBackbonePlausibility(bb_list,false));
}

BOOST_AUTO_TEST_CASE(test_transform){
  
  String seq = "HELLYEAH";
  
  BackboneList bb_list(seq);

  geom::Mat4 translation;
  geom::Vec3 t(1.0,3.0,5.0);
  translation.PasteTranslation(t);

  BackboneList transformed_list = bb_list;

  transformed_list.ApplyTransform(translation);

  for(uint i = 0; i < bb_list.size(); ++i){
    BOOST_CHECK(geom::Distance(bb_list.GetN(i)+t,transformed_list.GetN(i))<1e-3);
    BOOST_CHECK(geom::Distance(bb_list.GetCA(i)+t,transformed_list.GetCA(i))<1e-3);
    BOOST_CHECK(geom::Distance(bb_list.GetC(i)+t,transformed_list.GetC(i))<1e-3);
    BOOST_CHECK(geom::Distance(bb_list.GetCB(i)+t,transformed_list.GetCB(i))<1e-3);
    BOOST_CHECK(geom::Distance(bb_list.GetO(i)+t,transformed_list.GetO(i))<1e-3);
  }
}

BOOST_AUTO_TEST_CASE(test_replace_fragment) {
  
  // get dummy backbone and fragment
  String seq = "HELLYEAHEYECATCHER";
  BackboneList bb_list(seq);
  BackboneList frag = bb_list.Extract(4, 9);
  
  // vary fragment
  uint N = frag.size();
  for (uint res_index = 0; res_index < N; ++res_index) {
    if (res_index > 0) frag.RotateAroundPhiTorsion(res_index, 0.5);
    if (res_index < N-1) frag.RotateAroundPsiTorsion(res_index, 0.5);
  }

  // put back in
  bb_list.ReplaceFragment(frag, 4, true);

  // check for consistency
  BOOST_CHECK(CheckBackbonePlausibility(bb_list, false));
}

BOOST_AUTO_TEST_CASE(test_superpose) {
  
  // get dummy backbones
  String seq = "HELLYEAHEYECATCHER";
  BackboneList bb_list(seq);
  BackboneList bb_list2(seq);

  // sanity check distance metrics
  BOOST_CHECK(bb_list.MinCADistance(bb_list) == 0);
  BOOST_CHECK(bb_list.RMSD(bb_list, false) == 0);
  BOOST_CHECK(bb_list.CARMSD(bb_list, false) == 0);
  // check superpose, apply transform and get transform
  bb_list2.SuperposeOnto(bb_list);
  BOOST_CHECK(bb_list.MinCADistance(bb_list2) < 1e-4);
  BOOST_CHECK(bb_list.RMSD(bb_list2, false) < 1e-4);
  BOOST_CHECK(bb_list.CARMSD(bb_list2, false) < 1e-4);
  BOOST_CHECK(bb_list.RMSD(bb_list, true) < 1e-4);
  BOOST_CHECK(bb_list.CARMSD(bb_list, true) < 1e-4);
}

BOOST_AUTO_TEST_CASE(test_bb_to_ost) {
  
  // get dummy backbone and fragment
  String seq = "HELLYEAHEYECATCHER";
  BackboneList bb_list(seq);
  BackboneList frag = bb_list.Extract(4, 9);

  // extract OST entity
  ost::mol::EntityHandle ent = bb_list.ToEntity();
  CheckEntity(ent);
  CheckEntityBB(ent, bb_list);

  // insert fragment in position of 4th residue (fragment pos.)
  ost::mol::ResidueHandleList res_list = ent.GetResidueList();
  ost::mol::ResidueHandle res = res_list[4];
  ost::mol::ChainHandle chain = res.GetChain();
  frag.InsertInto(chain, res.GetNumber());
  CheckEntity(ent);
  CheckEntityBB(ent, bb_list);

  // trigger exceptions
  ost::mol::ChainHandle invalid_chain;
  BOOST_CHECK_THROW(frag.InsertInto(invalid_chain, res.GetNumber()),
                    promod3::Error);
}

BOOST_AUTO_TEST_SUITE_END();