local_exclusionmask.h

// -*-c++-*-
#ifndef STORMM_LOCAL_EXCLUSIONMASK_H
#define STORMM_LOCAL_EXCLUSIONMASK_H
 
#include "copyright.h"
#include "Accelerator/hybrid.h"
#include "DataTypes/common_types.h"
#include "DataTypes/stormm_vector_types.h"
#include "Math/formulas.h"
#include "Synthesis/atomgraph_synthesis.h"
#include "Topology/atomgraph.h"
#include "Topology/atomgraph_abstracts.h"
#include "energy_enumerators.h"
 
namespace stormm {
namespace energy {
 
using card::Hybrid;
using card::HybridTargetLevel;
using stmath::ipow;
using synthesis::AtomGraphSynthesis;
using topology::AtomGraph;
using topology::ChemicalDetailsKit;
using topology::NonbondedKit;

static const ullint lmask_mode_a       = 0x0LLU;
static const ullint lmask_mode_b       = 0x2000000000000000LLU;
static const ullint lmask_mode_c       = 0x4000000000000000LLU;
static const ullint lmask_mode_d       = 0x6000000000000000LLU;
static const ullint lmask_mode_e       = 0x8000000000000000LLU;
static const ullint lmask_mode_f       = 0xa000000000000000LLU;

static const ullint lmask_mode_bitmask = 0xe000000000000000LLU;
static const ullint lmask_a_excl       = 0x1fffffffffffffffLLU;
static const ullint lmask_b_excl       = 0x000000007fffffffLLU;
static const ullint lmask_b_lower_mask = 0x000001ff80000000LLU;
static const ullint lmask_b_upper_mask = 0x0007fe0000000000LLU;
static const ullint lmask_b_lower_shft = 0x00f8000000000000LLU;
static const ullint lmask_b_upper_shft = 0x1f00000000000000LLU;
static const ullint lmask_c_excl       = 0x000000007fffffffLLU;
static const ullint lmask_c_alt_mask   = 0x00001fff80000000LLU;
static const ullint lmask_c_shft       = 0x1fffe00000000000LLU;
static const ullint lmask_d_excl       = 0x000000007fffffffLLU;
static const ullint lmask_d_array_idx  = 0x003fffff80000000LLU;
static const ullint lmask_d_array_cnt  = 0x1fc0000000000000LLU;
static const ullint lmask_e_array_idx  = 0x000000ffffffffffLLU;
static const ullint lmask_e_array_cnt  = 0x1fffff0000000000LLU;
static const ullint lmask_f_idx        = 0x000000000000001fLLU;

static const int lmask_long_local_span   = 30;
static const int lmask_short_local_span  = 15;
static const int lmask_long_pivot_point  = lmask_long_local_span + 1;
static const int lmask_short_pivot_point = lmask_short_local_span + 1;
static const int lmask_long_extra_span   = 14;
static const int lmask_short_extra_span  = 10;
  
static const int lmask_mode_bits         =  3;
static const int lmask_b_lower_mask_bits = lmask_short_extra_span;
static const int lmask_b_upper_mask_bits = lmask_short_extra_span;
static const int lmask_b_shft_bits       =  5;
static const int lmask_b_max_reach       = lmask_short_local_span + 32 +
                                           lmask_short_extra_span - 1;
static const int lmask_c_alt_mask_bits   = lmask_long_extra_span;
static const int lmask_c_shft_bits       = 16;
static const int lmask_d_array_idx_bits  = 23;
static const int lmask_d_array_cnt_bits  =  7;
static const int lmask_e_array_idx_bits  = 40;
static const int lmask_e_array_cnt_bits  = 21;

static const int lmask_mode_pos          = (2 * lmask_long_local_span) + 1;
static const int lmask_b_lower_mask_pos  = (2 * lmask_short_local_span) + 1;
static const int lmask_b_upper_mask_pos  = lmask_b_lower_mask_pos + lmask_short_extra_span;
static const int lmask_b_lower_shft_pos  = lmask_b_upper_mask_pos + lmask_short_extra_span;
static const int lmask_b_upper_shft_pos  = lmask_b_lower_shft_pos + lmask_b_shft_bits;
static const int lmask_c_alt_mask_pos    = (2 * lmask_short_local_span) + 1;
static const int lmask_c_shft_pos        = lmask_c_alt_mask_pos + lmask_c_alt_mask_bits;
static const int lmask_d_array_idx_pos   = (2 * lmask_short_local_span) + 1;
static const int lmask_d_array_cnt_pos   = lmask_d_array_idx_pos + lmask_d_array_idx_bits;
static const int lmask_e_array_cnt_pos   = 40;

struct LocalExclusionMaskReader {

  LocalExclusionMaskReader(const int* prof_idx_in, const ullint* profiles_in,
                           const uint2* aux_masks_in);

  LocalExclusionMaskReader(const LocalExclusionMaskReader &original) = default;
  LocalExclusionMaskReader(LocalExclusionMaskReader &&original) = default;
  
