#include "GPSVolumeSampling.hh"
#include "BGMMetadata.hh"
#include "G4GeneralParticleSource.hh"
#include "G4GeneralParticleSourceData.hh"
#include "G4SingleParticleSource.hh"
#include "G4SPSPosDistribution.hh"
#include "G4SingleParticleSource.hh"
#include "G4SPSPosDistribution.hh"
#include "G4PhysicalVolumeStore.hh"
#include "G4VPhysicalVolume.hh"
#include "G4LogicalVolume.hh"
#include "G4VSolid.hh"
#include "G4ApplicationState.hh"
#include "G4UIcmdWithAString.hh"
#include "G4UIcmdWith3VectorAndUnit.hh"
#include "G4ThreeVector.hh"
#include "G4VisExtent.hh"
#include "G4TransportationManager.hh"
#include <algorithm>
#include "globals.hh"
#include "G4SystemOfUnits.hh"
#include "G4UnitsTable.hh"
#include "G4UImanager.hh"
#include "G4UIcommandTree.hh"
#include "G4NavigationHistory.hh"
#include "G4Tokenizer.hh"
#include "G4ThreeVector.hh"
#include "Randomize.hh"
#include "G4PhysicalVolumesSearchScene.hh"
#include <thread>
#include <mutex>

GPSVolumeSampling::GPSVolumeSampling(G4GeneralParticleSource* gps) :
  G4UImessenger(),  sampleVolumeCmd(0), sampleDirectionCmd(0),
  locateVolumeCmd(0),  centerVolumeCmd(0), offsetVolumeCmd(0),
  fGPS(gps)
{
  G4String path = "/bgm/generator/sampleVolume";
  sampleVolumeCmd = new G4UIcommand(path, this);
  sampleVolumeCmd->SetParameter(new G4UIparameter("volname", 's', false));
  G4UIparameter* copyno = new G4UIparameter("copyno",'i', true);
  copyno->SetDefaultValue(-1);
  sampleVolumeCmd->SetParameter(copyno);
  G4UIparameter* samplemode = new G4UIparameter("mode",'s',true);
  samplemode->SetDefaultValue("exclude");
  samplemode->SetParameterCandidates("exclude include only surface center bbox");
  sampleVolumeCmd->SetParameter(samplemode);
  sampleVolumeCmd->SetGuidance("Uniformly sample points from a volume");
  sampleVolumeCmd->SetGuidance("Usage: sampleVolume <volname [copyno] [mode]");
  sampleVolumeCmd->SetGuidance("If copyno==-1, sample all volumes with same name");
  sampleVolumeCmd->SetGuidance("mode determines how volumes with children are handled:");
  sampleVolumeCmd->SetGuidance("  exclude (default): only points not containing a child");
  sampleVolumeCmd->SetGuidance("  include: sample the whole mother's solid");
  sampleVolumeCmd->SetGuidance("  only: only select points from children");
  sampleVolumeCmd->SetGuidance("  surface: sample associated solid's surface");
  sampleVolumeCmd->SetGuidance("           Ignores daughter volumes");
  sampleVolumeCmd->SetGuidance("  center: offset by center of the volume(s)");
  sampleVolumeCmd->SetGuidance("  bbox: Use the solid's bounding box, translated");
  sampleVolumeCmd->SetGuidance("        to world coordinates, without checking volume");
  sampleVolumeCmd->SetGuidance("        If initial shape is 'gps', use that");

  initialShapeCmd = new G4UIcmdWithAString("/bgm/generator/initialSampleShape", this);
  initialShapeCmd->SetGuidance("Distribution of particles in solid's coordinates");
  initialShapeCmd->SetGuidance("before checking for volume in world coordinates.");
  initialShapeCmd->SetGuidance("  bbox: Use the solid's bounding box");
  initialShapeCmd->SetGuidance("  gps:  use the configured GeneralParticleSource");
  initialShapeCmd->SetGuidance("NB: for gps, use the solid's coordinates, not world!");
  initialShapeCmd->SetGuidance("Also 'confine' must OT be set");
  initialShapeCmd->SetParameterName("shape", false);
  initialShapeCmd->SetCandidates("bbox gps");


  locateVolumeCmd = new G4UIcmdWithAString("/geom/locateVolume", this);
  locateVolumeCmd->SetGuidance("Print the location of a volume in world coordinates");
  locateVolumeCmd->SetParameterName("volname", false);

  clearVolumeListCmd = new G4UIcommand("/bgm/generator/clearVolumeList", this);
  clearVolumeListCmd->SetGuidance("Clear the list of volumes to be sampled");

  sampleDirectionCmd = new G4UIcmdWithAString("/bgm/generator/sampleDirection", this);
  sampleDirectionCmd->SetGuidance("Determine the primary particle direction");
  sampleDirectionCmd->SetGuidance("  none: use the GPS chosen direction");
  sampleDirectionCmd->SetGuidance("  IsotopicOuter: Isotropic facing outward");
  sampleDirectionCmd->SetGuidance("  IsotropicInner: Isotropic facing inward");
  sampleDirectionCmd->SetGuidance("  LambertianOuter: Costheta from surface normal pointing outward");
  sampleDirectionCmd->SetGuidance("  LambertianInner: Costheta from surface normal pointing inward");
  sampleDirectionCmd->SetParameterName("dir", false);
  sampleDirectionCmd->SetCandidates("none IsotropicOuter IsotropicInner LambertianOuter LambertianInner");


  commandsShouldBeInMaster = true;

  centerVolumeCmd = new G4UIcmdWithAString("/gps/pos/centerOnVolume", this);
  centerVolumeCmd->SetGuidance("Center on a physical volume");
  centerVolumeCmd->SetParameterName("volname", false);
  centerVolumeCmd->SetToBeBroadcasted(false);

  offsetVolumeCmd = new G4UIcmdWith3VectorAndUnit("/gps/pos/offsetCenter", this);
  offsetVolumeCmd->SetGuidance("Shift the center point from it's current position");
  offsetVolumeCmd->SetUnitCategory("Length");
  offsetVolumeCmd->SetToBeBroadcasted(false);

  commandsShouldBeInMaster = false;
}

