/***************************************************************************
                          energy_mod.cpp  -  description
                             -------------------
    begin                : Tue Dec 23 08:37:05 2003
    copyright            : (C) 2002 by Cavalli Andrea
    email                : cavalli@bioc.unizh.ch
**************************************************************************/
/***************************************************************************
 *                                                                         *
 *   This program is free software; you can redistribute it and/or modify  *
 *   it under the terms of the GNU General Public License as published by  *
 *   the Free Software Foundation; either version 2 of the License, or     *
 *   (at your option) any later version.                                   *
 *                                                                         *
 ***************************************************************************/


#include <zeta/module.h>
#include <zeta/class.h>
#include <zeta/zextclass.h>
#include <zeta/zobj.hpp> 


#include <almost.h>
#include <docs.h>
#include <coor.h>
#include <mdb.h>
#include <molecules/molecules.h>
#include <ff.h>
#include <sstream>
#include <ffdoc.h>
#include <rundoc.h>
using namespace Almost;
#define ALMTOSTRING(s) \
template<> \
inline string to_string(const s & ){ \
  return "<" #s ">"; \
}

ALMTOSTRING(ConstraintCollection);

template<typename float_type>
inline string to_string(const EnergyData<float_type> & data){
  stringstream s;
  s<<data;
  return s.str();
}


template<typename float_type>
inline string to_string(const Coor<float_type> & data){
  return "Coor";
}



template<typename float_type>
inline string to_string(const EnergyOptions<float_type> & c19_opt){
  stringstream s;
  s<<c19_opt;
  return s.str();
}


template<typename float_type>
EnergyData<float_type> energy(const Molecules & molecules,
			      const EnergyOptions<float_type> & options,
			      const MDB & mdb){
  //create ff
  ForceField<float_type> ff(molecules,options,mdb);
  
  //coordinates
  Coor<float_type> atomCoor(molecules);

  EnergyData<float_type> ene  = ff.energy(atomCoor);
  
  return ene;
}




template<>
inline string to_string(const ForceField<double> &){
  return "ForceField<double>";
}

#if (__GNUC__ >= 3)
template<>
#endif
map<EnergySelector<double>,ForceField<double>::ff_factory> ForceField<double>::factory
#if (__GNUC__ >= 3)
= map<EnergySelector<double>,ForceField<double>::ff_factory>();
#endif
; 


#if (__GNUC__ >= 3)
template<>
#endif
map<EnergySelector<float>,ForceField<float>::ff_factory> ForceField<float>::factory
#if (__GNUC__ >= 3)
= map<EnergySelector<float>,ForceField<float>::ff_factory>();
#endif
; 

// #if (__GNUC__ >= 3)
// template<>
// #endif
// map<int,ForceField<float>::ff_factory> ForceField<float>::factory
// #if (__GNUC__ >= 3)
// = map<int,ForceField<float>::ff_factory>();
// #endif
// ; 


template<>
inline string to_string(const Molecules & molecules){
  return "<Molecules>";
}

template<>
inline string to_string(const MDB & mdb){
  return "<MDB>";
}


