// 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/structure_db.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/io/seq/load.hh>
#include <ost/mol/entity_view.hh>
#include <ost/mol/alg/accessibility.hh>
#include <ost/mol/alg/sec_struct.hh>

BOOST_AUTO_TEST_SUITE( loop );

namespace {

ost::mol::EntityView LoadTestStructure(const String& path) {
  ost::io::PDBReader reader(path, ost::io::IOProfile());
  ost::mol::EntityHandle test_ent = ost::mol::CreateEntity();
  reader.Import(test_ent);
  ost::conop::HeuristicProcessor heu_proc;
  heu_proc.Process(test_ent);
  return test_ent.CreateFullView();
}

ost::seq::ProfileHandlePtr LoadTestProfile(const String& path) {
  return ost::io::LoadSequenceProfile(path);
}

}

using namespace promod3::loop;

BOOST_AUTO_TEST_CASE(test_structuredb_initialization) {

  // default initialization is supposed to contain all data
  StructureDB db_one;

  BOOST_CHECK(db_one.HasData(StructureDB::Dihedrals));
  BOOST_CHECK(db_one.HasData(StructureDB::SolventAccessibilities));
  BOOST_CHECK(db_one.HasData(StructureDB::ResidueDepths));
  BOOST_CHECK(db_one.HasData(StructureDB::DSSP));
  BOOST_CHECK(db_one.HasData(StructureDB::AAFrequencies));
  BOOST_CHECK(db_one.HasData(StructureDB::AAFrequenciesStruct));


  // test of bitwise or operation 
  StructureDB db_two(StructureDB::Dihedrals | StructureDB::DSSP);

  BOOST_CHECK(db_two.HasData(StructureDB::Dihedrals));
  BOOST_CHECK(!db_two.HasData(StructureDB::SolventAccessibilities));
  BOOST_CHECK(!db_two.HasData(StructureDB::ResidueDepths));
  BOOST_CHECK(db_two.HasData(StructureDB::DSSP));
  BOOST_CHECK(!db_two.HasData(StructureDB::AAFrequencies));
  BOOST_CHECK(!db_two.HasData(StructureDB::AAFrequenciesStruct));
}


BOOST_AUTO_TEST_CASE(test_structuredb_errors_when_adding_coordinates) {

  ost::mol::EntityView test_ent = 
  LoadTestStructure("data/1A88A.pdb");
  
  ost::seq::ProfileHandlePtr profile = LoadTestProfile("data/1A88A.hhm");
  ost::seq::ProfileHandlePtr fake_profile = LoadTestProfile("data/1A88A.hhm");

  String raw_seqres = profile->GetSequence();
  String raw_seqres_fake = raw_seqres;
  raw_seqres_fake[5] = 'Q';
  fake_profile->SetSequence(raw_seqres_fake);

  ost::seq::SequenceHandle seqres = ost::seq::CreateSequence("seqres", 
                                                             raw_seqres);
  ost::seq::SequenceHandle fake_seqres = ost::seq::CreateSequence("fake_seqres", 
                                                              raw_seqres_fake);

  // some StructureDBs to play around with
  StructureDB db_full;
  StructureDB db_no_aa_frequencies(StructureDB::Dihedrals | 
                                   StructureDB::SolventAccessibilities |
                                   StructureDB::ResidueDepths |
                                   StructureDB::DSSP);

  // try to add chain that is not present
  BOOST_CHECK_THROW(db_full.AddCoordinates("1a88", "X", test_ent, 
                                            seqres, profile),
                                            promod3::Error);

  // try to provide no valid profile
  BOOST_CHECK_THROW(db_full.AddCoordinates("1a88", "A", test_ent, 
                                            seqres), promod3::Error);


  // db_full has AAFrequencies, it therefore checks consistency between 
  // seqres and profile... lets trigger that error
  BOOST_CHECK_THROW(db_full.AddCoordinates("1a88", "A", test_ent, 
                                            seqres, fake_profile),
                                            promod3::Error);

  // the above should work if we provide the right profile
  BOOST_CHECK_NO_THROW(db_full.AddCoordinates("1a88", "A", test_ent, 
                                              seqres, profile));

  // we check, whether we can add coordinates without valid profile
  // in the DB without AAFrequencies
  BOOST_CHECK_NO_THROW(db_no_aa_frequencies.AddCoordinates("1a88", "A", 
                                                           test_ent, seqres));

  // BUT an error is thrown if the seqres does not match with seqres!
  BOOST_CHECK_THROW(db_no_aa_frequencies.AddCoordinates("1a88", "A", 
                                                        test_ent, fake_seqres),
                                                        promod3::Error);


  // also check whether size is properly checked when adding structure profile
  ost::seq::ProfileHandlePtr short_profile = profile->Extract(10, 50);
  BOOST_CHECK_THROW(db_full.SetStructureProfile(0, short_profile), promod3::Error);
  BOOST_CHECK_NO_THROW(db_full.SetStructureProfile(0, profile));
}