GPSVolumeSampling::~GPSVolumeSampling()
{
  delete sampleVolumeCmd;
  delete initialShapeCmd;
  delete clearVolumeListCmd;
  delete sampleDirectionCmd;
  delete locateVolumeCmd;
  delete centerVolumeCmd;
  delete offsetVolumeCmd;
}


void GPSVolumeSampling::SetNewValue(G4UIcommand* command, G4String newValues)
{
  if(command == sampleVolumeCmd){
    PrepareSamplingList(command, newValues);
  }
  else if(command == initialShapeCmd){
    SetInitialShapeMode(newValues);
  }
  else if(command == clearVolumeListCmd){
    ClearTargetVolumes();
  }
  else if(command == sampleDirectionCmd){
    SetDirectionMode(newValues);
  }
  else if(command == locateVolumeCmd){
    LocateVolume(newValues);
  }
  else if(command == centerVolumeCmd){
    CenterGPSOnVolume(newValues);
  }
  else if(command == offsetVolumeCmd){
    ShiftGPSCenter(G4UIcommand::ConvertToDimensioned3Vector(newValues));
  }
  else {
    G4cout << "Error entering command" << G4endl;
  }
}

G4String GPSVolumeSampling::GetCurrentValue(G4UIcommand* )
{
  return "Not implemented";
}

const G4VPhysicalVolume* GPSVolumeSampling::FindVolume(const G4String& volname,
                                                       G4int copyno)
{
  //see if this volume actually exists
  G4PhysicalVolumeStore* pvStore = G4PhysicalVolumeStore::GetInstance();
  const auto& iter =
    std::find_if(pvStore->begin(), pvStore->end(),
     [&volname,copyno](const G4VPhysicalVolume* pv)
     { return pv->GetName() == volname &&
       ( copyno<0 || pv->GetCopyNo()==copyno); }

     );
  if(iter == pvStore->end()){
    G4ExceptionDescription msg;
    msg<<"No physical volume '"<<volname<<" (" << copyno << ")"<<"' found.";
    G4Exception("GPSVolumeSampling::SetNewValue","1", FatalException, msg);
  }
  return *iter;
}

