/*****************************************************
**
**  S-PRES Copyright (c) 2010
**
**  Joshua T Berryman
**
**  This file is part of S-PRES
**
**  S-PRES 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 3 of the License, or
**  (at your option) any later version.
**
**  S-PRES is distributed in the hope that it will be useful,
**  but WITHOUT ANY WARRANTY; without even the implied warranty of
**  MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
**  GNU General Public License for more details.
**
**  You should have received a copy of the GNU General Public License
**  along with this program.  If not, see <http://www.gnu.org/licenses/>.
**
**  If you make use of this program for research 
**  resulting in a publication, please cite:
**
**  Joshua T Berryman and Tanja Schilling, 
**  Sampling rare events in nonequilibrium and nonstationary systems. 
**  
**  J Chem Phys, 133(24):244101, Dec 2010.
**
******************************************************/
#ifdef PARALLEL_OMP
#include <omp.h>
#endif

#include <iostream>
#include <string>
#include <stdlib.h>
#include <stdio.h>
#include <math.h>
#include <cstring>
#include <time.h>
#include <assert.h>

#include "spres.h"


typedef enum {SHI_STATSOP_LARGEST_NUCLEUS    = 0,  
               SHI_STATSOP_ENERGY             = 1,  
               SHI_STATSOP_NUM_UPSPINS        = 2,  
               SHI_STATSOP_NUM_BURIED         = 3,  
               SHI_STATSOP_STRUCTURE_FACTOR   = 4,  
               SHI_STATSOP_NUM_OPS            = 5} shiOpId;


/** 
*   This is the main loop for the S-PRES algorithm
*
*/

