#include <cmath>
#include <algorithm>
#include <iostream>
#include <set>
#include <tut.h>

#include "TLorentzVector.h"
#include "TH1F.h"
#include "THnBase.h"
#include "THnSparse.h"

#include "HEPUnits.hxx"
#include "IRooTrackerHistogram.hxx"
#include "IRooTrackerFile.hxx"

namespace tut {
    struct baseRTHistGenerate {
        baseRTHistGenerate() {
            // Run before each test.
        }
        ~baseRTHistGenerate() {
            // Run after each test.
        }
    };

    // Declare the test
    typedef test_group<baseRTHistGenerate>::object testRTHistGenerate;
    test_group<baseRTHistGenerate> groupRTHistGenerate("RTHistGenerate");

    // Check that RT histogram distributions can correctly generate the 
    // frequency of particles
    template<> template<>
    void testRTHistGenerate::test<1> () {
        // Construct the output histogram
        std::string outputFileName = "test_RT_histIn_5.root";
        COMET::IOutputRooTrackerHistogram testOutput(outputFileName);
        bool init = testOutput.InitialiseForWriting("RECREATE");
        // Add a dummy proton
        int testProt = 2212;
        int testNeut = 2112;
        TLorentzVector zeroVector(0.0, 0.0, 0.0, 0.0);
        int nParts = 10;
        // Add 0 events with 0, 1 event with 1, 2 events with 2 etc.
        // etc.
        int nEvents = 0;
        int nSkipEvent = 0;
        int totalProt = 0;
        int totalNeut = 0;
        for (int i = 0; i < nParts; ++i){
            for (int j = 0; j < i; j++){
                testOutput.BeginEvent();
                bool skippedEvent = false;
                for (int k = 0; k < i; k++){
                    testOutput.AddParticle(testProt, zeroVector, zeroVector, 1.0);
                    totalProt++;
                    if (i%2 == 0){
                        testOutput.AddParticle(testNeut, 
                                               zeroVector,  
                                               zeroVector,
                                               1.0);
                        totalNeut++;
                    }
                    else skippedEvent = true;
                }
                testOutput.EndEvent();
                if (skippedEvent) nSkipEvent++;
                ++nEvents;
            }
        }
        // Number of total events read
        int emptyEvents = 5;
        int totalEvents = nEvents + emptyEvents;
        testOutput.Finalize(totalEvents);
        testOutput.Write();
        testOutput.Close();
        // Open as input rootracker histogram
        COMET::IRooTrackerHistogram* testInput = 
            COMET::IRooTrackerHistogram::Open(outputFileName);
        // Total input events is 50, sperad accross frequencies [0,9]
        // Multiply this by 10,000 to generate frequency distribution
        int sampleFactor = 10000;
        int totalSamples = totalEvents*sampleFactor;
        int freqNeut[10] = {0};
        int freqProt[10] = {0};
        for (int i = 0; i < totalSamples; i++){
            const COMET::IRooTrackerHistogram::PrimaryList& primaries = 
                testInput->GeneratePrimaries();
            int nNeut = 0;
            int nProt = 0;
            for (COMET::IRooTrackerHistogram::PrimaryList::const_iterator 
                    pP = primaries.begin();
                    pP != primaries.end();
                    pP++){
                if (pP->pid == testProt) ++nProt;
                if (pP->pid == testNeut) ++nNeut;
            }
            ensure("Freqency of protons and neutrons is not larger than allowed "
                   " by input", (nProt < nParts) && (nNeut < nParts));
            freqNeut[nNeut]++;
            freqProt[nProt]++;
        }
        // Check frequency recorded by 50,000 particles is proprotional to 
        // 5 empty events, 1 event with 1, 2 events with 2, ..., 9 events with 9
        for (int i = 0; i < nParts; ++i){
            int expectFreq = i;
            int expectNeutFreq = expectFreq;
            if (i==0) {
                // Had 5 empty eevents
                expectFreq = emptyEvents;
                // Neutrons not generated on (odd values of i + empty events) = 
                // 5 + 1 + 3 + 5+ 7 + 9 = 30
                expectNeutFreq = nSkipEvent + emptyEvents;
            }
            if (i%2 != 0) expectNeutFreq = 0;
            int nearFreq = round(float(freqProt[i])/sampleFactor); 
            ensure_equals("Frequency of protons matches the input", 
                          nearFreq, expectFreq);
            int nearNeutFreq = round(float(freqNeut[i])/sampleFactor); 
            ensure_equals("Frequency of neutrons matches the input", 
                          nearNeutFreq, expectNeutFreq);
        }
        // Check that generating one particle many many times respects the ratio 
        // of different particle production
        int nNeut = 0;
        int nProt = 0;
        for (int i = 0; i < sampleFactor*100; i++){
            COMET::IRooTrackerHistogram::PrimaryParticle_t aPart = 
                testInput->GenerateOnePrimary();
            if (aPart.pid == testProt) nProt++;
            else if (aPart.pid == testNeut) nNeut++;
        }
        // Callcualte the relative proportion of protons and neutrons
        double protToNeut = ((double)nProt)/nNeut;
        double protToNeutOrig = ((double)totalProt)/totalNeut;
        // TODO fix this test
        // Ensure the overall particle ratios are the same
        //ensure("Ratio of proton to neutron production correct for"
        //       " GenerateOnePrimary", protParts == neutParts);
    }