G4Transform3D GPSVolumeSampling::LocateVolume(const G4String& volname,
                                              G4int copyno)
{
  // get the list of all volumes and take the first one
  // this is inefficient, but prevents code duplication
  VolTransVec alltforms = LocateAllVolumes(volname, copyno);
  if(alltforms.empty()){
    G4ExceptionDescription msg;
    msg<<"Volume "<<volname<<" is not found within world geometry!";
    G4Exception("GPSVolumeSampling::LocateFolume","1", FatalException, msg);
  }
  G4cout<<"Volume "<<volname<<" located at"
        << (alltforms[0].transform.getTranslation())/cm<<" cm"
        << G4endl;
  return alltforms[0].transform;
}

G4Transform3D GPSVolumeSampling::LocateVolume(const G4VPhysicalVolume* vol)
{
  //todo: this will fail if multiple volumes have same names!
  return LocateVolume(vol->GetName(), vol->GetCopyNo());
}

std::vector<const G4VPhysicalVolume*>
GPSVolumeSampling::FindAllVolumes(const G4String& volname, G4int copyno)
{
  std::vector<const G4VPhysicalVolume*> result;
  G4PhysicalVolumeStore* pvStore = G4PhysicalVolumeStore::GetInstance();
  for(const G4VPhysicalVolume* pv : *pvStore){
    if(pv->GetName() == volname && (copyno < 0 || pv->GetCopyNo() == copyno)){
      result.push_back(pv);
    }
  }
  return result;
}

GPSVolumeSampling::VolTrans::VolTrans(const G4VPhysicalVolume* Pv,
                                      const G4Transform3D& Tform,
                                      ESamplingMode sampleMode,
                                      G4double Weight) :
  vol(Pv), transform(Tform), weight(Weight)
{
  vol->GetLogicalVolume()->GetSolid()->BoundingLimits(bbox_min, bbox_length);
  bbox_length -= bbox_min;
  SetSamplingMode(sampleMode);
}

void GPSVolumeSampling::VolTrans::SetSamplingMode(ESamplingMode sampleMode,
                                                  bool forcereweight)
{
  fSampleMode = sampleMode;
  //update weighting
  G4double weightunit = kg;
  if(weight <= 0 or forcereweight){
    switch(sampleMode){
      case kSurface:
        weight = vol->GetLogicalVolume()->GetSolid()->GetSurfaceArea();
        weightunit = cm2;
        break;
      case kExcludeChildren:
        weight = vol->GetLogicalVolume()->GetMass(true, false);
        break;
      case kIncludeChildren:
        weight = vol->GetLogicalVolume()->GetMass(true, true);
        break;
      case kOnlyChildren:
        weight = vol->GetLogicalVolume()->GetMass(true, true) -
                 vol->GetLogicalVolume()->GetMass(true, false);
        break;
      default:
        weight = 1;
        break;
    };
    if(weight <= 0){
      G4cerr<<"Error calculating weight of volume "<<vol->GetName()<<"; using 1\n";
      weight = 1;
    }
    //G4cout<<"Sampling volume "<<vol->GetName()<<" with weight "<<weight/weightunit<<G4endl;
  }
}

