/**@file GPSVolumeSampling.hh
   @author B. Loer
   @date 2018-03-29
   @brief Adds new macro optiones to G4GeneralParticleSource to simplify
   throwing sources from volumes.

*/

#ifndef GPSVolumeSampling_h
#define GPSVolumeSamplnig_h

#include "G4UImessenger.hh"
#include "globals.hh"
#include "G4AffineTransform.hh"
#include "G4Transform3D.hh"
#include "G4ThreeVector.hh"
#include "WeightedSampler.hh"
#include <vector>

class G4VPhysicalVolume;
class G4UIcommand;
class G4UIcmdWithAString;
class G4UIcmdWith3VectorAndUnit;
class G4GeneralParticleSource;
class G4NavigationHistory;
class G4GeneralParticleSource;

class GPSVolumeSampling : public G4UImessenger {
public:
  GPSVolumeSampling(G4GeneralParticleSource* gps=nullptr);
  virtual ~GPSVolumeSampling();

  /// different ways to treat children when sampling
  enum ESamplingMode { kExcludeChildren, kIncludeChildren, kOnlyChildren,
                       kSurface, kCenter, kBoundingBox };

  enum EInitialShapeMode { kInitBbox, kInitGps };

  /// different direction sampling
  enum EDirectionMode { kNoDirection, kIsotropicOuter, kIsotropicInner,
                        kLambertianOuter, kLambertianInner };

  /// simple struct to hold a volume and it's tranform
  struct VolTrans {
    const G4VPhysicalVolume* vol; //< physical volume
    G4Transform3D transform;      //< transform of volume in world coords
    ESamplingMode fSampleMode;    //< volume sampling algorithm
    G4ThreeVector bbox_min;       //< bounding box min in local coords
    G4ThreeVector bbox_length;    //< bounding box max in local coords
    G4double weight = 1;          //< random samppling weight

    VolTrans(const G4VPhysicalVolume* Pv, const G4Transform3D& Tform,
             ESamplingMode SampleMode=kExcludeChildren,  G4double Weight=1);

    G4ThreeVector SamplePositionAndDirection(G4ThreeVector& direction,
                                             EDirectionMode directionMode=kNoDirection,
                                             EInitialShapeMode initShape=kInitBbox,
                                             G4GeneralParticleSource* gps=nullptr);

    void SetSamplingMode(ESamplingMode sampleMode, bool forcereweight=false);
  };
  typedef std::vector<VolTrans> VolTransVec;


  void SetNewValue(G4UIcommand* command, G4String newValues);
  G4String GetCurrentValue(G4UIcommand* command);

  void SetGPS(G4GeneralParticleSource* gps){ fGPS = gps; }
  G4GeneralParticleSource* GetGPS() const { return fGPS; }

  ///utility: get a physical volume by name
  static const G4VPhysicalVolume* FindVolume(const G4String& volname,
               G4int copyno=-1);

  /// utility: get the transform for a volume by name
  static G4Transform3D LocateVolume(const G4String& volname, G4int copyno=-1);

  /// utility: get the transform for a volume pointer
  static G4Transform3D LocateVolume(const G4VPhysicalVolume* pvol);

  /// utility: get all volume pointers by name
  static std::vector<const G4VPhysicalVolume*>
    FindAllVolumes(const G4String& volname, G4int copyno=-1);

  /// utility: get all volumes and transforms by name
  static VolTransVec LocateAllVolumes(const G4String& volname, G4int copyno=-1);

  EDirectionMode GetDirectionMode() const { return fDirectionMode; }
  void SetDirectionMode(EDirectionMode mode){ fDirectionMode = mode; }
  void SetDirectionMode(const G4String& mode);

  EInitialShapeMode GetInitialShapeMode() const { return fInitialShapeMode; }
  void SetInitialShapeMode(EInitialShapeMode mode){ fInitialShapeMode = mode; }
  void SetInitialShapeMode(const G4String& mode);

  size_t GetNumTargetVolumes() const { return fTargetVolumes.size(); }
  
  void ClearTargetVolumes() { fTargetVolumes.clear(); fTargetVolumeWeights.clear();}

  /// sample the volume. Deprecated, use SampleVolumeAndDirection instead
  inline G4ThreeVector SampleVolume()
  {
    G4ThreeVector dummy(0,0,1);
    return SampleVolumeAndDirection(dummy);
  }

  /// Sample chosen volumes and optionally the primary direction
  G4ThreeVector SampleVolumeAndDirection(G4ThreeVector& direction);

  void CenterGPSOnVolume(const G4String& volname);
  void ShiftGPSCenter(const G4ThreeVector& offset);

private: 
  void PrepareSamplingList(G4UIcommand* command, const G4String& newValues);

  G4UIcommand* sampleVolumeCmd;
  G4UIcmdWithAString* initialShapeCmd;
  G4UIcommand* clearVolumeListCmd;
  G4UIcmdWithAString* sampleDirectionCmd;
  G4UIcmdWithAString* locateVolumeCmd;
  G4UIcmdWithAString* centerVolumeCmd;
  G4UIcmdWith3VectorAndUnit* offsetVolumeCmd;
  VolTransVec fTargetVolumes;
  WeightedSampler fTargetVolumeWeights;
  EDirectionMode fDirectionMode = kNoDirection;
  EInitialShapeMode fInitialShapeMode = kInitBbox;
  G4GeneralParticleSource* fGPS;
};

#endif //GPSVolumeSampling_h