//
// ********************************************************************
// * License and Disclaimer                                           *
// *                                                                  *
// * The  Geant4 software  is  copyright of the Copyright Holders  of *
// * the Geant4 Collaboration.  It is provided  under  the terms  and *
// * conditions of the Geant4 Software License,  included in the file *
// * LICENSE and available at  http://cern.ch/geant4/license .  These *
// * include a list of copyright holders.                             *
// *                                                                  *
// * Neither the authors of this software system, nor their employing *
// * institutes,nor the agencies providing financial support for this *
// * work  make  any representation or  warranty, express or implied, *
// * regarding  this  software system or assume any liability for its *
// * use.  Please see the license in the file  LICENSE  and URL above *
// * for the full disclaimer and the limitation of liability.         *
// *                                                                  *
// * This  code  implementation is the result of  the  scientific and *
// * technical work of the GEANT4 collaboration.                      *
// * By using,  copying,  modifying or  distributing the software (or *
// * any work based  on the software)  you  agree  to acknowledge its *
// * use  in  resulting  scientific  publications,  and indicate your *
// * acceptance of all terms of the Geant4 Software license.          *
// ********************************************************************
//
// G4VTHitsMap, G4THitsMultiMap, G4THitsUnorderedMap & G4THitsUnorderedMultiMap
//
// Class description:
//
// This is a template class of hits map and parametrised by
// The concrete class of G4VHit. This is a uniform collection for
// a particular concrete hit class objects.
// An intermediate layer class G4HitsMap appeared in this
// header file is used just for G4Allocator, because G4Allocator
// cannot be instantiated with a template class. Thus G4HitsMap
// class MUST NOT be directly used by the user.
//
// Authors: Makoto Asai, Jonathan Madsen
// --------------------------------------------------------------------
#ifndef G4THitsMap_h
#define G4THitsMap_h 1

#include "G4THitsCollection.hh"
#include "globals.hh"

#include <map>
#include <unordered_map>
#include <type_traits>

template <typename T, typename Map_t = std::map<G4int, T*>>
class G4VTHitsMap : public G4HitsCollection
{
 private:

  using mmap_t = std::multimap<G4int, T *>;
  using pair_t = std::pair<G4int, T *>;
  using uommap_t = std::unordered_multimap<G4int, T *>;

#define is_same_t(_Tp, _Up) std::is_same<_Tp, _Up>::value
#define is_multimap_t(_Mp) std::is_same<_Mp, mmap_t>::value
#define is_uommap_t(_Mp) std::is_same<_Mp, uommap_t>::value
#define is_mmap_t(_Mp) (is_multimap_t(_Mp) || is_uommap_t(_Mp))
#define is_fundamental_t(_Tp) std::is_fundamental<_Tp>::value

  template <bool _Bp, typename _Tp = void>
  using enable_if_t = typename std::enable_if<_Bp, _Tp>::type;

  // ensure fundamental types are initialized to zero
  template <typename U = T, enable_if_t<is_fundamental_t(U), G4int> = 0>
  T* allocate() const
  {
    return new T(0.);
  }
  // non-fundamental types should set values to appropriate values
  // and avoid issues such as:
  //   G4StatDouble stat(0.); stat += 1.0; gives n == 2;
  template <typename U = T, enable_if_t<! is_fundamental_t(U), G4int> = 0>
  T* allocate() const
  {
    return new T();
  }

 public:
  using this_type = G4VTHitsMap<T, Map_t>;
  using value_type = T;
  using map_type = Map_t;
  using iterator = typename map_type::iterator;
  using const_iterator = typename map_type::const_iterator;

  // generic constructor
  G4VTHitsMap();
  // det + collection description constructor
  G4VTHitsMap(const G4String& detName, const G4String& colNam);
  // destructor
  ~G4VTHitsMap() override;
  // equivalence operator
  G4bool operator==(const G4VTHitsMap<T, Map_t>& right) const;