    // Check that RT histogram distributions can correctly generate the 
    // distribution of particles
    template<> template<>
    void testRTHistGenerate::test<2> () {
        // Construct the output histogram
        std::string outputFileName = "test_RT_histIn_6.root";
        COMET::IOutputRooTrackerHistogram testOutput(outputFileName);
        bool init = testOutput.InitialiseForWriting("RECREATE");
        // Construct more convenient axes
        std::string setupString = " =(10000, -1, 10)";
        for (int d = 0; d < COMET::IRooTrackerHistogram::kNumDimensions; d++){
            COMET::IRooTrackerHistogram::Dimensions_t dim = 
                static_cast<COMET::IRooTrackerHistogram::Dimensions_t>(d); 
            testOutput.SetupAxis(COMET::IRooTrackerHistogram::DimensionStr(dim)+
                                 setupString);
        }
        // Add a dummy proton
        int testProt = 2212;
        TLorentzVector zeroVector(0.0, 0.0, 0.0, 0.0);
        TLorentzVector oneVector(1.0, 1.0, 1.0, 1.0);
        int nParts = 10;
        int totalProt = 0;
        testOutput.BeginEvent();
        // For 8 dimensions    (x, y, z, t, px, py, pz, E)
        //  * 0 particles with (0, 0, 0, 0,  0,  0,  0, 0)
        //  * 1 particle  with (1, 1, 1, 1,  1,  1,  1, 1)
        //  * ...
        //  * 9 particles with (9, 9, 9, 9,  9,  9,  9, 9)
        for (int i = 0; i < nParts; ++i){
            TLorentzVector thisVector = oneVector * static_cast<double>(i);
            for (int j = 0; j < i; j++){
                testOutput.AddParticle(testProt, thisVector, thisVector, 1.0);
                totalProt++;
            }
        }
        testOutput.EndEvent();
        // Number of total events read
        int totalEvents = 1;
        testOutput.Finalize(totalEvents);
        testOutput.Write();
        testOutput.Close();
        // Open as input rootracker histogram
        COMET::IRooTrackerHistogram* testInput = 
            COMET::IRooTrackerHistogram::Open(outputFileName);
        // Total number of particles is 45
        // Multiply this by 10,000 to generate output distribution
        int sampleFactor = 1000;
        int totalSamples = totalEvents*sampleFactor;
        int freqVect[10] = {0};
        for (int i = 0; i < totalSamples; i++){
            const COMET::IRooTrackerHistogram::PrimaryList& primaries = 
                testInput->GeneratePrimaries();
            for (COMET::IRooTrackerHistogram::PrimaryList::const_iterator 
                    pP = primaries.begin();
                    pP != primaries.end();
                    pP++){
                // Ensure this particle was described in the input, where each 
                // component had the same value
                bool okayP = true;
                okayP = okayP && TMath::AreEqualAbs(pP->posX, pP->posY,   5e-3);
                okayP = okayP && TMath::AreEqualAbs(pP->posX, pP->posZ,   5e-3);
                okayP = okayP && TMath::AreEqualAbs(pP->posX, pP->time,   5e-3);
                okayP = okayP && TMath::AreEqualAbs(pP->posX, pP->momX,   5e-3);
                okayP = okayP && TMath::AreEqualAbs(pP->posX, pP->momY,   5e-3);
                okayP = okayP && TMath::AreEqualAbs(pP->posX, pP->momZ,   5e-3);
                okayP = okayP && TMath::AreEqualAbs(pP->posX, pP->energy, 5e-3);
                ensure("Only particles described in distribution histogram "
                       "are generated", okayP);
                // Count this particle as having been generated
                int intPart = round(pP->posX);
                ++freqVect[intPart];
            }
        }
        // Check that generating one particle many times respects the 
        // description of the input particles
        // The total number of samples is the (sample factor * the total number 
        // of particles)
        int freqVectSingle[10] = {0};
        for (int i = 0; i < sampleFactor*totalProt; i++){
            COMET::IRooTrackerHistogram::PrimaryParticle_t aP = 
                testInput->GenerateOnePrimary();
            // Ensure this particle was described in the input, where each 
            // component had the same value
            bool okayP = true;
            okayP = okayP && TMath::AreEqualAbs(aP.posX, aP.posY,   5e-3);
            okayP = okayP && TMath::AreEqualAbs(aP.posX, aP.posZ,   5e-3);
            okayP = okayP && TMath::AreEqualAbs(aP.posX, aP.time,   5e-3);
            okayP = okayP && TMath::AreEqualAbs(aP.posX, aP.momX,   5e-3);
            okayP = okayP && TMath::AreEqualAbs(aP.posX, aP.momY,   5e-3);
            okayP = okayP && TMath::AreEqualAbs(aP.posX, aP.momZ,   5e-3);
            okayP = okayP && TMath::AreEqualAbs(aP.posX, aP.energy, 5e-3);
            ensure("Only particles described in distribution histogram "
                   "are generated in GenerateOnePrimary mode", okayP);
                // Count this particle as having been generated
            int intPart = round(aP.posX);
            ++freqVectSingle[intPart];
        }
        // Check the frequencies of the particles that were generated matched 
        // the frequncy of the input, i.e.  * 0 particles with (0, 0, 0, 0,  0,  
        // 0,  0, 0)
        //  * 1 particle  with (1, 1, 1, 1,  1,  1,  1, 1)
        //  * ...
        //  * 9 particles with (9, 9, 9, 9,  9,  9,  9, 9)
        for (int i = 0; i < nParts; ++i){
            int expectFreq = i;
            int nearFreq = round(float(freqVect[i])/sampleFactor); 
            int nearFreqSingle = round(float(freqVectSingle[i])/sampleFactor); 
            ensure_equals("Description of protons matches the input for "
                          "GeneratePrimaries method", 
                          nearFreq, expectFreq);
            ensure_equals("Description of protons matches the input for "
                          "GenerateOnePrimary method", 
                          nearFreqSingle, expectFreq);
        }

    }


};