void spres::spresRun( simulation *state, int numPathsPerBin, long numSpresBins ) {

   ofstream          spresRates;
   shotRecord       *shot, *parent;
   shotRecord      **parentArray;
   simulation      **stateArray;
   int               k, firstFullMatrix = 0;
   int               rcBin = 0, shotId = 0;
   int               highestHitBin, startBin;
   transitionMatrix *tmpMat;
   double            startFlux;
   

   //for parallelisation, it's convenient to have an array of states and shotpointers.
   parentArray=new shotRecord*[MAX_SAMPLES];
   stateArray= new simulation*[MAX_SAMPLES];
  
   //make duplicates of the initial data structure
   for ( int T = 0; T < MAX_SAMPLES; T++ ){
          stateArray[T]             = state->copySelf();
        *(stateArray[T]->get_iran()) *= (T + 1 ); //give each state its own rng conditions.
   }
   cerr << "Allocated extra copies of states... for easier OMP parallelisation." << endl;

   
   //initialisation
   pathsPerBin     = numPathsPerBin; 
   nRcBins         = numSpresBins;
   firstFullMatrix = 0; 
   highestHitBin   = 0;
   memBytes        = 0;
   shotMemSize     = 0;
   
   //init & seed the local RNG using the same seed as that passed into the simulation class.
   //We are assuming that the simulation class has its own thread-safe, industrial-strength RNG
   //however, that kind of thing isn't really needed for the wrapper S-PRES class.
   srand((unsigned int)*state->get_iran()); 
   
   
   //just test that everything makes sense
   sanitiseInputs(state, numPathsPerBin, numSpresBins );   // in spres.C

   stateVector     = new double[nRcBins];
   memset( stateVector, 0, nRcBins * sizeof(double) );
   stateVector[0]  = 1.0;
   pt_state = state;
   
   //setup file for output of rates
   spresRates.open( state->get_spresRatesFileName().c_str() );
   if( !spresRates.good() ) {
        cerr << "Warning: Could not open S-PRES sampling rates file for writing, name:\n\""
             <<  state->get_spresRatesFileName() << "\"\n";
   }
   spresLog.open( state->get_spresLogFileName().c_str()  );
   if( !spresLog.good() ) {
        cerr << "Error! Could not open S-PRES logfile, name:\n\""
             <<  state->get_spresLogFileName() << "\"\n";
        exit( EXIT_FAILURE );
   }
   //setup file for derived data
   dataLog.open( state->get_stateVectorFileName().c_str() );
   if( !dataLog.good() ) {
        cerr << "Error! Could not open outfile for derived data:\n\""
             <<  state->get_stateVectorFileName() << "\"\n";
        exit( EXIT_FAILURE );
   }
   

   cerr << " Entering S-PRES." << endl;

   //what is the reaction coordinate?
   state->calcMyStats();   
        
   //How many RC bins, therefore? 
   rcBinMax   = state->get_spresCriterionMax();
   rcBinMin   = state->get_spresCriterionMin();
   rcBinSize  = int( ( rcBinMax - rcBinMin ) / double(nRcBins) );

   //for each bin, will save the rates in a 'transfer matrix'
   transMat     = new transitionMatrix;
   transMat->init( nRcBins, pathsPerBin );
   transMat->reset();
   transMat->age = 0;

   //keep a pair of transition matrices... algorithm needs to look one step in the past.
   transMatPrev = new transitionMatrix;
   transMatPrev->init( nRcBins, pathsPerBin );
   transMatPrev->reset();
   transMatPrev->age = 0;

   initStateData();
   
   startFlux       = getStartFlux( state, transMat );
   cerr     << " #start flux out of A is: " << startFlux << endl;
   spresLog <<  " #start flux out of A is: " << startFlux << endl;

   tmpMat        = transMatPrev;
   transMatPrev  = transMat;
   transMat      = tmpMat;

   memBytes     -= transMat->reset();
   transMat->age = 1;
   startBin      = 0;
  
   transMat->print( &spresLog );
    

   cout<<"  ************************************************************** "<<endl;
   cerr << " Looping from " << startBin << " of " << nRcBins << " bins\n";
   cerr << " initially from matrix " << 1 << " to " << state->get_spresNumCompleteMatrices()  << endl; 


   //now fill each matrix - initially, only the 0,x row will contain any data 
   for(  timeStep = 1; timeStep  < state->get_spresNumCompleteMatrices(); timeStep++ ) 
   {
        bool runIsLive;

        runIsLive = false;

        writeDataFile();

        //clear the observables
        memset( dataVector, 0, dvCount * sizeof( double ) );

        for( workingBin = startBin; workingBin < nRcBins; workingBin++)
        { 
             // **************************************************************************
             highestHitBin = 0;
             // we probe if the "workingBin" is empty
             //test to see if there are any parent shots available 
                    
             parent = selectParentShot( state, transMatPrev, workingBin );
             if( parent == 0 ){      
                      continue; //no shots have reached this RC so cannot move on from it.
             }
             runIsLive = true;          
             
             //  *************************************************************************
             //  ***  Loop over shots we intend to make and prepare independent copies
             //  ***  of the start states.
             //  *************************************************************************
             for( int T = 0; T < transMatPrev->sampleRate[workingBin]; T++ ){
                      
                 parentArray[T] = selectParentShot( state, transMatPrev, workingBin );
                      
                 /* reload a state from the parent shot*/
                 if( parentArray[T]->inMemory == true ) {
                   
                   parentArray[T]->stateString.seekg( ios_base::beg );//set read pointer to the beginning
                   parentArray[T]->stateString.clear(); //reset EOF flags and similar.
                   
                   stateArray[T]->reLoad( &parentArray[T]->stateString );
                   
                 }else{ //if no state in memory, problem
                   
                    cerr << "Error : State not in memory " << endl;
                    exit( EXIT_FAILURE );
                 
                 }
             }
                
                
                  
//  ****************************************************************************************************************           
//  ****************************************************************************************************************            
//  ************************     Only this part is parallelized    *************************************************
//  ************************  ...the OMP pragma splits the 'for' loop up between the different threads. ************


#ifdef PARALLEL_OMP
#pragma omp parallel private(T)
{
            for( int T = omp_get_thread_num(); T < transMatPrev->sampleRate[workingBin]; T += omp_get_num_threads()){
                  stateArray[T]->stepTime();//run tau steps of dynamics
            }
}
#else
            for( int T = 0; T < transMatPrev->sampleRate[workingBin]; T++){
                  stateArray[T]->stepTime();//run tau steps of dynamics
            }
#endif
//  *****************************************************************************************************************
//  *****************************************************************************************************************


            //now collect and save the data from each thread
            for( int T = 0; T < transMatPrev->sampleRate[workingBin]; T++ )
            {
              //remember the chosen parent shot.
              parent = parentArray[T];

              /* create a new shot record and backup current state into it */
              shot = new shotRecord( parent->endRc );
              shot->myId     = shotId++;
              shot->state    = 0; //won't be reusing the state for a while 
              shot->endRc    = stateArray[T]->get_reactionCoordinate();
              shot->duration = 1;
              shot->start    = parent->start + parent->duration;

              /* write the shot endstate to memory, or to the `trajectory directory' */ 
              defineShotFilename( stateArray[T], shot, parent );
              //cerr<<"created new shot"<< shot->myFileName.c_str() << endl;

              //test for memory overflow
              if( memBytes + shotMemSize <= MAX_MEM_BYTES ) {
                
                shot->stateString.str().assign(""); //init the string
                shot->stateString.str().reserve( shotMemSize ); //request contiguous space
                shotMemSize     = stateArray[T]->save( &shot->stateString );
                shot->inMemory  = true;
                shot->myMemSize = shotMemSize;
                memBytes       += shotMemSize;
              } else{
                
                  cerr<<"Error: you have overrun the memory which you allowed to yourself. "<<endl;
                  cerr<<"Consider resetting MAX_MEM_BYTES in the Makefile. "<<endl;
                  exit(EXIT_FAILURE);
                  
              }
              
              /* record its effect in the transition matrix */
              rcBin = assignRcBin( shot->endRc );
              transMat->addShot( shot, workingBin, rcBin );
                      
              /* log the highest hit bin */
              if( rcBin > highestHitBin )
              {
                highestHitBin = rcBin;
              }
              
              /* if we are logging anything for output, it needs to be recorded */
              collectStateData(    stateVector[workingBin]/ double(transMatPrev->sampleRate[workingBin])   ); 

            }//finish loop over completed runs
              
            /* normalise the transfer probabilities */
            for( k = 0; k < nRcBins; k++ ){
                  transMat->transferRate[workingBin][k] /= double(transMatPrev->sampleRate[workingBin]);
            }
            
            /* consider increasing/decreasing the number of shots per bin for the next attempt
            ** in order to get good statistics */
            cerr << " Adjusting sample rate "  <<  workingBin << " from " << transMatPrev->sampleRate[workingBin];
            transMat->sampleRate[workingBin] = transMat->testDistribution(workingBin, transMatPrev->sampleRate[workingBin]);
            cerr << " to " <<  transMat->sampleRate[workingBin] << endl;

            //clean up the workspace as we go. Turn this off in order to recover trajectories.
            memBytes -= transMatPrev->deleteRow( workingBin );
            //cerr << " cleaned row " << workingBin <<  endl;
           
     }/*next row or end loop over rows */

     //save the transition matrix
     transMat->print( &spresLog );

     if( runIsLive == false )
     {
          cerr << "Error... system returned to absorbing boundary conditions without converging.\n";
          exit( 8 );
     }        

     /* require an unbroken series of complete matrices */
     cerr << "Testing convergence...\n";
     if( transMat->testRateNonZero() == false )
     { 
       firstFullMatrix = timeStep;
     }

     /* move the estimated probability distribution along */
     transMatPrev->advanceState( stateVector );

     /* shuffle the pointers and reinitialise */
     cerr << "Shuffling matrix pointers... currently using approx. "<< memBytes << " bytes for state.\n";
     tmpMat        = transMatPrev;
     transMatPrev  = transMat;
     transMat      = tmpMat;
     memBytes     -= transMat->reset();
     transMat->age = (timeStep + 1 );    

     cerr << "Shuffled.\n";

     /* only start from middle of matrix when loaded from filesystem 
      * ..subsequently startBin is 1 (absorbing boundaries at zero). */
     startBin = 0;

    }/* end loop over matrices: final convergence*/

    cerr << "writing data file...";
    writeDataFile();
    cerr << "...data file written\n";

    
}





