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

#include <ost/io/mol/pdb_reader.hh>
#include <ost/conop/heuristic.hh>
#include <ost/mol/builder.hh>

BOOST_AUTO_TEST_SUITE( modelling );

using namespace promod3;
using namespace promod3::modelling;

namespace {
// checks, wether bond lengths and angles are similar to those of the CHARMM
// forcefield. It's possible to skip the last backbone item. This is due to the
// fact, that CCD has to modify the last item for better convergence and
// therefore some deviations are possible.
bool CheckBackbonePlausibility(const promod3::loop::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
  Real d_n_ca = 1.453445;
  Real d_ca_c = 1.52119;
  Real d_c_n = 1.348435;
  Real d_c_o = 1.230;
  Real d_ca_cb = 1.5;//this value is not taken from CHARMM, but from the
                     //ost cb reconstruction algorithm (CHARMM would be: 1.538)
  Real a_c_n_ca = 2.16512457033;
  Real a_n_ca_c = 1.87199652246;
  Real a_ca_c_n = 2.04874601581;

  //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.02 || 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;
}
}

BOOST_AUTO_TEST_CASE(test_ccd_closing) {

  loop::TorsionSamplerPtr torsion_sampler = loop::LoadTorsionSamplerCoil();

  //load an entity to test...
  String pdb_name = "data/4nyk.pdb";

  ost::io::PDBReader reader(pdb_name, ost::io::IOProfile());
  ost::mol::EntityHandle test_ent = ost::mol::CreateEntity();
  reader.Import(test_ent);

  ost::conop::ProcessorPtr processor(new ost::conop::HeuristicProcessor(false,true,true,true,ost::conop::CONOP_SILENT));
  processor->Process(test_ent);


  String gap_sequence = "RIDCETRYL";

  // the defined loop already is rather helical, so a helical initialization is
  // already a good aproximation
  loop::BackboneList bb_list_one(gap_sequence);
  loop::BackboneList bb_list_two(gap_sequence);

  ost::mol::ResNum n_res_num(310);
  ost::mol::ResNum c_res_num(318);

  ost::mol::ResidueHandle n_stem = test_ent.FindResidue("A",n_res_num);
  ost::mol::ResidueHandle c_stem = test_ent.FindResidue("A",c_res_num);

  geom::Vec3 c_stem_n_pos = c_stem.FindAtom("N").GetPos();
  geom::Vec3 c_stem_ca_pos = c_stem.FindAtom("CA").GetPos();
  geom::Vec3 c_stem_c_pos = c_stem.FindAtom("C").GetPos();

  //test loop closing without torsion sampler
  CCD closer_one(n_stem, c_stem);
  bool converged_one = closer_one.Close(bb_list_one);
  //test loop closing with torsion sampler
  CCD closer_two(gap_sequence, n_stem, c_stem, torsion_sampler);
  bool converged_two = closer_two.Close(bb_list_two);

  //check, whether they converged
  BOOST_CHECK(converged_one);
  BOOST_CHECK(converged_two);

  //we still manually check, whether the c_stem rmsd really is below the threshold
  const uint last_idx = bb_list_one.size() - 1;
  Real rmsd_one = 0;
  Real rmsd_two = 0;
  rmsd_one += geom::Length2(c_stem_n_pos-bb_list_one.GetN(last_idx));
  rmsd_one += geom::Length2(c_stem_ca_pos-bb_list_one.GetCA(last_idx));
  rmsd_one += geom::Length2(c_stem_c_pos-bb_list_one.GetC(last_idx));
  rmsd_one /= 3;
  rmsd_one = std::sqrt(rmsd_one);

  rmsd_two += geom::Length2(c_stem_n_pos-bb_list_two.GetN(last_idx));
  rmsd_two += geom::Length2(c_stem_ca_pos-bb_list_two.GetCA(last_idx));
  rmsd_two += geom::Length2(c_stem_c_pos-bb_list_two.GetC(last_idx));
  rmsd_two /= 3;
  rmsd_two = std::sqrt(rmsd_two);

  BOOST_CHECK(rmsd_one < 0.1);
  BOOST_CHECK(rmsd_two < 0.1);

  //let's finally check whether the bond lenght and angles are in a plausible range
  BOOST_CHECK(CheckBackbonePlausibility(bb_list_one, true));
  BOOST_CHECK(CheckBackbonePlausibility(bb_list_two, true));
}

BOOST_AUTO_TEST_CASE(test_kic_closing) {
  // load an entity to test...
  String pdb_name = "data/4nyk.pdb";

  ost::io::PDBReader reader(pdb_name, ost::io::IOProfile());
  ost::mol::EntityHandle test_ent = ost::mol::CreateEntity();
  reader.Import(test_ent);

  ost::conop::ProcessorPtr processor(new ost::conop::HeuristicProcessor(false,true,true,true,ost::conop::CONOP_SILENT));
  processor->Process(test_ent);

  String gap_sequence = "PALDFLVEK";

  // note: the defined loop already is rather helical, so a helical
  // initialization (def. bb_list) is already a good aproximation
  loop::BackboneList bb_list(gap_sequence);

  ost::mol::ResNum n_res_num(347);
  ost::mol::ResNum c_res_num(355);

  ost::mol::ResidueHandle n_stem = test_ent.FindResidue("A",n_res_num);
  ost::mol::ResidueHandle c_stem = test_ent.FindResidue("A",c_res_num);

  std::vector<loop::BackboneList> kic_solutions;

  KIC closer(n_stem,c_stem);

  closer.Close(bb_list, 1, 4, 6, kic_solutions);

  // with this constellation, KIC should find 4 solutions
  BOOST_CHECK(kic_solutions.size() == 4);

  // check the plausibility of all kic solutions
  for (std::vector<loop::BackboneList>::iterator i = kic_solutions.begin();
       i != kic_solutions.end(); ++i) {
    BOOST_CHECK(CheckBackbonePlausibility(*i, false));
  }
}


BOOST_AUTO_TEST_SUITE_END();