#include <TMath.h>
#include "TF1.h"
#include "Math/WrappedTF1.h"
#include "Math/GSLIntegrator.h"
#include "Math/IntegrationTypes.h"
#include "G4PhysicalConstants.hh"
#include <RAT/Log.hh>
#include <RAT/PMTChgpdf.hh>

using namespace std;

namespace RAT {

PMTChgpdf* PMTChgpdf::ptrSelf=NULL;

PMTChgpdf::PMTChgpdf() :
    hhpLast(-100.), hhpLastNorm(-100.), thresholdLast(-100.), spreadLast(-100.)
{
    fPMTChargePtr = PMTCharge::Get();
    deltaQ = 1.;
    fMaxPE = dimPE;      // default value of fMaxPE
    for(int i=1; i<=dimPE; i++) {
        genChgCharFn[i] = new CharFn;
        truncChgCharFn[i] = new CharFn;
        genChgPDF[i] = new DiscretePDF;
        recChgPDF[i] = new DiscretePDF;
    }

    tfftFwd = new TFFTRealComplex(nPts,false);
    tfftFwd->Init("M", 0, 0);

    tfftRev = new TFFTComplexReal(nPts,false);
    tfftRev->Init("M", 0, 0);
}

PMTChgpdf* PMTChgpdf::Instance() {
    if(ptrSelf==NULL) { ptrSelf = new PMTChgpdf(); }
    return ptrSelf;
}

// Set maximum number of photoelectrons to be considered in the calculations
int PMTChgpdf::SetMaxPE( int _maxPE ) {
    fMaxPE = _maxPE;
    hhpLast = -100.;  // Force pdf recalculation at next entry to CalcPDFs
    if(fMaxPE>dimPE) {
        fMaxPE = dimPE;
        warn << "PMTChgpdf:  max number of PE to be treated reset from " <<
            _maxPE <<" "<< "to " << dimPE << endl;
    }
    return fMaxPE;      // Pass value back to PEnergy in case it was changed
}

// Check that number of photoelectrons for which a result is requested does
// not exceed the maximum number set previously
int PMTChgpdf::CheckNPE(int _nPE) {
    int nPE = _nPE;
    if(_nPE>fMaxPE) {
//         warn << "PMTChgpdf: Not initialized to calculate for more than " <<
//                                           fMaxPE<<" photoelectrons."<<endl;
        nPE = fMaxPE;
    }
    return nPE;
}

// Return the probability that the recorded charge value is q, assuming nPE
// photoelectrons were produced in the PMT
double PMTChgpdf::PdfRecordedVal(
        int nPE, double q, double hhp, double threshold, double pedSig) {
    nPE = CheckNPE(nPE);
    CalcPDFs( hhp, threshold, pedSig);

    double fBins = q/deltaQ;
    double value;
    if(fBins<0.) {  // guard against negative q, which one would think
                    // shouldn't occur but as of 5 July 2017 did sometimes
        value = recChgPDF[nPE]->p[0];
    }else{
        int idx = fBins;  // index of the charge bin whose center is below q
        // Linear interpolation in the recorded charge pdf:
        value = recChgPDF[nPE]->p[idx] + (fBins - idx)*
                        (recChgPDF[nPE]->p[idx+1] - recChgPDF[nPE]->p[idx]);
    }
    value = value/deltaQ;   // divide by deltaQ to get probability per unit
                            // change in charge
    if ( value<0. ) value = 0.;
    return value;
}

// Return the probability that the generated charge will be below the threshold
double PMTChgpdf::ProbBelowThreshold ( int nPE, double hhp, double threshold) {
    nPE = CheckNPE(nPE);
    CalcPDFs( hhp, threshold );
    double fBins =  threshold/deltaQ + 0.5;
    int idx = fBins;    // index of the charge bin in which the
                        // threshold value lies
    double sum = 0.;
    for(int j=0; j<idx; j++) {
        sum += genChgPDF[nPE]->p[j];
    }
    sum += (fBins-idx)*genChgPDF[nPE]->p[idx];
    if (sum < 0.)  sum = 0.;
    return sum;
}

void PMTChgpdf::CalcPDFs( double hhp, double threshold ) {
    // The function ProbBelowThreshold uses only the generated-charge pdf and
    // therefore calls this version of CalcPDF, which uses the same pedestal
    // spread as the last time CalcPDFs was called, perhaps avoiding
    // unnecessary recalculation of the pdfs.
    double spread;
    if (spreadLast == -100.) {
        spread=3.61;
    }else{
        spread = spreadLast;
    }
    CalcPDFs( hhp, threshold, spread) ;
    return;
}

// Calculate the various pdfs for the desired hhp, theshold, and pedSpread.
void PMTChgpdf::CalcPDFs( double hhp, double threshold, double pedSpread) {
    // Don't recalculate if the parameters have not changed
    if (hhp!=hhpLast || threshold!=thresholdLast || pedSpread!=spreadLast) {
        double q[nPts];  // Charge values at which discrete pdf values
                         // are input to the FFTs.
        for(int i=0; i<nPts; i++) { q[i] = (double)i*deltaQ; }
        double inChgDistn[nPts];  // pdf values input to the FFTs
        double fInvNPts = 1./(double)nPts;

        // Don't need to recalculate the generated charge pdfs if the high-half
        // point has not changed
        if( hhp!=hhpLast ) {
            SetSinglePENorm( hhp );
            // Non-zero values of the single-PE pdf for charge>200 are ignored
            int maxBinSingle = 200./deltaQ;
            for(int i=0; i<maxBinSingle; i++) {
                // This could be made more precise by making
                // inChgDistn[i] = the difference in the integrals of
                // DoublePolyaQSpectrum at the edges of the bin.
                inChgDistn[i] = deltaQ*DoublePolyaQSpectrum( q[i], hhp) ;
                genChgPDF[1]->p[i] = inChgDistn[i];
            }
            for(int i=maxBinSingle; i<nPts; i++) {
                inChgDistn[i] = 0.;
                genChgPDF[1]->p[i] = 0.;
            }
            // Calculate the discrete characteristic fn of the single-PE pdf:
            tfftFwd->SetPoints(inChgDistn);
            tfftFwd->Transform();
            tfftFwd->GetPointsComplex
                        (genChgCharFn[1]->real, genChgCharFn[1]->imag, false);
            for(int i=2; i<=fMaxPE; i++) {
                int j = i/2;
                int k = j;
                if( i%2==1) k++;
                // The characteristic function for nPE photoelectrons is the
                // nPE-th power of the single-PE characteristic function:
                MultiplyCharFn(
                        *genChgCharFn[j], *genChgCharFn[k], *genChgCharFn[i] );
                // Apply reverse FFT to obtain the generated-charge pdf for nPE
                // photoelectrons:
                tfftRev->SetPointsComplex(
                                genChgCharFn[i]->real, genChgCharFn[i]->imag );
                tfftRev->Transform();
                double* tempout = tfftRev->GetPointsReal(false);
                // Need to scale the result.  See, e.g., Numerical Recipes in
                // Fortran, equation 12.1.9.
                for(j=0; j<nPts; j++) {
                    genChgPDF[i]->p[j] = tempout[j]*fInvNPts;
                }
            }
        }

        if( hhp!=hhpLast || threshold!=thresholdLast ) {
            // Calculate the characteristic functions of the truncated
            // charge distributions
            if ( threshold>0. ) {   // sometimes the threshold is negative??
                for(int i=1; i<=fMaxPE; i++) {
                    // First truncate the charge pdf at the threshold
                    double fBins = threshold/deltaQ + 0.5;
                    int idx = fBins;    // index of the charge bin
                                        // in which the threshold value lies
                    for(int j=0; j<idx; j++) {
                        inChgDistn[j] = 0.;
                    }
                    // Interpolate for bin idx:
                    inChgDistn[idx] = (fBins-idx)*genChgPDF[i]->p[idx];
                    for(int j=idx+1; j<nPts; j++) {
                        inChgDistn[j] = genChgPDF[i]->p[j];
                    }
                    // Renormalize the pdf after truncation:
                    double normFactor = 1./PdfBinSum( inChgDistn, 0, nPts-1);
                    for(int j=0; j<nPts; j++) {
                        inChgDistn[j] = normFactor*inChgDistn[j];
                    }
                    // Apply FFT to get the characteristic function of the
                    // truncated distribution:
                    tfftFwd->SetPoints(inChgDistn);
                    tfftFwd->Transform();
                    tfftFwd->GetPointsComplex ( truncChgCharFn[i]->real,
                                            truncChgCharFn[i]->imag, false);
                }
            }else{
                for(int i=1; i<=fMaxPE; i++) {
                    *truncChgCharFn[i] = *genChgCharFn[i];
                }
            }
        }

        if( pedSpread != spreadLast) {
            // Calculate the characteristic function of the pedestal spread pdf
            double inSpreadDistn[nPts];
            for(int i=0; i<nPts; i++) {
                inSpreadDistn[i] = 0.;
            }
            // Pedestal spread pdf is a Gaussian:
            double finv = deltaQ/(pedSpread*sqrt(twopi));
            // Truncate the spread pdf at 4 times sigma:
            int maxIndex = 4.*pedSpread/deltaQ;
            double expfactor = 1./(2.*pedSpread*pedSpread);
            inSpreadDistn[0] = finv;
            // The pedestal spread can be negative as well as positive.  Pdf
            // values for negative spreads are input to the FFT at the upper
            // end of the distribution.  See Numerical Recipes in Fortran,
            // section 13.1.
            if(maxIndex>=nPts) {
                detail <<
                    "PMTChgdpf error: pedSpread, deltaQ, nPts, maxIndex = " <<
                   pedSpread<<" "<<  deltaQ<<" "<< nPts<<" "<< maxIndex <<endl;
            }
            for(int i=1; i<=maxIndex; i++) {
                inSpreadDistn[i] = finv*exp( -q[i]*q[i]*expfactor );
                inSpreadDistn[nPts-i] = inSpreadDistn[i] ;
            }
            // Apply FFT to get the characteristic function of the
            // pedestal spread pdf:
            tfftFwd->SetPoints(inSpreadDistn);
            tfftFwd->Transform();
            tfftFwd->GetPointsComplex(
                                spreadCharFn.real, spreadCharFn.imag, false);
        }

        CharFn tempChFn;
        for(int i=1; i<=fMaxPE; i++) {
            // Multiply the characteristic function of the truncated charge
            // distribution by the characteristic function of the pedestal
            // spread to get the characteristic fn of the recorded charge pdf.
            MultiplyCharFn ( *truncChgCharFn[i], spreadCharFn, tempChFn );
            // Apply the inverse transformation to get the pdf of the
            // recorded charge:
            tfftRev->SetPointsComplex( tempChFn.real, tempChFn.imag );
            tfftRev->Transform();
            double* tempout = tfftRev->GetPointsReal(false);
            for(int j=0; j<nPts; j++) {
                recChgPDF[i]->p[j] = tempout[j]*fInvNPts;
            }
        }

        hhpLast = hhp;
        thresholdLast = threshold;
        spreadLast = pedSpread;
    } // if (hhp!=hhpLast || threshold!=thresholdLast || pedSpread!=spreadLast)
}

void PMTChgpdf::MultiplyCharFn( CharFn& a, CharFn& b, CharFn& result) {
    // Multiplication of two characteristic functions
    for(int i=0; i<nOv2P1; i++) {
        result.real[i] = a.real[i]*b.real[i] - a.imag[i]*b.imag[i];
        result.imag[i] = a.imag[i]*b.real[i] + a.real[i]*b.imag[i];
    }
}

double PMTChgpdf::PdfBinSum( double* y, int lowLim, int hiLim) {
    // Sum bin values in a discrete dpf
    double sum = 0.;
    for(int i=lowLim; i<=hiLim; i++) {
        sum = sum + y[i];
    }
    return sum;
}

// Integrand for single-PE charge distribution
double PMTChgpdf::DPQSIntegrand (double* x, double* p) {
    return DoublePolyaQSpectrum(x[0], p[0]);
}

// Normalize the single-PE charge pdf if necessary
void PMTChgpdf::SetSinglePENorm( double hhp ) {
    if( hhp!=hhpLastNorm ) {
        fnorm = 1.;
        TF1* f = new TF1("fnorm", this, &PMTChgpdf::DPQSIntegrand,0.5,500.,1);
        f->SetParameter(0,hhp);
        ROOT::Math::WrappedTF1 wf(*f);
        ROOT::Math::GSLIntegrator ig(ROOT::Math::IntegrationOneDim::kADAPTIVE);
        ig.SetFunction(wf);
        // Set lower and upper limits for normalization consistent with
        // definition of input to discrete Fourier transform.
        fnorm = 1./ig.Integral(0.5, 199.5);
        hhpLastNorm = hhp;
        delete f;
    }
}

// Single PE generated-charge pdf values - probability per unit change in charge
double PMTChgpdf::DoublePolyaQSpectrum(const double Q, double hhp) {
    if (Q<0.5) return 0.;
    return fnorm*(double)fPMTChargePtr->DoublePolyaQSpectrum(Q, hhp);
}

}   // namespace RAT