#include <vector>
#include <random>
#include <numeric>
#include <pybind11/stl.h>
#include <pybind11/stl_bind.h>
#include <pybind11/pybind11.h>
#include <pybind11/numpy.h>
#include <Python.h>

using namespace std;

namespace py = pybind11;
#include "tcHeader.h"
PYBIND11_MAKE_OPAQUE(std::vector<double>);

#define EPSILON 1.0e-5

/// Returns the mean dfbf[ frame# ] for the specified shuffle[ trial# ]
/// using data[ trial# ][ frame# ]. Operates on data for a single cell.
static vector< double > aveOfTrials( 
				const vector< vector< double > >& data, 
				const vector< unsigned int >& shuff, 
				unsigned int numFrames )
{
	unsigned int numTrials = data.size();
	vector< double > aveBin( numFrames, 0.0 );
	assert( numTrials == data.size() && numTrials > 0  );
	for ( unsigned int ii = 0; ii < numTrials; ++ii ) {
		vector< double >::iterator aptr = aveBin.begin();
		auto dd = data[ii].begin();
		for ( unsigned int bb = shuff[ii]; bb < shuff[ii] + numFrames; bb++, aptr++) {
			*aptr += *( dd + bb%numFrames );
		}
	}
	for( vector< double >::iterator aptr = aveBin.begin(); aptr != aveBin.end(); aptr++ ) {
		*aptr /= numFrames;
	}
	return aveBin;
}


/// args: data[ frame# ][ trial# ][ cell# ]
/// Pass-back arg: ret[ cell# ][ trial# ][ frame# ]
/// Fills in only the data in the window around the stimuli, from 
/// CS_ONSET_FRAME - CIRC_PAD to US_ONSET_FRAME + CIRC_PAD.
static void reorderData( const double* data, unsigned int numCells, unsigned int numTrials, unsigned int numFrames, vector< vector< vector< double >>>& ret, const AnalysisParams& ap ) 
{
	ret.clear();
	ret.resize( numCells );
	unsigned int ncxnt = numCells * numTrials;
	unsigned int nf = ap.circShuffleFrames;
	for ( unsigned int cell = 0; cell < numCells; cell++ ) {
		vector< vector< double > >& rc = ret[cell];
		rc.resize( numTrials );
		for( unsigned int tt = 0; tt < numTrials; tt++ ) {
			vector< double >& rct = rc[ tt ];
			rct.resize( nf,  0.0 );
			for( unsigned int ff = 0; ff < nf; ++ff ) {
				rct[ff] = data[ cell + tt * numCells + (ff + ap.csOnsetFrame - ap.circPad) * ncxnt];
			}
		}
	}
}

// Returns mean and sdev across all frames and all trials for a given cell.
// data has already been sub-selected for a given cell.
pair< double, double > findCellStats(
		const vector< vector< double > >& data, unsigned int numFrames )
{
	double sum = 0.0;
	double sq = 0.0;
	for( auto trialD = data.begin(); trialD != data.end(); ++trialD ) {
		for( auto frameD = trialD->begin(); frameD != trialD->end(); ++frameD ) {
			sum += *frameD;
			sq += *frameD * *frameD;
		}
	}
	double numSamples = numFrames * data.size(); //numFrames * numTrials
	double mean = sum/numSamples;
	return pair<double, double>( mean, sqrt( sq/numSamples - mean*mean ) );
}

/// Returns time stamps of transients trial# ][ transient# ]. 
// given the reordered dfbf data[ trial# ][ frame# ]
// A given trial may have zero or more transients.
static vector< vector< unsigned int > > findTransients(
		const vector< vector< double > >& data, double cellThresh, unsigned int numFrames )
{
	unsigned int numTrials = data.size();
	vector< vector< unsigned int > > ret ( numTrials );
	vector< vector< unsigned int > >::iterator trialT = ret.begin();
	for( auto trialD = data.begin(); trialD != data.end(); ++trialD, ++trialT ) {
		trialT->clear();
		double lastFrame = 0.0;
		bool refractory = 0;
		for( unsigned int ii = 0; ii < numFrames; ++ii ) {
			double frameD = (*trialD)[ii];
			// Don't permit another transient till signal goes < cellThresh
			if (frameD > cellThresh && frameD > lastFrame && !refractory) {
				trialT->push_back( ii );
				refractory = 1;
			} else {
				refractory = (frameD > cellThresh);
			}
			lastFrame = frameD;
		}
	}
	return ret;
}