/** test if we have enough data to declare the calculation complete */
bool spres::spresRunComplete( int matricesSoFar, int firstFullMatrix, int maxMatrices, int minCompleteMatrices) {

    /* can run untill a fixed number of matrices */
    if( maxMatrices > 0 ){
        if( matricesSoFar < maxMatrices ){
            return( false );
        }
        else{
            return( true );
        }
    }
    else{ /* alternatively can run untill fixed number of complete matrices */
        if( matricesSoFar < firstFullMatrix + minCompleteMatrices ){
            return( false );
        } 
        else{
            return( true );
        }
    }
}


/** select a shot to start a new run from */
shotRecord* spres::selectParentShot( simulation *state, transitionMatrix *m, int row )
{

    int      fromRow, inQueue, candidates;
    double     totalTrans, scale, flip, *chanceFromHere;

    shotRecord *s;

    //catch an unusual case
    if( row < 0 || row >= m->nBins )
        return( 0 );
    
    //calculate chance to enter this bin from each other bin
    chanceFromHere = new double[m->nBins];
    scale = 0.0;
    for(int i = 0; i < m->nBins; i++ )
    {
             chanceFromHere[i] = stateVector[i] *  m->transferRate[i][row];
             scale            += chanceFromHere[i];
             if( chanceFromHere[i] !=  chanceFromHere[i] ) //test for NaN
             {
                  cerr << "Error: " << i << " has: " << chanceFromHere[i] << " transitions\n";
                  cerr << "state vector: " << stateVector[i] << " transfer count: " << m->transferRate[i][row] << endl;
                  exit( 8 );
             }
    }
             
    if( scale <= 0.0 )
    {
         delete[] chanceFromHere;
         return( 0 ); //there are no parent shots entering this bin, as yet.
    }
    
    scale = 1.0 / scale;
    for( int i = 0; i < m->nBins; i++ )
         chanceFromHere[i] *= scale;
           
  
    //select a bin that the parent shot has to come from
    flip       =  rand()/(double)RAND_MAX; 
    totalTrans =  0.0;	
    fromRow    =  -1;
    for(int i = 0; i < m->nBins; i++ )
    {
             totalTrans += chanceFromHere[i];
             if( totalTrans >= flip )
             {
                 fromRow = i;
                 break;
             } 
    }	  
    if( fromRow == -1 ){
             cerr << "error selecting shots to " << row << " from " << m->nBins << " bins.\n";
             cerr << "reached probability: " << totalTrans << " out of " << flip << endl;
             for(int i = 0; i < m->nBins; i++ )
             {
                 cerr << i << " " << chanceFromHere[i] << endl;
             } 
             exit( 8 );
    }
        

    //easy but inefficient count of number of usable candidate shots
    candidates = 0;
    s = m->shotsBetween[fromRow][row];
    while( s != 0 ){
        candidates++;
        s = s->fromNext;
    }

    if( candidates == 0 ){
        cerr << "error... no stored shots between " << fromRow << " - " << row << endl; 
        exit( 8 );
    }

              
    //select a shot from that bin
    inQueue    = candidates;
    while( inQueue >= candidates ){
            flip    = rand()/(double)RAND_MAX; 
            inQueue = int( flip * candidates );
    }

    int i = 0;
    s = m->shotsBetween[fromRow][row];
    while( i < inQueue )
    {
         s = s->fromNext;
         i++;
         if( s == 0 )
         {
              cerr << i << " :Error! list of shots taken from matrix is too short.\n";
              exit( 8 );
         }
    }


    delete[] chanceFromHere;
    return( s ); 
        
}