  //------------------------------------------------------------------------//
  // Generic operator += where add(...) overloads handle various
  //  U and MapU_t types
  //------------------------------------------------------------------------//
  template <typename U, typename MapU_t>
  this_type& operator+=(const G4VTHitsMap<U, MapU_t>& right) const
  {
    MapU_t* aHitsMap = right.GetMap();
    for (auto itr = aHitsMap->begin(); itr != aHitsMap->end(); ++itr)
      add<U, map_type>(itr->first, *(itr->second));
    return (this_type&)(*this);
  }
  //------------------------------------------------------------------------//

  void DrawAllHits() override;
  void PrintAllHits() override;
  //  These two methods invokes Draw() and Print() methods of all of
  //  hit objects stored in this map, respectively.

  // Generic iteration
  inline Map_t* GetContainer() const { return (Map_t*)theCollection; }

  inline typename Map_t::size_type size() { return GetContainer()->size(); }

  inline typename Map_t::size_type GetIndex(iterator itr) { return itr->first; }

  inline typename Map_t::size_type GetIndex(const_iterator itr) const { return itr->first; }

  template <typename MapU_t = Map_t, enable_if_t<! is_mmap_t(MapU_t), G4int> = 0>
  inline T* GetObject(G4int idx) const
  {
    return (GetContainer()->count(idx) != 0) ? (*GetContainer())[idx] : nullptr;
  }

  template <typename MapU_t = Map_t, enable_if_t<is_mmap_t(MapU_t), G4int> = 0>
  inline T* GetObject(G4int idx) const
  {
    return (GetContainer()->count(idx) != 0) ? GetContainer()->find(idx)->second : nullptr;
  }

  inline T* GetObject(iterator itr) const { return itr->second; }

  inline const T* GetObject(const_iterator itr) const { return itr->second; }

  iterator begin() { return GetContainer()->begin(); }
  iterator end() { return GetContainer()->end(); }
  const_iterator begin() const { return GetContainer()->begin(); }
  const_iterator end() const { return GetContainer()->end(); }
  const_iterator cbegin() const { return GetContainer()->cbegin(); }
  const_iterator cend() const { return GetContainer()->cend(); }

  //  Returns a pointer to a concrete hit object.
  inline Map_t* GetMap() const { return (Map_t*)theCollection; }
  //  Overwrite a hit object. Total number of hit objects stored in this
  // map is returned.
  inline std::size_t entries() const { return ((Map_t*)theCollection)->size(); }
  //  Returns the number of hit objects stored in this map
  inline void clear();

  G4VHit* GetHit(std::size_t) const override { return nullptr; }
  std::size_t GetSize() const override { return ((Map_t*)theCollection)->size(); }

  //------------------------------------------------------------------------//
  //  Add/Insert a hit object. Total number of hit objects stored in this
  //  map is returned.
  //------------------------------------------------------------------------//
  //  Standard map overload for same type
  //------------------------------------------------------------------------//
  // here we don't use allocate() since instances like G4Colour() == white
  // and += adds to white (not correct)
  template <typename U = T, typename MapU_t = Map_t,
    enable_if_t<is_same_t(U, T) && ! is_mmap_t(MapU_t), G4int> = 0>
  std::size_t add(const G4int& key, U*& aHit) const
  {
    map_type* theHitsMap = GetMap();
    if (theHitsMap->find(key) == theHitsMap->end())
      theHitsMap->insert(pair_t(key, new T(*aHit)));
    else
      *theHitsMap->find(key)->second += *aHit;
    return theHitsMap->size();
  }
  //------------------------------------------------------------------------//
  //  Multimap overload for same type T
  //------------------------------------------------------------------------//
  template <typename U = T, typename MapU_t = Map_t,
    enable_if_t<(is_same_t(U, T) && is_mmap_t(MapU_t)), G4int> = 0>
  std::size_t add(const G4int& key, U*& aHit) const
  {
    map_type* theHitsMap = GetMap();
    theHitsMap->insert(pair_t(key, aHit));
    return theHitsMap->size();
  }
  //------------------------------------------------------------------------//
  //  Multimap overload for different types
  //      assumes type T has overload of += operator for U
  //------------------------------------------------------------------------//
  template <typename U = T, typename MapU_t = Map_t,
    enable_if_t<(! is_same_t(U, T) && is_mmap_t(MapU_t)), G4int> = 0>
  std::size_t add(const G4int& key, U*& aHit) const
  {
    map_type* theHitsMap = GetMap();
    T* hit = allocate();
    *hit += *aHit;
    theHitsMap->insert(pair_t(key, hit));
    return theHitsMap->size();
  }

