// This file is part of VecGeom and is distributed under the
// conditions in the file LICENSE.txt in the top directory.
// For the full list of authors see CONTRIBUTORS.txt and `git log`.

/// \brief Declaration of the manager/registry class for VecGeom geometries.
/// \file management/RootGeoManager.h
/// \author created by Johannes de Fine Licht, Sandro Wenzel (CERN)

#ifndef VECGEOM_MANAGEMENT_ROOTMANAGER_H_
#define VECGEOM_MANAGEMENT_ROOTMANAGER_H_

#include "VecGeom/base/Global.h"
#include "VecGeom/base/TypeMap.h"
#include <vector>
#include "VecGeom/volumes/PlacedVolume.h"
#include "VecGeom/management/GeoManager.h"
#include <functional>

#include "TGeoNode.h"
class TGeoVolume;
class TGeoMatrix;
class TGeoTrap;
class TGeoArb8;

namespace vecgeom {
inline namespace VECGEOM_IMPL_NAMESPACE {

class LogicalVolume;
class Transformation3D;
class UnplacedBox;
class VUnplacedVolume;
class UnplacedGenTrap;

/// \brief Class allowing to construct a VecGeom geometry from a ROOT/TGeo one
///        and to lookup VecGeom instances from ROOT ones and vice versa.
///
/// \details Allows integration with ROOT geometries for compatibility reasons.
///          Is not necessary for VecGeom library, and will only have source
///          compiled if the VECGEOM_ROOT flag is set by the compiler, activated
///          with -DROOT=ON in CMake. When VecGeom is compiled with specialization
///          support, the conversion will return runtime-specialized versions.
class RootGeoManager {

private:
  /** Remember pointer to generated world from imported ROOT geometry. */
  VPlacedVolume const *fWorld;

  BidirectionalTypeMap<unsigned int, TGeoNode const *> fPlacedVolumeMap;
  BidirectionalTypeMap<VUnplacedVolume const *, TGeoShape const *> fUnplacedVolumeMap;
  BidirectionalTypeMap<LogicalVolume const *, TGeoVolume const *> fLogicalVolumeMap;
  BidirectionalTypeMap<Transformation3D const *, TGeoMatrix const *> fTransformationMap;

  std::vector<TGeoNode const *> fTGeoNodeVector;

  std::function<void *(TGeoMaterial const *)> fMaterialConversionLambda = [](TGeoMaterial const *) { return nullptr; };

  int fVerbose;

  bool fFlattenAssemblies; // flag deciding whether we flatten assemblies (like G4 does)

public:
  /// Access singleton instance.
  static RootGeoManager &Instance()
  {
    static RootGeoManager instance;
    return instance;
  }

  /**
   * @return Most recently generated world from ROOT geometry. Will return NULL
   *         if no ROOT geometry has been imported.
   * @sa LoadRootGeometry()
   */
  VPlacedVolume const *world() const { return fWorld; }

  int GetVerboseLevel() const { return fVerbose; }

  void SetFlattenAssemblies(bool b) { fFlattenAssemblies = b; }
  bool GetFlattenAssemblies() const { return fFlattenAssemblies; }

  /// Returns a TGeoNode corresponding to a given VecGeom placed volume p
  TGeoNode const *tgeonode(VPlacedVolume const *p) const
  {
    if (p == nullptr) return nullptr;
    return fTGeoNodeVector[p->id()];
  }

  /// Returns the VecGeom placed volume corresponding to a TGeoNode node
  VPlacedVolume const *Lookup(TGeoNode const *node) const;

  /// Get placed volume that corresponds to a TGeoNode
  VPlacedVolume const *GetPlacedVolume(TGeoNode const *n) const
  {
    if (n == NULL) return NULL;
    return (GeoManager::Instance().Convert(fPlacedVolumeMap[n]));
  }

  /// Returns the name of a VecGeom placed volume by querying the TGeo underlying object
  char const *GetName(VPlacedVolume const *p) const { return tgeonode(p)->GetName(); }

  /// Print a table of registered TGeoNodes (for debugging purposes)
  void PrintNodeTable() const;

  /// Sets a verbose level. Currently if > 0 some debug information will be printed.
  void set_verbose(const int verbose) { fVerbose = verbose; }

  /**
   * Queries the global ROOT GeoManager for the top volume and recursively
   * imports and converts to VecGeom geometry.
   * Will register the imported ROOT geometry as the new world of the VecGeom
   * GeoManager singleton.
   *
   *
   * Requires an already initialized gGeoManager object in ROOT.
   */
  void LoadRootGeometry();

  /**
   * Same but can take a root file as input.
   * Initializes gGeoManager object in ROOT at the same time.
   */
  void LoadRootGeometry(std::string);

  /**
   * Exports a VecGeom geometry to a ROOT geometry.
   */
  bool ExportToROOTGeometry(VPlacedVolume const *, std::string);

  /**
   * @brief Deletes the VecGeom geometry generated by this class.
   */
  void Clear();

  /// This sets the base units to G4 style (in mm)
  void EnableG4Units() { fUnitsInMM = true; }
  /// This sets the base units to TGeo style (in cm)
  void EnableTGeoUnits() { fUnitsInMM = false; }

  /// Converts a TGeoNode to a VPlacedVolume, recursively converting daughters.
  /// Will take care not to convert anything twice by checking the birectional
  /// map between ROOT and VecGeom geometry. Usually all other Convert functions
  /// are invoked from this one.
  VPlacedVolume *Convert(TGeoNode const *const node);

  /// Conversion function from a TGeoShape to a corresponding VecGeom VUnplacedVolume
  VUnplacedVolume *Convert(TGeoShape const *const shape);

  /// Special conversion function from a TGeoVolumeAssembly
  VUnplacedVolume *ConvertAssembly(TGeoVolume const *const shape);

  /// Conversion function from a TGeoVolume to a corresponding VecGeom LogicalVolume
  LogicalVolume *Convert(TGeoVolume const *const volume);

  /// Conversion function from a TGeoMatrix to a corresponding VecGeom Transformation3D
  Transformation3D *Convert(TGeoMatrix const *const trans);

  /**
   * This function allows to register a callback that is executed
   * whenever the geometry conversion process encounters a TGeoMaterial.
   *
   * We provide this as VecGeom does not have itself a material class, but
   * when the user still wants to associate a material to a logical volume.
   */
  void SetMaterialConversionHook(std::function<void *(TGeoMaterial const *)> &&f) { fMaterialConversionLambda = f; }

  /// Conversion function from a VecGeom placed volume to a corresponding TGeoNode
  TGeoNode *Convert(VPlacedVolume const *const node);

  /// Conversion function from a VecGeom placed volume + logical volume to a corresponding TGeoNode
  TGeoVolume *Convert(VPlacedVolume const *const, LogicalVolume const *const);

  /// Conversion function from a VecGeom transformation to a corresponding TGeoMatrix
  TGeoMatrix *Convert(Transformation3D const *const trans);

private:
  RootGeoManager()
      : fWorld(NULL), fPlacedVolumeMap(), fUnplacedVolumeMap(), fLogicalVolumeMap(), fTransformationMap(), fVerbose(0),
        fFlattenAssemblies(false)
  {
  }
  RootGeoManager(RootGeoManager const &);
  RootGeoManager &operator=(RootGeoManager const &);

  // helper function to check TGeo trap for consistency / trap conversion
  bool TGeoTrapIsDegenerate(TGeoTrap const *);
  UnplacedGenTrap *ToUnplacedGenTrap(TGeoArb8 const *);

  /// Internal helper function to post-adjust a converted transformation in certain cases
  bool PostAdjustTransformation(Transformation3D *, TGeoNode const *, Transformation3D *adjustment) const;

  // for the choice of units
  bool fUnitsInMM = false;
  // whether to instantiate VecGeom with mm as units (to be compliant with G4 for instance)
  // if false, we will use cm (to be compliant with TGeo)

  constexpr static double CMTOMM = 10.;

  // return length unit
  double LUnit() const { return fUnitsInMM ? CMTOMM : 1.; }
};
}
} // End global namespace

#endif // VECGEOM_MANAGEMENT_ROOTMANAGER_H_
