#include "IValidationAnalyzer.hxx"
#include "IHistogram.hxx"
#include "IPIDMultiHist.hxx"
#include "I2DParticleTally.hxx"
#include "IParticleTagger.hxx"
#include "IVolumePlot.hxx"
#include "IAnalyzerFactory.hxx"

#include <IHandle.hxx>
#include <IG4Trajectory.hxx>
#include <ICOMETLog.hxx>
#include <IOARuntimeParamBlock.hxx>

bool COMET::IValidationAnalyzer::Analyze(const IHandle<IDatum> data){

    COMET::IHandle<COMET::IG4TrajectoryContainer> trajectories =
        data->Get<COMET::IG4TrajectoryContainer>(".");
    if (trajectories) {
       AnalyzeContainer(*trajectories);
    } else return false;
    return true;
}

bool COMET::IValidationAnalyzer::SetupHistograms(){

    std::string title;
	
    COMET::IOARuntimeParamBlock params("all");
    COMET::IOARuntimeParamBlock target("target");
    double x_tgt_low=target.GetDouble("x_low",-8e3), x_tgt_high=target.GetDouble("x_high",8e3);
    double y_tgt_low=target.GetDouble("y_low",-2e3), y_tgt_high=target.GetDouble("y_high",2e3);
    double z_tgt_low=target.GetDouble("z_low",-1e4), z_tgt_high=target.GetDouble("z_high",1e4);
    double time_low= params.GetDouble("time_low" ,0), time_high =params.GetDouble("time_high" ,5e3);

    title="Final Position in X";
    hPositionX=new COMET::IPIDMultiHist("IH1F",GetDatumName(),"hPositionX",title+" for {particle}");
    hPositionX->Setup("Position in X (mm) (>2 MeV)",400,x_tgt_low,x_tgt_high);
    AddHistogram(hPositionX);

    title="Unrestricted X Position";
    hUnresPosX=new COMET::IPIDMultiHist("IH1F",GetDatumName(),"hUnresPosX",title+" for {particle}");
    hUnresPosX->Setup("Unrestricted X position (mm)",399,x_tgt_low,x_tgt_high);
    AddHistogram(hUnresPosX);

    title="Final Position in Y";
    hPositionY=new COMET::IPIDMultiHist("IH1F",GetDatumName(),"hPositionY",title+" for {particle}");
    hPositionY->Setup("Position in Y (mm)",400,y_tgt_low,y_tgt_high);
    AddHistogram(hPositionY);

    title="Final Position in Z";
    hPositionZ=new COMET::IPIDMultiHist("IH1F",GetDatumName(),"hPositionZ",title+" for {particle}");
    hPositionZ->Setup("Position in Z (mm)",400,z_tgt_low,z_tgt_high);
    AddHistogram(hPositionZ);

    title="Initial Momentum in X";
    hMomentumX=new COMET::IPIDMultiHist("IH1F",GetDatumName(),"hMomentumX",title+" for {particle}");
    hMomentumX->Setup("Momentum in X (MeV/c)",400,-2000,2000);
    AddHistogram(hMomentumX);

    title="Initial Momentum in Y";
    hMomentumY=new COMET::IPIDMultiHist("IH1F",GetDatumName(),"hMomentumY",title+" for {particle}");
    hMomentumY->Setup("Momentum in Y (MeV/c)",400,-10000,10000);
    AddHistogram(hMomentumY);

    title="Initial Momentum in Z";
    hMomentumZ=new COMET::IPIDMultiHist("IH1F",GetDatumName(),"hMomentumZ",title+" for {particle}");
    hMomentumZ->Setup("Momentum in Z (MeV/c)",400,-10000,10000);
    AddHistogram(hMomentumZ);

    title="Initial Phi angle";
    hPhi=new COMET::IPIDMultiHist("IH1F",GetDatumName(),"hPhi",title+" for {particle}");
    hPhi->Setup("Phi (rad)",400,-7,7);
    AddHistogram(hPhi);

    title="Initial Theta angle";
    hTheta=new COMET::IPIDMultiHist("IH1F",GetDatumName(),"hTheta",title+" for {particle}");
    hTheta->Setup("Theta (rad)",400,-7,7);
    AddHistogram(hTheta);

    title="Initial Total Momentum";
    hMomentum=new COMET::IPIDMultiHist("IH1F",GetDatumName(),"hMomentum",title+" for {particle}");
    hMomentum->Setup("Momentum (MeV/c)",400,0,10000);
    AddHistogram(hMomentum);

    title="Final Position";
    hPosition3D=new COMET::IPIDMultiHist("IVolumePlot",GetDatumName(),"hPosition3D",title+" for {particle}");
    hPosition3D->Setup("Position in X (mm)",400,x_tgt_low,x_tgt_high,
                       "Position in Y (mm)",400,y_tgt_low,y_tgt_high,
                       "Position in Z (mm)",400,z_tgt_low,z_tgt_high);
    AddHistogram(hPosition3D);

    title="Final Position 2D";
    hPosition2D=new COMET::IPIDMultiHist("IH2F",GetDatumName(),"hPosition2D",title+" for {particle}");
    hPosition2D->Setup("Position in Z (mm)",400,6000,9000,
                       "Position in Y (mm)",400,-1500,1500);
    AddHistogram(hPosition2D);

    title="Final Time";
    hTime=new COMET::IPIDMultiHist("IH1F",GetDatumName(),"hTime",title+" for {particle}");
    hTime->Setup("Time (ns)",1e3,time_low,time_high);
    AddHistogram(hTime);

    title="Start Time";
    hStartTime=new COMET::IPIDMultiHist("IH1F",GetDatumName(),"hStartTime",title+" for {particle}");
    hStartTime->Setup("Time (ns)",1e3,-100,1000);
    AddHistogram(hStartTime);

    return true;
}

bool COMET::IValidationAnalyzer::AnalyzeContainer(const COMET::IG4TrajectoryContainer& trajectories){
    for(IG4TrajectoryContainer::const_iterator i_traj=trajectories.begin();
            i_traj!=trajectories.end();++i_traj){

   const COMET::IG4Trajectory& trajectory=i_traj->second;

         TLorentzVector end_pos4=trajectory.GetFinalPosition();
         TLorentzVector start_pos4=trajectory.GetInitialPosition();
   	 TLorentzVector mom4=trajectory.GetInitialMomentum();
         double momentum=mom4.P();
         int pid=trajectory.GetPDGEncoding();
   	 if ( mom4.X() > 2 ){
             hPositionX->FillWith(pid,end_pos4.X());
         }
             hUnresPosX->FillWith(pid,end_pos4.X());
             hPositionY->FillWith(pid,end_pos4.Y());
             hPositionZ->FillWith(pid,end_pos4.Z());
             hMomentumX->FillWith(pid,mom4.X());
             hMomentumY->FillWith(pid,mom4.Y());
             hMomentumZ->FillWith(pid,mom4.Z());
             hMomentum->FillWith(pid,momentum);
             hPhi->FillWith(pid,mom4.Phi());
             hTheta->FillWith(pid,mom4.Theta());
             hPosition3D->FillWith(pid,end_pos4.X(),end_pos4.Y(),end_pos4.Z());
             hPosition2D->FillWith(pid,end_pos4.Z(),end_pos4.Y());
             hTime->FillWith(pid,end_pos4.T());
             hStartTime->FillWith(pid,start_pos4.T());
		
   	 
     }
     return true;
}

REGISTER_ANALYZER(IValidationAnalyzer)