SelectSubset.h

/**************************************************************************
 * basf2 (Belle II Analysis Software Framework)                           *
 * Author: The Belle II Collaboration                                     *
 *                                                                        *
 * See git log for contributors and copyright holders.                    *
 * This file is licensed under LGPL-3.0, see LICENSE.md.                  *
 **************************************************************************/
 
#pragma once
 
 
#include <framework/datastore/StoreArray.h>
#include <framework/datastore/RelationVector.h>
#include <framework/datastore/RelationsObject.h>
#include <framework/logging/Logger.h>
 
#include <string>
#include <functional>
#include <vector>
#include <type_traits>
 
 
namespace Belle2 {
  class SelectSubsetBase {
  public:
    std::vector<std::string> getInheritFromArrays() const { return m_inheritFromArrays; }
    std::vector<std::string> getInheritToArrays() const { return m_inheritToArrays; }
    bool getInheritToSelf() const { return m_inheritToSelf; }
 
    virtual StoreAccessorBase* getSet() = 0;
    virtual StoreAccessorBase* getSubSet() = 0;
 
    void swapSetsAndDestroyOriginal();
 
  protected:
    SelectSubsetBase() { }
    virtual ~SelectSubsetBase() { }
 
    std::vector<std::string> m_inheritFromArrays;
    std::vector<std::string> m_inheritToArrays;
    bool m_inheritToSelf = false;
    bool m_reduceExistingSet = false;
  };
 
  template < typename StoredClass >
  class SelectSubset : public SelectSubsetBase {
    static_assert(std::is_base_of<RelationsObject, StoredClass>::value,
                  "SelectSubset<T> only works with classes T inheriting from RelationsObject.");
  public:
    SelectSubset(): SelectSubsetBase() {};
 
    ~SelectSubset()
    {
      delete m_set;
      delete m_subset;
    }
 
    void registerSubset(const StoreArray< StoredClass >& set,
                        DataStore::EStoreFlags storeFlags = DataStore::c_ErrorIfAlreadyRegistered)
    {
      m_reduceExistingSet = true;
      registerSubset(set, set.getName() + "_tmpSubset", storeFlags | DataStore::c_DontWriteOut);
 
      inheritAllRelations();
    }
 
    void registerSubset(const StoreArray< StoredClass >& set, const std::string& subsetName,
                        DataStore::EStoreFlags storeFlags = DataStore::c_ErrorIfAlreadyRegistered)
    {
      if (m_set or m_subset) {
        B2FATAL("SelectSubset::registerSubset() can only be called once!");
      }
 
      m_set = new StoreArray<StoredClass>(set);
 
      m_subset = new StoreArray<StoredClass>(subsetName, m_set->getDurability());
      m_subset->registerInDataStore(storeFlags);
      m_subsetFlags = storeFlags;
 
      set.registerRelationTo(*m_subset, m_subset->getDurability(), storeFlags);
    }
 
    template<class T, class ... MoreArguments >
    void inheritRelationsFrom(const StoreArray<T>& array, MoreArguments... moreArgs)
    {
      if (array.getName() == m_set->getName()) {
        m_inheritToSelf = true;
        inheritRelationsFrom(*m_subset, moreArgs...);
      } else {
        const_cast<StoreArray<T>&>(array).isRequired();
 
        DataStore::EStoreFlags flags = m_subsetFlags;
        if (m_subset->notWrittenOut() or array.notWrittenOut())
          flags |= DataStore::c_DontWriteOut;
        array.registerRelationTo(*m_subset, m_subset->getDurability(), flags);
 
        if (array.getName() != m_subset->getName())
          m_inheritFromArrays.push_back(array.getName());
 
        inheritRelationsFrom(moreArgs ...);
      }
    }
 
    template<class T, class ... MoreArguments >
    void inheritRelationsTo(const StoreArray<T>& array, MoreArguments... moreArgs)
    {
      if (array.getName() == m_set->getName()) {
        m_inheritToSelf = true;
        inheritRelationsTo(*m_subset, moreArgs...);
      } else {
        const_cast<StoreArray<T>&>(array).isRequired();
 
        DataStore::EStoreFlags flags = m_subsetFlags;
        if (m_subset->notWrittenOut() or array.notWrittenOut())
          flags |= DataStore::c_DontWriteOut;
        m_subset->registerRelationTo(array, m_subset->getDurability(), flags);
 
        if (array.getName() != m_subset->getName())
          m_inheritToArrays.push_back(array.getName());
 
        inheritRelationsTo(moreArgs ...);
      }
    }
 