BOOST_AUTO_TEST_CASE(test_structuredb_longest_stretch_flag) {

  ost::mol::EntityView test_ent_with_gap = 
  LoadTestStructure("data/1A88A_with_gap.pdb");
  ost::seq::ProfileHandlePtr profile = LoadTestProfile("data/1A88A.hhm");
  String raw_seqres = profile->GetSequence();
  ost::seq::SequenceHandle seqres = ost::seq::CreateSequence("seqres", 
                                                             raw_seqres);


  StructureDB db;

  // we add coordinates with longest stretch flag set to true
  db.AddCoordinates("1a88", "A", test_ent_with_gap, 
                    seqres, profile, true);

  // we expect exactly one entry
  BOOST_CHECK(db.GetNumCoords() == 1);

  // with length 255
  BOOST_CHECK(db.GetCoordInfo(0).size == 255); 

  // we add coordinates with longest stretch flag set to false
  db.AddCoordinates("1a88", "A", test_ent_with_gap, 
                    seqres, profile, false);

  // we expect two more entries
  BOOST_CHECK(db.GetNumCoords() == 3);

  // with lengths 17 and 255
  BOOST_CHECK(db.GetCoordInfo(1).size == 17); 
  BOOST_CHECK(db.GetCoordInfo(2).size == 255); 
}

BOOST_AUTO_TEST_CASE(test_fragment_validity) {

  ost::mol::EntityView test_ent = LoadTestStructure("data/1A88A.pdb");
  ost::seq::ProfileHandlePtr profile = LoadTestProfile("data/1A88A.hhm");
  String raw_seqres = profile->GetSequence();
  ost::seq::SequenceHandle seqres = ost::seq::CreateSequence("seqres", 
                                                             raw_seqres);

  StructureDB db;

  // lets add some stuff
  db.AddCoordinates("1a88", "A", test_ent, seqres, profile, true);

  FragmentInfo valid_frag_info(0, 5, 12);
  FragmentInfo invalid_coord_idx(1, 5, 12);
  FragmentInfo invalid_offset(0, 1000, 12);
  FragmentInfo invalid_length(0, 5, 1000);

  BOOST_CHECK_NO_THROW(db.GetSequence(valid_frag_info));
  BOOST_CHECK_THROW(db.GetSequence(invalid_coord_idx), promod3::Error);
  BOOST_CHECK_THROW(db.GetSequence(invalid_offset), promod3::Error);
  BOOST_CHECK_THROW(db.GetSequence(invalid_length), promod3::Error);
}

