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

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

#include "analysis/TrackingNtupleObjs.hpp"

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

void register_objects(TrackingDataSet& tds){
    seeds.init(tds);
    sim_tracks.init(tds);
    sim_vertices.init(tds);
    gsf_tracks.init(tds);
    scls.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;

template <typename T>
struct KinColl {
    T* pt;
    T* eta;
    T* phi;
};


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();
    auto 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);
};

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

bool is_good_seed(const Seed& seed, float hoe_cut) {
    return seed.isECALDriven() and seed.hoe() < hoe_cut;
}

float reco_energy_rel_err(const SimTrack& sim_track, const Seed& seed) {
    return (sim_track.pt() - seed.Et()) / sim_track.pt() ;
}

float reco_energy_rel_err(const SimTrack& sim_track, const Track& track) {
    return (sim_track.pt() - track.pt()) / sim_track.pt() ;
}

template<typename TrackOrSeed>
bool reco_energy_consistent(const SimTrack& sim_track, const TrackOrSeed& track, float consistency_cut=0.1) {
    return fabs(reco_energy_rel_err(sim_track, track)) < consistency_cut;
}

double deltaR(const SimTrack& sim_track, const Track& gsf_track) {
    return sqrt(pow(sim_track.eta() - gsf_track.eta(), 2.0) +
                pow(sim_track.phi() - gsf_track.phi(), 2.0));
}

