#include <stdio.h>
#include <stdlib.h>
#include <iomanip>
#include <cstdlib>
#include <string>
#include <sstream>
#include <fstream>
#include <iostream>
#include <unistd.h>
#include "math.h"
#include <cassert>
#include <time.h>
#include <sys/time.h>
#include <cstdio> //timer stuff
#include <ctime>
#include <vector>
#include <map> //for map container //NOW NEED g++11
#include <algorithm> //sort from stl
#include <numeric> //for accumulate

#define NUMVECS int(1e8) //how many vectors to create
//#define NUMVECS int(1e9) //how many vectors to create
#define DENOM 10 //for ensemble construction. ratio of max to smallest prob

/*
DISCLAIMER: THIS SOFTWARE IS MADE FREELY AVAILABLE. 
NO GUARANTEES ARE MADE REGARDING ITS CORRECTNESS OR UTILITY.
PURPOSE: GENERATE AN ENSEMBLE OF PEPTIDE/PROTEIN BACKBONE CONFORMATIONS.
FOR USE OF THIS ENSEMBLE TO EVALUATE VARIOUS OBSERVABLES OF THE ENSEMBLE. 
(see program dihedtocc.cpp)
IN: condProbKmeans_TCB_ALL.dat  holds the pair conditional database.
see OUTS below.

SEE README.TXT FOR HELP AND AN EXAMPLE.

IF USED CITE XXXX 
 */

/*
  Need g++11 for map
  program analyticEns.cpp
  N.B. DUNBRACK prob1First[s]=1.0; //DUNBRACK - no info on first residue
  PURPOSE: Sample states of a Markov chain with probabilities
  W(x1,x2,....xn)=P(xn|xn-1)..P(x2|x1)W(x1)
  For given sequence, create the chain states and probabilities
  For each fixed s1 value generate a state with prob = product of condProb(each doublet) * prob(1) 
  in proportion to the P(s2|s1) values in the condProbKmeans_TCB_ALL.dat table.
  Note: Sum s2 P(s2|s1)=1 so normed already for each s1.

  PARAMETERS:  see begin of main 
  DENOM ratio of highest to lowest accepted  probability
  INPUT: (this dir) read condProbs DUNBRACK from condProbKmeans_ALL.dat file
  Line 1 XY 
  Line 2 condProbKmeans[k] k in K centroid[k] k in K (phi/psi each k)
  Line 3 ...
  Line K...
  METHOD
  construct [doublets][probs] for input sequence
  construct corresponding [doublets][numClus]  for the sequence
  The superstates in condProbKmeans_ALL.txt are ordered with s2 varying first DUNBRACK s1 =1 always. 
  (see N.B. below)
  generate  2^nD (nD =1,2, ..numDihed) succesive states using <vector>
  Each round, first cp the vecs from previous round then push_back possible rot values 
  from condProb and evaluate prodProb. 
  Descend sort (state, prodProb) on prodProb and delete if too small.
  DENOM controls relative prodProb to maxProb
  If accepted keep on list of currrent states. 
  if prob too small drop from next round construction. 
  Do until sequences are generated along with their probs. 
 
  OUTS: 
cout: to a check file for number sequence states as constructed in chain. 
Possible number of sequence states and number constructed based on DENOM.
Total kept probability.
mapPhiPsi.dat: the state phi psi correspondence  from read data 
StatesAndProbsDescend_*.dat  unique states and their probs in descend order.
Ensemble_*dat The ensemble of phi/psi dihedrals such that there are e.g. 9 (
for DENOM 10) for the highest prob down to e.g. 1 for the lowest accepted prob.
*/

/*
  N.B. for conditionals
  SUPERSTATE (SS) MATCH TO STATES S1 S2
  SS S1 S2
  0   0 0
  1   0 1
  2   1 0
  3   1 1
  4   2 0
  5   2 1
  where outer loop over Z1, inner loop over Z2 to match ss
  N.B. KEY POINT - don't need vecs 
  vecs1 run 0,1,..numVecs1 (same for 2) those
  integers are equally well covered by loops over s1 and s2.
  ss=numStates2*s1+IndexStates2[s2]; (with ss++) goes over the
  above superstates. 
*/

using namespace std;

