/*---------------------------------
* LoadTT_Axona
* MEX file
*
* input:
*    fn = file name string
*    records_to_get = an array that is either a range of values
*    record_units = 1: timestamp list.(a vector of timestamps to load (uses binsearch to find them in file))
*                   2: record number list (a vector of records to load)
*		    3: range of timestamps (a vector with 2 elements: a start and an end timestamp)
*		    4: range of records (a vector with 2 elements: a start and an end record number)
*		    5: return the count of spikes (records_to_get should be [] in this case)
*    if only fn is passed in then the entire file is opened.
*    if only fn is passed AND only t is provided as the output, then all 
*       of the timestamps for the entire file are returned.
*
* output:
*    [t, wv]
*    t = n x 1: timestamps of each spike in file
*    wv = n x 4 x 32 waveforms
*
* version 0.1
*
* Reads Axona spike data files (.[1-8]) generated by assorted Axona
* hardware (tested on "silver", "blue", and DacqUSB system generated
* files, with recording system running DOS/Win98/WinXP, MClust
* workstation running WinXP)
*
* mvdm Feb 2007 (based on adr's LoadTT_NeuralynxNT.cpp and Sturla
* Molden's importspikes.m) compiled with MATLAB 7.1's mex script and
* the OpenWatcom 1.6 compiler (www.openwatcom.org)
*
* NO WARRANTY given or implied (see MClust license for details)
*
* however, questions, comments, suggestions are welcome at 
* <M.vanderMeer@ed.ac.uk>
*
* Axona system files have a header ending with data_start, then some
* number of data records consisting of ch1 timestamp (4 bytes), 50x
* ch1 waveform sample (1 byte each) ... ch4 timestamp, 50x ch4
* waveform. For timestamps, a timebase of 96000 is assumed,
* i.e. each timestamp unit is 1/96000 s.
*
* Waveform samples are on a [-127,127] scale. How many uV that
* corresponds to can be worked out from the settings in the .SET file
* (each channel can have its own gain) but this is not done here since
* for cutting it doesn't matter anyway. HOWEVER, MClust's reported uV
* values will be wrong!
*
* NOTE: only the first 32 waveform samples are read in for easy
* compatibility with MClust components. They are scaled for visual
* display purposes.
*--------------------------------*/

#include "mex.h"
#include <stdlib.h>
#include <stdio.h>
#include <string.h>
#include <math.h>
#include <matrix.h>

// needed for structure to align correctly.
#pragma pack(1) 

#define BY_TIMESTAMP 1
#define BY_RECORD 2
#define BY_TIMESTAMP_RANGE 3
#define BY_RECORD_RANGE 4
#define COUNTSPIKES 5

#define TIMEBASE 96000
#define WAVEFORM_SCALE 15

//___________________________________________________________________________________
int FindDataStart(FILE *fp) {

  // Look for "data_start" in file header
  // returns 1 if found (and file is forwarded to directly after it as a side effect)
  //         0 if not found
  
  char headerflag[10];
  int i, file_len;
  
  fseek(fp,0,SEEK_END); // get file length
  file_len = ftell(fp);
  fseek(fp,0,SEEK_SET);

  for (i = 0; i < file_len-10; i++) {
    
    fseek(fp,i,SEEK_SET);
    fread(headerflag, sizeof(char), 10, fp);  
    //mexPrintf("Headerflag = %10s \n", headerflag);
    
    if (strncmp(headerflag,"data_start",10) == 0){
      mexPrintf("data_start found (at byte %d)\n",i);
      return 1;
    
    } 
  }
  
  return 0;

}

//___________________________________________________________________________________
int binsearch(int n, double* data, double* key)
{
  // n = number of records in data.
  // data is an array of records to search
  // key is the value to find.
  // 
  // binsearch returns the index of the closest item in data
  //
  // from ADR adapted by cowen. 
  // fixed to return always the lower integreal timestamp (floor(*data)) if key is not integral (PL)
  // changed end = n-1 to end = n to allow retrieval of last record in file (ncst)
  
  int start = 0;
  int end = 0; 
  int mid = 0;
  double tmp;
  
  // binary search 
  start = 0;
  end = n;                    
  while (start < (end-1))
    {
      mid = (int) floor((double)(start + end)/2);
      tmp = floor(data[mid]);
      if ((*key) == tmp)
	start = end = mid;
      if ((*key) < tmp)
	end = mid;
      if ((*key) > tmp)
	start = mid;
    }
  
  return start;
}

