/*****************************************************
**
**  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.
**
******************************************************/
#include <iostream>
#include <string>
#include <stdlib.h>
#include <stdio.h>
#include <math.h>
#include <cstring>

#include "spres.h"
 

//record a shot
void transitionMatrix::addShot( shotRecord *s, int startRcBin, int endRcBin ){

     if( startRcBin < 0 || endRcBin >= nBins ){
        cerr << "Adding shot out of bounds, between " 
             << startRcBin << " and " << endRcBin << endl;
        exit( 8 );
     }

     transferRate[startRcBin][endRcBin] += 1.0;
     shotsFrom[startRcBin]++;

     /* save the start and end points of this shot */
     s->fromNext = shotsBetween[startRcBin][endRcBin];
     shotsBetween[startRcBin][endRcBin] = s;


}



/** determine convergence for a matrix row */
int transitionMatrix::testDistribution(int row, int  shotsAttempted)
{

  int     i;
  double  gamma, delta, forwardShots, backwardShots, lowest;

  gamma = 0.3;  /* damping coefficient for adaptive sampling rate - should give numerical stability */

  forwardShots = 0.0;
  for( i = row + 1; i < nBins; i++ ){
    forwardShots += transferRate[row][i];
  }
  
  backwardShots = 0.0;
  for( i = 0; i < row; i++ ){
      backwardShots += transferRate[row][i];
  }
 
  if( forwardShots == 0.0 && backwardShots == 0.0 ){   
      if( gamma * 10 * shotsAttempted < MAX_SAMPLES ){
          return( int( gamma * 10 * shotsAttempted ) );
      }
      else return( MAX_SAMPLES );
  }   
  
  //try to have some backward as well as forward shots
  if( row == 0 ){
      lowest = forwardShots;
  }
  else if( row == nBins - 1 ){
      lowest = backwardShots;    
  }
  else if( forwardShots < backwardShots ){
      lowest = forwardShots;
  }
  else{
        //lowest = backwardShots;
      lowest = forwardShots;
  }
  if( lowest <= 0 ){
      if( gamma * 10 * shotsAttempted < MAX_SAMPLES )
        return( int( gamma * 10 * shotsAttempted ) );
      else return( MAX_SAMPLES );
  }


   /* proportionally increase/decrease the number of attempts to get minimum of targetShots forward */ 
  delta  = shotsAttempted * (targetShotsForward / (double)lowest);
  delta -= shotsAttempted;
  delta *= gamma;
  
  if( shotsAttempted + int(delta) > 0 ){
     if( shotsAttempted + int( delta ) < MAX_SAMPLES )
     {
        return( shotsAttempted + int( delta ) );
     }
     else return( MAX_SAMPLES );
  }
  else{
    return( targetShotsForward );
  }
}

/** determine convergence for the whole matrix - currently a bit preliminary. */
bool transitionMatrix::testRateNonZero(){

   int i;

   /* test for possibility of any path to final state */
   for( i = nBins - 2; i >= 0; i-- ) {
        if( transferRate[i][nBins-1] > 0.0 ) {
            return( true );
        }
   }

   return( false );
}

#ifdef HAVE_GSL 
void transitionMatrix::diagonalise( ofstream *logFile, int tau ){
  
    gsl_matrix                *m, *evec;
    gsl_vector                *eval;
    gsl_eigen_symmv_workspace *w;
    int                        i, j;
    
    /* copy into a matrix object compatible with the gnu scientific library */
    m = gsl_matrix_alloc( nBins, nBins );
    for( i = 0; i < nBins; i++ ){
        for( j = 0; j < nBins; j++ ) {
            gsl_matrix_set( m, i, j, transferRate[i][j]);
        }
    }

    /* define output objects */
    evec = gsl_matrix_alloc (nBins, nBins);
    eval = gsl_vector_alloc (nBins);

    /* allocate workspace and do the diagonalisation */
    w = gsl_eigen_symmv_alloc( nBins );
    gsl_eigen_symmv( m, eval, evec, w );
    gsl_eigen_symmv_free( w );
    
   
    gsl_vector_free (eval);
    gsl_matrix_free (evec);
    gsl_matrix_free(m);
    

}
#endif



