#include <iostream>
#include <vector>
#include <map>
#include <utility>
#include <numeric>
#include <limits>
#include <cmath>
#include <TSystem.h>

#include "filval/filval.hpp"
#include "filval/root/filval.hpp"

#include <boost/range/combine.hpp>
#include <boost/format.hpp>

#include "analysis/TrackingNtuple.h"
#include "analysis/TrackingNtupleObjs.hpp"
#include "analysis/common.hpp"

using namespace std;
using namespace fv;
using namespace fv::root;

enum class HitType {
  Pixel = 0,
  Strip = 1,
  Glued = 2,
  Invalid = 3,
  Phase2OT = 4,
  Unknown = 99
};

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_value();
    float dy = vertex.y() - bs.y->get_value();
    float dz = vertex.z() - bs.z->get_value();
    float radius = 5*sqrt(pow(bs.sigma_x->get_value(), 2) + pow(bs.sigma_y->get_value(), 2));
    float half_len = 5*bs.sigma_z->get_value();
    return sqrt(dx*dx + dy*dy) < radius and abs(dz) < half_len;
};

bool is_good_sim(size_t simTrkIdx) {
    const auto& track = sim_tracks[simTrkIdx];
    const auto& vertex = sim_vertices[track.parentVtxIdx()];
    return abs(track.pdgId()) == 11 and in_lum_region(vertex);
};