void run(){
    using namespace std;
    using namespace fv;
    using namespace fv_root;
    using namespace fv_util;
    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);

    auto hoe_cut = the_config->get("hoe_cut").as<float>(999);

    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")
    };

    enum TMType {
        NoMatch = 0,
        SeedMatched = 1,
        TrackMatched = 2,
        SeedAndTrackMatched = 3,
    };

    std::map<std::tuple<int, int>, ContainerTH2<float>*> residuals_dRz;
    std::map<std::tuple<int, int>, ContainerTH2<float>*> residuals_dPhi;

    std::map<std::tuple<int, int, int>, ContainerTH2<float>*> BPIX_residuals_dRz;
    std::map<std::tuple<int, int, int>, ContainerTH2<float>*> BPIX_residuals_dPhi;
    std::map<std::tuple<int, int, int>, ContainerTH2<float>*> FPIX_residuals_dRz;
    std::map<std::tuple<int, int, int>, ContainerTH2<float>*> FPIX_residuals_dPhi;

    stringstream name;
    auto set_name = [&name](const std::string& var, const std::string& region,
                            const int& layer, const int& hit, const int& tm_type) {
        name.str("");
        name << var << "_" << region;
        if (layer)
            name << "_L" << layer;
        name << "_H" << hit << "_v_Et";
        switch (tm_type) {
            case NoMatch:
                name << "_NoMatch";
                break;
            case SeedMatched:
                name << "_SeedMatched";
                break;
            case TrackMatched:
                name << "_TrackMatched";
                break;
            case SeedAndTrackMatched:
                name << "_SeedAndTrackMatched";
                break;
            default:
                break;
        }
    };

    THParams hist_params;
    for (int hit=1; hit<=3; hit++) {
        for (int tm_type = NoMatch; tm_type <= SeedAndTrackMatched; tm_type++) {
            for (int layer=1; layer<=4; layer++) {
                hist_params = (hit == 1) ? THParams::lookup("dRz_v_Et") : THParams::lookup("dRz_v_Et_outer_hits");
                set_name("dRz", "BPIX", layer, hit, tm_type);
                BPIX_residuals_dRz[make_tuple(layer, hit, tm_type)] = tds.register_container<ContainerTH2<float>>(name.str(), hist_params);
                set_name("dRz", "FPIX", layer, hit, tm_type);
                FPIX_residuals_dRz[make_tuple(layer, hit, tm_type)] = tds.register_container<ContainerTH2<float>>(name.str(), hist_params);

                hist_params = (hit == 1) ? THParams::lookup("dPhi_v_Et") : THParams::lookup("dPhi_v_Et_outer_hits");
                set_name("dPhi", "BPIX", layer, hit, tm_type);
                BPIX_residuals_dPhi[make_tuple(layer, hit, tm_type)] = tds.register_container<ContainerTH2<float>>(name.str(), hist_params);
                set_name("dPhi", "FPIX", layer, hit, tm_type);
                FPIX_residuals_dPhi[make_tuple(layer, hit, tm_type)] = tds.register_container<ContainerTH2<float>>(name.str(), hist_params);
            }
            hist_params = (hit == 1) ? THParams::lookup("dRz_v_Et") : THParams::lookup("dRz_v_Et_outer_hits");
            set_name("dRz", "ALL", 0, hit, tm_type);
            residuals_dRz[std::make_tuple(hit, tm_type)] = tds.register_container<ContainerTH2<float>>(name.str(), hist_params);

            hist_params = (hit == 1) ? THParams::lookup("dPhi_v_Et") : THParams::lookup("dPhi_v_Et_outer_hits");
            set_name("dPhi", "ALL", 0, hit, tm_type);
            residuals_dPhi[make_tuple(hit, tm_type)] = tds.register_container<ContainerTH2<float>>(name.str(), hist_params);
        }
    }

    std::map<int,vector<float>> eta_regions = {{1, {0.0, 1.1, 1.8, 3.0}},
                                               {2, {0.0, 1.4, 2.3, 3.0}},
                                               {3, {0.0, 1.0, 2.0, 3.0}}};

    auto get_region = [&eta_regions](int hit, float eta) {
        auto regions = eta_regions[hit];
        if (regions.size() <= 1) return 1;
        float abseta = abs(eta);
        for (int r_idx=1; r_idx < regions.size(); r_idx++) {
            if (abseta > regions[r_idx-1] and abseta < regions[r_idx]) {
                return r_idx;
            }
        }
        return static_cast<int>(regions.size());
    };
    std::map<int, ContainerTH2<float>*> dPhi_residuals_v_eta;
    std::map<int, ContainerTH2<float>*> dRz_residuals_v_eta;

    std::map<std::pair<int, int>, ContainerTH2<float>*> dPhi_residuals_v_region;
    std::map<std::pair<int, int>, ContainerTH2<float>*> dRz_residuals_v_region;
    for (int hit=1; hit<=3; hit++) {
        name.str("");
        name << "dPhi_residuals_v_eta_H" << hit;
        hist_params = (hit==1) ? THParams::lookup("dPhi_v_eta") : THParams::lookup("dPhi_v_eta_outer_hits");
        dPhi_residuals_v_eta[hit] = tds.register_container<ContainerTH2<float>>(name.str(), hist_params);

        name.str("");
        name << "dRz_residuals_v_eta_H" << hit;
        hist_params = (hit==1) ? THParams::lookup("dRz_v_eta") : THParams::lookup("dRz_v_eta_outer_hits");
        dRz_residuals_v_eta[hit] = tds.register_container<ContainerTH2<float>>(name.str(), hist_params);
        for (int region=1; region<=eta_regions[hit].size()-1; region++){
            hist_params = (hit==1) ? THParams::lookup("dRz_v_Et") : THParams::lookup("dRz_v_Et_outer_hits");
            name.str("");
            name << "dRz_residuals_H"  << hit << "_R" << region;
            dRz_residuals_v_region[{hit, region}] = tds.register_container<ContainerTH2<float>>(name.str(), hist_params);

            hist_params = (hit==1) ? THParams::lookup("dPhi_v_Et") : THParams::lookup("dPhi_v_Et_outer_hits");
            name.str("");
            name << "dPhi_residuals_H"  << hit << "_R" << region;
            dPhi_residuals_v_region[{hit, region}] = tds.register_container<ContainerTH2<float>>(name.str(), hist_params);
        }
    }

    KinColl<EfficiencyContainer<float>> seed_eff = {
            tds.register_container<EfficiencyContainer<float>>("seed_eff_v_pt",  THParams::lookup("eff_v_pt")),
            tds.register_container<EfficiencyContainer<float>>("seed_eff_v_eta", THParams::lookup("eff_v_eta")),
            tds.register_container<EfficiencyContainer<float>>("seed_eff_v_phi", THParams::lookup("eff_v_phi"))};
    KinColl<EfficiencyContainer<float>> seed_pur = {
            tds.register_container<EfficiencyContainer<float>>("seed_pur_v_pt",  THParams::lookup("pur_v_pt")),
            tds.register_container<EfficiencyContainer<float>>("seed_pur_v_eta", THParams::lookup("pur_v_eta")),
            tds.register_container<EfficiencyContainer<float>>("seed_pur_v_phi", THParams::lookup("pur_v_phi"))};
    KinColl<EfficiencyContainer<float>> tracking_eff = {
            tds.register_container<EfficiencyContainer<float>>("tracking_eff_v_pt",  THParams::lookup("eff_v_pt")),
            tds.register_container<EfficiencyContainer<float>>("tracking_eff_v_eta", THParams::lookup("eff_v_eta")),
            tds.register_container<EfficiencyContainer<float>>("tracking_eff_v_phi", THParams::lookup("eff_v_phi"))};
    KinColl<EfficiencyContainer<float>> tracking_pur = {
            tds.register_container<EfficiencyContainer<float>>("tracking_pur_v_pt",  THParams::lookup("pur_v_pt")),
            tds.register_container<EfficiencyContainer<float>>("tracking_pur_v_eta", THParams::lookup("pur_v_eta")),
            tds.register_container<EfficiencyContainer<float>>("tracking_pur_v_phi", THParams::lookup("pur_v_phi"))};

    KinColl<EfficiencyContainer<float>> tracking_eff_dR = {
            tds.register_container<EfficiencyContainer<float>>("tracking_eff_v_pt_dR",  THParams::lookup("eff_v_pt")),
            tds.register_container<EfficiencyContainer<float>>("tracking_eff_v_eta_dR", THParams::lookup("eff_v_eta")),
            tds.register_container<EfficiencyContainer<float>>("tracking_eff_v_phi_dR", THParams::lookup("eff_v_phi"))};
    KinColl<EfficiencyContainer<float>> tracking_pur_dR = {
            tds.register_container<EfficiencyContainer<float>>("tracking_pur_v_pt_dR",  THParams::lookup("pur_v_pt")),
            tds.register_container<EfficiencyContainer<float>>("tracking_pur_v_eta_dR", THParams::lookup("pur_v_eta")),
            tds.register_container<EfficiencyContainer<float>>("tracking_pur_v_phi_dR", THParams::lookup("pur_v_phi"))};

    KinColl<EfficiencyContainer<float>> fake_rate = {
            tds.register_container<EfficiencyContainer<float>>("fake_rate_v_pt",  THParams::lookup("eff_v_pt")),
            tds.register_container<EfficiencyContainer<float>>("fake_rate_v_eta", THParams::lookup("eff_v_eta")),
            tds.register_container<EfficiencyContainer<float>>("fake_rate_v_phi", THParams::lookup("pur_v_phi"))};

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

    auto& gsf_tracks_nmatch_sim_tracks = *tds.register_container<ContainerTH1<size_t>>("gsf_tracks_nmatch_sim_tracks ", THParams::lookup("gsf_tracks_nmatch_sim_tracks"));


    auto& ecal_energy_resolution = *tds.register_container<ContainerTH1<float>>("ecal_energy_resolution ", THParams::lookup("ecal_energy_resolution"));

    auto& n_seeds = *tds.register_container<ContainerTH1<size_t>>("n_seeds", THParams::lookup("n_seeds"));

    while (tds.next()) {

        for (const auto& sim_track : sim_tracks) {
            if (!is_good_sim(sim_track)) continue;
            if (seeds.size() == 0) continue;
            for (const auto &seed_idx : sim_track.seedIdx()) {
                const Seed& seed = seeds[seed_idx];
                if (!is_good_seed(seed, hoe_cut)) continue;
                ecal_energy_resolution.fill(reco_energy_rel_err(sim_track, seed));
            }
        }

        // Seeding Efficiency
        for (const auto& sim_track : sim_tracks) {
            if (!is_good_sim(sim_track)) continue;
            if (seeds.size() == 0) continue;

            bool matched = false;
            for (const auto& seed_idx : sim_track.seedIdx()) {
                const Seed& seed = seeds[seed_idx];
                if (is_good_seed(seed, hoe_cut)) {
                    matched=true;
                    break;
                }
            }
            if (abs(sim_track.eta()) < 2.4)
                seed_eff.pt->fill(sim_track.pt(), matched);
            if (sim_track.pt() > 20.0)
                seed_eff.eta->fill(sim_track.eta(), matched);
            if (abs(sim_track.eta()) < 2.4 and sim_track.pt() > 20)
                seed_eff.phi->fill(sim_track.phi(), matched);
        }

        // Tracking Efficiency
        for (const auto& sim_track : sim_tracks) {
            if (!is_good_sim(sim_track)) continue;
            if (gsf_tracks.size() == 0) continue;
            bool matched = false;
            for (const auto& track_idx : sim_track.trkIdx()) {
                const Seed& seed = seeds[gsf_tracks[track_idx].seedIdx()];
                if (is_good_seed(seed, hoe_cut)) {
                    matched=true;
                    break;
                }
            }
            if (abs(sim_track.eta()) < 2.4)
                tracking_eff.pt->fill(sim_track.pt(), matched);
            if (sim_track.pt() > 20.0)
                tracking_eff.eta->fill(sim_track.eta(), matched);
            if (abs(sim_track.eta()) < 2.4 and sim_track.pt() > 20)
                tracking_eff.phi->fill(sim_track.phi(), matched);

            // dR-matching
            matched = false;
            for (const auto& gsf_track : gsf_tracks) {
                const Seed& seed = seeds[gsf_track.seedIdx()];
                double dR = deltaR(sim_track, gsf_track);
                if (is_good_seed(seed, hoe_cut) and dR < 0.2) {
                    matched = true;
                    break;
                }
            }
            if (abs(sim_track.eta()) < 2.4)
                tracking_eff_dR.pt->fill(sim_track.pt(), matched);
            if (sim_track.pt() > 20.0)
                tracking_eff_dR.eta->fill(sim_track.eta(), matched);
            if (abs(sim_track.eta()) < 2.4 and sim_track.pt() > 20)
                tracking_eff_dR.phi->fill(sim_track.phi(), matched);
        }

        // Seeding Purity
        for (const auto& seed : seeds) {
            if (!is_good_seed(seed, hoe_cut)) continue;
            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.pt->fill(seed.pt(), match);
            if (seed.pt() > 20)
                seed_pur.eta->fill(seed.eta(), match);
            if (abs(seed.eta()) < 2.4 and seed.pt() > 20)
                seed_pur.phi->fill(seed.phi(), match);
        }

        // Tracking Purity
        for (const auto& gsf_track : gsf_tracks) {
            gsf_tracks_nmatch_sim_tracks.fill(gsf_track.simTrkIdx().size());
            const Seed& seed = seeds[gsf_track.seedIdx()];
            if (!is_good_seed(seed, hoe_cut)) continue;
            bool matched = false;
            for (const auto& sim_track_idx : gsf_track.simTrkIdx()) {
                if (is_good_sim(sim_tracks[sim_track_idx])) {
                    matched = true;
                    break;
                }
            }
            if (abs(gsf_track.eta()) < 2.4)
                tracking_pur.pt->fill(gsf_track.pt(), matched);
            if (gsf_track.pt() > 20)
                tracking_pur.eta->fill(gsf_track.eta(), matched);
            if (abs(gsf_track.eta()) < 2.4 and gsf_track.pt() > 20)
                tracking_pur.phi->fill(gsf_track.phi(), matched);


            // dR-matching
            matched = false;
            for (const auto& sim_track : sim_tracks) {
                double dR = deltaR(sim_track, gsf_track);
                if (is_good_sim(sim_track) and dR < 0.2) {
                    matched = true;
                    break;
                }
            }
            if (abs(gsf_track.eta()) < 2.4)
                tracking_pur_dR.pt->fill(gsf_track.pt(), matched);
            if (gsf_track.pt() > 20.0)
                tracking_pur_dR.eta->fill(gsf_track.eta(), matched);
            if (abs(gsf_track.eta()) < 2.4 and gsf_track.pt() > 20)
                tracking_pur_dR.phi->fill(gsf_track.phi(), matched);
        }

        // Fake Rate (#SC w/ 1 or more gsf-tracks / #SC) all w/ SC HOE>0.15
        std::map<size_t, size_t> scl2track;
        for (const auto &trk : gsf_tracks) {
            int seed_idx = trk.seedIdx();
            scl2track[seeds[seed_idx].sclIdx()] = trk.idx;
        }
        for (const auto &scl : scls) {
            if (scl.hoe() > hoe_cut) continue;
            float scl_pt = hypot(scl.px(), scl.py());
            float scl_eta = pseudorapidity(scl.x(), scl.y(), scl.z());
            float scl_phi = atan2(scl.py(), scl.px());

            bool has_track = scl2track.count(scl.idx) == 1;
            fake_rate.pt->fill(scl_pt, has_track);
            fake_rate.eta->fill(scl_eta, has_track);
            fake_rate.phi->fill(scl_phi, has_track);
        }

        // Hit Residuals
        for (const auto& seed : seeds) {
            if (seed.trkIdx() < 0) continue;  // require that seed produced gsf-track
            if (!is_good_seed(seed, hoe_cut)) continue;
            bool is_seed_sim_matched = false;
            for (const auto& sim_track_idx : seed.simTrkIdx()) {
                const SimTrack& sim_track = sim_tracks[sim_track_idx];
                if (is_good_sim(sim_track) and reco_energy_consistent(sim_track, seed)) {
                    is_seed_sim_matched = true;
                    break;
                }
            }

            bool is_track_sim_matched = false;
            const auto the_track = gsf_tracks[seed.trkIdx()];
            for (const auto& sim_track_idx : the_track.simTrkIdx()) {
                const SimTrack& sim_track = sim_tracks[sim_track_idx];
                if (is_good_sim(sim_track) and reco_energy_consistent(sim_track, the_track)) {
                    is_track_sim_matched = true;
                    break;
                }
            }

            std::vector<TMType> tm_types;
            if (is_seed_sim_matched and is_track_sim_matched) {
                tm_types = {SeedAndTrackMatched, SeedMatched, TrackMatched};
            } else if (is_seed_sim_matched) {
                tm_types = {SeedMatched};
            } else if (is_track_sim_matched) {
                tm_types = {TrackMatched};
            } else {
                tm_types = {NoMatch};
            }

            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();

            for (const auto& tm_type : tm_types) {
                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]);
                    int eta_region = get_region(static_cast<int>(hit_idx + 1), seed.eta());


                    if (is_seed_sim_matched) {
                        dRz_residuals_v_eta[hit_idx + 1]->fill(dRz, abs(seed.eta()));
                        dPhi_residuals_v_eta[hit_idx + 1]->fill(dPhi, abs(seed.eta()));

                        dRz_residuals_v_region[make_pair(hit_idx + 1, eta_region)]->fill(dRz, seed.Et());
                        dPhi_residuals_v_region[make_pair(hit_idx + 1, eta_region)]->fill(dPhi, seed.Et());
                    }

                    residuals_dRz[make_tuple(hit_idx + 1, tm_type)]->fill(dRz, seed.Et());
                    residuals_dPhi[make_tuple(hit_idx + 1, tm_type)]->fill(dPhi, seed.Et());
                    if (layerOrDiskNr == -1) continue;  // subDet is not set w/ old seeding

                    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));

                    if (isBarrel) {
                        BPIX_residuals_dRz[make_tuple(layerOrDiskNr, hit_idx + 1, tm_type)]->fill(dRz, seed.Et());
                        BPIX_residuals_dPhi[make_tuple(layerOrDiskNr, hit_idx + 1, tm_type)]->fill(dPhi, seed.Et());
                    } else {
                        FPIX_residuals_dRz[make_tuple(layerOrDiskNr, hit_idx + 1, tm_type)]->fill(dRz, seed.Et());
                        FPIX_residuals_dPhi[make_tuple(layerOrDiskNr, hit_idx + 1, tm_type)]->fill(dPhi, seed.Et());
                    }
                }
            }
        }
        n_seeds.fill(seeds.size());
    }

    tds.save_all();
}

int main(int argc, char * argv[]){
    using namespace fv_util;
    ArgParser args(argc, argv);
    if(args.cmd_option_exists("-c")) {
        init_config(args.get_cmd_option("-c"));
        args.update_config();
        if (args.cmd_option_exists("-s")) {
            the_config->update_key("silent", true);
        }
        if (args.cmd_option_exists("-b")) {
            the_config->update_key("batch", true);
        }
        init_log(LogPriority::kLogInfo);
//        gSystem->Load("libfilval.so");
        run();
    } else {
        cout << "Usage: ./" << argv[0] << " (-s) (-b) -c config_file.yaml" << endl;
    }
    return 0;
}