//	BLCMDtrace.cc
/*
This source file is part of G4beamline, http://g4beamline.muonsinc.com
Copyright (C) 2003,2004,2005,2006 by Tom Roberts, all rights reserved.

This program is free software; you can redistribute it and/or
modify it under the terms of the GNU General Public License
as published by the Free Software Foundation; either version 2
of the License, or (at your option) any later version.

This program is distributed in the hope that it will be useful,
but WITHOUT ANY WARRANTY; without even the implied warranty of
MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
GNU General Public License for more details.

http://www.gnu.org/copyleft/gpl.html
*/

#include <set>

#include "globals.hh"
#include "G4Track.hh"
#include "G4Step.hh"
#include "G4ParticleDefinition.hh"
#include "G4RunManager.hh"
#include "G4Event.hh"
#include "G4ThreeVector.hh"

#include "BLCoordinates.hh"
#include "BLEvaluator.hh"
#include "BLGlobalField.hh"
#include "BLCommand.hh"
#include "BLParam.hh"
#include "BLManager.hh"
#include "BLRunManager.hh"
#include "BLTrackNTuple.hh"
#include "BLCollectiveComputation.hh"
#include "mysnprintf.hh"

const char TraceFields[] =
    "x:y:z:Px:Py:Pz:t:PDGid:EventID:TrackID:ParentID:Weight:Bx:By:Bz:Ex:Ey:Ez";
const int NTraceFields = 18;

/**	class TraceNTuple implements an NTuple for tracing a track.
 **/
class TraceNTuple {
	BLTrackNTuple *ntuple;
	BLCoordinateType coordinateType;
	BLManager *manager;
	BLRunManager *runManager;
public:
	/// constructor.
	TraceNTuple(G4String format, G4String name, G4String filename,
					BLCoordinateType _coordinateType);

	/// destructor.
	virtual ~TraceNTuple();

	/// appendTrace() will append a track to the Trace NTuple.
	void appendTrace(G4Track *track);

	/// close() will close the NTuple.
	void close() { if(ntuple) ntuple->close(); }

	/// annotate() will add an annotation to ASCII NTuples.
	void annotate(G4String line) { if(ntuple) ntuple->annotate(line); }
};

/**	class BLCMDtrace implements the trace command.
 *
 **/
class BLCMDtrace : public BLCommand, public BLManager::SteppingAction,
	    public BLManager::TrackingAction, public BLCollectiveComputation {
	int nTrace;
	G4String format;
	G4int oneNTuple;
	G4int primaryOnly;
        G4int traceTune;
	G4String filename;
	G4String require;
	G4String coordinates;
	BLCoordinateType coordinateType;
	bool doTrace;
	TraceNTuple *trace;
	BLEvaluator *eval;
	bool init;
	BLManager *manager;
	BLRunManager *runManager;
	static TraceNTuple *allTrace;
	static TraceNTuple *tuneTrace;
	static TraceNTuple *referenceTrace;
	friend class TraceNTuple;
public:
	/// Constructor
	BLCMDtrace();

	BLCMDtrace(const BLCMDtrace &r);

	G4String commandName() { return "trace"; }
	
	int command(BLArgumentVector& argv, BLArgumentMap& namedArgs);

	void defineNamedArgs();

	/// help() prints help text.
	void help(bool detailed) {
		if(description[description.size()-2] == ':')
			description += BLTrackNTuple::getFormatList(); 
		BLCommand::help(detailed);
	}

	void newTrace(const G4Track *firstTrack);

	/// UserSteppingAction() from BLManager::SteppingAction.
	void UserSteppingAction(const G4Step *step);

	/// PreUserTrackingAction() from BLManager::TrackingAction.
	void PreUserTrackingAction(const G4Track *track);

	/// PostUserTrackingAction() from BLManager::TrackingAction.
	void PostUserTrackingAction(const G4Track *track);

	/// from BLCollectiveComputation
	virtual void beginCollectiveTracking(std::vector<BLTrackData>& v);
	virtual void collectiveStep(std::vector<BLTrackData>& v);
	virtual void endCollectiveTracking(std::vector<BLTrackData>& v);
};
TraceNTuple *BLCMDtrace::allTrace = 0;
TraceNTuple *BLCMDtrace::tuneTrace = 0;
TraceNTuple *BLCMDtrace::referenceTrace = 0;
BLCMDtrace defaultTrace;

