#include <iostream>
#include <vector>
#include <cmath>
#include <TSystem.h>

#include "filval.hpp"
#include "root_filval.hpp"

#include <boost/range/combine.hpp>

#include "analysis/TrackingNtupleObjs.hpp"


SeedCollection seeds;
SimTrackCollection sim_tracks;
SimVertexCollection sim_vertices;
TrackCollection gsf_tracks;

void register_objects(TrackingDataSet& tds){
    seeds.init(tds);
    sim_tracks.init(tds);
    sim_vertices.init(tds);
    gsf_tracks.init(tds);
}

struct {
    Value<float>* x;
    Value<float>* y;
    Value<float>* z;
    Value<float>* sigma_x;
    Value<float>* sigma_y;
    Value<float>* sigma_z;
} bs;

bool in_lum_region(const SimVertex& vertex) {
    /* Here approximate the luminous region as a cylinder with length 5*\sigma_z and radius
     * sqrt(\simga_x^2 + \sigma_y^2) centered as beamspot position.
     */
    float dx = vertex.x() - bs.x->get();
    float dy = vertex.y() - bs.y->get();
    float dz = vertex.z() - bs.z->get();
    float radius = static_cast<float>(5 * sqrt(pow(bs.sigma_x->get(), 2) + pow(bs.sigma_y->get(), 2)));
    float half_len = 5*bs.sigma_z->get();
    return sqrt(dx*dx + dy*dy) < radius and abs(dz) < half_len;
};

bool is_good_sim(const SimTrack& sim_track) {
    const auto& vertex = sim_vertices[sim_track.parentVtxIdx()];
    return abs(sim_track.pdgId()) == 11 and in_lum_region(vertex);
};