//___________________________________________________________________________________
//===================================================================================
//-----------------------------------------------------------------------------------
// Loading Functions

const int recSize = 216;

//___________________________________________________________________________________
void GetOneRecord(FILE *fp, double &T, double wv[128])
  // reads one record and goes to the next record
{
  
  char qwTimeStamp[4];
  char snData[4][50];
  long qwTimeStamp2;
  int c1;
  
  int pos = ftell(fp);
  //mexPrintf("GetOneRecord: start reading (at position %d)\n",pos);
  
  fread(&qwTimeStamp,  sizeof(char),   4, fp);
  fread(snData[0],      sizeof(char), 50, fp);
  fseek(fp,4,SEEK_CUR);
  fread(snData[1],      sizeof(char), 50, fp);
  fseek(fp,4,SEEK_CUR);
  fread(snData[2],      sizeof(char), 50, fp);
  fseek(fp,4,SEEK_CUR);
  fread(snData[3],      sizeof(char), 50, fp);
  
  //mexPrintf("GetOneRecord: reading done, starting timestamp byte swap (%d %d %d %d)\n",qwTimeStamp[0],qwTimeStamp[1],qwTimeStamp[2],qwTimeStamp[3]);
  
  // do reverse byte order manually
  qwTimeStamp2 = qwTimeStamp[0]*16777216+qwTimeStamp[1]*65536+qwTimeStamp[2]*256+qwTimeStamp[3];
  
  //mexPrintf("GetOneRecord: Casting timestamp (%d)\n",qwTimeStamp2);

  //T = (double) (qwTimeStamp2*10000)/TIMEBASE; // output in tenths of ms
  T = (double) qwTimeStamp2/9.6;

  //system("pause");

  //mexPrintf("GetOneRecord: Ordering waveform samples\n");

  for (int j = 0; j<4; j++) {
    for(int k = 0; k<32; k++) {
      
      if (snData[j][k] > 127) {
	wv[j + 4*k] = (double) snData[j][k] - 256;
      }
      else {
	wv[j + 4*k] = (double) snData[j][k];
      }

      //wv[j + 4*k] = wv[j+4*k] * WAVEFORM_SCALE;
      wv[j + 4*k] = wv[j+4*k] * 15;
      
      //mexPrintf("GetOneRecord: Processed snData[%d][%d], value %.2f\n",j,k,wv[j+4*k]);
    }
  }
  
  //mexPrintf("GetOneRecord: Done.\n");
  //system("pause");

}

//___________________________________________________________________________________
int GetNumberOfSpikes(char* fn){
  
  // open file 
  FILE* fp = fopen(fn, "rb");
  if (!fp)
    mexErrMsgTxt("Could not open file.");

  //skip header and determine file record size
  int start_found = FindDataStart(fp);  
  if (!start_found) {
    mexErrMsgTxt("data_start not found.");
    return(0);
  }

  // get filesize
  int postHeaderPos = ftell(fp);     // beginnig of file after header (if any)
  fseek(fp,0,2);                     // goto end of file
  int nbytes = ftell(fp) - postHeaderPos - 10; // 10 because of data_end0D0A
  
  int nSpikes = nbytes/recSize; // no need to skip last record for NT_cheetah files
  mexPrintf("Reading file %s:\nRecordSize = %d,  %d spikes, %d bytes.\n",
	    fn, recSize, nSpikes, nbytes);
  
  // cleanup
  fclose(fp);	
  
  return nSpikes;
}


//___________________________________________________________________________________
void ReadTT(char* fn, int nSpikes, double *t, double *wv){
  
  double t0; 
  double wv0[128];
  
  // open file 
  FILE *fp = fopen(fn, "rb");
  if (!fp)
    mexErrMsgTxt("ERROR: Could not open file.");
  
  // skip header 
  //int new_NT_format = SkipHeader(fp);  // flag for new NT_format TT files (0=old SUN, 1=new NT) 
  //if (!new_NT_format) 
  //  mexErrMsgTxt("Old sun format not supported by this loading engine.");
  //SkipCheetahNTHeader(fp);     // Skip standard Neuralynx header if present (Cheetah versions >= 1.3)
  //long postHeaderPos = ftell(fp);
  
  int start_found = FindDataStart(fp); 
  long postHeaderPos = ftell(fp);

  // read records and convert all to double
  fseek(fp, postHeaderPos, SEEK_SET);
  for (int i = 0; i < nSpikes; i++){
    GetOneRecord(fp,t0,wv0);
    t[i] = t0;
    for (int j = 0; j<4; j++)
      for(int k = 0; k<32; k++)
	wv[i + nSpikes*j + nSpikes*4*k] = (double) wv0[j + 4*k];  
  }
    
  fclose(fp);	
}