/** clear saved data, for reuse */
unsigned long int transitionMatrix::reset(){
       int                i, j;
       unsigned long int  freed;  
       shotRecord        *s;

       freed = 0;

       for( i = 0; i < nBins; i++ ) {
         memset(transferRate[i], 0, nBins * sizeof(double)); 
         shotsFrom[i]     = 0.0;

         for( j = 0; j < nBins; j++ ){
            while( shotsBetween[i][j] != 0 ) {
         
                s = shotsBetween[i][j];
                shotsBetween[i][j] = s->fromNext;

                if( s->inMemory == true ) { 
                    freed += s->myMemSize;
                } else{
                  if( remove( s->myFileName.c_str() ) != 0 ) {
                    cerr <<  "Error deleting file (in matrix reset): \n" 
                         <<  "between " << i << " and " << j << endl;

                    exit( EXIT_FAILURE );
                         
                  }
                }
                delete s;
            }
         }
      }

      return( freed );

}



/** clear saved data, for reuse */
long int transitionMatrix::deleteRow( int row ){
       int        i;
       shotRecord *s;
       long int   bytesFreed;       

       bytesFreed = 0; 

       for( i = 0; i < nBins; i++ ) {
           while( shotsBetween[i][row] != 0 ){
              if( shotsBetween[i][row]->inMemory == true ) {
                    bytesFreed += shotsBetween[i][row]->myMemSize;
              }
              else{
                if( remove( shotsBetween[i][row]->myFileName.c_str() ) != 0 ) {
                    cerr <<  "Error deleting file: \"" 
                         << shotsBetween[i][row]->myFileName << "\"\n";
                }
              }
              s = shotsBetween[i][row]->fromNext;
              delete shotsBetween[i][row]; 
              shotsBetween[i][row] = s;
           } 
       }

       return( bytesFreed );
         
}



/** destructor */
transitionMatrix::~transitionMatrix(){

   int        i, j;
   shotRecord *s;

   for( i = 0; i < nBins; i++ ) {
     delete transferRate[i];
        for( j = 0; j < nBins; j++ ) {
            while( shotsBetween[i][j] != 0 ){
                 s = shotsBetween[i][j];   
                 shotsBetween[i][j] = s->fromNext;
                 delete s;        
            }
        }
    }
    delete transferRate;
    delete shotsFrom;
    delete sampleRate;

}


void transitionMatrix::print( ofstream *logFile ){
        
     int  i, j;    

     for( i = 0; i < nBins; i++ ) {
       for( j = 0; j < nBins; j++ ) {
             if( transferRate[i][j] > 0.0 ){
                        *logFile << i << " " << j << " " 
                                 << ( transferRate[i][j] ) << " " << age << endl;                
            }
        }
     }

    *logFile << endl << endl;

}

void transitionMatrix::print( ofstream *logFile, int maxBin ){
        
     int  i, j;    

     for( i = 0; i < maxBin; i++ ) {
         for( j = 0; j < nBins; j++ ) {
             if( transferRate[i][j] > 0.0 ){
                        *logFile << i << " " << j << " " 
                                 << ( transferRate[i][j] ) << " " << age << endl;                
            }
        }
     }

    *logFile << endl << endl;

}

void transitionMatrix::printSampleRate( ofstream *logFile ){
        
     int  i;    

     for( i = 0; i < nBins; i++ ) {
        if( sampleRate[i] > 0 ){
               *logFile << i << " " << sampleRate[i]  << endl;                
        }
     }

    *logFile << endl << endl;

}


void transitionMatrix::advanceState( double *stateVector ){

  double *w, scale;
  int     from, to;

  w = new double[nBins];

  //matrix-vector multiply
  scale = 0.0; 
  for( to = 0; to < nBins; to++ ){
     w[to] = 0.0;
     for( from = 0; from < nBins; from++ ){
         w[to] += stateVector[from] * transferRate[from][to];
     }
     scale += w[to];
  }

#ifdef ALLEN_FRENKEL_FLUX_CALC
  //periodic boundaries
  stateVector[0]        += stateVector[nBins - 1];
  stateVector[nBins - 1] = 0.0;
#endif

  //quick renormalise for the sake of stability
  cerr << "state vector is now:\n";
  scale = 1.0 / scale;
  for( to = 0; to < nBins; to++ ){
        stateVector[to] = w[to] * scale;
        if( w[to] != 0.0 ) {
            cerr << to << " " << stateVector[to] << endl;
        }
  }

  delete[] w;

}














