<commit_before>/*!
  @file PDBReader.hpp
  @brief pdb reader.

  read pdb atoms from input stream and parse atoms as chains.

  @author <NAME> (<EMAIL>)
  @date 2016-06-10 13:00
  @copyright <NAME> 2016 on MIT License
*/

#ifndef COFFEE_MILL_PDB_READER_HPP
#define COFFEE_MILL_PDB_READER_HPP
#include <mill/util/logger.hpp>
#include "PDBChain.hpp"
#include <fstream>
#include <sstream>

namespace mill
{

//! PDBReader class
/*!
 * read atoms in pdb file until "ENDMDL" or "END".
 * and then parse the PDBAtoms into vector of PDBChains.
 * for movie file, parse only one model at once. you can specify the model
 * by model index.
 */
template<typename vectorT>
class PDBReader
{
  public:

    using vector_type  = vectorT;
    using atom_type    = PDBAtom<vector_type>;
    using residue_type = PDBResidue<vector_type>;
    using chain_type   = PDBChain<vector_type>;

  public:

    PDBReader() = default;
    ~PDBReader() = default;

    std::vector<chain_type> parse(const std::vector<atom_type>& atoms) const;

    std::vector<atom_type> read(const std::string& fname) const;
    std::vector<atom_type> read(const std::string& fname, const std::size_t model_idx) const;
    std::vector<atom_type> read(std::istream& fname) const;
    std::vector<atom_type> read(std::istream& fname, const std::size_t model_idx) const;
};

template<typename vectorT>
std::vector<typename PDBReader<vectorT>::atom_type>
PDBReader<vectorT>::read(const std::string& fname) const
{
    log::debug("reading PDB file ", fname);
    std::ifstream ifs(fname);
    if(not ifs.good()) {log::fatal("PDBReader: file open error: ", fname);}
    const auto data = this->read(ifs);
    return data;
}

template<typename vectorT>
std::vector<typename PDBReader<vectorT>::atom_type>
PDBReader<vectorT>::read(const std::string& fname, const std::size_t model_idx) const
{
    log::debug("reading PDB file ", fname, " model ", model_idx);
    std::ifstream ifs(fname);
    if(not ifs.good()) {log::fatal("PDBReader: file open error: ", fname);}
    const auto data = this->read(ifs, model_idx);
    return data;
}

template<typename vectorT>
std::vector<typename PDBReader<vectorT>::atom_type>
PDBReader<vectorT>::read(std::istream& ifs, const std::size_t model_idx) const
{
    while(not ifs.eof())
    {
        std::string line;
        std::getline(ifs, line);
        if(line.empty()) {continue;}
        if(line.substr(0, 5) == "MODEL")
        {
            std::size_t index;
            std::istringstream iss(line);
            std::string model;
            iss >> model >> index;
            log::debug("PDB model ", index, " found");
            if(index == model_idx) {break;}
        }
    }
    if(ifs.eof())
    {
        log::fatal("PDBReader: no model #", model_idx);
    }
    return this->read(ifs);
}

template<typename vectorT>
std::vector<typename PDBReader<vectorT>::atom_type>
PDBReader<vectorT>::read(std::istream& ifs) const
{
    std::vector<atom_type> atoms;
    while(not ifs.eof())
    {
        std::string line;
        std::getline(ifs, line);
        if(line.empty()) {continue;}
        if(line.substr(0, 3) == "END") {break;}

        atom_type atom;
        if(line >> atom)
        {
            log::debug("PDB reading `", line, "`");
            atoms.emplace_back(std::move(atom));
        }
    }
    return atoms;
}

template<typename vectorT>
std::vector<typename PDBReader<vectorT>::chain_type>
PDBReader<vectorT>::parse(const std::vector<atom_type>& atoms) const
{
    log::debug("PDB parsing atoms ...");
    std::vector<residue_type> residues;
    residue_type tmp_residue;
    int current_residue_id = std::numeric_limits<int>::min();
    for(const auto& atom : atoms)
    {
        if(atom.residue_id != current_residue_id)
        {
            if(not tmp_residue.empty())
            {
                residues.push_back(tmp_residue);
            }
            tmp_residue.clear();
            current_residue_id = atom.residue_id;
            log::debug("residue ", current_residue_id, " found.");
        }
        tmp_residue.push_back(atom);
    }
    if(not tmp_residue.empty())
    {
        residues.push_back(tmp_residue);
    }
    log::debug(residues.size(), " residues found.");

    std::vector<chain_type> chains;
    chain_type tmp_chain;
    std::string current_chain_id = "";
    for(const auto& residue : residues)
    {
        if(residue.chain_id() != current_chain_id)
        {
            if(not tmp_chain.empty())
            {
                chains.push_back(tmp_chain);
            }
            tmp_chain.clear();
            current_chain_id = residue.chain_id();
            log::debug("chain ", current_chain_id, " found.");
        }
        tmp_chain.push_back(residue);
    }
    if(not tmp_chain.empty())
    {
        chains.push_back(tmp_chain);
    }
    log::debug(chains.size(), " chains found.");
    return chains;
}

} // mill
#endif//COFFEE_MILL_PURE_PDB_READER<commit_msg>:recycle: use DeferedReader in PDBReader<commit_after>#ifndef COFFEE_MILL_PDB_READER_HPP
#define COFFEE_MILL_PDB_READER_HPP
#include <mill/util/logger.hpp>
#include <mill/common/Trajectory.hpp>
#include <mill/common/DeferedReader.hpp>
#include <fstream>
#include <sstream>