  //------------------------------------------------------------------------//
  //  Standard map overload for same type
  //      assumes type T has overload of += operator for U
  //------------------------------------------------------------------------//
  // here we don't use allocate() since instances like G4Colour() == white
  // and += adds to white (not correct)
  template <typename U = T, typename MapU_t = Map_t,
    enable_if_t<(is_same_t(U, T) && ! is_mmap_t(MapU_t)), G4int> = 0>
  std::size_t add(const G4int& key, U& aHit) const
  {
    map_type* theHitsMap = GetMap();
    if (theHitsMap->find(key) == theHitsMap->end())
      theHitsMap->insert(pair_t(key, new T(aHit)));
    else
      *theHitsMap->find(key)->second += aHit;
    return theHitsMap->size();
  }
  //------------------------------------------------------------------------//
  //  Standard map overload for different type
  //      assumes type T has overload of += operator for U
  //------------------------------------------------------------------------//
  template <typename U = T, typename MapU_t = Map_t,
    enable_if_t<(! is_same_t(U, T) && ! is_mmap_t(MapU_t)), G4int> = 0>
  std::size_t add(const G4int& key, U& aHit) const
  {
    map_type* theHitsMap = GetMap();
    if (theHitsMap->find(key) == theHitsMap->end()) theHitsMap->insert(pair_t(key, allocate()));
    *theHitsMap->find(key)->second += aHit;
    return theHitsMap->size();
  }
  //------------------------------------------------------------------------//
  //  Multimap overload for same type T
  //------------------------------------------------------------------------//
  template <typename U = T, typename MapU_t = Map_t,
    enable_if_t<(is_same_t(U, T) && is_mmap_t(MapU_t)), G4int> = 0>
  std::size_t add(const G4int& key, U& aHit) const
  {
    map_type* theHitsMap = GetMap();
    theHitsMap->insert(pair_t(key, new T(aHit)));
    return theHitsMap->size();
  }
  //------------------------------------------------------------------------//
  //  Multimap overload for different types
  //      assumes type T has overload of += operator for U
  //------------------------------------------------------------------------//
  template <typename U = T, typename MapU_t = Map_t,
    enable_if_t<(! is_same_t(U, T) && is_mmap_t(MapU_t)), G4int> = 0>
  std::size_t add(const G4int& key, U& aHit) const
  {
    map_type* theHitsMap = GetMap();
    T* hit = allocate();
    *hit += aHit;
    theHitsMap->insert(pair_t(key, hit));
    return theHitsMap->size();
  }