BLCMDtrace::BLCMDtrace() : BLCommand(), BLManager::SteppingAction(), 
            BLManager::TrackingAction(), BLCollectiveComputation(),
	    format()
{
	registerCommand(BLCMDTYPE_DATA);
	setSynopsis("Specifies tracing of tracks.");
	setDescription("Generates a separate NTuple for each track, with 1\n"
		"row per step, unless oneNTuple is nonzero (in which case "
		"all tracks are put into a single NTuple).\n"
		"So format=ascii generates one file per track with names\n"
		"generated by the pattern in filename (first %d is replaced "
		"by event #, second %d is replaced by trackId); for "
		"oneNTuple, the default filename is AllTracks.txt.\n"
		"\nNote that without a trace command no traces are generated, "
		"so to trace just the tune and reference particles include "
		"a trace command with no arguments.\n\n"
		"In collective tracking mode, oneNTuple must be nonzero, and "
		"the entries will be generated only at collective steps "
		"(usually at a specified deltaT).\n\n"
		"Unlike other NTuple commands, the require expression applies "
		"to entire tracks, not individual entries.\n\n"
		"The standard NTuple fields are:\n"
		"    x,y,z (mm)\n"
		"    Px,Py,Pz (MeV/c)\n"
		"    t (ns)\n"
		"    PDGid (11=e-, 13=mu-, 22=gamma, 211=pi+, 2212=proton, ...)\n"
		"    EventID (may be inexact above 16,777,215)\n"
		"    TrackID\n"
		"    ParentID (0 => primary particle)\n"
		"    Weight (defaults to 1.0)\n"
		"The trace includes the following fields:\n"
		"    Bx, By, Bz (Tesla)\n"
		"    Ex, Ey, Ez (Megavolts/meter)\n\n"
		"The following additional fields are appended for "
		"format=Extended, format=asciiExtended, and "
		"format=rootExtended:\n"
		"    ProperTime (ns)\n"
		"    PathLength (mm)\n"
		"    PolX, PolY, PolZ (polarization)\n"
		"    InitialKE (MeV when track was created)\n\n"
		"Valid Formats (ignore case): ");
	nTrace = 0;
	format = "";
	oneNTuple = 0;
	primaryOnly = 0;
        traceTune = 1;     //trace  tune track by default (previous behavior)
	filename = "";
	require = "";
	coordinates = "Centerline";
	coordinateType = BLCOORD_CENTERLINE;
	doTrace = true;
	trace = 0;
	eval = 0;
	init = false;
	manager = 0;	// too early to set these -- see command()
	runManager = 0;
}

BLCMDtrace::BLCMDtrace(const BLCMDtrace& r) : BLCommand(r), 
	    BLManager::SteppingAction(), BLManager::TrackingAction(),
	    BLCollectiveComputation(), format()
{
	nTrace = r.nTrace;
	format = r.format;
	oneNTuple = r.oneNTuple;
	primaryOnly = r.primaryOnly;
	traceTune= r.traceTune;
	filename = r.filename;
	require = r.require;
	coordinates = r.coordinates;
	coordinateType = r.coordinateType;
	doTrace = r.doTrace;
	trace = r.trace;
	eval = 0;
	init = r.init;
	manager = r.manager;
	runManager = r.runManager;
}

int BLCMDtrace::command(BLArgumentVector& argv, BLArgumentMap& namedArgs)
{
	manager = BLManager::getObject();
	runManager = BLRunManager::getObject();

	int retval = handleNamedArgs(namedArgs);

	if(!init) {
		init = true;
		manager->registerSteppingAction(this);
		manager->registerTrackingAction(this);
		runManager->registerCollectiveComputation(this);
	}

	if(filename == "")
		filename = (oneNTuple!=0 ? "AllTracks" : "Ev%dTrk%d");

	coordinateType = BLCoordinates::getCoordinateType(coordinates);
	// requir applies per track, not per step; handled in newTrace()
	if(require != "") eval = new BLEvaluator();

	// handle format
	for(G4String::size_type i=0; i<format.size(); ++i)
		format[i] = tolower(format[i]);
	G4String::size_type j=format.find("trace");
	if(j == format.npos) {
		j = format.find("extended");
		if(j == format.npos) format += "Trace";
	}

	print("");

	return retval;
}

void BLCMDtrace::defineNamedArgs()
{
	argInt(nTrace,"nTrace","Number of tracks to trace.");
	argString(format,"format","Format of the NTuple (see above for list).");
	argInt(oneNTuple,"oneNTuple","Nonzero to put all traces into a single NTuple.");
	argInt(primaryOnly,"primaryOnly","Nonzero to trace only primary tracks.");
	argInt(traceTune,"traceTune","Nonzero to trace tune tracks (default=1).");
	argString(filename,"filename","Filename (Ev%dTrk%d.txt or AllTracks.txt).");
	argString(filename,"file","synonym for filename.");
	argString(require,"require","Expression which must be nonzero to trace the track (default=1)",false);
	argString(coordinates,"coordinates","Coordinates: global, centerline, or reference (default=c).");
}

