/***************************************************************************
                          nevermind_mod.cpp  -  description
                             -------------------
    begin                : Tue Feb 10 09:39:51 2009
    copyright            : (C) 1999-2008 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 <molecules/molecules.h>
#include <mdb.h>

using namespace Almost;
typedef map<int,double> cs_map;
typedef map<int,double>::iterator cs_map_iter;
typedef pair<int,int> int_p;


struct ConnType {
  int first;
  int second;
  int spect;
  int peak;
  int kind;
  enum { ONE_TWO, ONE_THREE, ONE_FOUR, NOE};
  ConnType(int f, int s, int sp, int p):
    first(f),second(s),spect(sp),peak(p),kind(NOE){}
  ConnType(int f, int s,int k):
    first(f),second(s),spect(-1),peak(-1),kind(k){}
};

bool operator<(const ConnType &t1 , const ConnType & t2){
  if(t1.kind<t2.kind) return true;
  if(t1.kind>t2.kind) return false;
  if(t1.spect<t2.spect) return true;
  if(t1.spect>t2.spect) return false;
  if(t1.peak<t2.peak) return true;
  if(t1.peak>t2.peak) return false;
  if(t1.first<t2.first) return true;
  if(t1.first>t2.first) return false;
  if(t1.second<t2.second) return true;
  if(t1.second>t2.second) return false;
  return false;
}

struct Peak2D {
  double Vol;
  double Vol2;
  double W[2];
  map<int,double> cs_a;
  map<int,double> cs_b;

  //Flat assignment
  vector<int> a;
  vector<int> b;

  void generate_assignments(){
    a.clear(); b.clear();
    for(cs_map_iter a_iter = cs_a.begin(); a_iter!= cs_a.end(); ++a_iter){
      for(cs_map_iter b_iter = cs_b.begin(); b_iter!=cs_b.end(); ++b_iter){
	a.push_back(a_iter->first);
	b.push_back(b_iter->first);
      }
    }

    cout<<">> Generated "<<a.size()*b.size()<<" assignments\n";
    cout<<">> "<<cs_a.size()*cs_b.size()<<endl;
  }

  int assignment_size() const { return cs_a.size()*cs_b.size();}
};


struct Spect {
  double tol;
  cs_map cs;
  vector<Peak2D> peak;
};


struct NeverMind {
  vector<Spect> spect;


  int mode;
  Protein p;
  vector<string> seq;


  Protein p_ref;

  vector<vector<int> > mat_12;
  vector<vector<int> > mat_13;
  vector<vector<int> > mat_14;

  //Dist rest
  map<pair<int,int>, double>  Upper;
  map<pair<int,int>, double>  Lower;  

  enum { C19, C22, IUPAC };
  
  NeverMind(string name,int m):mode(m),p(name),p_ref(name){}
  
  void read_sequence(const char * file){
    seq.clear();
    ifstream in(file);
			
    if(!in){
      string error = string("\t>> FATAL ERROR: unable to open file: ")+file
	+string("\n");
      cout<<error;
      throw error;
    }
    
    
    while(in.good()){
      char line[256];
      in.getline(line,256);
      stringstream ss(line);
      string s,n;
      ss>>s;
      if(s.size()<3) continue;
      ss>>n;
      n="_"+n;
      if(s.size()==3) seq.push_back(s+n);
      else if(s.size()>3) seq.push_back(string(s.begin(),s.begin()+3)+n);
      else {
	cout<<"\t>> Error in reading sequence file: "<<s<<"\n";
	throw "\t>> Error in reading sequence";
      }
    }

    string mdb_file;
    if(mode==C19) mdb_file = "all19_eef1.mdb";
    if(mode==C22) mdb_file = "all22.mdb";
    if(mode==IUPAC) mdb_file = "a03.mdb";

    MDB mdb(mdb_file.c_str());
    p.build_sequ(seq,mdb,"DEFAULT","DEFAULT");


    //    mat_12.resize(pp.atom_size());
    for(int i=0;i<p.atom_size();i++){
      mat_12.push_back(p.bonds(i));
      //Upper lower
    }
			
    mat_13.resize(p.atom_size());
    mat_14.resize(p.atom_size());
    for(int i=0;i<mat_12.size();i++){
      int a1 = i;
      for(int j=0;j<mat_12[i].size();j++){
	int a2 = mat_12[i][j];
	for(int k=0;k<mat_12[a2].size();k++){
	  int a3 = mat_12[a2][k];
	  if(a1!=a3){
	    mat_13[a1].push_back(a3);
	    for(int l=0;l<mat_12[a3].size();l++){
	      int a4 = mat_12[a3][l];
	      if(a4!=a2&&a4!=a3&&a4!=a2){
		mat_14[a1].push_back(a4);
	      }
	    }
	  }
	}
      }
    }
			
			
    for(int i=0;i<mat_12.size();i++){
      for(int j=0;j<mat_12[i].size();j++){
	int_p pp(i,mat_12[i][j]);
	int_p pinv(mat_12[i][j],i);
	double d = p.distance(i,mat_12[i][j]);
	Upper[pp] = d;
	Lower[pp] = d;
	Upper[pinv] = d;
	Lower[pinv] = d;
      }
    }
			
    for(int i=0;i<mat_13.size();i++){
      for(int j=0;j<mat_13[i].size();j++){
	int_p pp(i,mat_13[i][j]);
	int_p pinv(mat_13[i][j],i);
	double d = p.distance(i,mat_13[i][j]);
	Upper[pp] = d;
	Lower[pp] = d;
	Upper[pinv] = d;
	Lower[pinv] = d;
      }
    }


  }

  void read_spect(string cs_file, string peaks_file, double tol){
    Spect s;
    read_shifts(cs_file.c_str(),s.cs);
    read_peaks(peaks_file.c_str(),s.peak);
    assign_peaks(s,tol);
    cout<<">> Read "<<s.peak.size()<<" peaks for spect "<<peaks_file<<"\n";
    spect.push_back(s);
  }


  void mk_pref(string pdb_file){
    string mdb_file;
    if(mode==C19) mdb_file = "all19_eef1.mdb";
    if(mode==C22) mdb_file = "all22.mdb";
    if(mode==IUPAC) mdb_file = "a03.mdb";

    MDB mdb(mdb_file.c_str());
    PDB pdb(pdb_file);
    p_ref.build_missing(pdb[0][0],mdb,"DEFAULT","DEFAULT");
  }



  void network(){


    map<int,set<int> > cmap;
    map<int_p,set<ConnType> > nw;
    map<int_p, set<pair<set<ConnType>,set<ConnType> > > > ind_nw;
    //1-2 bonds
    for(int a = 0; a<p.atom_size(); a++){
      for(int b = 0;b<p.atom_size();b++){
	if(a==b) continue;
	if(find(mat_12[a].begin(),
		mat_12[a].end(),b)!=mat_12[a].end()){
	  //are is 
	  ConnType t1(a,b,ConnType::ONE_TWO);
	  ConnType t2(b,a,ConnType::ONE_TWO);
       	  nw[int_p(a,b)].insert(t1);
	  nw[int_p(b,a)].insert(t2);
	  cmap[a].insert(b);
	  cmap[b].insert(a);
	}

	if(find(mat_12[b].begin(),
		mat_12[b].end(),b)!=mat_12[b].end()){
	  //are is 
	  ConnType t1(a,b,ConnType::ONE_TWO);
	  ConnType t2(b,a,ConnType::ONE_TWO);
       	  nw[int_p(a,b)].insert(t1);
	  nw[int_p(b,a)].insert(t2);
	  cmap[a].insert(b);
	  cmap[b].insert(a);
	}
      }
    }

    //1-3 bonds
    for(int a = 0; a<p.atom_size(); a++){
      for(int b = 0;b<p.atom_size();b++){
	if(a==b) continue;
	if(find(mat_13[a].begin(),
		mat_13[a].end(),b)!=mat_13[a].end()){
	  //are is 
	  ConnType t1(a,b,ConnType::ONE_THREE);
	  ConnType t2(b,a,ConnType::ONE_THREE);
       	  nw[int_p(a,b)].insert(t1);
	  nw[int_p(b,a)].insert(t2);
	  cmap[a].insert(b);
	  cmap[b].insert(a);

	}

	if(find(mat_13[b].begin(),
		mat_13[b].end(),b)!=mat_13[b].end()){
	  //are is 
	  ConnType t1(a,b,ConnType::ONE_THREE);
	  ConnType t2(b,a,ConnType::ONE_THREE);
       	  nw[int_p(a,b)].insert(t1);
	  nw[int_p(b,a)].insert(t2);
	  cmap[a].insert(b);
	  cmap[b].insert(a);

	}
      }
    }
    


    //1-4 bonds
    for(int a = 0; a<p.atom_size(); a++){
      for(int b = 0;b<p.atom_size();b++){
	if(a==b) continue;
	if(find(mat_14[a].begin(),
		mat_14[a].end(),b)!=mat_14[a].end()){
	  //are is 
	  ConnType t1(a,b,ConnType::ONE_FOUR);
	  ConnType t2(b,a,ConnType::ONE_FOUR);
       	  nw[int_p(a,b)].insert(t1);
	  nw[int_p(b,a)].insert(t2);
	  cmap[a].insert(b);
	  cmap[b].insert(a);
	}

	if(find(mat_14[b].begin(),
		mat_14[b].end(),b)!=mat_14[b].end()){
	  //are is 
	  ConnType t1(a,b,ConnType::ONE_FOUR);
	  ConnType t2(b,a,ConnType::ONE_FOUR);
       	  nw[int_p(a,b)].insert(t1);
	  nw[int_p(b,a)].insert(t2);
	  cmap[a].insert(b);
	  cmap[b].insert(a);
	}
      }
    }


    //Add peaks
    for(int s = 0; s<spect.size();s++){
      for(int p=0; p<spect[s].peak.size();p++){
	for(int i=0;i<spect[s].peak[p].a.size();i++){
	  
	  int_p ab(spect[s].peak[p].a[i],
		   spect[s].peak[p].b[i]);
	  int_p ba(spect[s].peak[p].b[i],
		   spect[s].peak[p].a[i]);

	  ConnType t(spect[s].peak[p].a[i],
		     spect[s].peak[p].b[i],s,p);
	  nw[ab].insert(t);
	  nw[ba].insert(t);
	  cmap[spect[s].peak[p].a[i]].insert(spect[s].peak[p].b[i]);
	  cmap[spect[s].peak[p].b[i]].insert(spect[s].peak[p].a[i]);

	}
      }
    }
    {
      map<int,set<int> >::iterator a_iter = cmap.begin(),
	a_end = cmap.end();
      while(a_iter!=a_end){
	int a = a_iter->first;
	set<int>::iterator b_iter = a_iter->second.begin(),
	  b_end = a_iter->second.end();
	while(b_iter!=b_end){
	  int b = *b_iter;
	  //now thak 
	  set<int>::iterator c_iter = cmap[b].begin(),
	    c_end = cmap[b].end();
	  while(c_iter!=c_end){
	    int c = *c_iter;
	    if(c==a){ ++c_iter; continue;}
	    if(c==b){ ++c_iter; continue;}
	    pair<set<ConnType>,set<ConnType> > pp(nw[int_p(a,b)],nw[int_p(b,c)]);
	    ind_nw[int_p(a,c)].insert(pp);
	    ++c_iter;
	  }
	  ++b_iter;
	}


	++a_iter;
      }
    }


    //Dump
    ofstream out; out.open("nevermind.dump");
    {
      for(int s = 0; s<spect.size();s++){
	for(int p=0; p<spect[s].peak.size();p++){
	  out<<"Peak "<<p+1<<" "<<s+1<<" "
	     <<spect[s].peak[p].W[0]<<" "
	     <<spect[s].peak[p].W[1]<<" "
	     <<spect[s].peak[p].Vol<<"\n";
	  for(int i=0;i<spect[s].peak[p].a.size();i++){
	    int_p ab(spect[s].peak[p].a[i],
		     spect[s].peak[p].b[i]);
	    set<ConnType>::iterator iter = nw[ab].begin(),
	      end = nw[ab].end();
	    out<<"\t"<<p_ref.atom_name(ab.first)<<" "
	       <<p_ref.atom_name(ab.second)<<" "
	       <<p_ref.distance(ab.first,ab.second)<<" "
	      ;
	    while(iter!=end){
	      out<<"( "<<iter->first<<" "<<iter->second<<" ";
	      out<<iter->kind<<" ";
	      out<<iter->peak<<" "<<iter->spect<<" ";
	      out<<") ";
	      ++iter;
	    }

	    //	    out<<"\n";

	    //Indirect 
	    {
	      set<pair<set<ConnType>,set<ConnType> > >::iterator 
		ind_iter = ind_nw[ab].begin(),
		ind_end = ind_nw[ab].end();
	      double P = 1;
	      while(ind_iter!=ind_end){
		//		cout<<"Indirect "<<ab.first<<" "<<ab.second<<endl;
		double pp = 0;
		if(ind_iter->first.begin()->kind==ConnType::NOE)
		  pp = 0.01;
		else 
		  pp = 0.5;

		if(ind_iter->second.begin()->kind==ConnType::NOE)
		  pp *= 0.01;
		else
		  pp *=0.5;
		P *= (1-pp);
// 		out<<"( "<<ind_iter->first.begin()->first<<" "
// 		   <<ind_iter->first.begin()->second<<" "
// 		   <<ind_iter->first.begin()->kind<<") "
// 		  ;
		
// 		out<<"( "<<ind_iter->second.begin()->first<<" "
// 		   <<ind_iter->second.begin()->second<<" "
// 		   <<ind_iter->second.begin()->kind<<") "
// 		  ;
// 		out<<pp<<" "<<1-P<<"\n";
// 		out<<"\n";
		++ind_iter;
	      }
	      out<<">> Prob: "<<1-P<<" "<<P<<endl;
	      if(0){
		while(ind_iter!=ind_end){
		  out<<"( "<<ind_iter->first.begin()->first<<" "
		     <<ind_iter->first.begin()->second<<" "
		     <<ind_iter->first.begin()->kind<<") "
		    ;
		  
		  out<<"( "<<ind_iter->second.begin()->first<<" "
		     <<ind_iter->second.begin()->second<<" "
		     <<ind_iter->second.begin()->kind<<") "
		    ;
		  out<<"\n";
		  ++ind_iter;
		  
		}
	      }
	    }
	  }
	}
      }
    }

  }
  


private:
  void read_shifts(const char * file, cs_map &cs){
    ifstream in(file);
			
    if(!in){
      string error = string("\t>> FATAL ERROR: unable to open file: ")+file
	+string("\n");
      cout<<error;
      throw error;
    }
			
    vector<string> lseq = p.sequence(Protein::LONG);
			
    while(in.good()){
      char line[256];
      in.getline(line,256);
      if(line[0]=='#'||in.gcount()<10) continue;
      stringstream ss(line);
				
      int id;
      double w;
      double err;
      string atm;
      int    nres;
      ss>>id;
      ss>>w;
      ss>>err;
      ss>>atm;
      ss>>nres;
      --nres;
				
      cout<<"\t>> Chemical shift for atom "<<atm<<" in res "<<lseq[nres]<<endl;
      int natm = p.find_atom(lseq[nres]+"/"+atm);
      {
	//Hack for CD1 ILE...
	if(natm==-1&&atm=="CD1"){
	  atm = "CD";
	  natm = p.find_atom(lseq[nres]+"/"+atm);
	}
					
      }
      if(natm==-1){
	cout<<"\t>> Atom not in protein\n";
      }
				
      if(natm!=-1){
	cs[natm] = w;
	cout<<"\t>> CS assignment "<<lseq[nres]+"/"+atm<<" "<<w<<"\n";
      }
				
    }
  }


  void read_peaks(const char * file, vector<Peak2D> & peak){
    ifstream in(file);
			
    if(!in){
      string error = string("\t>> FATAL ERROR: unable to open file: ")+file
	+string("\n");
      cout<<error;
      throw error;
    }
			
    while(in.good()){
      char line[256];
      in.getline(line,256);
      cout<<line<<"\n";
      if(line[0]=='#'||in.gcount()<50) continue;
				
      stringstream ss;
      ss<<line;
				
      //
      int id;
      char c;
      Peak2D p;
      //      p.Ulm = U;
      ss>>id;
      ss>> p.W[0];
      ss>> p.W[1];
      ss>>id;
      ss>> c;
      ss>> p.Vol;
      ss>> p.Vol2;
      peak.push_back(p);
    }
  }


  void assign_peaks(Spect & s, double tol){
    s.tol = tol;
    cs_map_iter iter= s.cs.begin(), end = s.cs.end();
    while(iter!=end){
      for(int i=0;i<s.peak.size();i++){
	if(abs(s.peak[i].W[0]-iter->second)<tol)
	  s.peak[i].cs_a[iter->first] = iter->second;

	if(abs(s.peak[i].W[1]-iter->second)<tol)
	  s.peak[i].cs_b[iter->first] = iter->second;

      }
      ++iter;
    }

    for(int i=0;i<s.peak.size();i++)
      s.peak[i].generate_assignments();

    {
      vector<Peak2D>::iterator iter= s.peak.begin(), end = s.peak.end();
    }
  }


};

ZETATOSTRING(NeverMind);


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

    Module mod = Module("nevermind","Nevermind");
    Class<NeverMind,ArgList2<string,int> >(mod.self(),"nevermind")
      .def_method("read_sequence",&NeverMind::read_sequence)
      .def_method("read_spect",&NeverMind::read_spect)
      .def_method("mk_pref",&NeverMind::mk_pref)
      .def_method("network",&NeverMind::network)
      ;

    mod
      .def_const("C19",(int)NeverMind::C19)	       
      .def_const("C22",(int)NeverMind::C22)
      .def_const("IUPAC",(int)NeverMind::IUPAC)
      ;
  }
}