#include <iostream>
#include <vector>
#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 "TrackingNtuple.h"
#include "analysis/obj_types.cpp"

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

using namespace std;

#define HIT_TYPE_PIXEL 0
#define HIT_TYPE_GLUED 1
#define HIT_TYPE_STRIP 2

Value<vector<Seed>>* seeds;
Value<vector<PixRecHit>>* pixrec_hits;
Value<vector<Track>>* tracks;
Value<vector<SimHit>>* sim_hits;
Value<vector<SimTrack>>* sim_tracks;

void register_objects(TrackingDataSet& tds){
    seeds = register_seeds(tds);

    pixrec_hits = register_pixrec_hits(tds);
    tracks = register_tracks(tds);

    sim_hits = register_sim_hits(tds);
    sim_tracks = register_sim_tracks(tds);
}

void setup_dphi_vs_eta(TrackingDataSet& tds){
    auto& calc_dphi_vs_eta =
    func<std::pair<vector<float>,vector<float>>(vector<Seed>,
                                                vector<PixRecHit>,
                                                vector<SimHit>, int)>("calc_dphi_vs_eta",
            FUNC(([](const vector<Seed>& seeds,
                     const vector<PixRecHit> pixrec_hits,
                     const vector<SimHit> sim_hits,
                     int layer){
                vector<float> dphi;
                vector<float> eta;
                for(const Seed &seed : seeds){
                    int hitIdx, hitType;
                    for(auto tup : boost::combine(seed.hitIdx, seed.hitType)){
                        boost::tie(hitIdx, hitType) = tup;
                        if(hitType == HIT_TYPE_PIXEL){
                            const PixRecHit &hit = pixrec_hits[hitIdx];
                            if(layer && hit.lay != layer) continue;
                            float hit_phi = atan2(hit.y, hit.x);
                            for(const int& simHitIdx : hit.simHitIdx){
                                const SimHit &sim_hit = sim_hits[simHitIdx];
                                dphi.push_back(hit_phi-atan2(sim_hit.y, sim_hit.x));
                                eta.push_back(seed.eta);
                            }
                        }
                    }
                }
                return std::make_pair(dphi,eta);
            })));

    auto dphi_vs_eta_B1 =
    fv::apply(calc_dphi_vs_eta, fv::tuple(seeds, pixrec_hits, sim_hits, constant<int>("1", 1)), "dphi_vs_eta_B1");
    auto dphi_vs_eta_B2 =
    fv::apply(calc_dphi_vs_eta, fv::tuple(seeds, pixrec_hits, sim_hits, constant<int>("2", 2)), "dphi_vs_eta_B2");
    auto dphi_vs_eta_B3 =
    fv::apply(calc_dphi_vs_eta, fv::tuple(seeds, pixrec_hits, sim_hits, constant<int>("3", 3)), "dphi_vs_eta_B3");
    auto dphi_vs_eta_B4 =
    fv::apply(calc_dphi_vs_eta, fv::tuple(seeds, pixrec_hits, sim_hits, constant<int>("4", 4)), "dphi_vs_eta_B4");


    tds.register_container<ContainerTH2Many<float>>("dphi_vs_eta_B1", "dphi vs eta (BPIX L1)",
                                                      dphi_vs_eta_B1, 100, -0.002, 0.002, 20, -3, 3,
                                                      "dphi",  "eta");
    tds.register_container<ContainerTH2Many<float>>("dphi_vs_eta_B2", "dphi vs eta (BPIX L2)",
                                                      dphi_vs_eta_B2, 100, -0.002, 0.002, 20, -3, 3,
                                                      "dphi",  "eta");
    tds.register_container<ContainerTH2Many<float>>("dphi_vs_eta_B3", "dphi vs eta (BPIX L3)",
                                                      dphi_vs_eta_B3, 100, -0.002, 0.002, 20, -3, 3,
                                                      "dphi",  "eta");
    tds.register_container<ContainerTH2Many<float>>("dphi_vs_eta_B4", "dphi vs eta (BPIX L4)",
                                                      dphi_vs_eta_B4, 100, -0.002, 0.002, 20, -3, 3,
                                                      "dphi",  "eta");

}