  //------------------------------------------------------------------------//
  //  Set a hit object. Total number of hit objects stored in this
  //  map is returned.
  //------------------------------------------------------------------------//
  //  Standard overload for same type T
  //------------------------------------------------------------------------//
  template <typename U = T, typename MapU_t = Map_t,
    enable_if_t<(is_same_t(U, T) && ! is_mmap_t(MapU_t)), G4int> = 0>
  inline std::size_t set(const G4int& key, U*& aHit) const
  {
    map_type* theHitsMap = GetMap();
    if (theHitsMap->find(key) != theHitsMap->end()) delete theHitsMap->find(key)->second;
    theHitsMap->find(key)->second = aHit;
    return theHitsMap->size();
  }
  //------------------------------------------------------------------------//
  //  Multimap overload for same type T
  //------------------------------------------------------------------------//
  template <typename U = T, typename MapU_t = Map_t,
    enable_if_t<(is_same_t(U, T) && is_mmap_t(MapU_t)), G4int> = 0>
  inline std::size_t set(const G4int& key, U*& aHit) const
  {
    map_type* theHitsMap = GetMap();
    if (theHitsMap->find(key) != theHitsMap->end())
      theHitsMap->insert(pair_t(key, aHit));
    else {
      delete theHitsMap->find(key)->second;
      theHitsMap->find(key)->second = aHit;
    }
    return theHitsMap->size();
  }
  //------------------------------------------------------------------------//
  //  Standard map overload for different types
  //------------------------------------------------------------------------//
  template <typename U = T, typename MapU_t = Map_t,
    enable_if_t<(! is_same_t(U, T) && ! is_mmap_t(MapU_t)), G4int> = 0>
  inline std::size_t set(const G4int& key, U*& aHit) const
  {
    map_type* theHitsMap = GetMap();
    T* hit = nullptr;
    if (theHitsMap->find(key) == theHitsMap->end())
      theHitsMap->insert(std::make_pair(key, hit = allocate()));
    else
      hit = theHitsMap->find(key)->second;
    *hit += *aHit;
    return theHitsMap->size();
  }
  //------------------------------------------------------------------------//
  //  Multimap overload for different types
  //------------------------------------------------------------------------//
  template <typename U = T, typename MapU_t = Map_t,
    enable_if_t<(! is_same_t(U, T) && is_mmap_t(MapU_t)), G4int> = 0>
  inline std::size_t set(const G4int& key, U*& aHit) const
  {
    map_type* theHitsMap = GetMap();
    T* hit = allocate();
    *hit += *aHit;
    if (theHitsMap->find(key) != theHitsMap->end())
      theHitsMap->insert(pair_t(key, hit));
    else {
      delete theHitsMap->find(key)->second;
      theHitsMap->find(key)->second = hit;
    }
    return theHitsMap->size();
  }

  //------------------------------------------------------------------------//
  //  Set a hit object. Total number of hit objects stored in this
  //  map is returned.
  //------------------------------------------------------------------------//
  //  Standard overload for same type T
  //------------------------------------------------------------------------//
  template <typename U = T, typename MapU_t = Map_t,
    enable_if_t<(is_same_t(U, T) && ! is_mmap_t(MapU_t)), G4int> = 0>
  inline std::size_t set(const G4int& key, U& aHit) const
  {
    map_type* theHitsMap = GetMap();
    T* hit = nullptr;
    if (theHitsMap->find(key) != theHitsMap->end())
      hit = theHitsMap->find(key)->second;
    else
      theHitsMap->insert(pair_t(key, hit = allocate()));
    *hit = aHit;
    return theHitsMap->size();
  }
  //------------------------------------------------------------------------//
  //  Multimap overload for same type T
  //------------------------------------------------------------------------//
  template <typename U = T, typename MapU_t = Map_t,
    enable_if_t<(is_same_t(U, T) && is_mmap_t(MapU_t)), G4int> = 0>
  inline std::size_t set(const G4int& key, U& aHit) const
  {
    map_type* theHitsMap = GetMap();
    if (theHitsMap->find(key) != theHitsMap->end())
      *theHitsMap->find(key)->second = aHit;
    else
      theHitsMap->insert(pair_t(key, new T(aHit)));
    return theHitsMap->size();
  }
  //------------------------------------------------------------------------//
  //  Standard map overload for different types
  //------------------------------------------------------------------------//
  template <typename U = T, typename MapU_t = Map_t,
    enable_if_t<(! is_same_t(U, T) && ! is_mmap_t(MapU_t)), G4int> = 0>
  inline std::size_t set(const G4int& key, U& aHit) const
  {
    map_type* theHitsMap = GetMap();
    T* hit = nullptr;
    if (theHitsMap->find(key) == theHitsMap->end())
      theHitsMap->insert(std::make_pair(key, hit = allocate()));
    else
      hit = theHitsMap->find(key)->second;
    *hit += aHit;
    return theHitsMap->size();
  }
  //------------------------------------------------------------------------//
  //  Multimap overload for different types
  //------------------------------------------------------------------------//
  template <typename U = T, typename MapU_t = Map_t,
    enable_if_t<(! is_same_t(U, T) && is_mmap_t(MapU_t)), G4int> = 0>
  inline std::size_t set(const G4int& key, U& aHit) const
  {
    map_type* theHitsMap = GetMap();
    T* hit = allocate();
    *hit += aHit;
    if (theHitsMap->find(key) != theHitsMap->end())
      *theHitsMap->find(key)->second = *hit;
    else
      theHitsMap->insert(pair_t(key, hit));
    return theHitsMap->size();
  }