  const int* prof_idx;     
  const ullint* profiles;  
  const uint2* aux_masks;  
};
  
class LocalExclusionMask {
public:

  LocalExclusionMask(const AtomGraph *ag_in = nullptr,
                     NonbondedTheme theme_in = NonbondedTheme::ALL);
 
  LocalExclusionMask(const AtomGraph &ag_in, NonbondedTheme theme_in = NonbondedTheme::ALL);
 
  LocalExclusionMask(const AtomGraphSynthesis *poly_ag_in,
                     NonbondedTheme theme_in = NonbondedTheme::ALL);
 
  LocalExclusionMask(const AtomGraphSynthesis &poly_ag_in,
                     NonbondedTheme theme_in = NonbondedTheme::ALL);

  LocalExclusionMask(const LocalExclusionMask &original) = default;
  LocalExclusionMask(LocalExclusionMask &&original) = default;
  LocalExclusionMask& operator=(const LocalExclusionMask &other) = default;
  LocalExclusionMask& operator=(LocalExclusionMask &&other) = default;

  int getAtomCount() const;

  int getProfileCount() const;
  int getProfileCount(int mode_index) const;

  int getMode(int atom_index, int system_index = 0) const;

  size_t getSecondaryMaskCount() const;

  ullint getProfile(int atom_index, int system_index = 0) const;

  const Hybrid<uint2> getSecondaryMaskView(int atom_index, int system_index = 0);

  std::vector<uint2> getSecondaryMasks(int atom_index, int system_index = 0) const;

  const AtomGraph* getTopologyPointer() const;

  const AtomGraphSynthesis* getTopologySynthesisPointer() const;

  bool testExclusion(int atom_i, int atom_j) const;
  bool testExclusion(int atom_i, int atom_j, int system_index) const;

  LocalExclusionMaskReader data(HybridTargetLevel tier = HybridTargetLevel::HOST) const;
 
#ifdef STORMM_USE_HPC
  void upload();
  
  void download();
#endif
  
  void setTopology(const AtomGraph *ag_in);
  void setTopology(const AtomGraph &ag_in);
  void setTopology(const AtomGraphSynthesis *poly_ag_in);
  void setTopology(const AtomGraphSynthesis &poly_ag_in);
  
private:

  int total_atom_count;

  NonbondedTheme theme;
  
  int mode_a_profile_count;
  int mode_b_profile_count;
  int mode_c_profile_count;
  int mode_d_profile_count;
  int mode_e_profile_count;
  int mode_f_profile_count;
  
  Hybrid<int> atom_profile_codes;
  
  Hybrid<ullint> atom_profiles;

  Hybrid<uint2> secondary_masks;

  AtomGraph *ag_pointer;

  AtomGraphSynthesis *poly_ag_pointer;

  void nullify();

  void validateAtomIndex(int atom_index, int system_index = 0, const char* caller = nullptr) const;
  void validateAtomIndex(int atom_index, const char* caller) const;

  void countProfileKinds();
  
  std::vector<int> extendMasks(const NonbondedKit<double> &nbk, const ChemicalDetailsKit &cdk);
  std::vector<int> extendMasks(const AtomGraph *ag);
};

std::vector<int> compileLocalExclusionList(const NonbondedKit<double> &nbk, int atom_index,
                                           NonbondedTheme theme, VdwCombiningRule lj_rule);

std::string lMaskModeToString(ullint mode);

void appendLMaskPrintout(std::string *result, const std::vector<int> &limits, ullint mask);
  
std::string lMaskToString(ullint mask);
  
bool lMaskFitsModeA(int atom_index, const std::vector<int> &excluded_atoms);

bool lMaskFitsModeB(int atom_index, const std::vector<int> &excluded_atoms);

bool lMaskFitsModeC(int atom_index, const std::vector<int> &excluded_atoms);

bool lMaskFitsModeD(int atom_index, const std::vector<int> &excluded_atoms,
                    const std::vector<uint2> &tmp_secondary_masks);

ullint lMaskCreateProfile(int atom_index, ullint mode, const std::vector<int> &excluded_atoms);
  
std::vector<uint2> lMaskCreateProfile(int atom_index, ullint mode,
                                      const std::vector<int> &excluded_atoms, ullint *t_profile);
  
void setProfile(int pos, ullint tprof, std::vector<int> *profile_indices,
                std::vector<ullint> *tmp_profiles);
 
void setProfile(int pos, ullint tprof, ullint mode,
                const std::vector<uint2> &tprof_secondary_masks, std::vector<int> *profile_indices,
                std::vector<ullint> *tmp_profiles, std::vector<uint2> *tmp_secondary_masks);

bool evaluateLocalMask(int atom_i, int atom_j, ullint prof, const uint2* secondary_ptr);

bool testExclusion(const LocalExclusionMaskReader &lemr, int atom_i, int atom_j);
 
} // namespace energy
} // namespace stormm
 
#endif