void BLCMDtrace::newTrace(const G4Track *firstTrack)
{
	int evNum = G4RunManager::GetRunManager()->GetCurrentEvent()->GetEventID();
	int trkId = manager->getExternalTrackID(firstTrack);

	if(oneNTuple != 0) {
		if(!allTrace) allTrace = new TraceNTuple(format,"AllTracks",
						filename, coordinateType);
	} else if(evNum == -2) {
		if(!tuneTrace) tuneTrace = new TraceNTuple(format,
			"TuneParticle","TuneParticle",coordinateType);
	} else if(evNum == -1) {
		if(!referenceTrace) referenceTrace = new TraceNTuple(format,
			"ReferenceParticle","ReferenceParticle",coordinateType);
	} else {
		if(trace) trace->close();
	}
	trace = 0;

	if(eval) {
		eval->setTrackVariables(firstTrack,coordinateType);
		G4double v = eval->evaluate(require);
		if(!eval->isOK()) {
			BLCommand::printError("trace: invalid expression require='%s'\n",
				require.c_str());
			delete eval;
			eval = 0;
			v = 1.0;
		}
		if(v == 0.0) return;
	}

	if(allTrace != 0) {
		trace = allTrace;
	} else if(evNum == -2) {
		trace = tuneTrace;
	} else if(evNum == -1) {
		trace = referenceTrace;
	} else if(evNum >= 0) {
		char tmp[128];
		snprintf(tmp,128,filename.c_str(),evNum,trkId);
		trace = new TraceNTuple(format,tmp,tmp,coordinateType);
	}
}


void BLCMDtrace::UserSteppingAction(const G4Step *step)
{
	if(manager->getState() == SPECIAL || runManager->getCollectiveMode())
		return;

	if(!trace || !doTrace) return;

	G4Track *track = step->GetTrack();
	trace->appendTrace(track);
}

void BLCMDtrace::PreUserTrackingAction(const G4Track *track)
{
	int evNum = runManager->GetCurrentEvent()->GetEventID();

	doTrace = false;

	if(evNum == -2) {
		doTrace = (traceTune && coordinateType != BLCOORD_REFERENCE);
		if(doTrace) newTrace(track);
	} else if(evNum < 0) {
		// cannot use reference coordinates for tune or reference
		doTrace = (coordinateType != BLCOORD_REFERENCE);
		if(doTrace) newTrace(track);
	} else if(runManager->getCollectiveMode() && nTrace > 0) {
		if(oneNTuple == 0) {
			G4Exception("trace","Collective mode, oneNTuple=0",
				FatalException,"oneNTuple must be nonzero");
		}
		newTrace(track);
	} else if(nTrace > 0) {
		doTrace = ((primaryOnly == 0) || (track->GetParentID() == 0));
		if(doTrace) {
			--nTrace;
			newTrace(track);
		}
	} else {
		if(trace) {
			trace->close();
			trace = 0;
		}
	}

	if(trace && doTrace && 
			!runManager->getCollectiveMode()) {
		char tmp[64];
		sprintf(tmp,"# Event %d Track %d",evNum,manager->
				getExternalTrackID(track));
		trace->annotate(tmp);
	}
}

void BLCMDtrace::PostUserTrackingAction(const G4Track *track)
{
	if(!trace) return;

	if(doTrace && !runManager->getCollectiveMode())
		trace->annotate("");

	int evNum = G4RunManager::GetRunManager()->GetCurrentEvent()->
								GetEventID();
	if(oneNTuple == 0 && evNum >= 0) {
		trace->close();
		trace = 0;
	}
}

void BLCMDtrace::beginCollectiveTracking(std::vector<BLTrackData>& v)
{
	(void)v;
}

void BLCMDtrace::collectiveStep(std::vector<BLTrackData>& v)
{
	if(trace) {
		for(unsigned i=0; i<v.size(); ++i) {
			G4Track *track = v[i].track;
			G4Event *event = v[i].event;
			runManager->setCurrentEvent(event);
			G4TrackStatus trackStatus = track->GetTrackStatus();
			if(trackStatus != fAlive && 
						trackStatus != fStopButAlive)
				continue;
			if(primaryOnly != 0 && track->GetParentID() > 0)
				continue;
			trace->appendTrace(track);
		}
	}
}

void BLCMDtrace::endCollectiveTracking(std::vector<BLTrackData>& v)
{
	(void)v;
}



TraceNTuple::TraceNTuple(G4String format, G4String name, G4String filename,
					BLCoordinateType _coordinateType)
{
	coordinateType = _coordinateType;
	manager = BLManager::getObject();
	runManager = BLRunManager::getObject();
	ntuple = BLTrackNTuple::create(format,"Trace",name,filename,
		coordinateType,"");
}

TraceNTuple::~TraceNTuple()
{
	if(ntuple) delete ntuple;
	ntuple = 0;
}

void TraceNTuple::appendTrace(G4Track *track)
{
	ntuple->appendTrack(track);
}