G4ThreeVector GPSVolumeSampling::VolTrans::
  SamplePositionAndDirection(G4ThreeVector& direction,
                             EDirectionMode directionMode,
                             EInitialShapeMode initialShapeMode,
                             G4GeneralParticleSource* gps)
{
  const G4VSolid* solid = vol->GetLogicalVolume()->GetSolid();

  // sample randomly in the bounding box and see if it matches criteria
  const int MAXTRIALS = 100000;
  int ntrials = 0;
  G4ThreeVector position;
  while(ntrials++ < MAXTRIALS){
    //first sample the solid in it's own coordinate system
    switch(fSampleMode){
      case kCenter:
        //nothing to do
        break;
      case kSurface:
        position = solid->GetPointOnSurface();
        break;
      case kExcludeChildren:
      case kIncludeChildren:
      case kOnlyChildren:
      default:
        if(initialShapeMode == kInitBbox){
          // random point in bounding box:
          for(size_t axis=0; axis<3; ++axis){
            position[axis] = G4UniformRand() * bbox_length[axis] + bbox_min[axis];
          }
        }else{ //initialShapeMode == kInitiGps
          position = gps->GetCurrentSource()->GetPosDist()->GenerateOne();
        }
        // is it within the solid?
        // should we test strictly inside, or allow surfaces?

        if(fSampleMode != kBoundingBox && solid->Inside(position) != kInside)
          continue;
        break;
    }
    if(directionMode != kNoDirection){
      G4ThreeVector normal = solid->SurfaceNormal(position);
      if(directionMode == kLambertianInner || directionMode == kIsotropicInner)
        normal = -normal;
      G4ThreeVector ortho = normal.orthogonal();
      double costheta = G4UniformRand();
      // lambertian distribution is proportional to costheta
      if(directionMode == kLambertianInner || directionMode == kLambertianOuter)
        costheta = std::sqrt(costheta);
      ortho = ortho.rotate(G4UniformRand()*CLHEP::twopi, normal);
      direction = normal.rotate(std::acos(costheta), ortho);
      direction = transform * (HepGeom::Vector3D<double>)direction;
    }

    // transform to the world coordinates
    position = transform * (HepGeom::Point3D<double>)position;
    if(fSampleMode == kExcludeChildren || fSampleMode == kOnlyChildren){
      // see what touchable volume we're in
      G4Navigator *gNavigator =G4TransportationManager::GetTransportationManager()
        ->GetNavigatorForTracking();
      const G4VPhysicalVolume* volAtPoint = gNavigator->LocateGlobalPointAndSetup(position);
      if(!volAtPoint // we're outside the tracking volume
         || (fSampleMode == kExcludeChildren && volAtPoint != vol) // we're in a child
         || (fSampleMode == kOnlyChildren && volAtPoint == vol)) // this is the mother
        continue;
    }
    // end of things to check. If we get here we have a good point
    break;
  }

  if(ntrials >= MAXTRIALS){
    G4ExceptionDescription msg;
    msg<<"Unable to sample volume with current settings.\n";
    msg<<"Switching to Bounding Box only mode\n";
    G4Exception("GPSVolumeSampling::SampleVolume()","", JustWarning, msg);
    fSampleMode = kBoundingBox;
    // bounding box only mode should always return so we won't inifinitely recurse
    position = SamplePositionAndDirection(direction, directionMode);
  }

  return position;
}

GPSVolumeSampling::VolTransVec
GPSVolumeSampling::LocateAllVolumes(const G4String& volname, G4int copyno)
{
  VolTransVec result;
  //get the world volume
  G4TransportationManager* tm = G4TransportationManager::GetTransportationManager();
  size_t nworlds = tm->GetNoWorlds();
  if(nworlds < 1){
    G4ExceptionDescription msg;
    msg<<"No world volume registered with transportation manager!";
    G4Exception("GPSVolumeSampling::SetNewValue","1", FatalException, msg);
  }
  size_t iworld = 0;
  auto worldsIter = tm->GetWorldsIterator();
  G4VPhysicalVolume* world = nullptr;
  while(iworld++ < nworlds && (world = *(worldsIter++))){
    //use vis tools to find the volumes
    G4PhysicalVolumeModel searchModel (world);  // Unlimited depth.
    G4ModelingParameters mp;  // Default - no culling.
    searchModel.SetModelingParameters (&mp);
    G4PhysicalVolumesSearchScene searchScene (&searchModel, volname, copyno);
    searchModel.DescribeYourselfTo (searchScene);  // Initiate search.
    for(const auto& finding : searchScene.GetFindings()){
      result.push_back(VolTrans(finding.fpFoundPV,
                                finding.fFoundObjectTransformation));
    }
  }
  return result;
}