//_________________________________________________________________________________________________
void ReadTTByRecord(char* fn, double *records_to_get, int n_records_to_get, double *t, double *wv)
  // Open the file and fseek to just those record number passed in 
  // in the array: records_to_get. The last record of records to get 
  // indicates the end of records. It's code is END_OF_RECORDS.
{
  int i = 0;
  double t0;
  double wv0[128];

  // open file 
  FILE *fp = fopen(fn, "rb");
  if (!fp)
    mexErrMsgTxt("ERROR: Could not open file.");
  
  // skip header 
  //int new_NT_format = SkipHeader(fp);  // flag for new NT_format TT files (0=old SUN, 1=new NT) 
  //if (!new_NT_format) 
  //  mexErrMsgTxt("Old sun format not supported by this loading engine.");
  //SkipCheetahNTHeader(fp);     // Skip standard Neuralynx header if present (Cheetah versions >= 1.3)
  //long postHeaderPos = ftell(fp);

  int start_found = FindDataStart(fp); 
  long postHeaderPos = ftell(fp);
	
  // read records and convert all to double
  while(i < n_records_to_get)
    {
      // Go directly to the record in question. Do not pass go. NO $200.
      fseek(fp, postHeaderPos+sizeof(char)*(recSize)*((long)records_to_get[i] - 1), SEEK_SET);
      GetOneRecord(fp,t0,wv0);
      t[i] = t0;
      for (int j = 0; j<4; j++)
	for(int k = 0; k<32; k++)
	  wv[i + n_records_to_get*j + n_records_to_get*4*k] = (double) wv0[j + 4*k];	
      i++;
    }
  fclose(fp);	
}	

//___________________________________________________________________________________
void ReadTT_timestamps(char* fn, int nSpikes, double *t)
{
  double t0;
  double wv0[128];
  
  // open file 
  FILE *fp = fopen(fn, "rb");
  if (!fp)
    mexErrMsgTxt("ERROR: Could not open file.");
  
  // skip header 
  //int new_NT_format = SkipHeader(fp);  // flag for new NT_format TT files (0=old SUN, 1=new NT) 
  //if (!new_NT_format) 
  //  mexErrMsgTxt("Old sun format not supported by this loading engine.");
  //SkipCheetahNTHeader(fp);     // Skip standard Neuralynx header if present (Cheetah versions >= 1.3)
  //long postHeaderPos = ftell(fp);

  int start_found = FindDataStart(fp); 
  long postHeaderPos = ftell(fp);
	
  // read records and convert all to double
  fseek(fp, postHeaderPos, SEEK_SET);
  for (int i = 0; i < nSpikes; i++)
    {
      mexPrintf("Attempting to get record %d\n",i);
      GetOneRecord(fp,t0,wv0);
      t[i] = t0;
    }
  fclose(fp);	
}	