namespace mill
{

class PDBReader final : public DeferedReaderBase
{
  public:
    using base_type = DeferedReaderBase;
    using trajectory_type          = base_type::trajectory_type         ;
    using snapshot_type            = base_type::snapshot_type           ;
    using particle_type            = base_type::particle_type           ;
    using attribute_container_type = base_type::attribute_container_type;

  public:

    explicit PDBReader(const std::string_view fname)
        : current_(0), file_name_(fname), pdb_(file_name_)
    {
        if(!pdb_.good())
        {
            log::fatal("PDBReader: file open error: ", fname);
        }
    }
    ~PDBReader() override = default;

    attribute_container_type read_header() override
    {
        log::debug("mill::PDBReader: reading header ...");
        //
        // TODO: REMARKs?
        std::string line;
        auto checkpoint = this->pdb_.tellg();
        while(line.substr(0, 6) != "MODEL " &&
              line.substr(0, 6) != "ATOM  " &&
              line.substr(0, 6) != "HETATM")
        {
            checkpoint = this->pdb_.tellg();
            std::getline(this->pdb_, line);
            log::debug("skipping ", line);
        }
        this->pdb_.seekg(checkpoint);
        log::debug("mill::PDBReader: the first line is ", line);
        log::debug("mill::PDBReader: done.");
        return attribute_container_type{};
    }
    trajectory_type read() override
    {
        log::debug("mill::PDBReader: reading the whole trajectory...");
        this->rewind();

        trajectory_type traj(this->read_header());
        while(!this->pdb_.eof())
        {
            if(auto frame = read_frame())
            {
                traj.snapshots().push_back(std::move(*frame));
                this->pdb_.peek();
            }
            else
            {
                log::debug("mill::PDBReader: frame read failed");
                break;
            }
        }
        log::debug("mill::PDBReader: done.");
        return traj;
    }

    std::optional<snapshot_type> read_frame() override
    {
        log::debug("mill::PDBReader: reading the next snapshot...");
        if(this->is_eof())
        {
            return std::nullopt;
        }
        snapshot_type frame;

        std::string line; // buffer for getline
        while(line.substr(0, 6) != "ENDMDL" && not this->is_eof())
        {
            std::getline(this->pdb_, line);

            const auto header = line.substr(0, 6);
            if(header == "ATOM  " || header == "HETATM")
            {
                Particle p;
                // according to the wwPDB field name
                p["record"  ]   = header;
                p["serial"  ]   = std::stoi(line.substr(6, 5));
                p["name"    ]   = line.substr(12, 4);
                p["alt_loc" ]   = line.substr(16, 1);
                p["res_name"]   = line.substr(17, 3);
                p["chain_id"]   = line.substr(21, 1);
                p["res_seq" ]   = std::stoi(line.substr(22, 4));
                p["i_code"  ]   = line.substr(26, 1);
                p.position()[0] = std::stod(line.substr(30, 8));
                p.position()[1] = std::stod(line.substr(38, 8));
                p.position()[2] = std::stod(line.substr(46, 8));

                if(60 <= line.size())
                {
                    p["occupancy"] = std::stod(line.substr(54, 6));
                }
                else
                {
                    p["occupancy"] = 0.0;
                }

                if(66 <= line.size())
                {
                    p["temp_factor"] = std::stod(line.substr(60, 6));
                }
                else
                {
                    p["temp_factor"] = 99.9;
                }

                if(78 <= line.size())
                {
                    p["element"] = line.substr(76, 2);
                }
                else
                {
                    p["element"] = "  ";
                }

                if(80 <= line.size())
                {
                    p["charge"] = line.substr(78, 2);
                }
                else
                {
                    p["charge"] = "  ";
                }
                frame.particles().push_back(std::move(p));
            }
            else if(header == "TER   ")
            {
                log::debug("PDBReader: `TER` found. Skipped.");
            }
            else if(header == "ANISOU")
            {
                log::debug("PDBReader: `ANISOU` is currently not supported.");
            }
        }

        current_ += 1;
        pdb_.peek();
        log::debug("mill::PDBReader: done.");
        return frame;
    }

    std::optional<snapshot_type> read_frame(const std::size_t idx) override
    {
        log::debug("mill::PDBReader: reading the ", idx, "-th snapshot");
        this->rewind();

        // skip n snapshots
        log::debug("mill::PDBReader: skipping ", idx, " snapshots...");
        std::string line;
        for(std::size_t i=0; i<idx; ++i)
        {
            if(const auto frame = this->read_frame(); not frame.has_value())
            {
                log::error("mill::PDBReader: ", idx, "-th snapshot does not exist");
                return std::nullopt;
            }
            pdb_.peek();
        }
        log::debug("mill::PDBReader: done.");
        return this->read_frame();
    }

    void rewind() override
    {
        log::debug("mill::PDBReader: rewinding the file");
        current_ = 0;
        pdb_.clear(); // clear failbit flags like the EOF flag
        pdb_.seekg(0, std::ios_base::beg);
        pdb_.peek();
        log::debug("mill::PDBReader: done.");
    }

    bool             is_eof()    const noexcept override {return pdb_.eof();}
    std::size_t      current()   const noexcept override {return current_;}
    std::string_view file_name() const noexcept override {return file_name_;}

  private:
    std::size_t   current_;
    std::string   file_name_;
    std::ifstream pdb_;
};

} // mill
#endif//COFFEE_MILL_PURE_PDB_READER
{'fraction_non_alphanumeric': 0.14225015490838275, 'fraction_numerical': 0.008143754979198017, 'mean_word_length': 2.0232807064490235, 'pattern_counts': {'":': 0, '<': 56, '<?xml version=': 0, '>': 74, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '792943', 'n_tokens_mistral': 3700, 'n_tokens_neox': 3515, 'n_words': 686}