BOOST_AUTO_TEST_CASE(test_data_extraction) {

  ost::mol::EntityView test_ent = LoadTestStructure("data/1A88A.pdb");
  ost::seq::ProfileHandlePtr profile = LoadTestProfile("data/1A88A.hhm");
  String raw_seqres = profile->GetSequence();
  ost::seq::SequenceHandle seqres = ost::seq::CreateSequence("seqres", 
                                                             raw_seqres);

  StructureDB db;

  // we add coordinates with longest stretch flag set to true,
  // the first entry therefore corresponds to residue with resnum 21
  db.AddCoordinates("1a88", "A", test_ent, seqres, profile, true);

  // we reload an entity in order to forget all generic properties that have 
  // been set in AddCoordinates
  test_ent = LoadTestStructure("data/1A88A.pdb");

  // we test the data for the residues in a range 25-50
  ost::mol::EntityView sel = test_ent.Select("rnum=25:50");
  ost::mol::ResidueViewList residues = sel.GetResidueList();
  ost::mol::ResidueHandleList residue_handles;
  for(ost::mol::ResidueViewList::iterator it = residues.begin();
      it != residues.end(); ++it) {
    residue_handles.push_back(it->GetHandle());
  }

  // lets assign secondary structures and solvent accessibilities
  ost::mol::alg::AssignSecStruct(test_ent);
  ost::mol::alg::Accessibility(test_ent);

  // lets extract data directly from the structure

  String expected_sequence(residues.size(), 'X');
  for(uint i = 0; i < residues.size(); ++i) {
    expected_sequence[i] = residues[i].GetOneLetterCode();
  }

  BackboneList expected_bb_list(expected_sequence, residue_handles);

  std::vector<Real> expected_sol_acc;
  for(ost::mol::ResidueViewList::iterator it = residues.begin();
      it != residues.end(); ++it) {
    expected_sol_acc.push_back(it->GetFloatProp("asaAbs"));
  }

  std::vector<char> expected_sec_struct;
  for(ost::mol::ResidueViewList::iterator it = residues.begin();
      it != residues.end(); ++it) {
    expected_sec_struct.push_back(char(it->GetSecStructure()));
  }

  std::vector<Real> expected_phi_torsions;
  std::vector<Real> expected_psi_torsions;
  for(ost::mol::ResidueViewList::iterator it = residues.begin();
      it != residues.end(); ++it) {
    expected_phi_torsions.push_back(it->GetPhiTorsion().GetAngle());
    expected_psi_torsions.push_back(it->GetPsiTorsion().GetAngle());
  }

  ost::seq::ProfileHandlePtr expected_prof = profile->Extract(24, 50);

  // and compare it with the data from the StructureDB
  FragmentInfo frag_info(0, 24, 26);

  String sequence = db.GetSequence(frag_info);
  BOOST_CHECK(sequence.size() == expected_sequence.size());
  for(uint i = 0; i < sequence.size(); ++i) {
    BOOST_CHECK(sequence[i] == expected_sequence[i]);
  }

  BackboneList bb_list;
  db.FillBackbone(bb_list, frag_info, sequence);
  BOOST_CHECK(bb_list.GetSequence() == expected_bb_list.GetSequence());
  BOOST_CHECK(bb_list.RMSD(expected_bb_list, true) < 0.1);

  // if we reverse the shift that has been applied in AddCoordinates,
  // the superposition flag in RMSD() should not be necessary anymore
  geom::Vec3 shift = db.GetCoordInfo(0).shift;
  geom::Mat4 transform;
  transform(0,3) = -shift[0];
  transform(1,3) = -shift[1];
  transform(2,3) = -shift[2];
  bb_list.ApplyTransform(transform);
  BOOST_CHECK(bb_list.RMSD(expected_bb_list, false) < 0.1);

  std::vector<int> sol_acc = db.GetSolventAccessibilities(frag_info);
  BOOST_CHECK(sol_acc.size() == expected_sol_acc.size());
  for(uint i = 0; i < sol_acc.size(); ++i) {
    BOOST_CHECK(static_cast<unsigned short>(expected_sol_acc[i]) == sol_acc[i]);
  }

  String sec_struct = db.GetDSSPStates(frag_info);
  BOOST_CHECK(sec_struct.size() == expected_sec_struct.size());
  for(uint i = 0; i < sec_struct.size(); ++i) {
    BOOST_CHECK(sec_struct[i] == expected_sec_struct[i]);
  }

  std::vector<std::pair<Real, Real> > torsion_angles = 
  db.GetDihedralAngles(frag_info);
  BOOST_CHECK(torsion_angles.size() == expected_phi_torsions.size());
  BOOST_CHECK(torsion_angles.size() == expected_psi_torsions.size());
  for(uint i = 0; i < torsion_angles.size(); ++i) {
    Real phi_d = std::abs(torsion_angles[i].first - expected_phi_torsions[i]);
    Real psi_d = std::abs(torsion_angles[i].second - expected_psi_torsions[i]);
    BOOST_CHECK(phi_d < 0.001);
    BOOST_CHECK(psi_d < 0.001);
  }

  ost::seq::ProfileHandlePtr prof = db.GetSequenceProfile(frag_info);
  BOOST_CHECK(prof->size() == expected_prof->size());
  BOOST_CHECK(prof->GetSequence() == expected_prof->GetSequence());

  for(uint i = 0; i < prof->size(); ++i) {
    const ost::seq::ProfileColumn& col = (*prof)[i];
    const ost::seq::ProfileColumn& expected_col = (*expected_prof)[i];

    const Real* freqs = col.freqs_begin();
    const Real* expected_freqs = expected_col.freqs_begin();

    for(uint j = 0; j < 20; ++j) {
      Real diff = std::abs(freqs[j] - expected_freqs[j]);
      BOOST_CHECK(diff < 0.001);
    } 
  }

  // the structure profile is expected to contain only zeros for now  
  ost::seq::ProfileHandlePtr structure_prof = db.GetStructureProfile(frag_info);
  for(uint i = 0; i < structure_prof->size(); ++i) {
    const ost::seq::ProfileColumn& col = (*structure_prof)[i];
    const Real* freqs = col.freqs_begin();
    for(uint j = 0; j < 20; ++j) {
      BOOST_CHECK(freqs[j] == 0.0);
    } 
  }

  // we set the sequence profile as structure profile
  db.SetStructureProfile(0, profile);
  structure_prof = db.GetStructureProfile(frag_info);

  // now it should be the same as expected_prof...
  BOOST_CHECK(structure_prof->size() == expected_prof->size());
  BOOST_CHECK(structure_prof->GetSequence() == expected_prof->GetSequence());

  for(uint i = 0; i < structure_prof->size(); ++i) {
    const ost::seq::ProfileColumn& col = (*structure_prof)[i];
    const ost::seq::ProfileColumn& expected_col = (*expected_prof)[i];

    const Real* freqs = col.freqs_begin();
    const Real* expected_freqs = expected_col.freqs_begin();

    for(uint j = 0; j < 20; ++j) {
      Real diff = std::abs(freqs[j] - expected_freqs[j]);
      BOOST_CHECK(diff < 0.001);
    } 
  }
}

