#include <iostream>
#include <vector>
#include <cmath>
#include <numeric>
#include <TSystem.h>
#include <Math/VectorUtil.h>

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

#include "analysis/TrackingNtupleObjs.hpp"

const double pi = 3.1415926535897932384626;
const double pi2 = 2*pi;

std::map<int, std::string> algoName = {
    {0, "undefAlgorithm"                },
    {1, "ctf"                           },
    {2, "duplicateMerge"                },
    {3, "cosmics"                       },
    {4, "initialStep"                   },
    {5, "lowPtTripletStep"              },
    {6, "pixelPairStep"                 },
    {7, "detachedTripletStep"           },
    {8, "mixedTripletStep"              },
    {9, "pixelLessStep"                 },
    {10, "tobTecStep"                   },
    {11, "jetCoreRegionalStep"          },
    {12, "conversionStep"               },
    {13, "muonSeededStepInOut"          },
    {14, "muonSeededStepOutIn"          },
    {15, "outInEcalSeededConv"          },
    {16, "inOutEcalSeededConv"          },
    {17, "nuclInter"                    },
    {18, "standAloneMuon"               },
    {19, "globalMuon"                   },
    {20, "cosmicStandAloneMuon"         },
    {21, "cosmicGlobalMuon"             },
    {22, "highPtTripletStep"            },
    {23, "lowPtQuadStep"                },
    {24, "detachedQuadStep"             },
    {25, "reservedForUpgrades1"         },
    {26, "reservedForUpgrades2"         },
    {27, "bTagGhostTracks"              },
    {28, "beamhalo"                     },
    {29, "gsf"                          },
    {30, "hltPixel"                     },
    {31, "hltIter0"                     },
    {32, "hltIter1"                     },
    {33, "hltIter2"                     },
    {34, "hltIter3"                     },
    {35, "hltIter4"                     },
    {36, "hltIterX"                     },
    {37, "hiRegitMuInitialStep"         },
    {38, "hiRegitMuLowPtTripletStep"    },
    {39, "hiRegitMuPixelPairStep"       },
    {40, "hiRegitMuDetachedTripletStep" },
    {41, "hiRegitMuMixedTripletStep"    },
    {42, "hiRegitMuPixelLessStep"       },
    {43, "hiRegitMuTobTecStep"          },
    {44, "hiRegitMuMuonSeededStepInOut" },
    {45, "hiRegitMuMuonSeededStepOutIn" }
};



typedef ROOT::Math::LorentzVector<ROOT::Math::PxPyPzE4D<float>> Vec4;

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

std::vector<SimTrack> sim_els;

Value<float>* truePileup;
Value<int>* inTimePileup;
Value<unsigned long>* event;
/* Value<vector<int>>* outOfTimePileup; */

/* Value<vector<unsigned int>>* totalSeedsAlgo; */
Value<unsigned int>* totalSeeds;

void register_objects(TrackingDataSet& tds){
    seeds.init(&tds);
    gens.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  // electrons
           in_lum_region(vertex) and         // from luminous region
           abs(sim_track.eta())<3.0;         // inside ECAL acceptance
};

bool is_good_sim(const std::vector<SimTrack> prompt_sims, const SimTrack& sim_track) {
    return find(prompt_sims.begin(), prompt_sims.end(), sim_track) != prompt_sims.end();
};

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

Vec4 scl_p4(const SuperCluster& scl) {
    return Vec4(scl.px(), scl.py(), scl.pz(), scl.e());
}

double dphi(double phi1, double phi2) {
    auto spec_mod = [](double a) {
        return a - floor(a/pi2)*pi2;
    };
    return spec_mod(phi2 - phi1 + pi) - pi;
}


double deltaR(double eta1, double phi1, double eta2, double phi2) {
    return sqrt(pow(eta1 - eta2, 2.0) +
                pow(dphi(phi1, phi2), 2.0));
}

struct Track_ {
    unsigned long idx;
    float pt;
    float eta;
    float phi;
};


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


    inTimePileup = tds.track_branch<int>("inTimePileup");
    truePileup = tds.track_branch<float>("truePileup");
    event = tds.track_branch<unsigned long>("event");

    while (tds.next()) {
        cout << "event: " << event << endl;

        // Find Z's Children
        /* for (const auto& gen : gens) { */
        /*     if (not gen.isPrompt()) continue; */
        /*     if (abs(gen.pdgId()) == 11 or */
        /*         abs(gen.pdgId()) == 13 or */
        /*         abs(gen.pdgId()) == 15) { */
        /*         cout << "  " << gen.pdgId() << endl; */
        /*     } */
        /* } */

        vector<Track_> tracks_;

        for (const auto& gsf_track : gsf_tracks) {
            const auto& seed = seeds[gsf_track.seedIdx()];
            tracks_.push_back({gsf_track.idx, gsf_track.pt(), gsf_track.eta(), gsf_track.phi()});
            cout << gsf_track.idx << ") " <<
                gsf_track.nPixel() << " : " << gsf_track.nPixelLay() << " | " <<
                gsf_track.nStrip() << " : " << gsf_track.nStripLay() << " | " <<
                algoName[gsf_track.originalAlgo()] <<  " | " <<
                algoName[gsf_track.algo()] <<  " | " <<
                algoName[seed.algo()] <<  " | " <<
                gsf_track.algoMask() <<  " | " <<
                endl;
        }
        std::sort(tracks_.begin(), tracks_.end(), [](Track_ a, Track_ b) {
            return a.phi < b.phi;
        });
        std::sort(tracks_.begin(), tracks_.end(), [](Track_ a, Track_ b) {
            return a.eta < b.eta;
        });
        std::sort(tracks_.begin(), tracks_.end(), [](Track_ a, Track_ b) {
            return a.pt < b.pt;
        });
        /* for (const auto& track_ : tracks_) { */
        /*     cout << track_.idx << ") " << */ 
        /*         track_.pt << " : " << */
        /*         track_.eta << " : " << */
        /*         track_.phi << endl; */
        /* } */
    }
    /* 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;
}