void check_const(ConstraintCollection & cc, 
		 Molecules & mols,
		 double dx, int steps){
  vector<Constraint *> xconst;
  xconst = cc.constraints();
  int asize= mols.atom_size();
  Coor<double> coor(mols), coor_buffer(mols),ff;
  ff.resize(asize);
  ff.clear();

  double W = 0;
  double Ene = 0;
  double tmp_ene = 0;
  for(int i=0;i<xconst.size();i++){
    Ene += xconst[i]->energy_force(coor,ff);
  }
  cout<<"Initial energy "<<Ene<<endl;
  for(int i = 0;i<steps;i++){
    coor_buffer.assign(coor);
    ff.clear();
    for(int j = 0;j<asize;j++){
      int pos = Coor<double>::DIM*j;
      coor.coor[pos] += dx*(drand48()-0.5);
      coor.coor[pos+1] += dx*(drand48()-0.5);
      coor.coor[pos+2] += dx*(drand48()-0.5);
    }
    tmp_ene = 0;
    for(int i=0;i<xconst.size();i++){
      tmp_ene += xconst[i]->energy_force(coor,ff);
    }
    cout<<"Current energy "<<tmp_ene<<endl;
    tmp_ene = tmp_ene - Ene;
    for(int j = 0;j<asize;j++){
      int pos = Coor<double>::DIM*j;
      W += 
	(coor.coor[pos]-coor_buffer.coor[pos])*ff.coor[pos];
      pos = pos + 1;
      W += 
	(coor.coor[pos]-coor_buffer.coor[pos])*ff.coor[pos];
      pos = pos + 1;
      W += 
	(coor.coor[pos]-coor_buffer.coor[pos])*ff.coor[pos];
    }
    cout<<"Current energy difference "<<tmp_ene<<" "<<W<<endl;
  }

}


double check_ff(Molecules & mols,
	      EnergyOptions<double> & opts,
	      MDB & mdb,
	      double dx, int steps){

  ForceField<double> force(mols,opts,mdb);
  int asize= mols.atom_size();
  Coor<double> coor(mols), coor_buffer(mols),ff;
  ff.resize(asize);
  ff.clear();

  double W = 0;
  double Ene0 = 0;
  double Ene1 = 0;
  double DeltaE = 0;

  double Einit = force.energy(coor)["energy"];
  for(int i = 0;i<steps;i++){
    coor_buffer.assign(coor);
    ff.clear();
    Ene0 = force.energy(coor_buffer)["energy"];
    force.force(coor_buffer,ff,true);
    //X1
    for(int j = 0;j<asize;j++){
      int pos = Coor<double>::DIM*j;
      coor.coor[pos] += dx*(drand48()-0.5);
      coor.coor[pos+1] += dx*(drand48()-0.5);
      coor.coor[pos+2] += dx*(drand48()-0.5);
    }
    force.force(coor,ff,true);
    Ene1 = force.energy(coor)["energy"];
    DeltaE += Ene1 - Ene0;
    for(int j = 0;j<asize;j++){
      int pos = Coor<double>::DIM*j;
      W += 
	0.5*(coor.coor[pos]-coor_buffer.coor[pos])*ff.coor[pos];
      pos = pos + 1;
      W += 
	0.5*(coor.coor[pos]-coor_buffer.coor[pos])*ff.coor[pos];
      pos = pos + 1;
      W += 
	0.5*(coor.coor[pos]-coor_buffer.coor[pos])*ff.coor[pos];
    }
    //    cout<<"Current energy difference "<<DeltaE<<" "<<W<<"\n";
    if(i%1000==0)
      cout<<"Diff "<<Ene1<<" "<<Einit+W<<" "<<W<<endl;
  }
  cout<<"End Diff "<<Einit<<" "<<Ene1<<" "<<Einit+W<<endl;
  return (Ene1-Einit-W);
}

