/*
 * coot-utils/bonded-atoms.cc
 *
 * Copyright 2018 by Medical Research Council
 * Author: Paul Emsley
 *
 * This file is part of Coot
 *
 * This program is free software; you can redistribute it and/or modify
 * it under the terms of the GNU Lesser General Public License as published
 * by the Free Software Foundation; either version 3 of the License, or (at
 * your option) any later version.
 *
 * This program is distributed in the hope that it will be useful, but
 * WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU
 * Lesser General Public License for more details.
 *
 * You should have received a copies of the GNU General Public License and
 * the GNU Lesser General Public License along with this program; if not,
 * write to the Free Software Foundation, Inc., 51 Franklin Street,
 * Fifth Floor, Boston, MA, 02110-1301, USA.
 * See http://www.gnu.org/licenses/
 *
 */

#include "compat/coot-sysdep.h"
#include "coot-coord-utils.hh"
#include "bonded-atoms.hh"

// called for atoms in different residues - we already know
// that at_2 will be in a different residue to at_1.
//
//
bool
coot::are_polymer_bonded(mmdb::Atom *at_1, mmdb::Atom *at_2) {

   bool state = false;

   std::string at_1_name(at_1->GetAtomName());
   std::string at_2_name(at_2->GetAtomName());

   if (at_1_name == " C  ")
      if (at_2_name == " N  ")
         state = true;

   if (at_1_name == " O3'")
      if (at_2_name == " P  ")
         state = true;

   if (! state) {
      if (at_1_name == " O4 " || at_1_name == " O3 " || at_1_name == " O2 ") {
         if (at_2_name == " C1 ") {
            // Let's add a distance check here
            clipper::Coord_orth pt_1 = co(at_1);
            clipper::Coord_orth pt_2 = co(at_2);
            float bond_dist_crit = 1.86; // catches SG-CG
            if ((pt_1-pt_2).lengthsq() < bond_dist_crit*bond_dist_crit) {
               state = true;
            }
         }
      }
   }

   return state;

}

