// GPL 3+
// Filippo Rusconi

#pragma once

/////////////////////// StdLib includes


/////////////////////// Qt includes


/////////////////////// Local includes
#include "msrundatasettreevisitor.h"
#include "msrundatasettreenode.h"
#include "msrunid.h"
#include "../exportinmportconfig.h"
#include "../processing/combiners/selectionpolygon.h"


namespace pappso
{

class MsRunDataSetTree;

typedef std::shared_ptr<MsRunDataSetTree> MsRunDataSetTreeSPtr;
typedef std::shared_ptr<const MsRunDataSetTree> MsRunDataSetTreeCstSPtr;

class PMSPP_LIB_DECL MsRunDataSetTree
{
  public:
  MsRunDataSetTree(MsRunIdCstSPtr ms_run_id_csp);
  virtual ~MsRunDataSetTree();

  MsRunDataSetTreeNode *
  addMassSpectrum(QualifiedMassSpectrumCstSPtr mass_spectrum);

  const std::map<std::size_t, MsRunDataSetTreeNode *> &getIndexNodeMap() const;
  const std::vector<MsRunDataSetTreeNode *> &getRootNodes() const;

  void accept(MsRunDataSetTreeNodeVisitorInterface &visitor);
  void accept(MsRunDataSetTreeNodeVisitorInterface &visitor,
              std::vector<pappso::MsRunDataSetTreeNode *>::const_iterator
                nodes_begin_iterator,
              std::vector<pappso::MsRunDataSetTreeNode *>::const_iterator
                nodes_end_iterator);

  MsRunDataSetTreeNode *
  findNode(QualifiedMassSpectrumCstSPtr mass_spectrum_csp) const;
  MsRunDataSetTreeNode *findNode(std::size_t spectrum_index) const;

  std::size_t massSpectrumIndex(const MsRunDataSetTreeNode *node) const;
  std::size_t massSpectrumIndex(
    QualifiedMassSpectrumCstSPtr qualified_mass_spectrum_csp) const;

  /****************** Flattened views ******************/

  std::vector<MsRunDataSetTreeNode *> flattenedView();

  std::vector<MsRunDataSetTreeNode *>
  flattenedViewMsLevel(std::size_t ms_level, bool with_descendants = false);

  /****************** Flattened views ******************/

  MsRunDataSetTreeNode *
  precursorNodeByProductSpectrumIndex(std::size_t product_spectrum_index);

  std::vector<MsRunDataSetTreeNode *>
  productNodesByPrecursorSpectrumIndex(std::size_t precursor_spectrum_index);

  std::vector<MsRunDataSetTreeNode *>
  precursorNodesByPrecursorMz(pappso_double mz, PrecisionPtr precision_ptr);

  // Utility functions.
  std::size_t depth() const;

  // The tree size as computed by going down the tree nodes.
  std::size_t size() const;

  // The size of the flat index/node map as filled in during file loading.
  std::size_t indexNodeMapSize() const;

  std::size_t getSpectrumCount() const;

  using NodeVector            = std::vector<MsRunDataSetTreeNode *>;
  using QualMassSpectraVector = std::vector<QualifiedMassSpectrumCstSPtr>;

  using DoubleNodeVectorMap = std::map<double, NodeVector>;

  std::size_t addDataSetTreeNodesInsideDtRtRange(double start,
                                                 double end,
                                                 NodeVector &nodes,
                                                 DataKind data_kind) const;

  std::size_t removeDataSetTreeNodesOutsideDtRtRange(double start,
                                                     double end,
                                                     NodeVector &nodes,
                                                     DataKind data_kind) const;

  std::size_t
  addDataSetQualMassSpectraInsideDtRtRange(double start,
                                           double end,
                                           QualMassSpectraVector &mass_spectra,
                                           DataKind data_kind) const;

  std::size_t removeDataSetQualMassSpectraOutsideDtRtRange(
    double start,
    double end,
    QualMassSpectraVector &mass_spectra,
    DataKind data_kind) const;

  private:
  MsRunIdCstSPtr mcsp_msRunId;

  std::size_t m_spectrumCount = std::numeric_limits<std::size_t>::min();

  std::vector<MsRunDataSetTreeNode *> m_rootNodes;
  std::map<std::size_t, MsRunDataSetTreeNode *> m_indexNodeMap;

  // We want to be able to list easily all the mass spectra that were acquired
  // at any given RT or DT.

  DoubleNodeVectorMap m_dtDoubleNodeVectorMap;
  DoubleNodeVectorMap m_rtDoubleNodeVectorMap;

  // Beware for data_kind == ion mobility to use the right
  // ion mobility value: 1/K0 for Bruker timsTOF data and
  // m_dt for Agilent and Synapt data.
  bool documentNodeInDtRtMap(double time,
                             MsRunDataSetTreeNode *node_p,
                             DataKind data_kind);

  // These addMassSpectrum functions are for forensic science.
  MsRunDataSetTreeNode *
  addMassSpectrum(QualifiedMassSpectrumCstSPtr mass_spectrum,
                  MsRunDataSetTreeNode *parent);

  MsRunDataSetTreeNode *
  addMassSpectrum(QualifiedMassSpectrumCstSPtr mass_spectrum,
                  std::size_t precursor_spectrum_index);
};

} // namespace pappso