void check_ff_pos(Molecules & mols,
		  EnergyOptions<double> & opts,
		  MDB & mdb,
		  double dx){

  ForceField<double> force(mols,opts,mdb);
  int asize= mols.atom_size();
  Coor<double> coor(mols), coor_buffer(mols),ff;
  ff.resize(asize);
  ff.clear();

  double Ene0  = 0;
  double Ene1  = 0;

  //Force
  force.energy_force(coor,ff,true);

  double err = 0;
  Ene0 = force.energy(coor)["energy"];
  for(int pos = 0;pos<mols.atom_size();pos++){
    int apos = pos*Coor<double>::DIM;
    {
      coor_buffer.assign(coor);
      //X 
      coor_buffer.coor[apos] += dx;
      Ene1 = force.energy(coor_buffer)["energy"];
      double terr = abs((Ene1-Ene0)-ff.coor[apos]*dx);
      if(terr>err) err = terr;
      cout<<"Grad: "<< (Ene1-Ene0)/dx<<" "<<ff.coor[apos]<<" ";
    }

    {
      coor_buffer.assign(coor);
      //Y 
      coor_buffer.coor[apos+1] += dx;
      Ene1 = force.energy(coor_buffer)["energy"];
      double terr = abs((Ene1-Ene0)-ff.coor[apos+1]*dx);
      if(terr>err) err = terr;
      cout<< (Ene1-Ene0)/dx<<" "<<ff.coor[apos+1]<<" ";
    }

    {
      coor_buffer.assign(coor);
      //Z 
      coor_buffer.coor[apos+2] += dx;
      Ene1 = force.energy(coor_buffer)["energy"];
      double terr = abs((Ene1-Ene0)-ff.coor[apos+2]*dx);
      if(terr>err) err = terr;
      cout<< (Ene1-Ene0)/dx<<" "<<ff.coor[apos+2]<<"\n";
    }
    cout<<"Maximal error: "<<err<<endl;
  }
}


template<typename float_type> 
EnergyOptions<float_type> amber_ene_defaults(){
  return EnergyOptions<float_type>(EnergyOptions<float_type>::AMBER);
}


template<typename float_type> 
EnergyOptions<float_type> c22_ene_defaults(){
  return EnergyOptions<float_type>(EnergyOptions<float_type>::C22);
}

template<typename float_type> 
EnergyOptions<float_type> c19_ene_defaults(int kind){
  return  c19_defaults<float_type>(kind);
  //  return EnergyOptions<float_type>(EnergyOptions<float_type>::C19);
}


