From e4b146e7d118277660fdfc3f5813eaebe61433ce Mon Sep 17 00:00:00 2001
From: Choi Shing Wan <choishingwan@gmail.com>
Date: Wed, 15 Jul 2020 22:18:53 +0800
Subject: [PATCH] [Fix] BGEN sample check when sample selection was performed

---
 inc/commander.hpp |  2 +-
 src/binarygen.cpp | 45 +++++++++++++++++++++++++++------------------
 src/genotype.cpp  |  4 +++-
 3 files changed, 31 insertions(+), 20 deletions(-)

diff --git a/inc/commander.hpp b/inc/commander.hpp
index 50c4e1ae..2f94380b 100644
--- a/inc/commander.hpp
+++ b/inc/commander.hpp
@@ -43,7 +43,7 @@
 #endif
 
 const std::string version = "2.3.2";
-const std::string date = "2020-06-23";
+const std::string date = "2020-07-10";
 class Commander
 {
 public:
diff --git a/src/binarygen.cpp b/src/binarygen.cpp
index 3ff6ccab..3afe9aee 100644
--- a/src/binarygen.cpp
+++ b/src/binarygen.cpp
@@ -260,26 +260,35 @@ void BinaryGen::check_sample_consistent(const genfile::bgen::Context& context,
         size_t sample_idx = 0;
         genfile::bgen::read_sample_identifier_block(
             bgen_file, tmp_context, [this, &sample_idx](const std::string& id) {
-                std::string ref_id =
-                    (m_ignore_fid ? "" : m_sample_id[sample_idx].FID + m_delim)
-                    + m_sample_id[sample_idx].IID;
-                if (ref_id != id && m_sample_id[sample_idx].IID != id)
+                auto&& find_id = m_sample_selection_list.find(id)
+                                 != m_sample_selection_list.end();
+                bool inclusion = m_remove_sample ^ find_id;
+                if (inclusion)
                 {
-                    throw std::runtime_error(
-                        "Error: Sample mismatch "
-                        "between bgen and phenotype file! Name in BGEN "
-                        "file is "
-                        ":"
-                        + id + " and in phentoype file is: " + ref_id
-                        + ". Please note that PRSice require the bgen file "
-                          "and "
-                          "the .sample (or phenotype file if sample file is "
-                          "not provided) to have sample in the same order. "
-                          "(We "
-                          "might be able to losen this requirement in future "
-                          "when we have more time)");
+                    std::string ref_id =
+                        (m_ignore_fid ? ""
+                                      : m_sample_id[sample_idx].FID + m_delim)
+                        + m_sample_id[sample_idx].IID;
+                    if (ref_id != id && m_sample_id[sample_idx].IID != id)
+                    {
+                        throw std::runtime_error(
+                            "Error: Sample mismatch "
+                            "between bgen and phenotype file! Name in BGEN "
+                            "file is "
+                            ":"
+                            + id + " and in phentoype file is: " + ref_id
+                            + ". Please note that PRSice require the bgen file "
+                              "and "
+                              "the .sample (or phenotype file if sample file "
+                              "is "
+                              "not provided) to have sample in the same order. "
+                              "(We "
+                              "might be able to losen this requirement in "
+                              "future "
+                              "when we have more time)");
+                    }
+                    ++sample_idx;
                 }
-                ++sample_idx;
             });
     }
 }
diff --git a/src/genotype.cpp b/src/genotype.cpp
index d919e10c..9c65f9e4 100644
--- a/src/genotype.cpp
+++ b/src/genotype.cpp
@@ -904,7 +904,6 @@ void Genotype::load_samples(bool verbose)
         BITCT_TO_WORDCT(m_unfiltered_sample_ct);
     m_exclude_from_std.resize(unfiltered_sample_ctl, 0);
     if (verbose) m_reporter->report(message);
-    m_sample_selection_list.clear();
 }
 
 bool Genotype::perform_freqs_and_inter(const QCFiltering& filter_info,
@@ -1064,6 +1063,9 @@ void Genotype::load_snps(
         message = "Error: No vairant remained!\n";
         throw std::runtime_error(message);
     }
+    // only remove selection list here, as we might need this for bgen file
+    // check
+    m_sample_selection_list.clear();
 }