void setup_dz_vs_eta(TrackingDataSet& tds){
    auto& calc_dz_vs_eta =
    func<std::pair<vector<float>,vector<float>>(vector<Seed>,
                                                vector<PixRecHit>,
                                                vector<SimHit>, int)>("calc_dz_vs_eta",
            FUNC(([](const vector<Seed>& seeds,
                     const vector<PixRecHit> pixrec_hits,
                     const vector<SimHit> sim_hits,
                     int layer){
                vector<float> dz;
                vector<float> eta;
                for(const Seed &seed : seeds){
                    int hitIdx, hitType;
                    for(auto tup : boost::combine(seed.hitIdx, seed.hitType)){
                        boost::tie(hitIdx, hitType) = tup;
                        if(hitType == HIT_TYPE_PIXEL){
                            const PixRecHit &hit = pixrec_hits[hitIdx];
                            if(layer && hit.lay != layer) continue;
                            for(const int& simHitIdx : hit.simHitIdx){
                                dz.push_back(hit.z-sim_hits[simHitIdx].z);
                                eta.push_back(seed.eta);
                            }
                        }
                    }
                }
                return std::make_pair(dz,eta);
            })));

    auto dz_vs_eta_B1 =
    fv::apply(calc_dz_vs_eta, fv::tuple(seeds, pixrec_hits, sim_hits, constant<int>("1", 1)), "dz_vs_eta_B1");
    auto dz_vs_eta_B2 =
    fv::apply(calc_dz_vs_eta, fv::tuple(seeds, pixrec_hits, sim_hits, constant<int>("2", 2)), "dz_vs_eta_B2");
    auto dz_vs_eta_B3 =
    fv::apply(calc_dz_vs_eta, fv::tuple(seeds, pixrec_hits, sim_hits, constant<int>("3", 3)), "dz_vs_eta_B3");
    auto dz_vs_eta_B4 =
    fv::apply(calc_dz_vs_eta, fv::tuple(seeds, pixrec_hits, sim_hits, constant<int>("4", 4)), "dz_vs_eta_B4");


    tds.register_container<ContainerTH2Many<float>>("dz_vs_eta_B1", "dz vs eta (BPIX L1)",
                                                      dz_vs_eta_B1, 20, -0.01, .01, 20, -3, 3,
                                                      "dz",  "eta");
    tds.register_container<ContainerTH2Many<float>>("dz_vs_eta_B2", "dz vs eta (BPIX L2)",
                                                      dz_vs_eta_B2, 20, -0.01, .01, 20, -3, 3,
                                                      "dz",  "eta");
    tds.register_container<ContainerTH2Many<float>>("dz_vs_eta_B3", "dz vs eta (BPIX L3)",
                                                      dz_vs_eta_B3, 20, -0.01, .01, 20, -3, 3,
                                                      "dz",  "eta");
    tds.register_container<ContainerTH2Many<float>>("dz_vs_eta_B4", "dz vs eta (BPIX L4)",
                                                      dz_vs_eta_B4, 20, -0.01, .01, 20, -3, 3,
                                                      "dz",  "eta");

}

void run_analysis(const std::string& input_filename, const std::string& data_label, bool silent){
    gSystem->Load("libfilval.so");
    auto replace_suffix = [](const std::string& input, const std::string& new_suffix){
        return input.substr(0, input.find_last_of(".")) + new_suffix;
    };
    string log_filename = replace_suffix(input_filename, "_result.log");
    fv::util::Log::init_logger(log_filename, fv::util::LogPriority::kLogDebug);

    string output_filename = replace_suffix(input_filename, "_result.root");
    TrackingDataSet tds(output_filename, input_filename, data_label, "trackingNtuple/tree");
    register_objects(tds);

    setup_dz_vs_eta(tds);
    setup_dphi_vs_eta(tds);


    tds.process(silent);
    tds.save_all();
}

int main(int argc, char * argv[])
{
    fv::util::ArgParser args(argc, argv);
    if(!args.cmdOptionExists("-l") || !args.cmdOptionExists("-f")) {
        cout << "Usage: ./main (-s) -l DATA_LABEL -f treefile.root" << endl;
        return -1;
    }
    bool silent = args.cmdOptionExists("-s");
    string input_filename = args.getCmdOption("-f");
    string data_label = args.getCmdOption("-l");
    run_analysis(input_filename, data_label, silent);
}