void GPSVolumeSampling::PrepareSamplingList(G4UIcommand*, const G4String& newValues)
{
  G4Tokenizer token(newValues);
  G4String volname = token();
  G4String copynostring = token();
  G4String modestring = token();
  G4int copyno = -1;
  ESamplingMode sampleMode = kExcludeChildren;


  if(!copynostring.empty() && copynostring != "*")
    copyno = std::stoi(copynostring);
  modestring.toLower();
  if(modestring == "exclude")
    sampleMode = kExcludeChildren;
  else if(modestring == "include")
    sampleMode = kIncludeChildren;
  else if(modestring == "only")
    sampleMode = kOnlyChildren;
  else if(modestring == "surface")
    sampleMode = kSurface;
  else if(modestring == "center")
    sampleMode = kCenter;
  else if(modestring == "bbox")
    sampleMode = kBoundingBox;
  else if(modestring != ""){
    G4ExceptionDescription msg;
    msg<<modestring<<" is not an allowed value for sampling mode";
    G4Exception("GPSVolumeSampling::PrepareSamplingList","1", FatalException, msg);
  }

  VolTransVec newvols = LocateAllVolumes(volname, copyno);
  if(newvols.size() == 0){
    G4ExceptionDescription msg;
    msg<<"No volumes found to sample for '"<<volname<<"', copy "<<copynostring;
    G4Exception("GPSVolumeSampling::PrepareSamplingList", "1", FatalException, msg);
  }
  for(auto& target : newvols)
    target.SetSamplingMode(sampleMode, true);
  // add newvols to the total list
  fTargetVolumes.insert(fTargetVolumes.end(), newvols.begin(), newvols.end());
  // and re-calculate the sampling weights
  auto getweight = [](const VolTrans& target){ return target.weight; };
  fTargetVolumeWeights.Fill(fTargetVolumes, getweight);

  if(sampleMode != kCenter){
    // set upseful normalization info
    auto meta = BGMMetadata::GetInstance();
    G4double unit = sampleMode == kSurface ? cm2 : kg;
    double sumweight = 0;
    for(const VolTrans& target : fTargetVolumes)
      sumweight += getweight(target);
    meta->Set("/bgm/generator/totalSampleWeight", sumweight / unit);
    G4String unitstr = sampleMode == kSurface ? "cm2" : "kg";
    meta->Set("/bgm/generator/totalSampleWeightUnit", unitstr);
  }
}

void GPSVolumeSampling::SetDirectionMode(const G4String& mode)
{
  if(mode == "none")
    fDirectionMode = kNoDirection;
  else if(mode == "IsotropicOuter")
    fDirectionMode = kIsotropicOuter;
  else if(mode == "IsotropicInner")
    fDirectionMode = kIsotropicInner;
  else if(mode == "LambertianOuter")
    fDirectionMode = kLambertianOuter;
  else if(mode == "LambertianInner")
    fDirectionMode = kLambertianInner;
  else{
    G4ExceptionDescription msg;
    msg<<"Unknown direction sampling mode string "<<mode;
    G4Exception("GPSVolumeSampling::SetDirectionMode", "1", FatalException, msg);
  }
}

void GPSVolumeSampling::SetInitialShapeMode(const G4String& mode)
{
  if(mode == "bbox")
    fInitialShapeMode = kInitBbox;
  else if(mode == "gps")
    fInitialShapeMode = kInitGps;
  else{
    G4ExceptionDescription msg;
    msg<<"Unknown initial shape mode string "<<mode;
    G4Exception("GPSVolumeSampling::SetInitialShapeMode", "1", FatalException, msg);
  }
}

G4ThreeVector GPSVolumeSampling::SampleVolumeAndDirection(G4ThreeVector& direction)
{
  if(fInitialShapeMode == kInitGps && !fGPS){
    G4ExceptionDescription msg;
    msg<<"GPS sampling requested, but gps is null";
    G4Exception("GPSVolumeSampling::SampleVolumeAndDirection", "1", FatalException, msg);
  }
  //first, randomly select a volume
  size_t volnum = fTargetVolumeWeights.Sample();
  return fTargetVolumes[volnum].SamplePositionAndDirection(direction, fDirectionMode,
                                                           fInitialShapeMode, fGPS);
}

void GPSVolumeSampling::CenterGPSOnVolume(const G4String& volname)
{
  G4ThreeVector pos = LocateVolume(volname, 0).getTranslation();
  G4GeneralParticleSourceData* gpsData = G4GeneralParticleSourceData::Instance();
  gpsData->Lock();
  gpsData->GetCurrentSource()->GetPosDist()->SetCentreCoords(pos);
  gpsData->Unlock();
  G4cout<<"GPS center set to "<<pos<<G4endl;
}

void GPSVolumeSampling::ShiftGPSCenter(const G4ThreeVector& offset)
{
  G4GeneralParticleSourceData* gpsData = G4GeneralParticleSourceData::Instance();
  gpsData->Lock();
  G4SPSPosDistribution* pdist = gpsData->GetCurrentSource()->GetPosDist();
  G4ThreeVector newpos = pdist->GetCentreCoords() + offset;
  pdist->SetCentreCoords(newpos);
  gpsData->Unlock();
  G4cout<<"GPS center is now set to "<<newpos<<G4endl;
}