/** select a shot to start a new run from */
shotRecord* spres::seekParentShot( simulation *state, transitionMatrix *m, int row, long int targetId ){

    int         fromRow;
    shotRecord *s;

    //catch an unusual case
    if( row < 0 || row >= m->nBins ){
        return( 0 );
    }

    for( fromRow = 0; fromRow < m->nBins; fromRow++ ){
      s = m->shotsBetween[fromRow][row];
      while( s ){
        if( s->myId == targetId ) {
          return( s );
        }
        s = s->fromNext;
      }
    }

    return( 0 ); 
        
}



//For comparability with the Allen and Frenkel (JChemPhys2006) FFS algorithm it is neccessary 
//to have strictly absorbing boundary conditions at A & B.
//
//To correct the state vector for this assumption, therefore, it is required to
//generate a Fudge, giving the number of states in A+1 which should in fact be
//deleted because they have crossed back to A in a zig-zag,
//as well as the number of transitions to A which should be added even though they were 
//not permanent.
//
//If you don't understand this, don't worry: just don't do an A&F flux calculation.
//
//
//
//Aside from the A&F initialisation activities, this function is also used to setup the
//transition matrix and the initial conditions.
double spres::getStartFlux( simulation *state, transitionMatrix *transMat ) 
{

    long       step, q;
    long       bin;
    long       lastBin;
    shotRecord *shot, *p;
    int        shotId = 0;
    double     officialStartRc;

    state->calcMyStats(); 

    //create a dummy shot 
    //cout<<"get reaction coordinate  "<<endl;	
    officialStartRc = state->get_reactionCoordinate();
  
    //cout<<"get reaction coordinate  "<<endl;
    p           = new shotRecord( officialStartRc );
    p->endRc    = officialStartRc;
    p->state    = state;
    p->myId     = shotId++;
    p->start    = 0;     
    p->duration = 0;

    //cerr << "start RC:: " << officialStartRc  << endl;
    shot           = new shotRecord( officialStartRc );
    
    //cout<<"nach shot"<<endl;    	
    shot->state    = state;
    shot->myId     = shotId++;
    shot->start    = 0;     
    shot->startRc  = 0.0;     
    shot->duration = 0;
    shot->parent   = p;

    q       = 0; 
    step    = 0;
    lastBin = 0;
    
    while( q < 10 )
    {   

  //cout<<"assignRcBin"<<endl;
       bin = assignRcBin( state->get_reactionCoordinate() );  // in spres_2.c    just finds out the bin  
       q++; 
       
  //cout<<"q  "<<q<< "bin "<<bin<<endl;
       /* write the shot endstate to the `trajectory directory' */
       shot->endRc = state->get_reactionCoordinate();// why do we calculate rc again ? no coarse graining 
       defineShotFilename( state, shot, p ); // in shotFiles.C    
  //cerr << "resetting..." << endl; 


        state->reset();	

  //cerr << "saving..." << endl; 
        shot->stateString.str().assign(""); //init the string
        shot->stateString.str().reserve( shotMemSize ); //request contiguous space
        shotMemSize     = state->save( &shot->stateString );
        shot->inMemory  = true;
        shot->myMemSize = shotMemSize;
        memBytes       += shotMemSize;
        
        collectStateData( 1.0 / 10.0 );
        
  
       /* record its existence in the transition matrix */
       transMat->addShot( shot, 0, bin ); 

       shot           = new shotRecord( officialStartRc );
       shot->state    = state;
       shot->start    = 0;     
       shot->duration = 0;
       shot->parent   = p;
       shot->myId     = shotId++;
       

    }   
    delete shot;

    for(int k = 0; k < transMat->nBins; k++ )
    {
      if( q > 0 )
      {
          transMat->transferRate[0][k] /= double(q);
      }
      else
      {
          transMat->transferRate[0][k] = 0.0;    
      }
    }
    return( 1.0 - transMat->transferRate[0][0] );
}












