#ifdef WIN32
#pragma warning(disable : 4786)
#endif

#include "AMBER_energy.h"

#include "Atom.h"
#include "Atom_group.h"
#include "Bond.h"
#include "Bond_group.h"
#include "Atom_kit.h"
#include "AMBER_Van_der_Waals_scaled_energy.h"
#include "Force_field.h"

#include <set>

namespace MM
{

AMBER_energy::
AMBER_energy (const Atom_group  & atoms,
              const Bond_group  & bonds,
              const Force_field & force_field)
//AMBER_energy::
//AMBER_energy (Array_of <Atom> & atoms, 
//              Array_of <Bond> & bonds,
//              Force_field     & force_field)
:
    force_field_(force_field)
{
    int i;

    //for (i=0;  i<bonds.size();  ++i)
    for (i=0;  i<bonds.bond_count();  ++i)
    {
        //setup_bond_energy    (bonds[i]);
        //setup_torsion_energy (bonds[i]);
        setup_bond_energy    (bonds.bond(i));
        setup_torsion_energy (bonds.bond(i));
    }

    //for (i=0;  i<atoms.size();  ++i)
    for (i=0;  i<atoms.atom_count();  ++i)
    {
        //setup_angle_energy (atoms[i]);
        setup_angle_energy (atoms.atom(i));
        setup_electrostatic_energy (i, atoms);
        setup_Van_der_Waals_energy (i, atoms);
    }
}

void AMBER_energy::
setup_bond_energy (const Bond & bond)
{
    double force_constant;
    double equilibrium_distance;

    const Atom & atom_1 = bond.atom(0);
    const Atom & atom_2 = bond.atom(1);

    force_field_.bond (atom_1.kit().mm_type(), atom_2.kit().mm_type(), 
                       force_constant, equilibrium_distance);

    bond_energy_.push_back (AMBER_bond_energy 
        (atom_1, atom_2, force_constant, equilibrium_distance));
}

void AMBER_energy::
setup_torsion_energy (const Bond & bond)
{
    REQUIRE ("Two centers bond", bond.atom_count() == 2);

    const Atom &  atom_2  = bond.atom (0);
    const Atom &  atom_3  = bond.atom (1);

    for (int i=0;  i<atom_2.bound_count();  ++i)
    {
        const Atom & atom_1 = atom_2.bound(i);
        if (&atom_1 == &atom_3) continue;

        for (int j=0;  j<atom_3.bound_count();  ++j)
        {
            const Atom & atom_4 = atom_3.bound(j);
            if (&atom_2 == &atom_4) continue;

            const Array <Force_field::Torsions> * 
                torsion_parameters = force_field_.torsion 
                    (atom_1.kit().mm_type(), 
                     atom_2.kit().mm_type(), 
                     atom_3.kit().mm_type(), 
                     atom_4.kit().mm_type());

            //if (torsion_parameters == 0)
            //    break; //fix

            torsion_energy_.push_back (AMBER_torsion_energy 
                (atom_1, atom_2, atom_3, atom_4, torsion_parameters));
        }
    }
}

void AMBER_energy::
setup_angle_energy (const Atom & atom_2)
{
    for (int j=0;  j<atom_2.bound_count();  ++j)
    {
        const Atom & atom_1 = atom_2.bound (j);

        for (int k=j+1;  k<atom_2.bound_count();  ++k)
        {
            const Atom & atom_3 = atom_2.bound (k);
            
            double force_constant, equilibrium_value;

            force_field_.angle (
                        atom_1.kit().mm_type(), 
                        atom_2.kit().mm_type(), 
                        atom_3.kit().mm_type(), 
                        force_constant, equilibrium_value);

            angle_energy_.push_back (AMBER_angle_energy 
                (atom_1, atom_2, atom_3, force_constant, equilibrium_value));
        }
    }
}

void AMBER_energy::
//setup_electrostatic_energy (int i, Array_of <Atom> & atoms)
//setup_electrostatic_energy (int i, const Atom_group & atoms)
setup_electrostatic_energy (int , const Atom_group & )
{
    //fix
}

namespace 
{

inline void 
find_neighbors (const Atom & atom_1, 
                std::set <const Atom *> & neighbors_1_2,
                std::set <const Atom *> & neighbors_1_3,
                std::set <const Atom *> & neighbors_1_4 )
{
    for (int i=0;  i<atom_1.bound_count();  ++i)
    {
        const Atom & atom_2 = atom_1.bound (i);
        neighbors_1_2.insert (&atom_2);

        for (int j=0;  j<atom_2.bound_count();  ++j)
        {
            const Atom & atom_3 = atom_2.bound (j);
            if (&atom_1 == &atom_3) continue;
            neighbors_1_3.insert (&atom_3);
            
            for (int k=0;  k<atom_3.bound_count();  ++k)
            {
                const Atom & atom_4 = atom_3.bound (k);
                if (&atom_2 == &atom_4) continue;
                neighbors_1_4.insert( &atom_4 );
            }
        }
    }
}

inline bool 
is_neighbor (const Atom & atom, const std::set<const Atom *> & neighbors )
{
    return neighbors.find (&atom) != neighbors.end();
}

}

void  AMBER_energy::
//setup_Van_der_Waals_energy (int i, Array_of <Atom> & atoms)
setup_Van_der_Waals_energy (int i, const Atom_group & atoms)
{
    //Atom & atom_1 = atoms[i];
    const Atom & atom_1 = atoms.atom(i);

    // Find neighbors of atom_1
    std::set <const Atom *> neighbors_1_2;
    std::set <const Atom *> neighbors_1_3;
    std::set <const Atom *> neighbors_1_4;
    find_neighbors (atom_1, neighbors_1_2, neighbors_1_3, neighbors_1_4);

    //for (int j=i+1;  j<atoms.size();  ++j)
    for (int j=i+1;  j<atoms.atom_count();  ++j)
    {
        //Atom & atom_2 = atoms[j];
        const Atom & atom_2 = atoms.atom(j);
        double A, B;
        force_field_.VdW (atom_1.kit().mm_type(), atom_2.kit().mm_type(), A,B);

        if      (is_neighbor (atom_2, neighbors_1_2));
        else if (is_neighbor (atom_2, neighbors_1_3));
        else if (is_neighbor (atom_2, neighbors_1_4))
        {
            Van_der_Waals_scaled_energy_.push_back (
                AMBER_Van_der_Waals_scaled_energy (atom_1, atom_2, A, B,
                    force_field_.VdW_Scale));
        }
        else
        {
            Van_der_Waals_energy_.push_back (
                AMBER_Van_der_Waals_energy (atom_1, atom_2, A, B));
        }
    }
}

double AMBER_energy::
handle_value () const
{
/*    REQUIRE( "Was initialized", electrostatic_energy_.was_initialized() &&
                                Van_der_Waals_Energy_.was_initialized() &&
                                bond_energy_         .was_initialized() &&
                                angle_energy_        .was_initialized() &&
                                dihedral_energy_     .was_initialized() );//*/

    return  //electrostatic_energy() +
            van_der_Waals_energy        () +
            bond_energy                 () +
            angle_energy                () +
            torsion_energy              ();
}

double AMBER_energy::
bond_energy () const
{
    double result = 0.;

    for (unsigned int i=0;  i<bond_energy_.size();  ++i)
        result += bond_energy_[i].value ();

    return result;
}

double AMBER_energy::
angle_energy () const
{
    double result = 0.;

    for (unsigned int i=0;  i<angle_energy_.size();  ++i)
        result += angle_energy_[i].value();

    return result;
}

double AMBER_energy::
torsion_energy () const
{
    double result = 0.;

    for (unsigned int i=0;  i<torsion_energy_.size();  ++i)
        result += torsion_energy_[i].value();

    return result;
}

//double AMBER_energy::electrostatic_energy() const;

double AMBER_energy::
van_der_Waals_energy () const
{
    double result = 0.;
    unsigned int i;

    for (i=0;  i<Van_der_Waals_energy_.size();  ++i)
        result += Van_der_Waals_energy_[i].value ();

    for (i=0;  i<Van_der_Waals_scaled_energy_.size();  ++i)
        result += Van_der_Waals_scaled_energy_[i].value ();

    return result;
}

}//MM