int main()
{
  //dirs
  //put results in named dir
  string sdirResults ="/path/analyticEnsOutDirectory";
  const int numDataBase=420; //the number of entries in condProbKmeans_ALL.dat  (exclude "X_*")
  //sequence to be simulated
  const string resName ("EGAAWAASS"); //fasta sequence //PEP_W
  int numRes=static_cast<int>(resName.length() );
 const int numDoub=numRes-1;
  int denom = DENOM; //for ensemble construction
  int numVecs = NUMVECS;
  //for fileNames
  string extnR=static_cast<ostringstream*>( &(ostringstream() << numRes) )->str();
  //for dataBase read
  const int numSupStatesOver=20; //overdim to use for variable numSupStates in nr loop
  double **probCondSeq; //[numDoub][numStates]
  probCondSeq=new double*[numDoub];
  for (int nd=0;nd<numDoub;nd++)
    probCondSeq[nd]=new double[numSupStatesOver];
  for (int nd=0;nd<numDoub;nd++)
    for(int s=0;s<numSupStatesOver;s++)
      probCondSeq[nd][s]=0.0;
  //for write to mapPhiPsi.dat
  double **phi, **psi;
  phi = new double*[numDoub];
  for (int nd=0;nd<numDoub;nd++)
    phi[nd]=new double[numSupStatesOver];
  for (int nd=0;nd<numDoub;nd++)
    for(int s=0;s<numSupStatesOver;s++)
      phi[nd][s]=0.0;
  psi = new double*[numDoub];
  for (int nd=0;nd<numDoub;nd++)
    psi[nd]=new double[numSupStatesOver];
  for (int nd=0;nd<numDoub;nd++)
    for(int s=0;s<numSupStatesOver;s++)
      psi[nd][s]=0.0;
  int *ns1, *ns2; //numStates1[numDoub] "2"for after doublet loop
  ns1 = new int[numDoub];
  ns2 = new int[numDoub];
  int *ns2Kept;
  ns2Kept = new int [numDoub]; //for count how many 

  int *numSupStates;
  numSupStates = new int[numDoub];
  numSupStates[numDoub]={0};
  //for (int nd=0;nd<numDoub;nd++)
  //numSupStates[nd]=0;
  int numStates1=0;
  int numStates2=0;
  //IFILE
  ifstream ifileDun;   
  for(int nd=0;nd<numDoub;nd++){ //seq
    string doubName = resName.substr(nd,2);
    ifileDun.open("condProbKmeans_TCB_ALL.dat"); //no the "X_" entrees
    assert(ifileDun.is_open());
    string doubNameData;
    int miss=0; //count dataBase
    for(int d=0;d<numDataBase;d++){  //over dataBase
      ifileDun>>doubNameData>>numStates2>>numStates1;
      double probTemp[numStates2], phiTemp[numStates2], psiTemp[numStates2];
      int numSupStatesTemp=numStates2*numStates1;
      for (int s=0;s<numSupStatesTemp;s++)
	ifileDun>>probTemp[s]>>phiTemp[s]>>psiTemp[s];
      if(doubName==doubNameData){
	ns1[nd]=numStates1;
	ns2[nd]=numStates2;
	numSupStates[nd]=numSupStatesTemp;
	for (int s=0;s<numSupStatesTemp;s++){
	  probCondSeq[nd][s]=probTemp[s];
	  phi[nd][s]=phiTemp[s];
	  psi[nd][s]=psiTemp[s];
	}
	ifileDun.close();
	break; //if found the doub in the seq
      } //end if
      miss++;
    }//end dataBase
    if(miss==numDataBase){
      cout<<"reached end of dataBase missing  seq doubName="<<doubName<<endl;
      exit(-1);
    }
  }  //end seq
  cout<<"Sequence: ";
  for(int nd=0;nd<numDoub;nd++)
    cout<<resName[nd];
  cout<<resName[numRes-1]<<endl;

  //get the max prodCond
  double maxProdCondProb=1.0;
  double minProdCondProb=1.0;
  for(int nd=0;nd<numDoub;nd++){
    vector<double> temp(numSupStates[nd]);
    for (int s=0;s<numSupStates[nd];s++)
      temp[s]=probCondSeq[nd][s];
    std::vector<double>::iterator result;
    std::vector<double>::iterator result_min;
    result = std::max_element(temp.begin(), temp.end());
    result_min = std::min_element(temp.begin(), temp.end());
    double maxVal=temp[std::distance(temp.begin(), result)];
    double minVal=temp[std::distance(temp.begin(), result_min)];
    maxProdCondProb *= maxVal;
    minProdCondProb *= minVal;
  }
  double probCutOff=maxProdCondProb/denom; //Keep e.g. first decade of probs
  cout<<"maxProdCondProb="<<maxProdCondProb<<" minProdCondProb="<<
    minProdCondProb<<" proCutOff="<<probCutOff<<endl;
  int unsigned long long numChainStates=1;
  numChainStates *= ns1[0];//first one s1
  for (int nr=0;nr<numDoub;nr++){//the rest are s2
    numChainStates *= ns2[nr];
    cout<<"for doublet="<<nr<<" and ns2="<<ns2[nr]<<" there are "<<numChainStates<<
      " subsequence states"<<endl;
  }
  cout<<"number of possible sequence states="<<numChainStates<<endl; 
 
  //for write OUTs
  chdir(sdirResults.c_str());
  //OFILE
  string ofileName;
  ofileName="mapPhiPsi_nR_"+extnR+".dat";
  ofstream ofMap(ofileName.data());
  ofMap<<numRes<<"   ";
  for(int nd=0;nd<numDoub;nd++)
    ofMap<<resName[nd]<<" ";
  ofMap<<resName[numRes-1]<<endl;
  for(int nd=0;nd<numDoub;nd++)
    ofMap<<resName[nd]<<resName[nd+1]<<" "<<numSupStates[nd]<<" ";
  ofMap<<endl;
  for(int nd=0;nd<numDoub;nd++){
    ofMap<<numSupStates[nd]<<"  "<<endl;
    for(int s=0;s<numSupStates[nd];s++)
      ofMap<<phi[nd][s]<<" ";
    ofMap<<endl;
    for(int s=0;s<numSupStates[nd];s++)
      ofMap<<psi[nd][s]<<" ";
    ofMap<<endl;
  }
  ofMap.close();

  // outfiles
  string   fnameS="StatesAndProbsDescend_"+extnR+"_TCB.dat";
  ofstream ofSPD;
  ofSPD.open(fnameS);
  string fnameE="Ensemble_nR_"+extnR+"_TCB.dat";
  ofstream ofEns;
  ofEns.open(fnameE);

  int cumSamps=0; //find total number of Samps to write at end of Ens 
  double *condProb; //current cond prob
  condProb = new double [numVecs];
  double *prodCondProb;   //to track product of condProb
  prodCondProb = new double[numVecs];
  for (int nv=0;nv<numVecs;nv++){
    condProb[nv]=0.0;
    prodCondProb[nv]=1.0;
  }
  //timer start
  clock_t start = clock();

  int nvTop=1;
  int nvTopPrev=nvTop;
  int vTop=0;
  std::vector <std::vector<int> > Vecs(numVecs);
  for(int i=0;i<numVecs;i++)
    {
      vector<int> row;
      Vecs.push_back(row);
    }
  //set base vecs with one vector of size 1 and prob 1
  vTop++;
  for (int nv=0;nv<nvTop;nv++){
    Vecs[nv].push_back(nv);
    condProb[nv]=1.0;
    prodCondProb[nv] *= condProb[nv];
  }
  //begin construction of new vectors
  double temp=0.0;
  int nD=0;
  while(nD<numDoub)
    {
      nvTopPrev=nvTop;
      nvTop *= ns2[nD];

      if(nvTop>numVecs-100){
	cout<<"for nD+1="<<nD+1<<endl;
	cout<<"nvTop="<<nvTop<<" and numVecs="<<numVecs<<" : you need more vecs"<<endl;
	exit(-1);
      }
      for(int nv=nvTopPrev;nv<nvTop;nv++) {//cp (only for new ones
	Vecs[nv]=Vecs[nv%nvTopPrev];
	prodCondProb[nv]=prodCondProb[nv%nvTopPrev];
      }
      vTop++; //increment vTop

      //get conditionals pC[s2|s1] for nr from read probCondSeq
      double *probCond;
      probCond = new double[numSupStates[nD]];
      for(int s=0;s<numSupStates[nD];s++)
	probCond[s]=-1.0;
      int ss=0;
      for(int s1=0;s1<ns1[nD];s1++) //over s1 then over s2
	for(int s2=0;s2<ns2[nD];s2++)
	  {
	    probCond[ss]=probCondSeq[nD][ss]; //from read data
	    ss++;
	  }
      for(int nv=0;nv<nvTop;nv++)//create
	{
	  Vecs[nv].push_back(nv/nvTopPrev);
	  prodCondProb[nv] *=  probCond[Vecs[nv][vTop-1]];
	}

      //sort on prodProb
      vector<pair<double,std::vector<int>>> A; //for sort of vectors by prodCondProb
      for(int nv=0;nv<nvTop;nv++){
	A.push_back(make_pair(prodCondProb[nv],Vecs[nv]));
      }
      std::sort(A.rbegin(), A.rend()); //sort on prodProb
      int cntCut=0;
      for(int nv=0;nv<nvTop;nv++){
	if(A[nv].first<probCutOff)
	  cntCut++;
      }
      cout<<"for nD= "<<nD<<endl;
      cout<<"nvTop before cut="<<nvTop;
      nvTop -= cntCut; //decrement nvTop
      //set Vecs to the remaining ones
      for(int nv=0;nv<nvTop;nv++){
	for (int v=0;v<vTop;v++)
	  Vecs[nv][v]=A[nv].second[v];
	prodCondProb[nv]=A[nv].first;
      }
      cout<<" after cut="<<nvTop<<endl;
      nD++;

      //use A to build ensemble
      if(nD==numDoub)
	{
	  //create ensemble
	  int Asize = A.size();
	  Asize -= cntCut; //set size to the reduce size
	  int fEtot=0;
	  vector<int> fE(Asize);
	  double Asum=0.0;
	  for(int p=0;p<Asize;p++)
	    Asum += A[p].first;
	  for(int p=0;p<Asize;p++)
	    fE[p] = A[p].first/Asum; //norm
	  for(int p=0;p<Asize;p++){
	    fE[p] = A[p].first/A[Asize-1].first; //convert to int number in prop to probs
	    fE[p] = int(fE[p]);
	    fEtot += fE[p];
	  }
	  ofSPD<<"numStates "<<Asize<<endl;
	  ofSPD<<"state\tprob\tsamples"<<endl;
	  for(int p=0;p<Asize;p++){
	    for(int v=0;v<A[0].second.size();v++)
	      ofSPD<<A[p].second[v]<<" ";
	  ofSPD<<" "<<A[p].first<<" "<<fE[p]<<endl;
	  }
	  cumSamps += fEtot; //cume samps in Ens
	  //write ensemble
	  ofEns<<" numRes "<<numRes<<" seq ";
	  for(int nr=0;nr<numDoub;nr++)
	    ofEns<<resName[nr];
	  ofEns<<resName[numRes-1]<<endl;

	  //convert from states to corresponding phi/psi s
	  for(int p=0;p<Asize;p++){
	    //write for doublets (not residues)
	    //over the fE number of times
	    for (int f=0;f<fE[p];f++){
	      for(int nd=0;nd<numDoub;nd++)
		ofEns<<phi[nd][A[p].second[nd+1]]<<" ";
	      ofEns<<endl;
	      for(int nd=0;nd<numDoub;nd++)
		ofEns<<psi[nd][A[p].second[nd+1]]<<" ";
	      ofEns<<endl;
	    }
	  } //end convert from states ...
	  //write to end of Ens*.dat 
	  ofEns<<cumSamps<<endl;
	} //end nD==numDoub
    }
  //end vector construction

  //sum the probs to see how much we get
  double sumProb=0.0;
  for (int nv=0;nv<nvTop;nv++)
    sumProb += prodCondProb[nv];
  cout<<"for "<<numChainStates<<" possible sequence states and "
      <<nvTop <<" states kept, the kept probability is "<<sumProb<<endl;

  clock_t end = clock();
  clock_t clockTicksTaken = end - start;
  double timeInSeconds = clockTicksTaken / (double) CLOCKS_PER_SEC;
  cout<<"time to create states="<<timeInSeconds<<" seconds"<<endl;

  start = clock(); //time the rest 
  //close OFILES
  ofSPD.close();
  ofEns.close();
  end = clock();
  clockTicksTaken = end - start;
  timeInSeconds = clockTicksTaken / (double) CLOCKS_PER_SEC;
  cout<<"time for the rest="<<timeInSeconds<<" seconds"<<endl;
  cout<<"normal termination"<<endl;

  delete [] prodCondProb;
  delete [] numSupStates;

  for (int nd=0;nd<numDoub;nd++) {
    delete [] probCondSeq[nd];
    delete [] phi[nd];
    delete [] psi[nd];
  }
  delete[] probCondSeq;
  delete[] phi;
  delete[] psi;
  delete[] ns1;
  delete[] ns2;
  
  return 0;
}