  //------------------------------------------------------------------------//
  //  Enable bracket operator. Return pointer to data indexed by key or
  //      last occurring instance of pointer to data index by key in the
  //      case of a multimap
  //------------------------------------------------------------------------//
  template <typename MapU_t = Map_t, enable_if_t<! is_mmap_t(MapU_t), G4int> = 0>
  T* operator[](G4int key) const
  {
    map_type* theHitsMap = GetMap();
    if (theHitsMap->find(key) != theHitsMap->end()) return theHitsMap->find(key)->second;
    return nullptr;
  }
  //------------------------------------------------------------------------//
  template <typename MapU_t = Map_t, enable_if_t<is_mmap_t(MapU_t), G4int> = 0>
  T* operator[](G4int key) const
  {
#ifdef G4VERBOSE
    static G4bool _first = true;
    if (_first) {
      _first = false;
      G4Exception("G4THitsMap operator[]", "calling [] on multimap", JustWarning,
        "Returning the last matching entry");
    }
#endif
    map_type* theHitsMap = GetMap();
    iterator itr = theHitsMap->find(key);
    if (itr != theHitsMap->end()) {
      std::advance(itr, theHitsMap->count(key) - 1);
      return itr->second;
    }
    return nullptr;
  }
  //------------------------------------------------------------------------//

#undef is_same_t
#undef is_multimap_t
#undef is_uommap_t
#undef is_mmap_t
#undef is_fundamental_t
};

//============================================================================//

template <typename T, typename Map_t>
G4VTHitsMap<T, Map_t>::G4VTHitsMap()
{
  theCollection = (void*)new Map_t;
}

//============================================================================//

template <typename T, typename Map_t>
G4VTHitsMap<T, Map_t>::G4VTHitsMap(const G4String& detName, const G4String& colNam)
  : G4HitsCollection(detName, colNam)
{
  theCollection = (void*)new Map_t;
}

//============================================================================//

template <typename T, typename Map_t>
G4VTHitsMap<T, Map_t>::~G4VTHitsMap()
{
  map_type* theHitsMap = GetMap();
  for (auto itr = theHitsMap->begin(); itr != theHitsMap->end(); ++itr)
    delete itr->second;
  delete theHitsMap;
}

//============================================================================//

template <typename T, typename Map_t>
G4bool G4VTHitsMap<T, Map_t>::operator==(const G4VTHitsMap<T, Map_t>& right) const
{
  return (collectionName == right.collectionName);
}

//============================================================================//

template <typename T, typename Map_t>
void G4VTHitsMap<T, Map_t>::DrawAllHits()
{
  ;
}

//============================================================================//

template <typename T, typename Map_t>
void G4VTHitsMap<T, Map_t>::PrintAllHits()
{
  G4cout << "G4THitsMap " << SDname << " / " << collectionName << " --- " << entries() << " entries"
         << G4endl;
  /*----- commented out for the use-case where <T> cannot be initialized
          to be zero or does not support += operator.
   Map_t * theHitsMap = GetMap();
   typename Map_t::iterator itr = theHitsMap->begin();
   T sum = 0.;
   for(; itr != theHitsMap->end(); ++itr) {
    G4cout << "  " << itr->first << " : "
           << *(itr->second) << G4endl;
    sum += *(itr->second);
   }
   G4cout << "             Total : " << sum << G4endl;
  ----------------------------------------------------------------------*/
}

//============================================================================//

template <typename T, typename Map_t>
void G4VTHitsMap<T, Map_t>::clear()
{
  Map_t* theHitsMap = GetMap();
  for (iterator itr = theHitsMap->begin(); itr != theHitsMap->end(); ++itr)
    delete itr->second;
  theHitsMap->clear();
}

