#include <iostream>
#include <fstream>

#include <TStopwatch.h>
#include <TString.h>
#include <TChain.h>
#include <TH2F.h>
#include <TProfile.h>

#include "fmcTypes.h"
#include "fmcReader.h"
#include "fmcEvent.h"
#include "fmcEventHeader.h"
#include "fmcParticle.h"
#include "fmcParticleLoopInfo.h"
#include "fmcLoopProcess.h"
#include "fmcLoopOption.h"
#include "fmcQvector.h"
#include "fmcConstants.h"
#include "fmcFunctions.cxx"

#ifndef _MCINI_
const char *_TREENAME_ = "mctree";
#endif
#ifdef _MCINI_
const char *_TREENAME_ = "events";
#endif

int main(int argc, char **argv)
{
    TString inFilelistName, outFileName, refMultFileName, histMultName = "hRefMultSTAR";
    if (argc < 7)
    {
        std::cerr << "./readRes -i filelist -o output -refMult refMultFile" << std::endl;
        return 10;
    }
    for (int i = 1; i < argc; i++)
    {
        if (std::string(argv[i]) != "-i" &&
            std::string(argv[i]) != "-o" &&
            std::string(argv[i]) != "-refMult")
        {
            std::cerr << "\nreadRes: Unknown parameter " << i
                      << ": " << argv[i] << std::endl;
            return 11;
        }
        else
        {
            if (std::string(argv[i]) == "-i" && i != argc - 1)
            {
                inFilelistName = argv[++i];
            }
            if (std::string(argv[i]) == "-i" && i == argc - 1)
            {
                std::cerr << "\nreadRes: Input file name was not specified!" << std::endl;
                return 20;
            }
            if (std::string(argv[i]) == "-o" && i != argc - 1)
            {
                outFileName = argv[++i];
            }
            if (std::string(argv[i]) == "-o" && i == argc - 1)
            {
                std::cerr << "\nreadRes: Output file name was not specified!" << std::endl;
                return 21;
            }
            if (std::string(argv[i]) == "-refMult" && i != argc - 1)
            {
                refMultFileName = argv[++i];
            }
            if (std::string(argv[i]) == "-refMult" && i == argc - 1)
            {
                std::cerr << "\nreadRes: Output file name was not specified!" << std::endl;
                return 22;
            }
        }
    }

    TChain *chain = new TChain(_TREENAME_);
    std::fstream file(inFilelistName.Data());
    std::string str;
    int nfiles = 0;

    while (std::getline(file, str))
    {
        if (chain->Add(str.c_str()))
            nfiles++;
    }

    std::cout << "Processing readRes..." << std::endl;
    std::cout << "Found " << nfiles << " files in the filelist." << std::endl;
    fmcReader *reader = new fmcReader();
    reader->Init(chain);

    fmcLong Nevents = reader->GetEntries();
    std::cout << "readRes: found " << Nevents << " events." << std::endl;

    // SetOptions for particle loop
    fmcLoopOption *loopOption = new fmcLoopOption();
    loopOption->ProcessRefMult();
    loopOption->ProcessQv();

    // Init particle loop info
    fmcParticleLoopInfo *loopInfo = nullptr;

    // Init particle loop process
    fmcLoopProcess *process = new fmcLoopProcess();
    process->SetLoopOptions(loopOption);
    process->SetReader(reader);

    // Prepare centrality selection
    fmc::Lim lim;
    fmcInt centrality;
    TFile *fiMult = new TFile(refMultFileName.Data(), "read");
    TH1I *hMult = (TH1I *)fiMult->Get(histMultName.Data());
    if (!hMult)
        return 30;
    fmcDouble CentEff = (histMultName == "hRefMultSTAR") ? 1. : 0.93;
    lim = fmc::CentralityBorders(hMult, CentEff);
    delete hMult;
    fiMult->Close();
    delete fiMult;

    // Init output
    TH2F *hQxEastTPC[fmc::nTpcEtaGaps];
    TH2F *hQxWestTPC[fmc::nTpcEtaGaps];
    TH2F *hQyEastTPC[fmc::nTpcEtaGaps];
    TH2F *hQyWestTPC[fmc::nTpcEtaGaps];
    TProfile *presTPC[fmc::nTpcEtaGaps];
    TH2F *hQxEastRXN, *hQxEastFHCal;
    TH2F *hQxWestRXN, *hQxWestFHCal;
    TH2F *hQyEastRXN, *hQyEastFHCal;
    TH2F *hQyWestRXN, *hQyWestFHCal;
    TProfile *presRXN, *presFHCal;
    for (int i = 0; i < fmc::nTpcEtaGaps; i++)
    {
        hQxEastTPC[i] = new TH2F(Form("hQxEastTPC_gap%i", i),
                                 Form("hQxEastTPC_gap%i", i), 500, -50., 50., 10, 0., 100.);
        hQxWestTPC[i] = new TH2F(Form("hQxWestTPC_gap%i", i),
                                 Form("hQxWestTPC_gap%i", i), 500, -50., 50., 10, 0., 100.);
        hQyEastTPC[i] = new TH2F(Form("hQyEastTPC_gap%i", i),
                                 Form("hQyEastTPC_gap%i", i), 500, -50., 50., 10, 0., 100.);
        hQyWestTPC[i] = new TH2F(Form("hQyWestTPC_gap%i", i),
                                 Form("hQyWestTPC_gap%i", i), 500, -50., 50., 10, 0., 100.);

        presTPC[i] = new TProfile(Form("presTPC_gap%i", i),
                                  Form("presTPC_gap%i", i), 10, 0., 100.);
    }
    hQxEastRXN = new TH2F(Form("hQxEastRXN"),
                          Form("hQxEastRXN"), 500, -50., 50., 10, 0., 100.);
    hQxWestRXN = new TH2F(Form("hQxWestRXN"),
                          Form("hQxWestRXN"), 500, -50., 50., 10, 0., 100.);
    hQyEastRXN = new TH2F(Form("hQyEastRXN"),
                          Form("hQyEastRXN"), 500, -50., 50., 10, 0., 100.);
    hQyWestRXN = new TH2F(Form("hQyWestRXN"),
                          Form("hQyWestRXN"), 500, -50., 50., 10, 0., 100.);

    presRXN = new TProfile(Form("presRXN"),
                           Form("presRXN"), 10, 0., 100.);

    hQxEastFHCal = new TH2F(Form("hQxEastFHCal"),
                            Form("hQxEastFHCal"), 500, -50., 50., 10, 0., 100.);
    hQxWestFHCal = new TH2F(Form("hQxWestFHCal"),
                            Form("hQxWestFHCal"), 500, -50., 50., 10, 0., 100.);
    hQyEastFHCal = new TH2F(Form("hQyEastFHCal"),
                            Form("hQyEastFHCal"), 500, -50., 50., 10, 0., 100.);
    hQyWestFHCal = new TH2F(Form("hQyWestFHCal"),
                            Form("hQyWestFHCal"), 500, -50., 50., 10, 0., 100.);

    presFHCal = new TProfile(Form("presFHCal"),
                             Form("presFHCal"), 10, 0., 100.);

    TStopwatch timer;
    timer.Start();

    fmcInt k;

    for (fmcLong iEv = 0; iEv < Nevents; iEv++)
    {
        if (iEv % 1000 == 0)
            std::cout << "readRes: Event [" << iEv << "/"
                      << Nevents << "]" << std::endl;

        loopInfo = process->ProcessParticleLoop(iEv);
        if (!loopInfo)
            continue;

        // Define centrality bin
        if (loopInfo->GetMultSTAR() == 0)
        {
            delete loopInfo;
            continue;
        }
        centrality = fmc::GetCentrality(lim, loopInfo->GetMultSTAR());
        if ( (centrality/10.) > lim.L.size() ||
           (centrality/10.) > lim.H.size()) continue;

        fmcQvector *qvEastTPC[fmc::nTpcEtaGaps], *qvWestTPC[fmc::nTpcEtaGaps];
        fmcQvector *qvEastRXN, *qvWestRXN;
        fmcQvector *qvEastFHCal, *qvWestFHCal;
        // Define Qvectors
        k = 0;
        for (int i = 0; i < fmc::nTpcEtaGaps; i++)
        {
            qvEastTPC[i] = loopInfo->GetQv(k);
            k++;
            qvWestTPC[i] = loopInfo->GetQv(k);
            k++;
        }
        qvEastRXN = loopInfo->GetQv(k);
        k++;
        qvWestRXN = loopInfo->GetQv(k);
        k++;
        qvEastFHCal = loopInfo->GetQv(k);
        k++;
        qvWestFHCal = loopInfo->GetQv(k);

        for (int i = 0; i < fmc::nTpcEtaGaps; i++)
        {
            hQxEastTPC[i]->Fill(qvEastTPC[i]->GetX(), centrality);
            hQyEastTPC[i]->Fill(qvEastTPC[i]->GetY(), centrality);
            hQxWestTPC[i]->Fill(qvWestTPC[i]->GetX(), centrality);
            hQyWestTPC[i]->Fill(qvWestTPC[i]->GetY(), centrality);
            presTPC[i]->Fill(centrality, TMath::Cos(2 * (qvEastTPC[i]->GetPhi() - qvWestTPC[i]->GetPhi())));
        }
        hQxEastRXN->Fill(qvEastRXN->GetX(), centrality);
        hQyEastRXN->Fill(qvEastRXN->GetY(), centrality);
        hQxWestRXN->Fill(qvWestRXN->GetX(), centrality);
        hQyWestRXN->Fill(qvWestRXN->GetY(), centrality);
        presRXN->Fill(centrality, TMath::Cos(2 * (qvEastRXN->GetPhi() - qvWestRXN->GetPhi())));
        hQxEastFHCal->Fill(qvEastFHCal->GetX(), centrality);
        hQyEastFHCal->Fill(qvEastFHCal->GetY(), centrality);
        hQxWestFHCal->Fill(qvWestFHCal->GetX(), centrality);
        hQyWestFHCal->Fill(qvWestFHCal->GetY(), centrality);
        presFHCal->Fill(centrality, TMath::Cos(1 * (qvEastFHCal->GetPhi() - qvWestFHCal->GetPhi())));

        delete loopInfo;
    }

    // Write output
    TFile *fo = new TFile(outFileName.Data(), "recreate");
    fo->mkdir("Qv");
    fo->mkdir("Res2");
    fo->cd("Qv");
    for (int i = 0; i < fmc::nTpcEtaGaps; i++)
    {
        hQxEastTPC[i]->Write();
        hQyEastTPC[i]->Write();
        hQxWestTPC[i]->Write();
        hQyWestTPC[i]->Write();
    }
    hQxEastRXN->Write();
    hQyEastRXN->Write();
    hQxWestRXN->Write();
    hQyWestRXN->Write();
    hQxEastFHCal->Write();
    hQyEastFHCal->Write();
    hQxWestFHCal->Write();
    hQyWestFHCal->Write();

    fo->cd("Res2");
    for (int i = 0; i < fmc::nTpcEtaGaps; i++)
    {
        presTPC[i]->Write();
    }
    presRXN->Write();
    presFHCal->Write();
    fo->Close();

    timer.Stop();
    timer.Print();
    return 0;
}