void run(bool silent){
    using namespace std;
    using namespace fv;
    using namespace fv_root;
    auto file_list = the_config->get_source_files();
    string output_filename = the_config->get_output_filename();
    TrackingDataSet tds(output_filename, file_list, "trackingNtuple/tree");
    register_objects(tds);

    bs = {
        tds.track_branch<float>("bsp_x"),
        tds.track_branch<float>("bsp_y"),
        tds.track_branch<float>("bsp_z"),
        tds.track_branch<float>("bsp_sigmax"),
        tds.track_branch<float>("bsp_sigmay"),
        tds.track_branch<float>("bsp_sigmaz")
    };

    auto& dRz_BPIX_L1_H1  = *tds.register_container<ContainerTH1<float>>("dRz_BPIX_L1_H1",  TH1Params::lookup("dRz"));
    auto& dPhi_BPIX_L1_H1 = *tds.register_container<ContainerTH1<float>>("dPhi_BPIX_L1_H1", TH1Params::lookup("dPhi"));
    auto& dRz_BPIX_L1_H1_v_Et  = *tds.register_container<ContainerTH2<float>>("dRz_BPIX_L1_H1_v_Et",  TH2Params::lookup("dRz_v_Et"));
    auto& dPhi_BPIX_L1_H1_v_Et = *tds.register_container<ContainerTH2<float>>("dPhi_BPIX_L1_H1_v_Et", TH2Params::lookup("dPhi_v_Et"));

    auto& dRz_BPIX_L2_H2  = *tds.register_container<ContainerTH1<float>>("dRz_BPIX_L2_H2",  TH1Params::lookup("dRz_outer_hits"));
    auto& dPhi_BPIX_L2_H2 = *tds.register_container<ContainerTH1<float>>("dPhi_BPIX_L2_H2", TH1Params::lookup("dPhi_outer_hits"));
    auto& dRz_BPIX_L2_H2_v_Et  = *tds.register_container<ContainerTH2<float>>("dRz_BPIX_L2_H2_v_Et",  TH2Params::lookup("dRz_v_Et_outer_hits"));
    auto& dPhi_BPIX_L2_H2_v_Et = *tds.register_container<ContainerTH2<float>>("dPhi_BPIX_L2_H2_v_Et", TH2Params::lookup("dPhi_v_Et_outer_hits"));

    auto& dRz_BPIX_L3_H3  = *tds.register_container<ContainerTH1<float>>("dRz_BPIX_L3_H3",  TH1Params::lookup("dRz_outer_hits"));
    auto& dPhi_BPIX_L3_H3 = *tds.register_container<ContainerTH1<float>>("dPhi_BPIX_L3_H3", TH1Params::lookup("dPhi_outer_hits"));
    auto& dRz_BPIX_L3_H3_v_Et  = *tds.register_container<ContainerTH2<float>>("dRz_BPIX_L3_H3_v_Et",  TH2Params::lookup("dRz_v_Et_outer_hits"));
    auto& dPhi_BPIX_L3_H3_v_Et = *tds.register_container<ContainerTH2<float>>("dPhi_BPIX_L3_H3_v_Et", TH2Params::lookup("dPhi_v_Et_outer_hits"));

    auto& seed_eff_v_pt  = *tds.register_container<EfficiencyContainer<float>>("seed_eff_v_pt",  TH1Params::lookup("eff_v_pt"));
    auto& seed_eff_v_eta = *tds.register_container<EfficiencyContainer<float>>("seed_eff_v_eta", TH1Params::lookup("eff_v_eta"));
    auto& seed_eff_v_phi = *tds.register_container<EfficiencyContainer<float>>("seed_eff_v_phi", TH1Params::lookup("eff_v_phi"));

    auto& tracking_eff_v_pt  = *tds.register_container<EfficiencyContainer<float>>("tracking_eff_v_pt",  TH1Params::lookup("eff_v_pt"));
    auto& tracking_eff_v_eta = *tds.register_container<EfficiencyContainer<float>>("tracking_eff_v_eta", TH1Params::lookup("eff_v_eta"));
    auto& tracking_eff_v_phi = *tds.register_container<EfficiencyContainer<float>>("tracking_eff_v_phi", TH1Params::lookup("eff_v_phi"));

    auto& seed_pur_v_pt  = *tds.register_container<EfficiencyContainer<float>>("seed_pur_v_pt",  TH1Params::lookup("pur_v_pt"));
    auto& seed_pur_v_eta = *tds.register_container<EfficiencyContainer<float>>("seed_pur_v_eta", TH1Params::lookup("pur_v_eta"));
    auto& seed_pur_v_phi = *tds.register_container<EfficiencyContainer<float>>("seed_pur_v_phi", TH1Params::lookup("pur_v_phi"));

    auto& tracking_pur_v_pt  = *tds.register_container<EfficiencyContainer<float>>("tracking_pur_v_pt",  TH1Params::lookup("pur_v_pt"));
    auto& tracking_pur_v_eta = *tds.register_container<EfficiencyContainer<float>>("tracking_pur_v_eta", TH1Params::lookup("pur_v_eta"));
    auto& tracking_pur_v_phi = *tds.register_container<EfficiencyContainer<float>>("tracking_pur_v_phi", TH1Params::lookup("pur_v_phi"));


    auto& hit_vs_layer_barrel = *tds.register_container<ContainerTH2<int>>("hit_vs_layer_barrel", TH2Params::lookup("hit_vs_layer"));
    auto& hit_vs_layer_forward = *tds.register_container<ContainerTH2<int>>("hit_vs_layer_forward", TH2Params::lookup("hit_vs_layer"));

    while (tds.next(!silent)) {

        for (const auto& sim_track : sim_tracks) {
            if (!is_good_sim(sim_track)) continue;
            if (seeds.size() == 0) continue;
            bool matched = !sim_track.seedIdx().empty();
            if (abs(sim_track.eta()) < 2.4)
                seed_eff_v_pt.fill(sim_track.pt(), matched);
            if (abs(sim_track.pt()) > 20.0)
                seed_eff_v_eta.fill(sim_track.eta(), matched);
            if (abs(sim_track.eta()) < 2.4 and abs(sim_track.pt()) < 20)
                seed_eff_v_phi.fill(sim_track.phi(), matched);
        }

        for (const auto& sim_track : sim_tracks) {
            if (!is_good_sim(sim_track)) continue;
            if (gsf_tracks.size() == 0) continue;
            bool matched = !sim_track.trkIdx().empty();
            if (abs(sim_track.eta()) < 2.4)
                tracking_eff_v_pt.fill(sim_track.pt(), matched);
            if (sim_track.pt() > 20.0)
                tracking_eff_v_eta.fill(sim_track.eta(), matched);
            if (abs(sim_track.eta()) < 2.4 and sim_track.pt() > 20)
                tracking_eff_v_phi.fill(sim_track.phi(), matched);
        }

        for (const auto& seed : seeds) {
            bool match = false;
            for (const auto& sim_track_idx : seed.simTrkIdx()) {
                if(is_good_sim(sim_tracks[sim_track_idx])) {
                    match = true;
                    break;
                }
            }
            if (abs(seed.eta()) < 2.4)
                seed_pur_v_pt.fill(seed.pt(), match);
            if (seed.pt() > 20)
                seed_pur_v_eta.fill(seed.eta(), match);
            if (abs(seed.eta()) < 2.4 and seed.pt() > 20)
                seed_pur_v_phi.fill(seed.phi(), match);
        }

        for (const auto& gsf_track : gsf_tracks) {
            bool match = false;
            for (const auto& sim_track_idx : gsf_track.simTrkIdx()) {
                if(is_good_sim(sim_tracks[sim_track_idx])) {
                    match = true;
                    break;
                }
            }
            if (abs(gsf_track.eta()) < 2.4)
                tracking_pur_v_pt.fill(gsf_track.pt(), match);
            if (gsf_track.pt() > 20)
                tracking_pur_v_eta.fill(gsf_track.eta(), match);
            if (abs(gsf_track.eta()) < 2.4 and gsf_track.pt() > 20)
                tracking_pur_v_phi.fill(gsf_track.phi(), match);

        }

        for (const auto& seed : seeds) {
            if (seed.trkIdx() < 0) continue;
            const auto the_track = gsf_tracks[seed.trkIdx()];
            vector<int>   hits_valid;
            vector<float> hits_dRz;
            vector<float> hits_dPhi;
            if (the_track.q() == 1) {
                hits_valid = seed.isValidPos();
                hits_dRz   = seed.dRZPos();
                hits_dPhi  = seed.dPhiPos();
            } else {
                hits_valid = seed.isValidNeg();
                hits_dRz   = seed.dRZNeg();
                hits_dPhi  = seed.dPhiNeg();
            }
            const vector<int>& hits_isBarrel = seed.isBarrel();
            const vector<int>& hits_layerOrDiskNr = seed.layerOrDiskNr();

            size_t nHits = hits_valid.size();
            for(size_t hit_idx=0; hit_idx<nHits; hit_idx++) {
                if (!hits_valid[hit_idx]) continue;
                bool isBarrel = hits_isBarrel[hit_idx] == 1;
                int layerOrDiskNr = hits_layerOrDiskNr[hit_idx];
                float dRz = abs(hits_dRz[hit_idx]);
                float dPhi = abs(hits_dPhi[hit_idx]);
                if (isBarrel and layerOrDiskNr == 1 and hit_idx == 0 ) {
                    dRz_BPIX_L1_H1.fill(dRz);
                    dPhi_BPIX_L1_H1.fill(dPhi);
                    dRz_BPIX_L1_H1_v_Et.fill(dRz, seed.pt());
                    dPhi_BPIX_L1_H1_v_Et.fill(dPhi, seed.pt());
                }
                if (isBarrel and layerOrDiskNr == 2 and hit_idx == 1 ) {
                    dRz_BPIX_L2_H2.fill(dRz);
                    dPhi_BPIX_L2_H2.fill(dPhi);
                    dRz_BPIX_L2_H2_v_Et.fill(dRz, seed.pt());
                    dPhi_BPIX_L2_H2_v_Et.fill(dPhi, seed.pt());
                }
                if (isBarrel and layerOrDiskNr == 3 and hit_idx == 2 ) {
                    dRz_BPIX_L3_H3.fill(dRz);
                    dPhi_BPIX_L3_H3.fill(dPhi);
                    dRz_BPIX_L3_H3_v_Et.fill(dRz, seed.pt());
                    dPhi_BPIX_L3_H3_v_Et.fill(dPhi, seed.pt());
                }
                if (isBarrel)
                    hit_vs_layer_barrel.fill(layerOrDiskNr, static_cast<const int &>(hit_idx + 1));
                else
                    hit_vs_layer_forward.fill(layerOrDiskNr, static_cast<const int &>(hit_idx + 1));
            }
        }
    }

    tds.save_all();
}

int main(int argc, char * argv[]){
    using namespace fv_util;
    ArgParser args(argc, argv);
    bool silent = args.cmdOptionExists("-s");
    if(args.cmdOptionExists("-c")) {
        init_config(args.getCmdOption("-c"));
        init_log(LogPriority::kLogInfo);
//        gSystem->Load("libfilval.so");
        run(silent);
    } else {
        cout << "Usage: ./" << argv[0] << " (-s) -c config_file.yaml" << endl;
    }
    return 0;
}