#include "RooRealVar.h"
#include "RooDataSet.h"
#include "RooGaussian.h"
#include "RooLandau.h"
#include "RooFFTConvPdf.h"
#include "RooPlot.h"
#include "TCanvas.h"
#include "TAxis.h"
#include "TH1.h"
#include "TFile.h"
#include <vector>
#include <string>
#include <iostream>
#include "TChain.h"
#include "TTree.h"
#include "TH2D.h"

using namespace RooFit;

void troll1_hl()
{
    Int_t hit0, hit1, hit2, hit3, hit4, hit5, hit6, hit7, hit8, hit9, hit10, hit11, hit12, hit13, hit14, hit15;
    double grain1X, grain1Y, muonDZ;
    int eventType;
    double diffTrack2;
    // set up some constants
    const double    coeff[16] = {0.83313962, 0.88315463, 1.39515245, 1.17506533, 1.107663,   0.37426376,
    0.73992582, 1.04746587, 1.33358453, 1.18417938, 1.18210636, 0.86438165,
    0.98166011, 1.21217571, 0.44520008, 1.20872314};

    const double posXNref[16] = {0.32694611 , 1.04790419 ,-1.01437126, -0.36047904 , 0.32694611,  1.01437126,
                                 -1.01437126 ,-0.34371257,  0.32694611 , 1.01437126 ,-1.03113772, -0.36047904,
                                 0.32694611 ,  0.98083832 ,-1.03113772 ,-0.37724551 };

    const double posYNref[16] = { 0.99760479 , 0.99760479 , 1.01437126 , 0.98083832  ,0.31017964 , 0.29341317,
                                  0.31017964 , 0.34371257 ,-0.42754491 ,-0.42754491, -0.42754491 ,-0.41077844,
                                  -1.11497006, -1.11497006, -1.13173653 ,-1.16526946};

    const double listFibre[16] = {13, 9, 5, 1, 12, 8, 4, 0, 15, 11, 7, 3, 14, 10, 6, 2};
    const double nQ[16] = {1, 2, 2, 1,  2, 4, 4, 2,  2,  4, 4, 2,  1,  2, 2, 1};
    const double mlistQ[4][16] = {{1,  1, 2, 2,  1, 0, 0, 2,  0,  0, 0, 3,  0,  0, 3, 3},
                                  {1,  2, 1, 2,  0, 1, 1, 3,  1,  1, 1, 2,  0,  3, 0, 3},
                                  {1,  1, 2, 2,  1, 2, 2, 2,  0,  2, 2, 3,  0, 0, 3, 3},
                                  {1,  2, 1, 2,  0, 3, 3, 3,  1,  3, 3, 2,  0, 3, 0, 3}};

    const double rmin = 0.1;
    const double rmax = 0.35;

    const double mVmax = 2000;
    const double mVmin = 0;

    const int runNumber = 60;
    std::string fileName ="/Users/chanal/Documents/grainita/test beam/Run_"+std::to_string(runNumber)+".root";
    TChain *tChain = new TChain("DATA");
    tChain->Add(fileName.c_str());
    std::cout << "File " << fileName << " added to the chain" << std::endl;
    tChain->SetBranchAddress("hit0Cor", &hit0);
    tChain->SetBranchAddress("hit1Cor", &hit1);
    tChain->SetBranchAddress("hit2Cor", &hit2);
    tChain->SetBranchAddress("hit3Cor", &hit3);
    tChain->SetBranchAddress("hit4Cor", &hit4);
    tChain->SetBranchAddress("hit5Cor", &hit5);
    tChain->SetBranchAddress("hit6Cor", &hit6);
    tChain->SetBranchAddress("hit7Cor", &hit7);
    tChain->SetBranchAddress("hit8Cor", &hit8);
    tChain->SetBranchAddress("hit9Cor", &hit9);
    tChain->SetBranchAddress("hit10Cor", &hit10);
    tChain->SetBranchAddress("hit11Cor", &hit11);
    tChain->SetBranchAddress("hit12Cor", &hit12);
    tChain->SetBranchAddress("hit13Cor", &hit13);
    tChain->SetBranchAddress("hit14Cor", &hit14);
    tChain->SetBranchAddress("hit15Cor", &hit15);
    tChain->SetBranchAddress("xM", &grain1X);
    tChain->SetBranchAddress("yM", &grain1Y);
    tChain->SetBranchAddress("eventType", &eventType);
    tChain->SetBranchAddress("diffTrack2", &diffTrack2);
    tChain->SetBranchAddress("muonDZ", &muonDZ);

    int nEntries = tChain->GetEntries();
    std::cout << "Number of entries: " << nEntries << std::endl;

    RooRealVar hTot("hTot", "hTot", 0, 2000);
    RooDataSet *data[16];
    for (int i=0;i<16;i++)
    {
        data[i] = new RooDataSet("hTot", "hTot", hTot);
    }
    TH2D *h2[16];
    for (int i=0;i<16;i++)
    {
        std::string tname = "h2_" + std::to_string(i);
        h2[i] = new TH2D(tname.c_str(), tname.c_str(), 100, -1.3, 1.3, 100, -1.3, 1.3);
    }
    int iFibre = 0;
    long nEvent[16] = {0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0};
    for (long i=0;i<nEntries;i++)
    {
        tChain->GetEntry(i);
        double hit[16];
        hit[0] = hit0; hit[1] = hit1; hit[2] = hit2; hit[3] = hit3; hit[4] = hit4; hit[5] = hit5; hit[6] = hit6; hit[7] = hit7;
        hit[8] = hit8; hit[9] = hit9; hit[10] = hit10; hit[11] = hit11; hit[12] = hit12; hit[13] = hit13; hit[14] = hit14; hit[15] = hit15;

       double hTotVal = 0;
        bool isVmin = false;
        for (int iFibre=0;iFibre<16;iFibre++) {
            if (hit[iFibre]/coeff[iFibre] > 50) {
                isVmin = true;
            }
            hTotVal += hit[iFibre] / coeff[iFibre];
        }

        for (int iFibre=0;iFibre<16;iFibre++) {
            int nFibre = listFibre[iFibre];
            double dist2 = sqrt((grain1X - posXNref[nFibre]) * (grain1X - posXNref[nFibre]) +
                                (grain1Y - posYNref[nFibre]) * (grain1Y - posYNref[nFibre]));
            double dx = grain1X - posXNref[nFibre];
            double dy = grain1Y - posYNref[nFibre];

            bool cut1 = eventType == 4 && hTotVal > 0 && diffTrack2 < 1. && muonDZ > 0;
            bool cut2 = dist2 > rmin && dist2 < rmax && isVmin && hTotVal > 0 && hTotVal < mVmax && grain1X > -1.3 && grain1X < 1.3 && grain1Y > -1.3 && grain1Y < 1.3;
            int nQt = nQ[iFibre];
            bool cut3 = false;
            if (i==0){
                std::cout << "Fibre " << nFibre << " has " << nQt << " quadrants, iFibre =" << iFibre << std::endl;
            }

            for (int iQ = 0; iQ < nQt; iQ++) {
                int mQ = mlistQ[iQ][iFibre];
                if (mlistQ[iQ][iFibre] == 0)
                    cut3 = cut3 || (dx > 0 && dy > 0);
                if (mlistQ[iQ][iFibre] == 1)
                    cut3 = cut3 || (dx < 0 && dy > 0);
                if (mlistQ[iQ][iFibre] == 2)
                    cut3 = cut3 || (dx < 0 && dy < 0);
                if (mlistQ[iQ][iFibre] == 3)
                    cut3 = cut3 || (dx > 0 && dy < 0);
            }


            if (cut1 && cut2 && cut3) {
                hTot.setVal(hTotVal);
                data[nFibre]->add(hTot);
                h2[nFibre]->Fill(grain1X, grain1Y);
                nEvent[nFibre]++;
            }
        }
    }
    std::cout << "Number of events: " << nEvent[iFibre] << std::endl;
    for (int iFibre=0;iFibre<16;iFibre++) {
        std::string tname = "h2" + std::to_string(iFibre);
        TCanvas *myCanvas = new TCanvas(tname.c_str(), tname.c_str(), 600, 600);
        h2[iFibre]->Draw("colz");
        std::string myFileName = "h2_" + std::to_string(iFibre) + ".png";
        myCanvas->SaveAs(myFileName.c_str());
        delete myCanvas;
    }

    // ---------------------------------------
    // S e t u p   c o m p o n e n t   p d f s
    // ---------------------------------------

    // Construct observable

    // Construct landau(t,ml,sl) ;
    RooRealVar ml("ml", "mean landau", 500., 300, 700);
    RooRealVar sl("sl", "sigma landau", 25, 0.1, 100);
    RooLandau landau("lx", "lx", hTot, ml, sl);

    // Construct gauss(t,mg,sg)
    RooRealVar mg("mg", "mg", 0);
    RooRealVar sg("sg", "sg", 70, 0.1, 100);
    RooGaussian gauss("gauss", "gauss", hTot, mg, sg);

    // C o n s t r u c t   c o n v o l u t i o n   p d f
    // ---------------------------------------

    // Set #bins to be used for FFT sampling to 10000
    hTot.setBins(10000, "cache");

    // Construct landau (x) gauss
    RooFFTConvPdf lxg("lxg", "landau (X) gauss", hTot, landau, gauss);

    // S a m p l e ,   f i t   a n d   p l o t   c o n v o l u t e d   p d f
    // ----------------------------------------------------------------------

    // Sample 1000 events in x from gxlx
    //std::unique_ptr<RooDataSet> data{lxg.generate(t, 10000)};

    // Fit gxlx to data
    for(int iFibre=0;iFibre<16;iFibre++) {
        lxg.fitTo(*data[iFibre], PrintLevel(-1));

        // Plot data, landau pdf, landau (X) gauss pdf
        std::string myTitle = "Fibre " + std::to_string(iFibre);
        RooPlot *frame = hTot.frame(Title(myTitle.c_str()));
        data[iFibre]->plotOn(frame);
        lxg.plotOn(frame);
        lxg.paramOn(frame);
        //landau.plotOn(frame, LineStyle(kDashed));

        // Draw frame on canvas
        std::string cname = "landau" + std::to_string(iFibre);
        TCanvas *myCanvas = new TCanvas(cname.c_str(), cname.c_str(), 600, 600);
        gPad->SetLeftMargin(0.15);
        frame->GetYaxis()->SetTitleOffset(1.4);
        std::string myFileName = "landau" + std::to_string(iFibre) + ".png";
        frame->Draw();
        //frame->SaveAs(myFileName.c_str());
        myCanvas->SaveAs(myFileName.c_str());
        delete myCanvas;
        delete frame;

    }


}