    void inheritAllRelations()
    {
      auto arrays = DataStore::Instance().getListOfRelatedArrays(*m_set);
 
      for (std::string arrayName : arrays) {
        StoreArray<TObject> array(arrayName, m_set->getDurability());
        if (array == *m_subset)
          continue; // from registerSubset(), ignore
 
        if (array.optionalRelationTo(*m_set, m_set->getDurability()))
          inheritRelationsFrom(array);
        if (m_set->optionalRelationTo(array, m_set->getDurability()))
          inheritRelationsTo(array);
      }
    }
 
    void select(const std::function<bool (const StoredClass*)>& f);
 
    StoreAccessorBase* getSet() override { return m_set; }
    StoreAccessorBase* getSubSet() override { return m_subset; }
 
  protected:
    std::map<int, int> copySetWithRelations(std::function<bool (const StoredClass*)> f);
 
    void copyRelationsToSelf();
 
    void inheritRelationsFrom() { }
    void inheritRelationsTo() { }
 
    StoreArray<StoredClass>* m_set = nullptr;
    StoreArray<StoredClass>* m_subset = nullptr;
    DataStore::EStoreFlags m_subsetFlags = DataStore::c_WriteOut;
  };
 
  template < class StoredClass>
  std::map<int, int>
  SelectSubset< StoredClass >::copySetWithRelations(std::function<bool (const StoredClass*)> f)
  {
    std::map<int, int> oldToNew;
    for (const StoredClass& setObject : *m_set) {
      if (!f(&setObject))
        continue;
 
      oldToNew[setObject.getArrayIndex()] = m_subset->getEntries();
      const StoredClass* subsetObject = m_subset->appendNew(setObject);
      if (!m_reduceExistingSet)
        setObject.addRelationTo(subsetObject);
    }
 
 
    //TODO this is the slow bit, can probably be improved by directly dealing with indices
    for (const auto& oldToNewPair : oldToNew) {
      const StoredClass* setObject = (*m_set)[oldToNewPair.first];
      const StoredClass* subsetObject = (*m_subset)[oldToNewPair.second];
 
      for (std::string fromArray : m_inheritFromArrays) {
        const RelationVector<RelationsObject>& relations = setObject->template getRelationsFrom<RelationsObject>(fromArray);
        for (unsigned int iRel = 0; iRel < relations.size(); iRel++) {
          relations.object(iRel)->addRelationTo(subsetObject, relations.weight(iRel));
        }
      }
      for (std::string toArray : m_inheritToArrays) {
        const RelationVector<RelationsObject>& relations = setObject->template getRelationsTo<RelationsObject>(toArray);
        for (unsigned int iRel = 0; iRel < relations.size(); iRel++) {
          subsetObject->addRelationTo(relations.object(iRel), relations.weight(iRel));
        }
      }
    }
 
    return oldToNew;
  }
 
 
  template < class StoredClass>
  void
  SelectSubset< StoredClass >::copyRelationsToSelf()
  {
    for (const StoredClass& subsetObject1 : *m_subset) {
      //TODO: change relation direction to set -> subset?
      const StoredClass* setObject1 = subsetObject1.template getRelatedFrom<StoredClass>(m_set->getName());
 
      if (setObject1 != nullptr) {
        //get all objects in original set related to setObject1
        const RelationVector<StoredClass>& relations = setObject1->template getRelationsTo<StoredClass>(m_set->getName());
        for (unsigned int iRel = 0; iRel < relations.size(); iRel++) {
          const StoredClass* setObject2 = relations.object(iRel);
          const double weight = relations.weight(iRel);
          //if setObject2 was selected into subset, inherit relation
          const StoredClass* subsetObject2 = setObject2->template getRelatedTo<StoredClass>(m_subset->getName());
          if (subsetObject2) {
            subsetObject1.addRelationTo(subsetObject2, weight);
          }
        }
      }
    }
  }
 
  template < class StoredClass>
  void
  SelectSubset< StoredClass >::select(const std::function<bool (const StoredClass*)>& f)
  {
    copySetWithRelations(f);
 
    if (m_inheritToSelf) {
      copyRelationsToSelf();
    }
 
    if (m_reduceExistingSet) {
      swapSetsAndDestroyOriginal();
    }
  }
}