//___________________________________________________________________________________
void mexFunction(int nOUT, mxArray *pOUT[],
				 int nINP, const mxArray *pINP[])
{
	
	int i;
	double *records_to_get, *range_to_get,*t,*wv, *all_timestamps, *record_units_ptr;
	int n_records_to_get = 0;
	int record_units;
	int nSpikesInFile = 0;
	int length_records_to_get = 0;	
	int length_record_units = -1; //?
	int start_idx=0;
	int end_idx = 0;
	int idx = 0;	
	/* check number of arguments: expects 1 input */
	if (nINP != 3 && nINP != 1)
				mexErrMsgTxt("Call with fn or fn and array of recs to load(vector), and 1(timestamp) or 2(record number	) as inputs.");
	if (nOUT > 2)
				mexErrMsgTxt("Requires two outputs (t, wv).");

	/* read inputs */
	int fnlen = (mxGetM(pINP[0]) * mxGetN(pINP[0])) + 1;
	//mexPrintf("fnlen = %d\n", fnlen);
	char *fn = (char *) mxCalloc(fnlen, sizeof(char)); 
	if (!fn)
		mexErrMsgTxt("Not enough heap space to hold converted string.");
	int errorstatus = mxGetString(pINP[0], fn,fnlen);    
	if (errorstatus)
		mexErrMsgTxt("Could not convert string data.");

	nSpikesInFile = GetNumberOfSpikes(fn);
	////////////////////////////////////////////////////////
	// If only one input is passed, assume the whole file is
	// to be loaded. If only one output is present, assume it is
	// only timestamps.
	////////////////////////////////////////////////////////

	if(nINP == 1 && nOUT == 1)
	{
		/* Return the timestamps of all of the records */
			// create outputs 
		pOUT[0] = mxCreateDoubleMatrix(nSpikesInFile, 1, mxREAL);
		t = mxGetPr(pOUT[0]);

		// load tt file fn into t and wv arrays 
		ReadTT_timestamps(fn,nSpikesInFile,t);

		// cleanup
		mxFree(fn);
		return;

	}else if(nINP == 1 && nOUT == 2)
	{
		////////////////////////////////////////////////////////
		// create outputs 
		////////////////////////////////////////////////////////
		mexPrintf("Getting %i records.\n",nSpikesInFile);
		pOUT[0] = mxCreateDoubleMatrix(nSpikesInFile , 1, mxREAL);
		t = mxGetPr(pOUT[0]);
		int wvDims[] = {nSpikesInFile , 4, 32};    // wv = nSpikes x 4 x 32   FORTRAN (column major) array
		pOUT[1] = mxCreateNumericArray(3, wvDims, mxDOUBLE_CLASS, mxREAL);
		wv = mxGetPr(pOUT[1]);

		////////////////////////////////////////////////////////
		// load tt file fn into t and wv arrays 
		////////////////////////////////////////////////////////

		ReadTT(fn,nSpikesInFile ,t,wv);

		////////////////////////////////////////////////////////
		// cleanup
		////////////////////////////////////////////////////////

		mxFree(fn);
		return;
	}
	////////////////////////////////////////////////////////
	// unpack inputs 
	////////////////////////////////////////////////////////
	length_records_to_get = mxGetM(pINP[1]) * mxGetN(pINP[1]);
	records_to_get = mxGetPr(pINP[1]);

	length_record_units = mxGetM(pINP[2]) * mxGetN(pINP[2]);
	record_units_ptr = (double *) mxGetPr(pINP[2]);
	record_units = (int) *record_units_ptr;
	
	//mexPrintf("length_record_units = %d\n",length_record_units);
	//mexPrintf("record_units_ptr = %.2f\n", *record_units_ptr);
	//mexPrintf("record_units = %d\n", record_units);
	//mexPrintf("length_records_to_get = %d\n",length_records_to_get);
	
	switch(record_units)
	
	{
		case BY_TIMESTAMP:
			////////////////////////////////////////////////////////
			// Convert the timestamps into record numbers. This will
			// make loading these records easy.
			////////////////////////////////////////////////////////
			////////////////////////////////////////////////////////
			// Create a very large array of all of the timestamps
			////////////////////////////////////////////////////////
			n_records_to_get = length_records_to_get;
			all_timestamps = (double *)calloc( nSpikesInFile, sizeof( double ) );
			if (all_timestamps == NULL)
				mexErrMsgTxt("NOT ENOUGH MEMORY");
			range_to_get = (double *) calloc( n_records_to_get, sizeof( double ) );
			if (range_to_get == NULL)
				mexErrMsgTxt("NOT ENOUGH MEMORY");

			ReadTT_timestamps(fn,nSpikesInFile,all_timestamps);

			for (i = 0;i<n_records_to_get;i++)
			{
				idx = binsearch(nSpikesInFile, all_timestamps, &records_to_get[i]);
				range_to_get[i] = idx + 1; // Add one since records are assumed
											 // to be from 1 to end.
			}

			free(all_timestamps);
			break;

		case BY_RECORD:
			////////////////////////////////////////////////////////
			// Get records asked for. First, subract one since matlab
			// users typically think records as being ordered from 
			// 1 to ....
			////////////////////////////////////////////////////////

			n_records_to_get = length_records_to_get;
			range_to_get = records_to_get;
			break;

		case BY_TIMESTAMP_RANGE:

			////////////////////////////////////////////////////////
			// Error checking
			////////////////////////////////////////////////////////

			if(length_records_to_get != 2)
			{
				mexErrMsgTxt("Must pass in two arguements for parameter 2.");
				return;
			}

			////////////////////////////////////////////////////////
			// Create a very large array of all of the timestamps
			////////////////////////////////////////////////////////
			all_timestamps = (double *)calloc( nSpikesInFile, sizeof( double ) );
			if (all_timestamps == NULL)
				mexErrMsgTxt("NOT ENOUGH MEMORY");

			ReadTT_timestamps(fn,nSpikesInFile,all_timestamps);
			////////////////////////////////////////////////////////
			// Find the index in all_timestamps of the start record.
			////////////////////////////////////////////////////////

			start_idx = binsearch(nSpikesInFile, all_timestamps, &records_to_get[0]) + 1; // Add one since records are assumed
																					  	  // to be from 1 to end.

			////////////////////////////////////////////////////////
			// Find the index in all_timestamps of the end record.
			////////////////////////////////////////////////////////
			end_idx = binsearch(nSpikesInFile, all_timestamps, &records_to_get[1]) + 1;  // Add one since records are assumed
																				  	 // to be from 1 to end.


			free(all_timestamps);
			
			n_records_to_get = end_idx - start_idx + 1;
			////////////////////////////////////////////////////////
			// Allocate the space
			// Using ints would be much more efficient, but it would make it necessary
			// to rewrite the ReadTT program. I don't want to do that.
			////////////////////////////////////////////////////////

			range_to_get = (double *) calloc( n_records_to_get, sizeof( double ) );
			if (range_to_get == NULL)
			{
				mexErrMsgTxt("NOT ENOUGH MEMORY");
				return;
			}

			for (i = 0; i<n_records_to_get;i++)
				range_to_get[i] = start_idx + i;

			break;

		case BY_RECORD_RANGE:
			////////////////////////////////////////////////////////
			// This is easy. First check to make sure the user passed in
			// the proper number of arguements.
			////////////////////////////////////////////////////////
			if(length_records_to_get != 2)
			{
				mexErrMsgTxt("Must pass in a start and end record number for argument 2.");
				return;
			}
			start_idx = (int) records_to_get[0] ;
			end_idx   = (int) records_to_get[1] ;

			////////////////////////////////////////////////////////
			n_records_to_get = end_idx - start_idx + 1;
			////////////////////////////////////////////////////////
			// Allocate the space
			////////////////////////////////////////////////////////
			range_to_get = (double  *)calloc( n_records_to_get, sizeof( double  ) );
			if (range_to_get == NULL)
			{
				mexErrMsgTxt("NOT ENOUGH MEMORY");
				return;
			}
			for (i = 0; i<n_records_to_get;i++)
			{
				range_to_get[i] = start_idx + i;
			}
			break;
			
		case COUNTSPIKES: // ADR/NCST
		  mexPrintf("Counting spikes...\n");
		  pOUT[0] = mxCreateDoubleMatrix(1,1,mxREAL);
		  t = mxGetPr(pOUT[0]);
		  t[0] = nSpikesInFile;		
		  break;

		default:
			mexErrMsgTxt("Incorrect parameter 3.");
	}
	
	if (!(record_units==COUNTSPIKES))
	
	{
		////////////////////////////////////////////////////////
		// Allocate the space
		////////////////////////////////////////////////////////
		printf("Getting %i records.\n",n_records_to_get);
		pOUT[0] = mxCreateDoubleMatrix(n_records_to_get, 1, mxREAL);
		t = mxGetPr(pOUT[0]);
		int wvDims[] = {n_records_to_get, 4, 32};    // wv = nSpikes x 4 x 32   FORTRAN (column major) array
		pOUT[1] = mxCreateNumericArray(3, wvDims, mxDOUBLE_CLASS, mxREAL);
		wv = mxGetPr(pOUT[1]);
		////////////////////////////////////////////////////////
		// Get the data. Record_units will whether to get by 
		// timestamp or by record number.
		////////////////////////////////////////////////////////
		ReadTTByRecord(fn,range_to_get,n_records_to_get,t,wv);

		// Free this variable only if it was not a matlab variable.
		if (record_units == BY_TIMESTAMP_RANGE || record_units == BY_RECORD_RANGE )
					free(range_to_get);
	}
	// cleanup
	mxFree(fn);

}