#include "Event.hh"
#include "BaselineFinder.hh"
#include "ConvertData.hh"
#include "SumChannels.hh"
#include "intarray.hh"
#include "RootWriter.hh"
#include "TGraph.h"
#include <vector>
#include <cmath>
#include <string>
#include <iomanip>
#include <iostream>
#include <functional>
#include <algorithm>
#include <numeric>

BaselineFinder::BaselineFinder():
  ChannelModule(GetDefaultName(), "Find the baseline (zero) of the channel in the samples read before the trigger"),
  fixed_params("fixed_params","Parameters for fixed baseline search mode"),
  interp_params("interp_params","Parameters for linear interpolation mode"),
  drifting_params("drifting params","Parameters for drifting baseline search")
{
  AddDependency<ConvertData>();
  
  //Register all the config handler parameters
  RegisterParameter("fixed_baseline", fixed_baseline = false,
		    "Search for a flat baseline in pre-trigger window, otherwise search for a drifting baseline");
  RegisterParameter("signal_begin_time", signal_begin_time = 0,
		    "Search for baseline before this time [us] ");
  
  //parameters for fixed baseline
  RegisterParameter("fixed_params",fixed_params);

  fixed_params.RegisterParameter("segment_samps", segment_samps = 15,
		    "Samples in each baseline segment");
  fixed_params.RegisterParameter("min_valid_samps", min_valid_samps = 50,
		    "Minimum samples for a baseline to be valid");
  fixed_params.RegisterParameter("max_sigma", max_sigma = 1.2,
		    "Maximum standard deviation for a baseline segment to be accepted");
  fixed_params.RegisterParameter("max_sigma_diff", max_sigma_diff = 1,
		    "Maximum difference between baseline sigma and the sigma of a new valid segment");
  fixed_params.RegisterParameter("max_mean_diff", max_mean_diff = 1,
		    "Maximum difference between baseline mean and the mean of a new valid segment");
	
  //parameters for linear baseline with few baseline points
  fixed_params.RegisterParameter("linear_interpolation", 
				 linear_interpolation = false,
				 "Compute further baseline estimates throughout trigger window and perform interpolation");
  fixed_params.RegisterParameter("interp_params",interp_params);
  interp_params.RegisterParameter("avg_samps", avg_samps = 1000,
					"Number of contiguous samples to use to perform new baseline estimate");
  interp_params.RegisterParameter("max_sigma_factor", max_sigma_factor = 3,
					"Samples withing this number of flat baseline sigmas will be considered for baseline estimate");
  interp_params.RegisterParameter("pulse_threshold", pulse_threshold = 4,
					"Threshold (in number of max_sigmas) to consider a pulse and ignore cooldown samples from baseline estimate");
  interp_params.RegisterParameter("cooldown", cooldown = 2050,
					"Total number of samples around pulses to ignore for baseline estimate");
  interp_params.RegisterParameter("pre_cooldown", pre_cooldown = 50,
					"Number of samples before pulse to ignore for baseline estimate");

  //parameters for drifting baseline
  RegisterParameter("drifting_params",drifting_params);
  drifting_params.RegisterParameter("max_amplitude", max_amplitude = 5,
		    "Max amplitude for sample to be considered in baseline");
  drifting_params.RegisterParameter("max_return_amplitude", max_return_amplitude = 5,
		    "Max amplitude for sample to be considered in baseline after an excluded region");
  drifting_params.RegisterParameter("max_sum_amplitude", 
				    max_sum_amplitude = 0.2,
		    "max_amplitude for sum channel");
  drifting_params.RegisterParameter("pre_samps", pre_samps = 5,
		    "Samples before to include in moving average");
  drifting_params.RegisterParameter("post_samps", post_samps = 5,
		    "Samples after to include in average");
  drifting_params.RegisterParameter("laserwindow_begin_time", 
				    laserwindow_begin_time = -999,
		    "Specify the beginning time of a laser window.");		
  drifting_params.RegisterParameter("laserwindow_end_time", 
				    laserwindow_end_time = -999,
		    "Specify the ending time of a laser window.");	    
  drifting_params.RegisterParameter("laserwindow_freeze", 
				    laserwindow_freeze = false,
   "Specify whether to always interpolate the baseline in the laser window.");
  drifting_params.RegisterParameter("save_interpolations", 
				    save_interpolations = false, 
		     "Save identified interpolation regions as spe");
}