extern "C" {
  void init_energy(){
    //declarations here

    Module mod = Module("energy","Energy Tools");
    Class<EnergyData<double> >(mod.self(),"data");

    ZObj::BinaryOperator<INDEX>::BOM()[pair<TypeInfo,TypeInfo>
				       (typeid(ZType< EnergyData<double> >),
					typeid(ZString))] =
      index_operator<EnergyData<double>,const char *,double>;


    Class<EnergyOptions<double> >(mod.self(),"energy_options")
      .def_attribute("kind",&EnergyOptions< double >::kind,FFDoc::kind)
      .def_attribute("elec",&EnergyOptions<double>::elec,FFDoc::elec)
      .def_attribute("elec_cutoff",&EnergyOptions<double>::elec_cutoff,FFDoc::elec_cutoff)
      .def_attribute("vdw",&EnergyOptions<double>::vdw,FFDoc::vdw)
      .def_attribute("vdw_cutoff",&EnergyOptions<double>::vdw_cutoff,FFDoc::vdw_cutoff)
      .def_attribute("solv",&EnergyOptions<double>::solv,FFDoc::solv)
      .def_attribute("solv_ignore_h",&EnergyOptions<double>::solv_ignore_h,FFDoc::solv_ignore_h)
      .def_attribute("geometry",&EnergyOptions<double>::geometry,RunDoc::geometry)
      .def_attribute("Lx",&EnergyOptions<double>::Lx,RunDoc::Lx)
      .def_attribute("Ly",&EnergyOptions<double>::Ly,RunDoc::Ly)
      .def_attribute("Lz",&EnergyOptions<double>::Lz,RunDoc::Lz)
      .def_attribute("vect",&EnergyOptions<double>::vect,RunDoc::vect)
      .def_attribute("cut_on",&EnergyOptions<double>::cut_on,FFDoc::cut_on)
      .def_attribute("cut_off",&EnergyOptions<double>::cut_off,FFDoc::cut_off)
      .def_attribute("cut_nb",&EnergyOptions<double>::cut_nb,FFDoc::cut_nb)
      .def_attribute("alpha",&EnergyOptions<double>::alpha,FFDoc::alpha)
      .def_attribute("e_zero",&EnergyOptions<double>::e_zero,FFDoc::e_zero)
      .def_attribute("e_14",&EnergyOptions<double>::e_14,FFDoc::e_14)
      .def_attribute("eps",&EnergyOptions<double>::eps,FFDoc::eps)
      .def_attribute("r_water",&EnergyOptions<double>::r_water,FFDoc::r_water)
      .def_attribute("rgbmax",&EnergyOptions<double>::rgbmax,"Undocumented")
      .def_attribute("fsmax",&EnergyOptions<double>::fsmax,"Undocumented")
      .def_attribute("eps_water",&EnergyOptions<double>::eps_water,"Undocumented")
      .def_attribute("offset",&EnergyOptions<double>::offset,"Undocumented")
      .def_attribute("surftens",&EnergyOptions<double>::surftens,"Undocumented")
      .def_attribute("nsmooth",&EnergyOptions<double>::nsmooth,"Undocumented")
      .def_attribute("memb_width",&EnergyOptions<double>::memb_width,"Undocumented")
      .def_attribute("elec_strength",&EnergyOptions<double>::elec_strength,"Undocumented")
      ;

    Class<ForceField<double>,ArgList3<Molecules,EnergyOptions<double>,MDB> >
      (mod.self(),"force_field")
      .def_method("energy",&ForceField<double>::energy)
      .def_method("intenergy",&ForceField<double>::intenergy)
      .def_method("energy_force",&ForceField<double>::energy_force)
      .def_method("force",&ForceField<double>::force)
      .def_method("intforce",&ForceField<double>::intforce)
      ;
    


    mod.def_function("energy","Computes C19 energy for molecules",
		     energy<double>)
      .def_function("check_const","Check that MD constraints work",check_const)
      .def_function("check_ff","Check that FF are correct",check_ff)
      .def_function("check_ff_pos","Check that FF are correct",check_ff_pos);
    ;

    mod
      .def_const("C19",(int)EnergyOptions<double>::C19)
      .def_const("C22",(int)EnergyOptions<double>::C22)
      .def_const("AMBER",(int)EnergyOptions<double>::AMBER)
      .def_const("NONE",(int)EnergyOptions<double>::NONE)
      .def_const("SASA",(int)EnergyOptions<double>::SASA)
      .def_const("EEF1",(int)EnergyOptions<double>::EEF1)
      .def_const("IMM1",(int)EnergyOptions<double>::IMM1)
      .def_const("GB",(int)EnergyOptions<double>::GB)
      .def_const("GB2",(int)EnergyOptions<double>::GB2)
      .def_const("OFF",(int)EnergyOptions<double>::OFF)
      .def_const("RDIE",(int)EnergyOptions<double>::RDIE)
      .def_const("CDIE",(int)EnergyOptions<double>::CDIE)
      .def_const("LJ",(int)EnergyOptions<double>::LJ)
      .def_const("SHIFT",(int)EnergyOptions<double>::SHIFT)
      .def_const("SWITCH",(int)EnergyOptions<double>::SWITCH)
      .def_const("TRUNC",(int)EnergyOptions<double>::TRUNC)
      .def_const("PLAIN",(int)EnergyOptions<double>::PLAIN)
      .def_const("CUBIC",(int)EnergyOptions<double>::CUBIC)
      .def_const("SSE",(int)EnergyOptions<double>::SSE)
      //      .def_const("NONE",(int)Shake::NONE)
      //      .def_const("BONDH",(int)Shake::BONDH)
      //      .def_const("BONDALL",(int)Shake::BONDALL)
      ;

    mod
      //       .def_function("c19_options","Return options with right defaults",c19_md_options<double>)
      .def_function("amber_defaults","Return options with right defaults",amber_ene_defaults<double>)
      .def_function("c22_defaults","Return options with right defaults",c22_ene_defaults<double>)
      .def_function("c19_defaults","Return options with right defaults",c19_ene_defaults<double>);

  }
}