/// Returns mean and SD of event widths for this cell,
// given the reordered dfbf data[ trial# ][ frame# ]
// Closely based on findTransients()
// A given trial may have zero or more events, it looks at all events.
pair< double, double > findEventWidthStats(
		const vector< vector< double > >& data, double cellThresh, unsigned int numFrames )
{
	double sum = 0.0;
	double sumSq = 0.0;
	double num = 0.0;

	for( auto trialD = data.begin(); trialD != data.end(); ++trialD ) {
		double lastFrame = 0.0;
		unsigned int startFrameOfEvent = 0;
		bool refractory = 0;
		for( unsigned int ii = 0; ii < numFrames; ++ii ) {
			double frameD = (*trialD)[ii];
			// Don't permit another transient till signal goes < cellThresh
			if (frameD > cellThresh && frameD > lastFrame && !refractory) {
				startFrameOfEvent = ii;
				refractory = 1;
			} else if ( refractory ) {
				refractory = (frameD > cellThresh);
				unsigned int dt = ii - startFrameOfEvent;
				if ( dt > 1 && !refractory ) { // End of event. get its stats.
					sum += dt;
					sumSq += dt*dt;
					num += 1.0;
				}
			}
			lastFrame = frameD;
		}
	}

	if ( num > 0.5 ) {
		double mean = sum/num;
		return pair< double, double >( mean, sqrt( sumSq/num - mean*mean) );
	}

	return pair< double, double >( -1.0, -1.0 );
}

// Returns hitTrialRation and the sdev of frames diff between meanPkIdx 
// and nearest transient.
pair< double, double > findHitRatioAndImprecision( unsigned int meanPkIdx, vector< vector< unsigned int > > transients, double hitWindow )
{
	double sumImp = 0.0;
	double sqImp = 0.0;
	double numImp = 0.0;
	double numHits = 0.0;
	for ( auto tt : transients ) {
		unsigned int numT = tt.size();
		if( numT > 0 ) {
			double minDt = 10000000;
			for (auto jj : tt ) {
				double dt = double( jj ) - meanPkIdx;
				if ( abs(minDt) > abs(dt) )
					minDt = dt;
			}
			sumImp += minDt;
			sqImp += minDt*minDt;
			numImp += 1.0;
			numHits += double( abs(minDt) < hitWindow );
		}
	}
	if ( numImp > 0.1 ) {
		double mean = sumImp / numImp;
		double sdev = sqrt( sqImp/numImp - mean*mean);
		return pair< double, double >( numHits/transients.size(), sdev );
	}
	return pair< double, double >( -1.0, -1.0 );
}