BaselineFinder::~BaselineFinder()
{
  Finalize();
}

int BaselineFinder::Initialize()
{
  return 0;
}

int BaselineFinder::Finalize()
{   
  return 0;
}

int BaselineFinder::Process(ChannelData* chdata)
{
  if(fixed_baseline) return FixedBaseline(chdata);
  else return DriftingBaseline(chdata);
}

int BaselineFinder::DriftingBaseline(ChannelData* chdata)
{
  Baseline& baseline = chdata->baseline;
  double* wave = chdata->GetWaveform();
  std::vector<double>& baseform = chdata->subtracted_waveform;
  const int nsamps = chdata->nsamps;
  baseform.resize(nsamps);
  
  //find the maximum sample value within the pre_trigger area
  int pre_trig_samp = chdata->TimeToSample(signal_begin_time);
  if(pre_trig_samp < 0) pre_trig_samp = pre_samps+post_samps;
  if(pre_trig_samp >= nsamps) pre_trig_samp = nsamps-1;
  double max_pre_trig = *std::max_element(wave,wave+pre_trig_samp);
  double mean_pre_trig = std::accumulate(wave,wave+pre_trig_samp,0.) /
    pre_trig_samp;
  if( std::abs(chdata->GetVerticalRange() - max_pre_trig ) < 0.01)
    baseline.saturated = true;
  //loop through the data, calculating a moving average as we go,
  //but only for samples < 2* max_amplitude away from max_pre_trig
  //until we find a good baseline
  double var = max_amplitude;
  double var2 = max_return_amplitude;
  if(chdata->channel_id == ChannelData::CH_SUM){
    var = max_sum_amplitude;
    var2 *= max_sum_amplitude / max_amplitude;
  }
  
  
  double sum = 0;
  int sum_samps = 0;
  int samp = -1;
  int window_samps = pre_samps + post_samps + 1;
  int last_good_samp = -1;
  double moving_base = -1;
  baseline.found_baseline = false;
  int laserwindow_begin_samp = chdata->TimeToSample(laserwindow_begin_time);
  int laserwindow_end_samp = chdata->TimeToSample(laserwindow_end_time);
  while(++samp < chdata->nsamps){
    bool pass_amp = false;
    if(!baseline.found_baseline && max_pre_trig - wave[samp] < 2*var)
      pass_amp = true;
    //if(!baseline.found_baseline && std::abs(mean_pre_trig-wave[samp]) < var)
    //pass_amp = true;
    else if(baseline.found_baseline){
      if(sum_samps > 0 && std::abs(wave[samp]-moving_base)<var)
	pass_amp = true;
      else if(sum_samps == 0 && std::abs(wave[samp]-moving_base)<var2)
	pass_amp = true; 
    }
    else if(samp>=laserwindow_begin_samp && samp <= laserwindow_end_samp)
      baseline.laserskip = true;
    if(laserwindow_freeze && samp>=laserwindow_begin_samp && samp <= laserwindow_end_samp ){
      pass_amp = false;
      }
  
    if(pass_amp){
      //this sample is part of the baseline
      sum += wave[samp];
      sum_samps++;
      if(sum_samps > window_samps){
	//we have collected too many samples, so remove the earliest one
	sum -= wave[samp-window_samps];
	sum_samps--;
      }
      if(sum_samps == window_samps){
	//we have a validly averaged sample
	double mean = sum/sum_samps;
	moving_base = mean;
	baseform[samp - post_samps] = mean;
	if(last_good_samp < samp-post_samps-1){
	  //linearly interpolate the baseline to fill this region
	  double preval = mean;
	  if(last_good_samp >=0)
	    preval = baseform[last_good_samp];
	  double slope = (mean-preval)/((samp-post_samps)-last_good_samp);
	  for(int backsamp = last_good_samp+1; backsamp < samp; backsamp++){
	    baseform[backsamp] = preval + slope*(backsamp-last_good_samp);
	  }
	  if(save_interpolations && last_good_samp >=0){
	    // locate pe region
	    Spe pe;
	    pe.start_time = chdata->SampleToTime(last_good_samp);
	    pe.length = chdata->SampleToTime(samp-post_samps)-pe.start_time;
	    int peak_samp = std::min_element(wave+last_good_samp, wave+samp-post_samps)-wave;
	    pe.peak_time = chdata->SampleToTime(peak_samp);
	    pe.amplitude = baseform[peak_samp]-wave[peak_samp];
	    baseline.interpolations.push_back(pe);
	    baseline.ninterpolations++;
	  }
	}
	last_good_samp = samp-post_samps;
	if(!baseline.found_baseline){
	  baseline.found_baseline = true;
	  baseline.mean = mean;
	  baseline.search_start_index = samp - window_samps;
	  baseline.length = window_samps;
	  //calculate the variance
	  double sum2 = 0;
	  for(int backsamp = samp-window_samps+1; backsamp<=samp; backsamp++)
	    sum2 += wave[backsamp]*wave[backsamp];
	  baseline.variance = sum2/window_samps - mean*mean;
	}
      }
    }
    else{  //!pass_amp
      //we are part of a real signal now
      sum = 0; 
      sum_samps = 0;
      if(!baseline.found_baseline && samp > pre_trig_samp){
	//can't find baseline in pre-trigger region! abort!
	return 0;
      }
      //continue;
    }
  } // end loop over samples
  //fill in the missing part at the end
  for(samp = last_good_samp+1; samp<nsamps; samp++){
    baseform[samp] = baseform[last_good_samp];
  }
  //subtract off the baseline
  for(samp=0; samp<nsamps; samp++){
    baseform[samp] = wave[samp]-baseform[samp];
  }
      
  return 0;
}