//============================================================================//
//                                                                            //
//                                                                            //
//                      Helpers for different map types                       //
//                                                                            //
//                                                                            //
//============================================================================//

template <typename _Tp>
class G4THitsMap : public G4VTHitsMap<_Tp, std::map<G4int, _Tp*>>
{
 public:
  using parent_type = G4VTHitsMap<_Tp, std::map<G4int, _Tp *>>;

 public:
  G4THitsMap() : parent_type() {}
  G4THitsMap(G4String detName, G4String colName) : parent_type(detName, colName) {}

  using parent_type::operator+=;
  using parent_type::operator==;
  using parent_type::operator[];
  using parent_type::add;
  using parent_type::begin;
  using parent_type::cbegin;
  using parent_type::cend;
  using parent_type::clear;
  using parent_type::DrawAllHits;
  using parent_type::end;
  using parent_type::entries;
  using parent_type::GetHit;
  using parent_type::GetMap;
  using parent_type::GetSize;
  using parent_type::PrintAllHits;
  using parent_type::set;
};

//============================================================================//

template <typename _Tp>
class G4THitsMultiMap : public G4VTHitsMap<_Tp, std::multimap<G4int, _Tp*>>
{
 public:
  using parent_type = G4VTHitsMap<_Tp, std::multimap<G4int, _Tp *>>;

 public:
  G4THitsMultiMap() : parent_type() {}
  G4THitsMultiMap(const G4String& detName, const G4String& colName)
    : parent_type(detName, colName) {}

  using parent_type::operator+=;
  using parent_type::operator==;
  using parent_type::operator[];
  using parent_type::add;
  using parent_type::begin;
  using parent_type::cbegin;
  using parent_type::cend;
  using parent_type::clear;
  using parent_type::DrawAllHits;
  using parent_type::end;
  using parent_type::entries;
  using parent_type::GetHit;
  using parent_type::GetMap;
  using parent_type::GetSize;
  using parent_type::PrintAllHits;
  using parent_type::set;
};

//============================================================================//

template <typename _Tp>
class G4THitsUnorderedMap : public G4VTHitsMap<_Tp, std::unordered_map<G4int, _Tp*>>
{
 public:
  using parent_type = G4VTHitsMap<_Tp, std::unordered_map<G4int, _Tp *>>;

 public:
  G4THitsUnorderedMap() : parent_type() {}
  G4THitsUnorderedMap(const G4String& detName, const G4String& colName)
    : parent_type(detName, colName) {}

  using parent_type::operator+=;
  using parent_type::operator==;
  using parent_type::operator[];
  using parent_type::add;
  using parent_type::begin;
  using parent_type::cbegin;
  using parent_type::cend;
  using parent_type::clear;
  using parent_type::DrawAllHits;
  using parent_type::end;
  using parent_type::entries;
  using parent_type::GetHit;
  using parent_type::GetMap;
  using parent_type::GetSize;
  using parent_type::PrintAllHits;
  using parent_type::set;
};

//============================================================================//

template <typename _Tp>
class G4THitsUnorderedMultiMap : public G4VTHitsMap<_Tp, std::unordered_multimap<G4int, _Tp*>>
{
 public:
  using parent_type = G4VTHitsMap<_Tp, std::unordered_multimap<G4int, _Tp *>>;

 public:
  G4THitsUnorderedMultiMap() : parent_type() {}
  G4THitsUnorderedMultiMap(const G4String& detName, const G4String& colName)
    : parent_type(detName, colName) {}

  using parent_type::operator+=;
  using parent_type::operator==;
  using parent_type::operator[];
  using parent_type::add;
  using parent_type::begin;
  using parent_type::cbegin;
  using parent_type::cend;
  using parent_type::clear;
  using parent_type::DrawAllHits;
  using parent_type::end;
  using parent_type::entries;
  using parent_type::GetHit;
  using parent_type::GetMap;
  using parent_type::GetSize;
  using parent_type::PrintAllHits;
  using parent_type::set;
};

//============================================================================//

#endif