BOOST_AUTO_TEST_CASE(test_coordinate_removal) {

  ost::mol::EntityView test_ent_with_gap = 
  LoadTestStructure("data/1A88A_with_gap.pdb");
  ost::seq::ProfileHandlePtr profile = LoadTestProfile("data/1A88A.hhm");
  String raw_seqres = profile->GetSequence();
  ost::seq::SequenceHandle seqres = ost::seq::CreateSequence("seqres", 
                                                             raw_seqres);


  StructureDB db_one;
  StructureDB db_two;

  // we add coordinates with longest stretch flag set to false
  db_one.AddCoordinates("1a88", "A", test_ent_with_gap, 
                        seqres, profile, false);

  db_two.AddCoordinates("1a88", "A", test_ent_with_gap, 
                        seqres, profile, false);

  // we're expecting two entries in the database
  BOOST_CHECK(db_one.GetNumCoords() == 2);
  BOOST_CHECK(db_two.GetNumCoords() == 2);

  db_two.RemoveCoordinates(0);

  // now its one...
  BOOST_CHECK(db_two.GetNumCoords() == 1);

  // check the CoordInfoObjects
  BOOST_CHECK(db_one.GetCoordInfo(1).id == 
              db_two.GetCoordInfo(0).id);
  BOOST_CHECK(db_one.GetCoordInfo(1).chain_name == 
              db_two.GetCoordInfo(0).chain_name);
  // offset requires some special treatment
  BOOST_CHECK(db_one.GetCoordInfo(1).offset - db_one.GetCoordInfo(0).size == 
              db_two.GetCoordInfo(0).offset);
  BOOST_CHECK(db_one.GetCoordInfo(1).size == 
              db_two.GetCoordInfo(0).size);
  BOOST_CHECK(db_one.GetCoordInfo(1).start_resnum == 
              db_two.GetCoordInfo(0).start_resnum);
  BOOST_CHECK(db_one.GetCoordInfo(1).shift == 
              db_two.GetCoordInfo(0).shift);

  // the FragmentInfo objects defined here should point at the same information
  // in their corresponding databases
  FragmentInfo frag_info_one(1, 24, 10);
  FragmentInfo frag_info_two(0, 24, 10);

  BOOST_CHECK(db_one.GetSequence(frag_info_one) == 
              db_two.GetSequence(frag_info_two));

  BackboneList bb_list_one;
  db_one.FillBackbone(bb_list_one, frag_info_one, 
                      db_one.GetSequence(frag_info_one));

  BackboneList bb_list_two;
  db_two.FillBackbone(bb_list_two, frag_info_two, 
                      db_two.GetSequence(frag_info_two));

  BOOST_CHECK(bb_list_one.RMSD(bb_list_two) < 0.001);

  BOOST_CHECK(db_one.GetDSSPStates(frag_info_one) == 
              db_two.GetDSSPStates(frag_info_two));

  BOOST_CHECK(db_one.GetDihedralAngles(frag_info_one) == 
              db_two.GetDihedralAngles(frag_info_two));

  BOOST_CHECK(db_one.GetResidueDepths(frag_info_one) == 
              db_two.GetResidueDepths(frag_info_two));

  BOOST_CHECK(db_one.GetSolventAccessibilities(frag_info_one) == 
              db_two.GetSolventAccessibilities(frag_info_two));

  BOOST_CHECK((*db_one.GetSequenceProfile(frag_info_one)) == 
              (*db_two.GetSequenceProfile(frag_info_two)));

  BOOST_CHECK((*db_one.GetStructureProfile(frag_info_one)) == 
              (*db_two.GetStructureProfile(frag_info_two)));
}




BOOST_AUTO_TEST_SUITE_END();