/**
 * @file
 * @brief Implementation of ADTreeWriter module
 * @copyright Copyright (c) 2019 CERN and the Allpix Squared authors.
 * This software is distributed under the terms of the MIT License, copied verbatim in the file "LICENSE.md".
 * In applying this license, CERN does not waive the privileges and immunities granted to it by virtue of its status as an
 * Intergovernmental Organization or submit itself to any jurisdiction.
 */

#include "ADTreeWriterModule.hpp"

#include <string>
#include <utility>

#include "core/utils/log.h"


using namespace allpix;

ADTreeWriterModule::ADTreeWriterModule(Configuration& config,
                                       Messenger* messenger,
                                       const std::shared_ptr<Detector>& detector)
    : SequentialModule(config, detector), detector_(detector), messenger_(messenger) {

    // Enable parallelization of this module if multithreading is enabled
    allow_multithreading();

    messenger_->bindSingle<PixelHitMessage>(this, MsgFlags::REQUIRED);
    config_.setDefault("file_name", detector_->getName());
}

void ADTreeWriterModule::initialize() {
    output_path_ = createOutputFile(config_.get<std::string>("file_name"), "root", true);
    output_file_ = std::make_unique<TFile>(output_path_.c_str(), "RECREATE");
    output_file_->cd();

    // Initialize the events tree
    data_.data_tree_ = new TTree("rawtree", "Allpix Squared simulation");
    data_.data_tree_->Branch("run", &data_.run_, "Run/I");
    data_.data_tree_->Branch("track", &data_.track_, "Track/I");
    data_.data_tree_->Branch("nhits", &data_.nhits_, "Nhits/I");
   // data_.data_tree_->Branch("clToT", &data_.cluster_tot_, "clToT/I");
   // data_.data_tree_->Branch("clX", &data_.cluster_x_, "clX/F");
   // data_.data_tree_->Branch("clY", &data_.cluster_y_, "clY/F");
    data_.data_tree_->Branch("row", &data_.rows_, "Row/I");
    data_.data_tree_->Branch("col", &data_.columns_, "Col/I");
    data_.data_tree_->Branch("tot", &data_.value_, "ToT/I");
    data_.data_tree_->Branch("toa", &data_.toa_, "ToA/I");
    data_.data_tree_->Branch("ftoa_rise", &data_.ftoa_rise_, "FToA_rise/I");
    data_.data_tree_->Branch("ftoa_fall", &data_.ftoa_fall_, "FToA_fall/I");
    data_.data_tree_->Branch("uftoa_start", &data_.uftoa_start_, "UFToA_start/I");
    data_.data_tree_->Branch("uftoa_stop", &data_.uftoa_stop_, "UFToA_stop/I");
    data_.data_tree_->Branch("toa_hd", &data_.toa_hd_, "ToA HD/F");
    data_.data_tree_->Branch("tot_hd", &data_.tot_hd_, "ToT HD/F");
   // data_.data_tree_->Branch("SpidrTime", &data_.time_, "SpidrTime[Nhits]/I");
   // data_.data_tree_->Branch("GlobalTime", &data_.global_time_, "GlobalTime[Nhits]/l");
    data_.data_tree_->Branch("particleID", &data_.pdg_codes_, "ParticleID/I");

    data_.time_tree_ = new TTree("timetree", "Allpix Squared simulation");
    data_.time_tree_->Branch("TrigTimeGlobal", &data_.trig_time_, "TrigTimeGlobal/l");
}

void ADTreeWriterModule::run(Event* event) {
    auto detector_name = detector_->getName();
    auto root_lock = root_process_lock();

    // Reset data for this event:
    data_.nhits_ = 0;

    LOG(TRACE) << "Writing event data for detector " << detector_name;

    // We are using the track number as event ID here
    data_.track_ = event->number;
    data_.trig_time_ = 0;

    // Loop over all hits
    auto message = messenger_->fetchMessage<PixelHitMessage>(this, event);

    for(const auto& hit : message->getData()) {
        if(allpix::ADTreeWriterModule::sensor_data::MAXHITS <= data_.nhits_) {
            LOG(ERROR) << "More than " << allpix::ADTreeWriterModule::sensor_data::MAXHITS << " in detector "
                       << detector_name;
            continue;
        }

        // Fill the tree with received messages
        auto i = data_.nhits_;
        data_.columns_ = hit.getPixel().getIndex().x();      // NOLINT
        data_.rows_ = hit.getPixel().getIndex().y();         // NOLINT
        data_.value_ = static_cast<UInt_t>(hit.getSignal()); // NOLINT
        data_.time_ = static_cast<UInt_t>(hit.getLocalTime());    // NOLINT
	data_.global_time_ = static_cast<UInt_t>(hit.getGlobalTime());
	data_.toa_ = static_cast<UInt_t>(hit.getLocalTime());

        // Only select primary particles
        auto particles = hit.getPrimaryMCParticles();
        if(particles.size() > 1) {
            LOG(WARNING) << "More than one primary particle detected for pixel hit - using first entry";
        }
        auto mcparticle = (particles.empty() ? nullptr : particles.front());
        data_.pdg_codes_ = (mcparticle != nullptr ? mcparticle->getParticleID() : 0); // NOLINT

        data_.nhits_ += 1;

        LOG(TRACE) << detector_name << " i=" << i << " x=" << data_.columns_ << " y=" << data_.rows_
                   << " t=" << data_.time_ << " v=" << data_.value_ << " mcp=" << data_.pdg_codes_;
    

        data_.data_tree_->Fill();
        data_.time_tree_->Fill();
    }
    LOG(DEBUG) << "Wrote data for detector " << detector_name << " and event " << event->number;
}

void ADTreeWriterModule::finalize() {
    output_file_->Write();
    LOG(STATUS) << "Wrote all events to file:" << std::endl << output_path_;
}