atommask.h

// -*-c++-*-
#ifndef STORMM_AMBMASK_H
#define STORMM_AMBMASK_H
 
#include <string>
#include <vector>
#include <climits>
#include "copyright.h"
#include "Constants/behavior.h"
#include "Topology/atomgraph.h"
#include "Trajectory/coordinateframe.h"
#include "Trajectory/phasespace.h"
#include "chemical_features.h"
 
namespace stormm {
namespace chemistry {
 
using constants::ExceptionResponse;
using parse::WildCardKind;
using topology::AtomGraph;
using trajectory::CoordinateFrame;
using trajectory::PhaseSpace;

enum class MaskOperator {
  NONE,            
  ATOM_LT,         
  ATOM_LE,         
  ATOM_GT,         
  ATOM_GE,         
  RESIDUE_LT,      
  RESIDUE_GT,      
  RESIDUE_LE,      
  RESIDUE_GE,      
  AND,             
  OR,              
  NOT              
};

enum class SelectionItemKind {
  ATOM_NAME,      
  ATOM_NUMBER,    
  ELEMENT_NAME,   
  ELEMENT_NUMBER, 
  ATOM_TYPE,      
  RESIDUE_NAME,   
  RESIDUE_NUMBER, 
#if 0
  RING_SIZE,      
#endif
  NONE            
};

class SelectionItem {
public:

  SelectionItem(const SelectionItemKind kind_in, const char4 name_in,
                const std::vector<WildCardKind> &wildcards_in);
  SelectionItem(const SelectionItemKind kind_in, const int begin_in, const int end_in = INT_MIN);
  
  SelectionItemKind kind;               
  char4 name;                           
  std::vector<WildCardKind> wildcards;  
  int scan_begin;                       
  int scan_end;                         
};

enum class MaskComponentKind {
  OPERATOR,  
  MASK       
};

class MaskComponent {
public:

  MaskComponent(MaskOperator op_in, const std::string &basis_in, int start_idx, int end_idx);
  MaskComponent(MaskOperator op_in, double range, const std::string &basis_in, int start_idx,
                int end_idx);
  MaskComponent(const std::vector<SelectionItem> &parts_in, const AtomGraph *ag,
                const ChemicalFeatures &chemfe, const std::string &basis_in, int start_idx,
                int end_idx);
  MaskComponent(const std::vector<uint> &primitive_mask_in, const std::string &basis_in,
                int start_idx, int end_idx);

  MaskComponentKind getKind() const;

  MaskOperator getOperator() const;

  std::vector<uint> getMask() const;

  std::string getText() const;

  int2 getTextLimits() const;

  void applyNotOperator();

  void applyAndOperator(const std::vector<uint> &other);

  void applyOrOperator(const std::vector<uint> &other);

  MaskComponent applyRangeOperator(const std::vector<uint> &other, const AtomGraph *ag,
                                   const CoordinateFrameReader &cfr);
 
private:
  MaskComponentKind kind;
  MaskOperator op;
  double range;
  std::vector<uint> primitive_mask;
  std::string text_basis;
  int2 text_limits;
};

enum class MaskTraversalMode {
  COMPLETE,    
  PARTIAL,     
  SEGMENTED    
};

enum class MaskInputMode {
  AMBMASK,  
  VMD       
};
  
class AtomMask {
public:

  AtomMask(const AtomGraph *ag_in = nullptr);
 
  AtomMask(const std::string &input_text_in, const AtomGraph *ag_in,
           const ChemicalFeatures &chemfe, const CoordinateFrameReader &cfr,
           MaskInputMode mode = MaskInputMode::AMBMASK,
           const std::string &description_in = std::string("No description provided"));
 
  AtomMask(const std::string &input_text_in, const AtomGraph *ag_in,
           const ChemicalFeatures &chemfe, const CoordinateFrame &cf,
           MaskInputMode mode = MaskInputMode::AMBMASK,
           const std::string &description_in = std::string("No description provided"));
 
  AtomMask(const std::string &input_text_in, const AtomGraph *ag_in,
           const ChemicalFeatures &chemfe, const PhaseSpace &ps,
           MaskInputMode mode = MaskInputMode::AMBMASK,
           const std::string &description_in = std::string("No description provided"));

  AtomMask(const AtomMask &original) = default;
  AtomMask(AtomMask &&original) = default;
  
  MaskTraversalMode getRecommendation() const;

  const std::vector<uint>& getRawMask() const;

  int getMaskedAtomCount() const;

  std::vector<bool> getMask() const;
 
  // \brief Get a list of the masked atoms
  std::vector<int> getMaskedAtomList() const;

  std::vector<int2> getSegments() const;

  std::string getInputText() const;

  MaskInputMode getInputKind() const;
  
  std::string getDescription() const;

  const AtomGraph* getTopologyPointer() const;

  bool isAtomInMask(const char4 atom_name) const;
  bool isAtomInMask(const char4 res_name, const char4 atom_name) const;
  bool isAtomInMask(int atom_index) const;
  
  void addAtoms(const std::vector<int> &new_indices,
                ExceptionResponse policy = ExceptionResponse::DIE);
 
  void addAtoms(const std::vector<char4> &new_names);
 
  void addAtoms(const AtomMask &new_mask, const CoordinateFrame &cf,
                const ChemicalFeatures &chemfe);
 
  void addAtoms(const std::string &new_mask, const CoordinateFrame &cf,
        const ChemicalFeatures &chemfe);
 
  void addAtoms(const AtomMask &new_mask, const CoordinateFrame &cf);
  
  void addAtoms(const std::string &new_mask, const CoordinateFrame &cf);
  
private:
  MaskTraversalMode recommended_scan;      
  MaskInputMode style;                     
  int masked_atom_count;                   
  std::vector<uint> raw_mask;              
  std::vector<int2> segments;              
  std::string input_text;                  
  std::string description;                 
  const AtomGraph *ag_pointer;             

  double extractRangeValue(int start_char, int final_char, int *position);

  std::vector<SelectionItem> evaluateInclusions(const std::string &inclusion_list);
  
  std::vector<uint> parseMask(const std::vector<int> &scope_levels, int *position,
                              const CoordinateFrameReader &cfr, const ChemicalFeatures &chemfe);
};
 
} // namespace chemistry
} // namespace stormm
 
#endif