/// Computes PEQ score for a given neuron using peeling for multi-peak detection.
// For each significant peak: Q score uses that peak as the timing reference.
// Significance is determined by a shuffle test on the mean trace (same criterion
// as TI's tuningCurve). allPkScores holds per-peak Q; allPkPvalues from shuffle.
CellScore cellPeqScore( const vector< vector< double > >& data, const AnalysisParams& ap, const PeqAnalysisParams& pep )
{
	CellScore cs;
	cs.meanScore = 0.0;
	cs.baseScore = 0.0;
	cs.percentileScore = 0.0;
	cs.eventWidthMean = 0.0;
	cs.eventWidthSdev = 0.0;
	cs.imprecision = 0.0;
	unsigned int numTrials = data.size();
	assert( numTrials > 0 );
	unsigned int numFrames = ap.circShuffleFrames;

	vector< unsigned int > nonShuff( numTrials, 0 );
	cs.meanTrace = aveOfTrials( data, nonShuff, numFrames );

	// Pre-compute shuffled mean traces for significance testing.
	std::mt19937 rng( 1234 );
	std::uniform_int_distribution<std::mt19937::result_type> shuffler( 0, numFrames - 1 );
	vector< vector< double > > shuffMeans( ap.numShuffle );
	for ( unsigned int ii = 0; ii < ap.numShuffle; ii++ ) {
		vector< unsigned int > shuff( numTrials );
		for ( unsigned int tt = 0; tt < numTrials; tt++ )
			shuff[tt] = shuffler( rng );
		shuffMeans[ii] = aveOfTrials( data, shuff, numFrames );
	}

	// Global cell stats: used in Q formula, same for all peaks.
	pair<double, double> stats = findCellStats( data, numFrames );
	double mean = cs.meanScore = stats.first;
	double sdev = cs.sdev = stats.second;
	double cellThresh = mean + pep.transientThresh * sdev;
	vector< vector< unsigned int > > transients = findTransients( data, cellThresh, numFrames );
	pair<double, double> ewStats = findEventWidthStats( data, cellThresh, numFrames );
	cs.eventWidthMean = ewStats.first;
	cs.eventWidthSdev = ewStats.second;

	unsigned int minSepFrames = (unsigned int)round( ap.minPeakSep / ap.frameDt );
	if ( minSepFrames < 1 ) minSepFrames = 1;

	// Dip threshold: mean + dipSdev * sdev of the mean trace.
	double dipThresh;
	{
		double sum = 0.0, sumSq = 0.0;
		for ( auto v : cs.meanTrace ) { sum += v; sumSq += v * v; }
		double mn = sum / numFrames;
		double sd = sqrt( sumSq / numFrames - mn * mn );
		dipThresh = mn + ap.dipSdev * sd;
	}

	vector< bool > available( numFrames, true );
	bool firstPeak = true;

	while ( true ) {
		// Find the highest available frame in the original mean trace.
		unsigned int pkFrame = numFrames; // sentinel
		double pkVal = -1e30;
		for ( unsigned int ff = 0; ff < numFrames; ff++ )
			if ( available[ff] && cs.meanTrace[ff] > pkVal ) {
				pkVal = cs.meanTrace[ff]; pkFrame = ff;
			}
		if ( pkFrame == numFrames || pkVal <= 0.0 ) break;

		unsigned int adjFrame;
		if ( pkFrame == 0 ) adjFrame = 1;
		else if ( pkFrame == numFrames - 1 ) adjFrame = numFrames - 2;
		else adjFrame = (cs.meanTrace[pkFrame-1] > cs.meanTrace[pkFrame+1]) ? pkFrame-1 : pkFrame+1;

		unsigned int numOK = 0;
		for ( unsigned int ii = 0; ii < ap.numShuffle; ii++ )
			numOK += ( cs.meanTrace[pkFrame]  > shuffMeans[ii][pkFrame]  ) &&
			         ( cs.meanTrace[adjFrame] > shuffMeans[ii][adjFrame] );
		if ( (numOK * 100) <= (99 * ap.numShuffle) ) break;

		// Dip check: for 2nd+ peaks the trace must drop below dipThresh for
		// at least dipFrames consecutive frames between this and the nearest accepted peak.
		unsigned int lo = (pkFrame >= minSepFrames) ? pkFrame - minSepFrames : 0;
		unsigned int hi = min( pkFrame + minSepFrames, numFrames - 1 );
		if ( !firstPeak ) {
			unsigned int nearest = cs.allPkIndices[0];
			unsigned int minDist = (pkFrame > nearest) ? pkFrame - nearest : nearest - pkFrame;
			for ( auto prevIdx : cs.allPkIndices ) {
				unsigned int d = (pkFrame > prevIdx) ? pkFrame - prevIdx : prevIdx - pkFrame;
				if ( d < minDist ) { minDist = d; nearest = prevIdx; }
			}
			if ( !hasDipBetween( cs.meanTrace, nearest, pkFrame, dipThresh, ap.dipFrames ) ) {
				for ( unsigned int ff = lo; ff <= hi; ff++ ) available[ff] = false;
				continue;
			}
		}

		// Per-peak Q: use this peak as the timing reference for hit/imprecision.
		pair<double, double> hiStats = findHitRatioAndImprecision( pkFrame, transients, pep.hitWindow );
		double fracFired = hiStats.first;
		double sdevImp   = hiStats.second;
		double signal = -1.0;
		if ( fracFired > 0.0 )
			signal = sdev * fracFired / cs.meanTrace[pkFrame];
		double Q = 0.0;
		if ( !(sdevImp < 0.0 || ewStats.first < 0.0 || signal < 0.0) )
			Q = fracFired * exp( -pep.alpha * signal
			                   + pep.beta * ewStats.second / ewStats.first
			                   + sdevImp / double(numFrames) );

		double pval = 1.0 - (double)numOK / ap.numShuffle;
		cs.allPkIndices.push_back( pkFrame );
		cs.allPkScores.push_back( Q );
		cs.allPkPvalues.push_back( pval );

		if ( firstPeak ) {
			cs.meanPkIdx       = pkFrame;
			cs.baseScore       = Q;
			cs.fracTrialsFired = fracFired;
			cs.imprecision     = sdevImp;
			firstPeak          = false;
		}

		// Mark the separation window as unavailable.
		for ( unsigned int ff = lo; ff <= hi; ff++ )
			available[ff] = false;
	}

	if ( firstPeak ) {
		cs.meanPkIdx = (unsigned int)(
			max_element( cs.meanTrace.begin(), cs.meanTrace.end() ) - cs.meanTrace.begin() );
		cs.baseScore = 0.0;
	}

	return cs;
}


// Returns the reliability index.
vector< CellScore > peqScore( py::array_t<double> xs, const AnalysisParams& ap, const PeqAnalysisParams& pep )
{
	py::buffer_info info = xs.request();
	auto data = static_cast< double* >( info.ptr);

	vector< vector< vector< double >>> reorderedData;
	reorderData( data, ap.numCells, ap.numTrials, ap.numFrames, reorderedData, ap );

	vector< CellScore > ret( ap.numCells );
	for( unsigned int cellIdx = 0; cellIdx < ap.numCells; cellIdx++ ) {
		ret[cellIdx] = cellPeqScore( reorderedData[cellIdx], ap, pep );
	}

	return ret;
}