//	BLCMDtotalenergy.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 <stdio.h>
#include <map>
#include <vector>
#include "fnmatch.h"

#include "BLManager.hh"
#include "BLCommand.hh"
#include "BLAsciiFile.hh"

/**	class BLCMDtotalenergy prints the total energy deposited into each
 *	selected volume, at the end of the run.
 **/
class BLCMDtotalenergy : public BLCommand, public BLCallback, 
			public BLManager::SteppingAction,
			public BLManager::TrackingAction {
	G4String volumes;
	G4int ancestors;
	G4String filename;
	bool alreadyRegistered;
	struct Totalizer {
		double total;
		Totalizer() { total=0.0; }
	};
	std::map<G4VPhysicalVolume*,Totalizer> total;
	std::vector<G4VPhysicalVolume*> printOrder;
	std::vector<G4String> patterns;
	G4VPhysicalVolume *world;
public:
	/// Constructor.
	BLCMDtotalenergy();

	/// commandName() returns "totalenergy".
	virtual G4String commandName() { return "totalenergy"; }
	
	/// command() implements the totalenergy command.
	virtual int command(BLArgumentVector& argv, BLArgumentMap& namedArgs);

	/// defineNamedArgs() defines the named arguments for this command.
	virtual void defineNamedArgs();

	/// callback() from BLCallback.
	void callback(int type);

	/// walk the tree of physical volumes
	void walkPVTree(G4VPhysicalVolume *pv, G4String indent="");

	/// isMatch() returns true if the name of this PV matches some pattern
	bool isMatch(G4VPhysicalVolume *pv);

	/// sum() will sum the energy into the relevant volumes.
	void sum(const G4TouchableHandle &th, double energy);

	/// 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);
};
BLCMDtotalenergy defaultTotalenergy;


BLCMDtotalenergy::BLCMDtotalenergy() : volumes("*"), total(), printOrder(), 
							patterns()
{
	registerCommand(BLCMDTYPE_DATA);
	setSynopsis("Print total energy deposited in selected volumes.");
	setDescription("At end of run, prints the total energy deposited in "
		"the selected volumes.\n\n"
		"Volume-name patterns are like UNIX filename patterns: "
		"e.g. '*[AB]*' matches any name containing an A or a B.\n\n"
		"Tracks that are killed have their kinetic energy summed into "
		"the volume where they were killed, unless they are killed "
		"because they leave the World.\n\n"
		"With ancestors=1, energy deposited in matching volumes is "
		"added into their ancestors; energy deposited directly into "
		"those ancestors is not summed into them unless their names "
		"also match. "
		"That is, if A is a child of B, but only A matches the list "
		"of volume-names, energy "
		"deposited into A will be reported in both A and B, but energy "
		"deposited directly into B is ignored.");

	ancestors = 0;
	filename = "";
	alreadyRegistered = false;
	world = 0; // initialized in callback()
}

int BLCMDtotalenergy::command(BLArgumentVector& argv, BLArgumentMap& namedArgs)
{
	int retval = handleNamedArgs(namedArgs);

	if(alreadyRegistered) {
		printError("Duplicate totalenergy command");
		return -1;
	}

	BLManager::getObject()->registerCallback(this,0);
	BLManager::getObject()->registerCallback(this,2);
	BLManager::getObject()->registerBeamStep(0,this);
	BLManager::getObject()->registerTrackingAction(this);
	alreadyRegistered = true;

	print("");

	patterns = splitString(volumes, ",", true);

	return retval;
}

void BLCMDtotalenergy::defineNamedArgs()
{
	argString(volumes,"volumes","Comma-separated list of Volume-Name patterns (*)");
	argInt(ancestors,"ancestors","Set nonzero to sum energy into all ancestor (enclosing) volumess (0).");
	argInt(ancestors,"enclosing","Synonym for ancestors.");
	argString(filename,"filename","Filename to write summary (stdout).");
	argString(filename,"file","Synonym for filename.");
}

void BLCMDtotalenergy::callback(int type)
{
	if(type == 0) { 
		// Walk the tree of physical volumes
		printf("\nPhysical volumes for totalenergy (* = totaled):\n");
		world = BLManager::getObject()->getWorldPhysicalVolume();
		walkPVTree(world);
	} else if(type == 2) {
		FILE *out=stdout;
		if(filename.size() > 0) {
			out = BLAsciiFile::fopen(filename);
			if(!out) {
				printError("Cannot write to '%s' - stdout used\n",
					filename.c_str());
				out = stdout;
			}
		}
		fprintf(out,"# Total energy deposited in selected volumes (MeV):\n");
		for(unsigned i=0; i<printOrder.size(); ++i) {
			G4VPhysicalVolume *pv = printOrder[i];
			double t=total[pv].total/MeV;
			if(t < 10.0)
			    fprintf(out,"%s\t%.6f\n",pv->GetName().c_str(),t);
			else if(t < 10000.0)
			    fprintf(out,"%s\t%.3f\n",pv->GetName().c_str(),t);
			else if(t < 100000.0)
			    fprintf(out,"%s\t%.1f\n",pv->GetName().c_str(),t);
			else
			    fprintf(out,"%s\t%.4g\n",pv->GetName().c_str(),t);
		}
		if(out != stdout) BLAsciiFile::fclose(out);
	}
}

void BLCMDtotalenergy::walkPVTree(G4VPhysicalVolume *pv, G4String indent)
{
	printf("%s%s",indent.c_str(),pv->GetName().c_str());
	if(isMatch(pv)) {
		printf(" *");
		if(total.count(pv) != 0)
			G4Exception("totalenergy command",
				"Duplicate PhysicalVolume cut",JustWarning, "");
		total[pv] = Totalizer();
		printOrder.push_back(pv);
	}
	printf("\n");

	G4LogicalVolume *lv = pv->GetLogicalVolume();
	int n=lv->GetNoDaughters();
	for(int i=0; i<n; ++i)
		walkPVTree(lv->GetDaughter(i),indent+"    ");
}

bool BLCMDtotalenergy::isMatch(G4VPhysicalVolume *pv)
{
	const char *name = pv->GetName().c_str();
	for(unsigned i=0; i<patterns.size(); ++i) {
		if(patterns[i].size() == 0) continue;
		if(fnmatch(patterns[i].c_str(),name,0) == 0)
			return true;
	}
	return false;
}

void BLCMDtotalenergy::sum(const G4TouchableHandle &th, double energy)
{
	if(!th) return;

	for(int depth=0; depth<=th->GetHistoryDepth(); ++depth) {
		G4VPhysicalVolume *pv = th->GetVolume(depth);
		if(!pv) break;
		if(depth == 0 && total.count(pv) == 0) break;
		total[pv].total += energy; // create if needed
		if(pv == world) break;
		if(!ancestors) break;
	}
}

void BLCMDtotalenergy::UserSteppingAction(const G4Step *step)
{
	G4StepPoint *prePoint = step->GetPreStepPoint();
	if(!prePoint) return;
	G4double energy = step->GetTotalEnergyDeposit();
	sum(prePoint->GetTouchableHandle(),energy);
}

void BLCMDtotalenergy::PreUserTrackingAction(const G4Track *track)
{
}

void BLCMDtotalenergy::PostUserTrackingAction(const G4Track *track)
{
	if(track->GetTrackStatus() == fSuspend) return;
	if(track->GetNextVolume() == 0) return; // don't sum if leaving World
	sum(track->GetTouchableHandle(),track->GetKineticEnergy());
}