//search for a flat baseline in the pre-trigger window
int BaselineFinder::FixedBaseline(ChannelData* chdata){
	
	Baseline & baseline = chdata->baseline;
	double * wave = chdata->GetWaveform();
	const int nsamps = chdata->nsamps;
	
	//find the maximum sample value within the pre_trigger area
	int pre_trig_samp = chdata->TimeToSample(signal_begin_time);
	if(pre_trig_samp <= 0 || pre_trig_samp >= nsamps) return 0;
	double max_pre_trig = *std::max_element(wave,wave+pre_trig_samp);
	double min_pre_trig = *std::min_element(wave,wave+pre_trig_samp);
	if( std::abs(chdata->GetVerticalRange() - max_pre_trig ) < 0.01
	   || min_pre_trig < 0.01)
		baseline.saturated = true;
	
	double sum=0, sum2=0, mean=0, sigma=0;
	double seg_sum=0, seg_sum2=0, seg_mean=0, seg_sigma=0;
	int sum_samps = 0;
	
	baseline.found_baseline = false;
	int samp = -1;
	while(++samp < chdata->nsamps){
		if(samp>=pre_trig_samp) break;
		
		if(!(samp % segment_samps) && samp){ // another segment recorded
			seg_mean = seg_sum/segment_samps;
			seg_sigma = std::sqrt(seg_sum2/segment_samps-seg_mean*seg_mean);
			
			if(seg_sigma>max_sigma){ //baseline variance is too big, skip this segment
				// do nothing
			}
			else if(sum<1e-6 && sum2<1e-6) { //initialize the baseline
				sum  = seg_sum;
				sum2 = seg_sum2;
				mean = seg_mean;
				sigma = seg_sigma;
				sum_samps = segment_samps;
				baseline.search_start_index = samp - segment_samps;
			}
			else if(std::fabs(sigma-seg_sigma)<max_sigma_diff && std::fabs(mean-seg_mean)<max_mean_diff){
				//add new segment to baseline calculateion
				sum += seg_sum;
				sum2 += seg_sum2;
				sum_samps += segment_samps;
				mean = sum/sum_samps;
				sigma = sqrt(sum2/sum_samps-mean*mean);
			}
			else if(sum_samps>=min_valid_samps) break; //baseline is valid, ignore the remaining part
			else if(sigma>seg_sigma){ //start new baseline if the new segment has much smaller sigma
				sum = seg_sum;
				sum2 = seg_sum2;
				mean = seg_mean;
				sigma = seg_sigma;
				sum_samps = segment_samps;
				baseline.search_start_index = samp - segment_samps;
			}//end new baseline
			else{ //reset baseline and start over if everything is a mess
				sum = 0;
				sum2 = 0;
				mean = 0;
				sigma = 0;
				sum_samps = 0;
			}
			seg_sum = 0;
			seg_sum2 = 0;
		}
		seg_sum  += wave[samp];
		seg_sum2 += wave[samp]*wave[samp];
	}//end while loop
	
	if(sum_samps>=min_valid_samps){
		baseline.found_baseline = true;
		baseline.mean = mean;
		baseline.variance = sigma*sigma;
		baseline.length = sum_samps;
		
		std::vector<double>& baseform = chdata->subtracted_waveform;
		baseform.resize(nsamps);
		
		if(linear_interpolation){
			//new vector to hold events to evaluate baseline
			double max_sigma_lin = max_sigma_factor*1.4142*sigma;
			int cooldown_timer = 0;
			
			std::vector<double> vx;
			std::vector<double> vy;
			double totx=0, toty=0, asamps=0;
			
			//subtract off the baseline
			for(samp=0; samp<nsamps-1; samp++){
				
				bool reset = false;
				
				if(std::fabs(wave[samp]-wave[samp+1]) < max_sigma_lin && cooldown_timer==0){
					asamps++;
					totx+=samp;
					toty+=wave[samp];
				}
				
				if(asamps>=avg_samps){
					vx.push_back(totx/asamps);
					vy.push_back(toty/asamps);
					reset = true;
				}
				
				if(samp<nsamps-pre_cooldown && std::fabs(wave[samp+pre_cooldown]-mean)>pulse_threshold*max_sigma_lin){
					cooldown_timer = cooldown;
					reset = true;
				}
				else if(cooldown_timer>0) cooldown_timer--;
				
				if(reset){
					asamps=0;
					totx=0;
					toty=0;
				}
			}
			
			TGraph* gbase = new TGraph( vx.size(), &vx[0], &vy[0]);
			
			for(samp=0; samp<nsamps; samp++){
				
				if(samp<nsamps-pre_cooldown && std::fabs(wave[samp+pre_cooldown]-mean)>pulse_threshold*max_sigma_lin) cooldown_timer = cooldown;
				else if(cooldown_timer>0) cooldown_timer--;
				
				double meanp = mean;
				if(gbase->GetN()>1) meanp = gbase->Eval(samp);
				
				if(std::fabs(wave[samp]-meanp) < max_sigma_lin && cooldown_timer == 0) baseform[samp] = 0;
				else {
					
					baseform[samp] = wave[samp]-meanp;
					
					if(samp>0 && std::fabs(wave[samp-1]-meanp) < max_sigma_lin) baseform[samp-1] = wave[samp-1]-meanp;
					if(samp<nsamps-1 && std::fabs(wave[samp+1]-meanp) < max_sigma_lin){
						baseform[samp+1] = wave[samp+1]-meanp;
						samp++;
						if(cooldown_timer>0) cooldown_timer--;
					}
				}
			}
			
			gbase->Delete();
		}
		
		else{
			//subtract off the baseline
			for(samp=0; samp<nsamps; samp++){
				baseform[samp] = wave[samp]-mean;
			}
		}
		
		return 0;
	}
	else return 0; 
}