// std::vector<std::set> > ?
std::vector<std::vector<unsigned int> >
coot::make_bonds(mmdb::Manager *mol, int n_selected_atoms, int udd_atom_index_handle) {

   std::vector<std::vector<unsigned int> > v(n_selected_atoms);
   for (int i=0; i<n_selected_atoms; i++)
      v[i].reserve(4);

   bool calc_only = true; // doesn't change anything

   mol->MakeBonds(calc_only); // doesn't do a good job for alt confs, it seems

   for(int imod = 1; imod<=mol->GetNumberOfModels(); imod++) {
      mmdb::Model *model_p = mol->GetModel(imod);
      if (model_p) {

         // make a bond tree for this model

         int n_chains = model_p->GetNumberOfChains();
         for (int ichain=0; ichain<n_chains; ichain++) {
            mmdb::Chain *chain_p = model_p->GetChain(ichain);
            int nres = chain_p->GetNumberOfResidues();
            if (nres > 1) {
               for (int ires=0; ires<(nres-1); ires++) {
                  mmdb::Residue *residue_1_p = chain_p->GetResidue(ires);
                  mmdb::Residue *residue_2_p = chain_p->GetResidue(ires+1);
                  int n_atoms_res_1 = residue_1_p->GetNumberOfAtoms();
                  int n_atoms_res_2 = residue_2_p->GetNumberOfAtoms();
                  for (int iat=0; iat<n_atoms_res_1; iat++) {
                     mmdb::Atom *at_i = residue_1_p->GetAtom(iat);
                     for (int jat=0; jat<n_atoms_res_2; jat++) {
                        mmdb::Atom *at_j = residue_2_p->GetAtom(jat);
                        // std::cout << "iat " << iat << " jat " << jat << std::endl;
                        bool b = are_polymer_bonded(at_i, at_j);
                        if (b) {
                           int idx_1;
                           int idx_2;
                           at_i->GetUDData(udd_atom_index_handle, idx_1);
                           at_j->GetUDData(udd_atom_index_handle, idx_2);
                           if ((idx_1 < 0) || (idx_1 >= n_selected_atoms)) {
                              std::cout << "atom index problem " << idx_1 << " "
                                        << n_selected_atoms << std::endl;
                           } else {
                              if ((idx_2 < 0) || (idx_2 >= n_selected_atoms)) {
                                 std::cout << "atom index problem " << idx_2 << " "
                                           << n_selected_atoms << std::endl;
                              } else {

                                 // Happy path

                                 v[idx_1].push_back(idx_2);
                                 v[idx_2].push_back(idx_1);
                              }
                           }
                        }
                     }
                  }
               }
            }
         }

         // --- make residue internal bonds

         for (int ichain=0; ichain<n_chains; ichain++) {
            mmdb::Chain *chain_p = model_p->GetChain(ichain);
            int nres = chain_p->GetNumberOfResidues();
            for (int ires=0; ires<nres; ires++) {
               mmdb::Residue *residue_p = chain_p->GetResidue(ires);
               int n_atoms = residue_p->GetNumberOfAtoms();
               for (int iat=0; iat<n_atoms; iat++) {
                  mmdb::Atom *at = residue_p->GetAtom(iat);
                  int nb = at->GetNBonds();
                  mmdb::AtomBond *atom_bonds = 0;
                  int n_atom_bonds = 0;
                  at->GetBonds(atom_bonds, n_atom_bonds);
                  if (false)
                     std::cout << "at " << atom_spec_t(at) << " has " << n_atom_bonds
                               << " bonds " << std::endl;
                  if (n_atom_bonds > 0) {
                     for (int ib=0; ib<n_atom_bonds; ib++) {
                        mmdb::Atom *at_bonded = atom_bonds[ib].atom;
                        int idx_1;
                        int idx_2;
                        at->GetUDData(udd_atom_index_handle, idx_1);
                        at_bonded->GetUDData(udd_atom_index_handle, idx_2);
                        if ((idx_1 < 0) || (idx_1 >= n_selected_atoms)) {
                           std::cout << "internal bonds error idx_1 " << idx_1 << std::endl;
                        } else {
                           if ((idx_2 < 0) || (idx_2 >= n_selected_atoms)) {
                              std::cout << "internal bonds error idx_2 " << idx_2 << std::endl;
                           } else {
                              v[idx_1].push_back(idx_2);
                           }
                        }
                        // v[idx_2].push_back(idx_1); // MakeBonds does both ways
                     }
                  }
               }
            }
         }
      }
   }

   mol->RemoveBonds();

   return v;
}

std::vector<std::vector<unsigned int> >
coot::find_1_4_connections(const std::vector<std::vector<unsigned int> > &bonds_vec) {

   std::vector<std::vector<unsigned int> > v(bonds_vec.size());
   for (std::size_t i=0; i<bonds_vec.size(); i++)
      v[i].reserve(4);

   // needs backtracking check

   for (std::size_t i=0; i<bonds_vec.size(); i++) {
      // std::cout << "i " << i << std::endl;
      const std::vector<unsigned int> &v1 = bonds_vec[i];
      for (std::size_t j=0; j<v1.size(); j++) {
         const std::vector<unsigned int> &v2 = bonds_vec[v1[j]];
         // std::cout << "j " << v1[j] << std::endl;
         for (std::size_t k=0; k<v2.size(); k++) {
            if (v2[k] != i) {
               // std::cout << "k " << v2[k] << std::endl;
               const std::vector<unsigned int> &v3 = bonds_vec[v2[k]];
               for (std::size_t l=0; l<v3.size(); l++) {
                  if (v3[l] != i) {
                     if (v3[l] != v1[j]) {
                        // std::cout << "1-4: " << i << " " << v3[l] << std::endl;
                        v[i].push_back(v3[l]);
                     }
                  }
               }
            }
         }
      }
   }

   return v;

}