void run_analysis(const vector<fv::util::DataFileDescriptor>& dfds, const string output_filename, bool silent){
    TrackingDataSet tds(output_filename, dfds, "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")
    };

    // Indices of prompt-ish sim electrons
    auto good_sim_electrons = func_value<vector<size_t>>("good_sim_electrons",
        FUNC(([](){
            vector<size_t> idxs;
            for(const auto& sim_track : sim_tracks) {
                if (is_good_sim(sim_track.idx)) idxs.push_back(sim_track.idx);
            }
            return idxs;
        })));

    // Indices of prompt-ish sim electrons that have a matched seed
    function<bool(size_t)> with_seed = [](const size_t& e_idx) {
        return seeds.size() != 0 and sim_tracks[e_idx].seedIdx().size() != 0;
    };

    // Indices of prompt-ish sim electrons that have a matched gsftrack
    function<bool(size_t)> with_track = [](const size_t& e_idx) { return sim_tracks[e_idx].trkIdx().size() > 0; };

    function<float(size_t)> sim_pt = [](const size_t& sim_idx) { return sim_tracks[sim_idx].pt(); };

    function<float(size_t)> sim_eta = [](const size_t& sim_idx) { return sim_tracks[sim_idx].eta(); };

    function<float(size_t)> sim_phi = [](const size_t& sim_idx) { return sim_tracks[sim_idx].phi(); };

    float pt_cut = 20; // GeV
    float eta_cut = 2.4;
    function<bool(size_t)> sim_in_pt = [pt_cut](size_t sim_idx) {
        return sim_tracks[sim_idx].pt() > pt_cut;
    };

    function<bool(size_t)> sim_in_eta = [eta_cut](size_t sim_idx) {
        return abs(sim_tracks[sim_idx].eta()) < eta_cut;
    };

    function<bool(size_t)> sim_in_eta_and_pt = [eta_cut, pt_cut](size_t sim_idx) {
        return abs(sim_tracks[sim_idx].eta()) < eta_cut and
               sim_tracks[sim_idx].pt() > pt_cut;
    };

    tds.register_container<EfficiencyContainer<size_t>>("seed_eff_v_pt",
        good_sim_electrons, TH1Params::lookup("eff_v_pt"), &sim_in_eta, &with_seed, &sim_pt);

    tds.register_container<EfficiencyContainer<size_t>>("seed_eff_v_eta",
        good_sim_electrons, TH1Params::lookup("eff_v_eta"), &sim_in_pt, &with_seed, &sim_eta);

    tds.register_container<EfficiencyContainer<size_t>>("seed_eff_v_phi",
        good_sim_electrons, TH1Params::lookup("eff_v_phi"), &sim_in_eta_and_pt, &with_seed, &sim_phi);

    tds.register_container<EfficiencyContainer<size_t>>("track_eff_v_pt",
        good_sim_electrons, TH1Params::lookup("eff_v_pt"), &sim_in_eta, &with_track, &sim_pt);

    tds.register_container<EfficiencyContainer<size_t>>("track_eff_v_eta",
        good_sim_electrons, TH1Params::lookup("eff_v_eta"), &sim_in_pt, &with_track, &sim_eta);

    tds.register_container<EfficiencyContainer<size_t>>("track_eff_v_phi",
        good_sim_electrons, TH1Params::lookup("eff_v_phi"), &sim_in_eta_and_pt, &with_track, &sim_phi);

    // Indices of ecal-driven seeds
    auto ecal_seeds = func_value<vector<size_t>>("ecal_seeds",
        FUNC(([](){
            vector<size_t> idxs;
            for (const auto& seed : seeds) {
                idxs.push_back(seed.idx);
            }
            return idxs;
        })));

    function<bool(size_t)> seed_with_sim_electron = [](const size_t& seed_idx) {
        for (const int& simTrkIdx : seeds[seed_idx].simTrkIdx()) {
            if (is_good_sim(simTrkIdx)) return true;
        }
        return false;
    };

    function<float(size_t)> seed_pt = [](const size_t& seed_idx) { return seeds[seed_idx].pt(); };

    function<float(size_t)> seed_eta = [](const size_t& seed_idx) { return seeds[seed_idx].eta(); };

    function<float(size_t)> seed_phi = [](const size_t& seed_idx) { return seeds[seed_idx].phi(); };

    function<bool(size_t)> seed_in_pt = [pt_cut](size_t seed_idx) { return seeds[seed_idx].pt() > pt_cut; };

    function<bool(size_t)> seed_in_eta = [eta_cut](size_t seed_idx) { return abs(seeds[seed_idx].eta()) < eta_cut; };

    function<bool(size_t)> seed_in_eta_and_pt = [eta_cut, pt_cut](size_t seed_idx) {
        return abs(seeds[seed_idx].eta()) < eta_cut and
               seeds[seed_idx].pt() > pt_cut;
    };

    tds.register_container<EfficiencyContainer<size_t>>("seed_pur_v_pt",
        ecal_seeds, TH1Params::lookup("pur_v_pt"), &seed_in_eta, &seed_with_sim_electron, &seed_pt, true);

    tds.register_container<EfficiencyContainer<size_t>>("seed_pur_v_eta",
        ecal_seeds, TH1Params::lookup("pur_v_eta"), &seed_in_pt, &seed_with_sim_electron, &seed_eta, true);

    tds.register_container<EfficiencyContainer<size_t>>("seed_pur_v_phi",
        ecal_seeds, TH1Params::lookup("pur_v_phi"), &seed_in_eta_and_pt, &seed_with_sim_electron, &seed_phi, true);


    // Indices of gsf-tracks from ECAL-Driven Seeds
    auto ecal_tracks = func_value<vector<size_t>>("ecal_tracks",
        FUNC(([]() {
            vector<size_t> idxs;
            for (const auto& seed : seeds) {
                const int& trk_idx = seed.trkIdx();
                if (trk_idx >= 0) idxs.push_back(trk_idx);
            }
            return idxs;
        })));

    function<bool(size_t)> gsf_track_with_sim_electron = [](const size_t& track_idx) {
        for (const int& simTrkIdx : gsf_tracks[track_idx].simTrkIdx()) {
            if (is_good_sim(simTrkIdx)) return true;
        }
        return false;
    };

    function<float(size_t)> track_pt = [](const size_t& track_idx) { return gsf_tracks[track_idx].pt(); };
    function<float(size_t)> track_eta = [](const size_t& track_idx) { return gsf_tracks[track_idx].eta(); };
    function<float(size_t)> track_phi = [](const size_t& track_idx) { return gsf_tracks[track_idx].phi(); };

    function<bool(size_t)> track_in_pt = [pt_cut](size_t track_idx) { return gsf_tracks[track_idx].pt() > pt_cut; };

    function<bool(size_t)> track_in_eta = [eta_cut](size_t track_idx) { return abs(gsf_tracks[track_idx].eta()) < eta_cut; };

    function<bool(size_t)> track_in_eta_and_pt = [eta_cut, pt_cut](size_t track_idx) {
        return abs(gsf_tracks[track_idx].eta()) < eta_cut and gsf_tracks[track_idx].pt() > pt_cut;
    };

    tds.register_container<EfficiencyContainer<size_t>>("track_pur_v_pt",
        ecal_tracks, TH1Params::lookup("pur_v_pt"), &track_in_eta, &gsf_track_with_sim_electron, &track_pt, true);

    tds.register_container<EfficiencyContainer<size_t>>("track_pur_v_eta",
        ecal_tracks, TH1Params::lookup("pur_v_eta"), &track_in_pt, &gsf_track_with_sim_electron, &track_eta, true);

    tds.register_container<EfficiencyContainer<size_t>>("track_pur_v_phi",
        ecal_tracks, TH1Params::lookup("pur_v_phi"), &track_in_eta_and_pt, &gsf_track_with_sim_electron, &track_phi, true);

    auto residuals_rPhi = func<vector<float>(int, int, int)>("residuals_rPhi",
        FUNC(([](int detSel, int layerSel, int hitSel) {
        std::vector<float> residuals;
        for (const Seed& seed : seeds) {
            int nHits = seed.isBarrel().size();
            for (int idx_hit=0; idx_hit<nHits; idx_hit++) {
                if ((detSel==0 or detSel==seed.isBarrel()[idx_hit]) and
                    (layerSel==0 or layerSel==seed.layerOrDiskNr()[idx_hit]) and
                    (hitSel==0 or hitSel==(idx_hit+1))) {
                    if (seed.isValidPos()[idx_hit]) {
                        residuals.push_back(seed.dRZPos()[idx_hit]);
                    }
                    if (seed.isValidNeg()[idx_hit]) {
                        residuals.push_back(seed.dRZNeg()[idx_hit]);
                    }
                }
            }
        }
        return residuals;
    })));

    tds.register_container<ContainerTH1Many<float>>("residuals_all",
                                                    tup_apply(residuals_rPhi, fv::tuple(constant("BPIX", 1), constant("ALL", 0), constant("ALL", 0))),
                                                    TH1Params::lookup("residuals_all"));

    tds.process(silent);
    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"));
        auto output_filename = the_config->get_output_filename();
        if (output_filename == "") return -1;

        init_log(fv::util::LogPriority::kLogInfo);
        gSystem->Load("libfilval.so");

        auto file_list = the_config->get_source_files();
        run_analysis(file_list, output_filename, silent);
    } else {
        cout << "Usage: ./tracking_eff (-s) -c config_file.yaml" << endl;
    }
    return 0;
}