/***************************************************************************
 tmd.cpp  -  description
 -------------------
 begin                : Tue Nov 18 01:45:35 2003
 copyright            : (C) 1999-2008 by Cavalli Andrea
 author               : : cavalli $
 date                 : : 2003/12/09 11:10:33 $
 id                   : : coor.h,v 1.1.2.2 2003/12/09 11:10:33 cavalli Exp $
 email                : cavalli@bioc.unizh.ch amc82@cam.ac.uk
**************************************************************************/
/***************************************************************************
 *                                                                         *
 *   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 <almost.h>
#include <tmd/tmd.h>
extern void prot2pdb(const Almost::Protein & prot, string ofile);


namespace Almost {
  namespace TMD {
    MoleculesTree::MoleculesTree(const Protein & protein){
      tau = 100;//is 10 succesfull was 20 !!
      dof = 0;
      cc = NULL;
      int begin = protein.atom_offset();
      int end   = protein.atom_size()+protein.atom_offset();
		
		
      for(int i=begin;i<end;i++){
	int nbonds = protein.bonds(i).size();
	for(int j=0;j<nbonds;j++){
	  if(i>protein.bonds(i)[j]) continue;
	  if(protein.is_rotatable(i,protein.bonds(i)[j])){
	    if(protein.atom_name(i,Protein::SHORT)=="C"&&protein.atom_name(protein.bonds(i)[j],Protein::SHORT)=="N"){
	      cout<<">> WARNING omega rotatable "<<protein.atom_name(i,Protein::BASE)
		  <<"-"<<protein.atom_name(protein.bonds(i)[j],Protein::BASE)<<" skipped\n";
	      continue;
	    }
	    rot_bond.push_back(pair<int,int>(i,protein.bonds(i)[j]));
	    cout<<">> Detected rotatable bond "<<protein.atom_name(i,Protein::BASE)<<"-"<<protein.atom_name(protein.bonds(i)[j],Protein::BASE)<<"\n";
	  }
	}
      }
		
      split(protein,rot_bond);
		
      init(protein);
      MoleculesNode::id_counter = 0;
    }
	
    MoleculesTree::~MoleculesTree(){
      delete root;
    }
	
    void MoleculesTree::split(const Protein & p,vector<pair<int,int> > &rot_bond){
      //    visited.insert(rot_bond[0].first);
      vector<MoleculesNode* > node;
      if(!p.atom_size()) return;
      if(!rot_bond.size()) return;
      {
	set<int> visited;
	visited.insert(0);
	__visit(0,
		visited, 
		rot_bond,
		p);     
			
	root = new MoleculesNode();
	node.push_back(root);
	set<int>::iterator iter = visited.begin(),
	  end =visited.end();
	cout<<">> Node: ";
	while(iter!=end){
	  cout<<p.atom_name(*iter,Protein::BASE)<<" ";
	  root->atom.push_back(*iter++);	 
	}
	cout<<" <\n";
	root->atm1=-1;
	root->atm2=-1;
	root->parent = NULL;
      }
      for(int i=0;i<rot_bond.size();i++){
	set<int> visited;
	visited.insert(rot_bond[i].second);
	__visit(rot_bond[i].second,
		visited, 
		rot_bond,
		p);
	MoleculesNode * new_node = new MoleculesNode();
	++dof;
	bond_node.push_back(new_node);
	node.push_back(new_node);
	set<int>::iterator iter = visited.begin(),
	  end =visited.end();
	cout<<">> Node: ";
	while(iter!=end){
	  cout<<p.atom_name(*iter,Protein::BASE)<<" ";
	  new_node->atom.push_back(*iter++);
	}
	cout<<" <\n";
	//       MoleculesNode::find_parent(root,new_node,rot_bond[i].first);
	new_node->atm1 = rot_bond[i].first;
	new_node->atm2 = rot_bond[i].second;
	for(int d=0;d<3;d++){
	  new_node->coor_atm1.d[d] = p[rot_bond[i].first][d];
	  new_node->coor_atm2.d[d] = p[rot_bond[i].second][d];
	}
      }
		
      for(int i=0;i<node.size();i++){
	//Find parent
	if(node[i]->atm2==-1) continue;
	MoleculesNode * par = NULL;
	for(int j=0;j<node.size();j++){
	  if(find(node[j]->atom.begin(),
		  node[j]->atom.end(),node[i]->atm1)!=node[j]->atom.end()){
	    par = node[j];
	    break;
	  }
	}
	if(par){
	  par->child.push_back(node[i]);
	  node[i]->parent = par;
	}
	else {
	  aout<<" Warning unable to find node parent...\n";
	  aout<<p.atom_name(node[i]->atm1,Protein::BASE)<<" "
	      <<p.atom_name(node[i]->atm2,Protein::BASE)<<"\n";
	}
      }
    }
	
    Protein MoleculesTree::mini(const Protein & protein, int steps,
				const TmdOptions & opts,
				const MDB & mdb){
      Protein p = protein;
      Molecules mols;
      mols.add_protein(p);
      Coor<double> coor(mols);
      Coor<double> force(mols);
      ff = new ForceField<double>(mols,opts,mdb);
      EnergyData<double> data, data_old;
      double dt = 0.01;
      double f=1;
      data_old = ff->energy_force(coor,force,true);
      for(int i=0;i<steps; i++){
	root->forward_loop();
	force.clear(0);
	data = ff->energy_force(coor,force,true);
	root->compute_theta_force(coor,force);
	f = root->min_update_theta(dt);
	// 	if(data["energy"]<=data_old["energy"]){
	// 	  dt = 1.2*dt;	  
	// 	}
	// 	else dt = 0.5*dt;
	data_old = data;
	//	dt = f*dt;
	root->compute_new_coor(coor);
	if(i%100==0){
	  cout<<data<<"\n";
	  cout<<"DEBUG Step "<<i+1<<" "<<f<<" "<<dt<<"\n";
	}
      }
		
		
      delete ff;
      assign(mols,coor);
      return mols.protein(0);
    }
	
    Protein MoleculesTree::const_mini(const Protein & protein, int steps,
				      const TmdOptions & opts,
				      const ConstraintCollection & cc,
				      const MDB & mdb){
      Protein p = protein;
      Molecules mols;
      mols.add_protein(p);
      Coor<double> coor(mols);
      Coor<double> force(mols);
      ff = new ForceField<double>(mols,
				  opts,
				  const_cast<ConstraintCollection &>(cc),
				  mdb);
      EnergyData<double> data, data_old;
      double dt = 0.00001;
      double f=1;
      data_old = ff->energy_force(coor,force,true);
      for(int i=0;i<steps; i++){
	root->forward_loop();
	force.clear(0);
	data = ff->energy_force(coor,force,true);
	root->compute_theta_force(coor,force);
	f = root->min_update_theta(dt);
	if(data["energy"]<data_old["energy"]){
	  dt = 1.01*dt;	  
	}
	else dt = 0.9*dt;
	if(dt<0.000001) dt = 0.000001;
	data_old = data;
	//	dt = f*dt;
	root->compute_new_coor(coor);
	if(i%100==0){
	  cout<<data<<"\n";
	  cout<<"DEBUG Step "<<i+1<<" "<<f<<" "<<dt<<"\n";
	}
      }
		
      delete ff;
      assign(mols,coor);
      return mols.protein(0);
    }

    Protein MoleculesTree::run(const Protein & protein, 
			       const TmdOptions & opts,
			       const MDB & mdb){ 
      DCDOStream dcd(opts.trj);
      dcd.init(opts.steps/opts.nsave,protein.atom_size());
      ofstream rst; rst.open(opts.rst.c_str());
      Protein p = protein;
      Molecules mols;
      mols.add_protein(p);
      Coor<double> coor(mols);
      Coor<double> force(mols);
      if(cc)
	ff = new ForceField<double>(mols,opts,*cc,mdb);
      else 
	ff = new ForceField<double>(mols,opts,mdb);

      double T = opts.temp;
      int nprint = opts.nprint;
      int nsave = opts.nsave;
      int steps = opts.steps;
      int nupdate = opts.nupdate;
      tau = opts.tau;
      //Data
      double dt1;
      double dt2;
      double time = 0;
      EnergyData<double> data;
      double E_current;
      double E_back;
      double ene_accuracy = opts.ene_accuracy;//0.025;
      double lambda_energy_max = 1.025;
      double E_pot = 0;
      double E_kin_current = 0;
      double E_kin_back = 0;
      dt1 = dt2 = opts.time_step;
      E_current = 0;
      E_back = 0;
      srand48(opts.seed);
      {
	double mult = 0.01;
	for(int i=0;i<1000;i++){
	  root->init_vel(mult);
	  root->forward_loop();
	  double et = 0;
	  root->kin_energy(et);
	  double temp  =2*et/(dof*Almost::kbol);
	  cout<<"KINETIC ENEGY: "<<i+1<<" "<<et<<" "<<temp<<" "<<T<<" "
	      <<abs(T-temp)/T<<" "
	      <<"\n";
	  if(abs(T-temp)/T<0.1)
	    break;
	  else if(temp > T)
	    mult *=0.75;
	  else {
	    mult *= 1.01;
	  }
	}
      }


      {
	force.clear(0);
	data = ff->energy_force(coor,force,true);
	root->compute_theta_force(coor,force);
	E_pot = data["energy"];
	//E_pot = 0;
	root->forward_loop();
	root->backward_loop();
	root->end_loop();
			
	double e= 0;
	root->kin_energy(e);
	cout<<"KINETIC ENEGY: "<<e<<" "<<2*e/(dof*Almost::kbol)<<"\n";
			
      }
		
		
      E_kin_current = 0;
      root->kin_energy(E_kin_current);
      E_kin_back = E_kin_current;
      E_current = E_kin_current+E_pot;
      E_back = E_kin_back+E_pot;
      //
      for(int i=0;i<steps;i++){
			
	{
	  //	  compute_torsional_acceleration;
	  force.clear(0);
	  data = ff->energy_force(coor,force,i%nupdate==0);
	  E_pot = data["energy"];
	  root->compute_theta_force(coor,force);
				
	}
	E_kin_current = 0;
	root->kin_energy(E_kin_current);
			
	E_current = E_pot + E_kin_current;
	if(i==0) E_back = E_current;
	dt2 = dt1*compute_dt2(E_current,E_back,ene_accuracy,
			      tau,lambda_energy_max);
			
	if(!compute_scaled_theta(T,tau,E_kin_current))
 	  dt2 = 0.1*dt2;	
	//	compute_scaled_theta(T,tau,E_kin_current);
	E_kin_current = 0;
	root->kin_energy(E_kin_current); //Scaled!!!!
			
	root->forward_loop();
	root->backward_loop();
	root->end_loop();
			
			
	{
	  //	  compute_velocities_at_half_step
	  root->new_velocity(dt1,dt2);
	}
	{
	  //	  compute_torsional_increment
	  root->torsion_increment(dt2);
	}
	{
				
	  //	  root->shift_vel();
	}
	{
	  //New coor
	  root->compute_new_coor(coor);
	  if(i%nsave==0){
	    dcd.write(coor);
	    root->write_restart(rst);
	    // 	    cout<<"DEBUG EPOT "<<E_pot<<" ";
	    // 	    cout<<"DEBUG kin "<<E_kin_current<<"\n";
	    //	    prot2pdb(p,"dump.pdb");
	  }
	  if(i%nprint==0){
	    //	    cout<<">> Step : "<<i<<"\n";
	    cout<<data<<"\n";
	    cout<<"DEBUG TIME "<<" "<<i+1<<" "<<time <<" T "
		<< 2*E_kin_current/(Almost::kbol*dof)<<" "
		<<E_kin_current<<" "
		<< dt1 <<" "<<dt2 <<" "
		<<"\n";
	  }
	}
	time += dt2;
			
	dt1 = dt2;
	E_back = E_pot + E_kin_current;
	E_kin_current = 0;
      }

      root->write_restart(rst);
      rst.close();
      root->print_theta();
      {//Dump
	for(int i=0;i<p.atom_size();i++){
	  int ipos = Coor<double>::DIM*i;
	  p[i][0] = coor.coor[ipos];
	  p[i][1] = coor.coor[ipos+1];
	  p[i][2] = coor.coor[ipos+2];
	}
			
			
      }
		
      delete ff;
      return p;
    }

    Protein MoleculesTree::cont(const Protein & protein, 
				const TmdOptions & opts,
				const MDB & mdb,
				string rst_file){ 
      Protein p = protein;
      Molecules mols;
      mols.add_protein(p);
      Coor<double> coor(mols);
      Coor<double> force(mols);
      EnergyData<double> data;
      if(cc)
	ff = new ForceField<double>(mols,opts,*cc,mdb);
      else 
	ff = new ForceField<double>(mols,opts,mdb);

      {
	ifstream rst_in; rst_in.open(rst_file.c_str());
	root->read_restart(rst_in);
	root->forward_loop();
	double e= 0;
	root->kin_energy(e);
      }

      DCDOStream dcd(opts.trj);
      dcd.init(opts.steps/opts.nsave,protein.atom_size());
      ofstream rst; rst.open(opts.rst.c_str());


      double T = opts.temp;
      int nprint = opts.nprint;
      int nsave = opts.nsave;
      int steps = opts.steps;
      int nupdate = opts.nupdate;
      tau = opts.tau;
      //Data
      double dt1;
      double dt2;
      double time = 0;

      double E_current;
      double E_back;
      double ene_accuracy = opts.ene_accuracy;//0.025;
      double lambda_energy_max = 1.025;
      double E_pot = 0;
      double E_kin_current = 0;
      double E_kin_back = 0;
      dt1 = dt2 = opts.time_step;
      E_current = 0;
      E_back = 0;
      srand48(opts.seed);

      {
	force.clear(0);
	data = ff->energy_force(coor,force,true);
	root->compute_theta_force(coor,force);
	E_pot = data["energy"];
	//E_pot = 0;
	root->forward_loop();
	root->backward_loop();
	root->end_loop();
			
	double e= 0;
	root->kin_energy(e);
	cout<<"KINETIC ENEGY: "<<e<<" "<<2*e/(dof*Almost::kbol)<<"\n";
			
      }
		
		
      E_kin_current = 0;
      root->kin_energy(E_kin_current);
      E_kin_back = E_kin_current;
      E_current = E_kin_current+E_pot;
      E_back = E_kin_back+E_pot;
      //
      for(int i=0;i<steps;i++){
			
	{
	  //	  compute_torsional_acceleration;
	  force.clear(0);
	  data = ff->energy_force(coor,force,i%nupdate==0);
	  E_pot = data["energy"];
	  root->compute_theta_force(coor,force);
				
	}
	E_kin_current = 0;
	root->kin_energy(E_kin_current);
			
	E_current = E_pot + E_kin_current;
	if(i==0) E_back = E_current;
	dt2 = dt1*compute_dt2(E_current,E_back,ene_accuracy,
			      tau,lambda_energy_max);
			
	if(!compute_scaled_theta(T,tau,E_kin_current))
 	  dt2 = 0.1*dt2;	
	//	compute_scaled_theta(T,tau,E_kin_current);
	E_kin_current = 0;
	root->kin_energy(E_kin_current); //Scaled!!!!
			
	root->forward_loop();
	root->backward_loop();
	root->end_loop();
			
			
	{
	  //	  compute_velocities_at_half_step
	  root->new_velocity(dt1,dt2);
	}
	{
	  //	  compute_torsional_increment
	  root->torsion_increment(dt2);
	}
	{
				
	  //	  root->shift_vel();
	}
	{
	  //New coor
	  root->compute_new_coor(coor);
	  if(i%nsave==0){
	    dcd.write(coor);
	    root->write_restart(rst);
	    // 	    cout<<"DEBUG EPOT "<<E_pot<<" ";
	    // 	    cout<<"DEBUG kin "<<E_kin_current<<"\n";
	    //	    prot2pdb(p,"dump.pdb");
	  }
	  if(i%nprint==0){
	    //	    cout<<">> Step : "<<i<<"\n";
	    cout<<data<<"\n";
	    cout<<"DEBUG TIME "<<" "<<i+1<<" "<<time <<" T "
		<< 2*E_kin_current/(Almost::kbol*dof)<<" "
		<<E_kin_current<<" "
		<< dt1 <<" "<<dt2 <<" "
		<<"\n";
	  }
	}
	time += dt2;
			
	dt1 = dt2;
	E_back = E_pot + E_kin_current;
	E_kin_current = 0;
      }

      root->write_restart(rst);
      rst.close();
      root->print_theta();
      {//Dump
	for(int i=0;i<p.atom_size();i++){
	  int ipos = Coor<double>::DIM*i;
	  p[i][0] = coor.coor[ipos];
	  p[i][1] = coor.coor[ipos+1];
	  p[i][2] = coor.coor[ipos+2];
	}
			
			
      }
		
      delete ff;
      return p;
    }
	
    Protein MoleculesTree::anneal(const Protein & protein, 
				  double T0, 
				  double T1,				 
				  const TmdOptions & opts,
				  const MDB & mdb){ 
      DCDOStream dcd("dump.dcd");
      dcd.init(opts.steps/opts.nsave,protein.atom_size());
      Protein p = protein;
      Molecules mols;
      mols.add_protein(p);
      Coor<double> coor(mols);
      Coor<double> force(mols);
      if(cc)
	ff = new ForceField<double>(mols,opts,*cc,mdb);
      else 
	ff = new ForceField<double>(mols,opts,mdb);
      //      T=T*Almost::kbol;
      //      int nprint = opts.nprint;
      //      int nsave = opts.nsave;
      int steps = opts.steps;
      int nupdate = opts.nupdate;
      tau = opts.tau;
      double s = 1 - pow(T1/T0,1.0/((double)steps));
      double T = T0;

      //Data
      double dt1;
      double dt2;
      double time = 0;
      EnergyData<double> data;
      double E_current;
      double E_back;
      double ene_accuracy = opts.ene_accuracy;//0.005;
      double lambda_energy_max = 1.025;
      double E_pot = 0;
      double E_kin_current = 0;
      double E_kin_back = 0;
      dt1 = dt2 = opts.time_step;;
      E_current = 0;
      E_back = 0;

      {
	double mult = 0.01;
	for(int i=0;i<1000;i++){
	  root->init_vel(mult);
	  root->forward_loop();
	  double et = 0;
	  root->kin_energy(et);
	  double temp  =2*et/(dof*Almost::kbol);
	  cout<<"KINETIC ENEGY: "<<i+1<<" "<<et<<" "<<temp<<" "<<T<<" "
	      <<abs(T-temp)/T<<" "
	      <<"\n";
	  if(abs(T-temp)/T<0.1)
	    break;
	  else if(temp > T)
	    mult *=0.75;
	  else {
	    mult *= 1.01;
	  }
	}
      }
		
		
      {
	force.clear(0);
	data = ff->energy_force(coor,force,true);
	root->compute_theta_force(coor,force);
	E_pot = data["energy"];
	//E_pot = 0;
	root->forward_loop();
	root->backward_loop();
	root->end_loop();
			
	double e= 0;
	root->kin_energy(e);
	cout<<"KINETIC ENEGY: "<<e<<" "<<2*e/(dof*Almost::kbol)<<"\n";
			
      }
		
		
      E_kin_current = 0;
      root->kin_energy(E_kin_current);
      E_kin_back = E_kin_current;
      E_current = E_kin_current+E_pot;
      E_back = E_kin_back+E_pot;
      //
      for(int i=0;i<steps;i++){
			
	{
	  //	  compute_torsional_acceleration;
	  force.clear(0);
	  data = ff->energy_force(coor,force,i%nupdate==0);
	  E_pot = data["energy"];
	  root->compute_theta_force(coor,force);
				
	}
	E_kin_current = 0;
	root->kin_energy(E_kin_current);
			
	E_current = E_pot + E_kin_current;
	if(i==0) E_back = E_current;
	dt2 = dt1*compute_dt2(E_current,E_back,ene_accuracy,
			      tau,lambda_energy_max);
			
			
 	if(!compute_scaled_theta(T,tau,E_kin_current)){
 	  dt2 = 0.1*dt2;
	  //	  tau *= 2; 
	}
	//	compute_scaled_theta(T,tau,E_kin_current);
	E_kin_current = 0;
	root->kin_energy(E_kin_current); //Scaled!!!!
			
	root->forward_loop();
	root->backward_loop();
	root->end_loop();
			
			
	{
	  //	  compute_velocities_at_half_step
	  root->new_velocity(dt1,dt2);
	}
	{
	  //	  compute_torsional_increment
	  root->torsion_increment(dt2);
	}
	{
				
	  //	  root->shift_vel();
	}
	{
	  //New coor
	  root->compute_new_coor(coor);
	  if(i%1000==0){
	    dcd.write(coor);
	    // 	    cout<<"DEBUG EPOT "<<E_pot<<" ";
	    // 	    cout<<"DEBUG kin "<<E_kin_current<<"\n";
	    //	    prot2pdb(p,"dump.pdb");
	  }
	  if(i%100==0){
	    //	    cout<<">> Step : "<<i<<"\n";
	    cout<<data<<"\n";
	    cout<<"DEBUG TIME "<<" "<<i+1<<" "<<time <<" T "
		<< 2*E_kin_current/(Almost::kbol*dof)<<" "
		<<E_kin_current<<" "
		<< dt1<<" "<<dt2<<" "
		<<"\n";
	  }
	}
	time += dt2;
	T = (1-s)*T;	
	dt1 = dt2;
	E_back = E_pot + E_kin_current;
	E_kin_current = 0;
      }
		
      root->print_theta();
      {//Dump
	for(int i=0;i<p.atom_size();i++){
	  int ipos = Coor<double>::DIM*i;
	  p[i][0] = coor.coor[ipos];
	  p[i][1] = coor.coor[ipos+1];
	  p[i][2] = coor.coor[ipos+2];
	}
			
			
      }
		
      delete ff;
      return p;

    }


    void MoleculesTree::scan(const Protein & prot, double dOmega){
      typedef pair<int,int> ipair;
      map<pair<int,int>,double > upl;
      map<pair<int,int>,double > lol;
      Protein p = prot;
      Molecules mols;
      mols.add_protein(p);
      Coor<double> coor(mols);
      Coor<double> force(mols);
      cout<<">> Initi scan: \n";

      vector<double> T;
      T.resize(rot_bond.size());

      vector<int> C;
      for(int i=0;i<p.atom_size();i++){
	if(p[i].name()[0]=='C'){
	  C.push_back(i);
	}
      }


      bool done = false;
      cout<<"Rotats "<<bond_node.size()<<endl;
      while(!done){
	//	cout<<"Theta: ";
	for(int a=0;a<T.size();a++){
	  //	  cout<<T[a]<<" ";
	  bond_node[a]->theta = M_PI*(T[a]/180.0);
	}
	//	cout<<"\n";
	//	root->print_theta();
	root->compute_new_coor(coor);
	
	for(int aa=0;aa<C.size();aa++){
	  for(int bb=aa+1;bb<C.size();bb++){
	    double d = 0;
	    for(int dd = 0;dd<3;dd++){
	      double dx = coor.coor[Coor<double>::DIM*C[aa]+dd]-coor.coor[Coor<double>::DIM*C[bb]+dd];	    
	      d += dx*dx;
	    }
	    d = sqrt(d);
	    ipair pp(C[aa],C[bb]);
	    if(upl[pp]<d) upl[pp] = d;
	    if(lol[pp]>d||lol[pp]==0) lol[pp]=d;
	  }
	}


	done = true;
	//First and last fixed
	for(int q=T.size()-2;q>=1;q--){
	  //Can inc q?
	  if(T[q]+dOmega<360){
	    T[q] += dOmega;
	    for(int a=q+1;a<T.size();a++)
	      T[a] = 0;
	    done = false;
	    if(q==0){
	      cout<<"Andgle "<<T[q]<<endl;
	    }
	    break;
	  }
	}

      }


      map<ipair, double>::iterator iter = upl.begin(), end = upl.end();
      while(iter!=end){
	ipair pp(iter->first.first,iter->first.second);
	cout<<p.atom_name(iter->first.first)<<" "<<p.atom_name(iter->first.second)<<" "
	    <<lol[pp]<<" "
	    <<" "<<iter->second<<"\n";
	++iter;
      }

    }



    //Name spaces
  }
}
// Local Variables:
// mode:C++
// End: