#include "IAnalysisUtilities.hxx"
#include "IChannelId.hxx"
#include "IGeometryId.hxx"
#include "IOADatabase.hxx"
#include "IGeometryDatabase.hxx"
#include "IGeomIdManager.hxx"
#include <IG4VHit.hxx>
#include "IEventFolder.hxx"
#include "ICOMETEvent.hxx"

#include <TGeoNode.h>
#include <TGeoManager.h>
#include <TMath.h>
#include <TDirectory.h>

int COMET::IAnalysisUtilities::GetPID(const COMET::IG4VHit& hit){
    COMET::IHandle<COMET::IG4Trajectory> traj=hit.GetTrajectory();
    if(!traj){
        COMETError("Cannot get PID from hit, No trajectory...");
        return 0;
    }

    return traj->GetPDGEncoding();
}

TDirectory* COMET::IAnalysisUtilities::GetOrMakeDirectory(const std::string& name, TDirectory* parentDir){
    if(!parentDir){
        parentDir=TDirectory::CurrentDirectory();
    }
    TDirectory* dir=parentDir->GetDirectory(name.c_str());
    if(!dir) dir=parentDir->mkdir(name.c_str());
    return dir;
}

TGeoNode* COMET::IAnalysisUtilities::GetNode(const TVector3& position){
    return COMET::IOADatabase::Get().GeomId().FindNode(position);
}

std::string COMET::IAnalysisUtilities::FullNodePath(const TVector3& position){
    std::string fullPath = COMET::IOADatabase::Get().GeomId().FullNodePath(position);
    return fullPath;
}

void COMET::IAnalysisUtilities::ConvertToPolarAngles(
        const TVector3& beamLineDirection,
        const TVector3& direction,
        double& theta, double& phi){
    static const double degrees=180/TMath::Pi();
    TVector3 unitDir=direction.Unit();
    theta=unitDir.Angle(beamLineDirection)*degrees;
    unitDir-=unitDir.Dot(beamLineDirection)*beamLineDirection;
    phi=unitDir.Angle(TVector3(0,1,0))*degrees;
    if(unitDir.X() <0 ) phi=360-phi;
}

COMET::IChannelId COMET::IAnalysisUtilities::GetChannelID(const double pos_x, const double pos_y, const double pos_z){

    COMET::IGeometryId geomId=GetGeometryID(pos_x,pos_y,pos_z);
    COMET::IChannelId channelId;
    if(!geomId.IsValid()){
        COMET::IGeometryDatabase::Get().GetChannelId(channelId, geomId);
    }
    return channelId;
}

COMET::IGeometryId COMET::IAnalysisUtilities::GetGeometryID(const double pos_x, const double pos_y, const double pos_z){
    COMET::IGeometryId geomId;
    bool foundGeom= COMET::IOADatabase::Get().GeomId().GetGeometryId( pos_x, pos_y, pos_z, geomId);
    if(!foundGeom){
        COMETError("Couldn't find geometry for hit at: ("<<pos_x<<", "<<pos_y<<", "<<pos_z<<")");
    }
    return geomId;
}

bool COMET::IAnalysisUtilities::CheckVolumePath(
        const char* volumePath,
        const VolumeNames* matchNames,
        const VolumeNames* antiMatch){
    // Check that at least the match or anti-match name exists
    if ((!matchNames) && (!antiMatch)) return false;
    // Get the volume name as a tstring
    TString volPath = TString(volumePath);
    // Check the match names
    if (matchNames) {
        // Iterate through ALL match names
        for (auto a_name = matchNames->begin();
                  a_name != matchNames->end();
                  a_name++){
            // Check the volume has ALL match names in its name
            if (!volPath.Contains(a_name->c_str())) return false;
        }
    }
    // Check the anti-match names
    if (antiMatch) {
        // Iterate through ALL anti-match names
        for (auto a_name = antiMatch->begin();
                  a_name != antiMatch->end();
                  a_name++){
            // Check the volume has NONE of the anti matches names in its name
            if (volPath.Contains(a_name->c_str())) return false;
        }
    }
    // If it all checks out, return true
    return true;
}

bool COMET::IAnalysisUtilities::CheckVolumePath(
        const TVector3& pos,
        const VolumeNames* matchNames,
        const VolumeNames* antiMatch){
    // Check if the names match up
    return CheckVolumePath(FullNodePath(pos).c_str(), matchNames, antiMatch);
}

bool COMET::IAnalysisUtilities::ParticleStoppedInStopTgt(const COMET::IG4Trajectory& traj){
    // Did the particle come to a stop?
    if(traj.GetTrajectoryPoints().back().GetMomentum().Mag()>0.00001) return false;
    // Is the last point in the stopping target?
    const TVector3 end_pos3 = traj.GetFinalPosition().Vect();
    return IsInStoppingTarget(end_pos3);
}

bool COMET::IAnalysisUtilities::IsInStoppingTarget(const char* volumePath){
    // Match names
    VolumeNames toMatch;
    toMatch.push_back("MuonStoppingTarget");
    toMatch.push_back("TargetDisk");
    return CheckVolumePath(volumePath, &toMatch);
}

bool COMET::IAnalysisUtilities::IsInStoppingTarget(const TVector3& pos3){
    // Get the volume from this position
    const char* volumePath = FullNodePath(pos3).c_str();
    return IsInStoppingTarget(volumePath);
}

bool COMET::IAnalysisUtilities::CheckParticleFlux(
        const COMET::IG4Trajectory& traj,
        const bool skipStopped,
        const VolumeNames* matchNames,
        const VolumeNames* antiMatch){
    // Loop over the traj points to find if the trajectory passed through th
    // needed volume
    const COMET::IG4Trajectory::Points& points = traj.GetTrajectoryPoints();
    for(COMET::IG4Trajectory::Points::const_iterator i_point=points.begin(); 
                                                     i_point!=points.end(); 
                                                   ++i_point){
        // Skip the last, stopped point
        if (skipStopped && 
            i_point == --points.end() && 
            i_point->GetMomentum().Mag() > 0.0001) continue;
        // Get the position of the trajectory point
        const TLorentzVector pos4 = i_point->GetPosition();
        // Convert to ROOT units
        TVector3 pos3 = pos4.Vect();
        if(CheckVolumePath(pos3, matchNames, antiMatch)) return true;
    }
    return false;
}

void COMET::IAnalysisUtilities::GetParticleFluxVolumes(
        const COMET::IG4Trajectory& traj,
        NameSet& allNames,
        const bool skipStopped,
        const VolumeNames* matchNames,
        const VolumeNames* antiMatch){
    // Loop over the traj points to find if the trajectory passed through th
    // needed volume
    const COMET::IG4Trajectory::Points& points = traj.GetTrajectoryPoints();
    for(COMET::IG4Trajectory::Points::const_iterator i_point=points.begin(); 
                                                     i_point!=points.end(); 
                                                   ++i_point){
        // Skip the last, stopped point
        if (skipStopped && 
            i_point == --points.end() && 
            i_point->GetMomentum().Mag() > 0.0001) continue;
        // Get the position of the trajectory point
        const TLorentzVector pos4 = i_point->GetPosition();
        // Convert to ROOT units
        TVector3 pos3 = pos4.Vect();
        std::string nodeName = FullNodePath(pos3);
        // Check if we want the volume
        bool addVolume = true;
        if (matchNames || antiMatch) 
            addVolume = CheckVolumePath(nodeName.c_str(), matchNames, antiMatch);
        if (addVolume) allNames.insert(nodeName);
    }
}



const COMET::IHandle<COMET::IG4Trajectory>
COMET::IAnalysisUtilities::GetTrajectory(int trackID){
    const COMET::ICOMETEvent* event = COMET::IEventFolder::GetCurrentEvent();
    COMET::IHandle<COMET::IG4TrajectoryContainer> trajectories =
        event->Get<COMET::IG4TrajectoryContainer>("truth/G4Trajectories");
    if(!trajectories) return COMET::IHandle<COMET::IG4Trajectory>(NULL);
    return trajectories->GetTrajectory(trackID);
}