#include <algorithm>
#include <iostream>
#include "TF1.h"
#include "TH1.h"
#include "TH2.h"
#include "TTree.h"
#include "TFile.h"
#include "TProfile.h"
#include "TRandom3.h"
#include "TMath.h"

void g2gen() {

  //define a few constants

  // Momentum
  double pmumean = 3.094349631;
  double pmusd = 0.001 * pmumean; //0.1%
  double pmurang = pmumean + pmusd;
  double pmulo = pmumean-3*pmusd;
  double pmuhi = pmumean+3*pmusd;

  // Time & Freq
  double rMagic = 711.214761;
  double trevav = 2.*TMath::Pi()*rMagic/(0.99942*29.979246); //1.0e9*2.*TMath::Pi()*711.215/(0.99942*2.9979246e10);
  double freqav = 1.0/trevav;
  double freqrang = 0.020794769/pmurang;
  double freqsd = freqav - freqrang;
  double freqlo = freqav-freqsd*3.;
  double freqhi = freqav+freqsd*3.;

  // Muon pars
  double taumu = 2196.9803; //MuLan lifetime in ns
  double gamma = 29.4;

  // T0 Pars
  double t0mean = 0.0;
  double t0sd = 15;

  // Initialise histograms
  TH1D* hFreqSpec = new TH1D("hFreqSpec","Frequency Spectrum (GHz)",101,freqlo,freqhi);
  TH1D* hMomSpec = new TH1D("hMomSpec","Momentum Spectrum (GeV/c)",101,pmulo,pmuhi);
  TH1D* hT0Pulse = new TH1D("hT0Pulse","T0Pulse (ns) ",201,-100.5,100.5);
  const int ndet = 1;
  TH1D* hSignalRaw[ndet];
  TH1D* hSignalWeighted[ndet];
  TH1D* hSignalRatio_U[ndet];
  TH1D* hSignalRatio_V[ndet];
  TH1D* hSignalRatio[ndet];
  for(int idet=0;idet<ndet;idet++){
    hSignalRaw[idet] = new TH1D(Form("hSignalRaw%02d",idet),"Raw Signal;Time [#mus]",150000,-0.0005,149.9995); //only signal 00 is used
    hSignalWeighted[idet] = new TH1D(Form("hSignalWeighted%02d",idet),"Background-Removed Signal;Time [#mus]",150000,-0.0005,149.9995); //Holder for background removed (with wiggle & lifetime fit removed)
    hSignalRatio_U[idet] = new TH1D(Form("hSignalRatio_U%02d",idet),"Ratio Signal: U",150000,-0.0005,149.9995); // 5 histograms for calculating ratio
    hSignalRatio_V[idet] = new TH1D(Form("hSignalRatio_V%02d",idet),"Ratio Signal: V",150000,-0.0005,149.9995);
    hSignalRatio[idet] = new TH1D(Form("hSignalRatio%02d",idet),"Ratio Signal: U-V/U+V",150000,-0.0005,149.9995);
  }
  TH1D* hTimeDecay = new TH1D("hTimeDecay","Muon Decay Time ",150000,0.0,150.0);
  TH1D* hDetector = new TH1D("hDetector","Detector Number ",24,-0.5,23.5);
  TH1D* hnTurns = new TH1D("hnTurns","Number of turns",10000,0.0,6000.0);
  TH2D* htrevt0 = new TH2D("htrevt0","Time of revolution vs t0", 150, -75.0, 75.0, 120, 149.26-0.6, 149.26+0.6);
  TH1D* hRadius = new TH1D("hRadius","Radius of muon orbit",650,711.2-6.5,711.2+6.5); 

  // Initialise 5-parameter wiggle function
  TF1* fWiggle = new TF1("fWiggle","[0]*exp(-x/[1])*(1+[2]*cos([3]*x + [4]))",0,150000);
  //  fWiggle->SetParameters(1, 64372, 0.4, 0.00143944, 0.0);
  fWiggle->SetParameters(1, 64372, 0, 0.00143944, 0.0);  // Turn off wiggle for exponential decay only
  fWiggle->SetNpx(150000);
  
  // Main event loop
  TRandom3* re = new TRandom3(6280225);
  double nEvents = 1e10;
  for(long long iev = 0; iev < nEvents; iev++){
  
    if(iev%1000000 ==0) cout << "Muon number " << iev << endl;

    //Select a T0 for the muon (ns)
    double t0 = -1;

    //Ordinary Gaussian time distribution
    if(true){
      t0 = re->Gaus(t0mean,t0sd);
    }

    // Select from uniform distribution
    if(false){
      t0 = re->Uniform(-5.5,5.5);
    }

    // Select from histogram of t0 pulse
    if(false){
      //      t0 = ht0->GetRandom() * tbinh;
    }

    //Sawtooth t0 pulse
    if(false){
      t0 = re->Uniform(t0mean-t0sd, t0mean+t0sd);
      double t1 = re->Uniform(-1.0, 1.0);
      while ( t0 / t0sd < t1){
	t0 = re->Uniform(t0mean-t0sd, t0mean+t0sd);
	t1 = re->Uniform(-1.0, 1.0);
      }
    }
    
    //Now select "momentum", actually a revolution time
    double pmu = -1;

    //correlated momentum selection
    if(false){
      double pmm = (t0 > 0) ? pmulo : pmuhi;
      pmu = re->Gaus(pmm,pmusd);
    }

    // Momentum distribution will be an ordinary gaussian
    if(true){
      pmu = re->Gaus(pmumean,pmusd);
      while (pmu < pmulo or pmu > pmuhi) pmu = re->Gaus(pmumean,pmusd);
    }

    //The momentum distribution will be a half-gaussian
    if(false){
      pmu = re->Gaus(pmumean,pmusd);
      while( pmu-pmumean < 0 or pmu > pmuhi) pmu = re->Gaus(pmumean,pmusd);
    }

    //The momentum distribution will be an ordinary gaussian
    //but with a width that depends on t0. As a first try
    //we'll go with an n percent change in width with abs(t0) 
    if(false){
      double pmusdt0 = (fabs(t0) < 3*t0sd) ? pmusd*(1.0-0.1*fabs(t0/t0sd)) : pmusd*(1.0-0.3);
      pmu = re->Gaus(pmumean,pmusdt0);
      while(pmu < pmulo or pmu > pmuhi) pmu = re->Gaus(pmumean,pmusdt0);
    }			   
      
    // Convert to time 
    double trev = trevav*pmu/pmumean; //This is incorrect!! - Comment is inherited (I don't know why it's incorrect)

    //Now choose decay time
    double tdec = -1;

    //Use exponential decay
    if(false){
      double randb = re->Uniform(1.0);
      double rnmut = -log(1.0-randb);
      double taudil = gamma*taumu*pmu/pmumean;
      tdec = taudil*rnmut;
    }

    // Use wiggle decay
    if(true){
      tdec = fWiggle->GetRandom();
    }

    //Use uniform decay
    //This will not neceearily work for very early times
    //when tdec could be smaller than t0, for example
    //so fit for times > 0.5 turns
    if(false){
      tdec = 140000.0*re->Uniform(1.0);
    }      


    //With t0 and trev, grab the momentum (not relatistic) and frequency
    double frev = 1./trev;
    hMomSpec->Fill(pmu);
    hFreqSpec->Fill(frev);

    //TDec is the time of decay for this muon
    hTimeDecay->Fill(tdec / 1e3); // us

    //and determine the radius
    double murad = rMagic*trev/trevav;
    hRadius->Fill((float)murad);
      
    //Fill t0 pulse histogram
    double t0fill = t0;
    hT0Pulse->Fill(-t0fill);  //For use with fasrof

    //Calculate number of turns, detector
    double nturns = (tdec+t0)/trev; //number of turns; this assumes that positive t0 is already past detector 0
    hnTurns->Fill(nturns);
    double ttt = nturns-(int)nturns; // this is distance past detector 0 in turns. But the corresponding time is tdec
    double tfill = tdec-(ttt*trev); //this is the detector 0 signal

    int cdet = (int)(ndet*ttt); 

    hDetector->Fill((float)cdet);
    if(cdet > -1 && cdet < ndet) {
      hSignalRaw[cdet]->Fill(tfill/1e3);   //filling only detector 0  (us)

      // Ratio histograms
      double r = re->Uniform();
      if (r < 0.5){
       	hSignalRatio_U[cdet]->Fill(tfill/1e3);
      } else {
	hSignalRatio_V[cdet]->Fill(tfill/1e3 + trevav*(re->Uniform()-0.5)/1e3);
      }
    }

  }

  // Loop over detectors
  for(int idet=0;idet<ndet;idet++) {

    //////////////////////////////////////////////////////////////////////////////////////////////
    // Fit wiggle and remove from weighted histogram
    //////////////////////////////////////////////////////////////////////////////////////////////

    TH1D* hSignalWiggle = (TH1D*)hSignalRaw[idet]->Clone(Form("hSignalWiggle_%d",idet));

    // Bin out fast rotation (incoming bins were 1 ns, so rebin by 150)
    hSignalWiggle->Rebin(150);

    // Fit function (normal 5 parameter wiggle)
    TF1 *fitWiggle = new TF1(Form("fitWiggle_%d",idet), "[0]*exp(-x/[1])*(1 + [2]*cos([3]*x + [4]))", 30.0, 95.0);
    fitWiggle->SetParameters(1, 64.4, 0.4, 1.4, 0);
    fitWiggle->FixParameter(2,0);
    fitWiggle->FixParameter(3,0);
    fitWiggle->FixParameter(4,0);

    // Choose good starting guess by comparing integral of function and histogram
    double normGuess = hSignalWiggle->Integral(hSignalWiggle->GetXaxis()->FindBin(30), hSignalWiggle->GetXaxis()->FindBin(95), "WIDTH") / fitWiggle->Integral(30,95);
    fitWiggle->SetParameter(0,normGuess);

    // Fit function
    hSignalWiggle->Fit(fitWiggle,"NREML");

    // Set normalisation to 1 to just keep function without normalisation
    TF1* fitWiggleNorm = (TF1*)fitWiggle->Clone(Form("fitWiggleNorm_%d",idet));
    fitWiggleNorm->SetParameter(0,1);

    // Make histogram divided by wiggle and lifetime
    for(int j = 1; j <= hSignalRaw[idet]->GetNbinsX(); j++) {
      hSignalWeighted[idet]->SetBinContent(j, hSignalRaw[idet]->GetBinContent(j) / fitWiggleNorm->Eval(hSignalRaw[idet]->GetBinCenter(j)) );
      hSignalWeighted[idet]->SetBinError(j, hSignalRaw[idet]->GetBinError(j) / fitWiggleNorm->Eval(hSignalRaw[idet]->GetBinCenter(j)) );
    }

    //////////////////////////////////////////////////////////////////////////////////////////////
    // Make Ratio Histogram
    //////////////////////////////////////////////////////////////////////////////////////////////

    // Scale denominator to 1 at t = 0 (same as fit function);
    hSignalRatio_V[idet]->Scale(fitWiggleNorm->Integral(hSignalRatio_V[idet]->GetXaxis()->GetXmin(), hSignalRatio_V[idet]->GetXaxis()->GetXmax()) / hSignalRatio_V[idet]->Integral("WIDTH"));

    for(int iBin = 1; iBin <= hSignalRatio_V[idet]->GetNbinsX(); iBin++){
      if(hSignalRatio_V[idet]->GetBinContent(iBin) > 0){
	double U = hSignalRatio_U[idet]->GetBinContent(iBin);
	double V = hSignalRatio_V[idet]->GetBinContent(iBin);
	double UErr = hSignalRatio_U[idet]->GetBinError(iBin);
	double VErr = hSignalRatio_V[idet]->GetBinError(iBin);
	double ratio = U/V;
	double ratioErr = U > 0 ? ratio * sqrt(pow(UErr/U,2) + pow(VErr/V,2)) : 0;
	hSignalRatio[idet]->SetBinContent(iBin,ratio);
	hSignalRatio[idet]->SetBinError(iBin,ratioErr);
      }
    }
    hSignalRatio[idet]->ResetStats();

  }


  // Write out histograms to file
  string outfile = "g2genExpDecay.root";
  cout << "file is " << outfile << endl;
  TFile *fout = new TFile(outfile.c_str(),"RECREATE");
  hDetector->Write();
  hMomSpec->Write();
  hnTurns->Write();
  hFreqSpec->Write();
  hT0Pulse->Write();
  //  ht0->Write();
  hTimeDecay->Write();
  for(int idet=0;idet<ndet;idet++) {
    hSignalRaw[idet]->Write();
    hSignalWeighted[idet]->Write();
    hSignalRatio_U[idet]->Write();
    hSignalRatio_V[idet]->Write();
    hSignalRatio[idet]->Write();
  }
  htrevt0->Write();
  hRadius->Write();
  fout->Write();
  fout->Close();
}