1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133

//! Provides functionality for working with residue atoms in molecular structures.
//!
//! This module contains the `ResidueAtoms` struct and its implementation, which provides methods for
//! accessing and manipulating atoms within a single residue of a molecular structure. A residue is
//! a structural unit in a protein or other macromolecule consisting of multiple atoms.
//!
use super::info::AtomInfo;
use crate::info::constants::{is_amino_acid, is_carbohydrate, is_nucleotide};
use crate::selection::{AtomView, Selection};
use crate::AtomCollection;
use pdbtbx::Element;

pub struct ResidueAtoms<'a> {
    pub start_idx: usize,
    pub end_idx: usize,
    pub res_id: i32,
    pub res_name: String,
    pub chain_id: String,
    pub atoms: Selection,
    pub parent: &'a AtomCollection,
}

impl<'a> ResidueAtoms<'a> {
    // Get all atom coordinates for this residue
    pub fn coords(&self) -> Vec<&[f32; 3]> {
        (self.start_idx..self.end_idx)
            .map(|i| self.parent.get_coord(i))
            .collect()
    }

    // Get all atom names for this residue
    pub fn atom_names(&self) -> Vec<&String> {
        (self.start_idx..self.end_idx)
            .map(|i| self.parent.get_atom_name(i))
            .collect()
    }

    // Get all elements for this residue
    pub fn elements(&self) -> Vec<&Element> {
        (self.start_idx..self.end_idx)
            .map(|i| self.parent.get_element(i))
            .collect()
    }

    // Get atom view for this residue
    pub fn view(&self) -> AtomView {
        self.parent.view(self.atoms.clone())
    }

    // Get number of atoms in this residue
    pub fn atom_count(&self) -> usize {
        self.end_idx - self.start_idx
    }

    // Iterator over atoms in this residue
    pub fn iter_atoms(&self) -> impl Iterator<Item = AtomInfo> + '_ {
        (self.start_idx..self.end_idx).map(|i| AtomInfo {
            index: i,
            coords: self.parent.get_coord(i),
            element: self.parent.get_element(i),
            atom_name: self.parent.get_atom_name(i),
            is_hetero: self.parent.get_is_hetero(i),
        })
    }

    // Get specific atom by index within residue
    pub fn get_atom(&self, residue_atom_idx: usize) -> Option<AtomInfo> {
        let abs_idx = self.start_idx + residue_atom_idx;
        if abs_idx < self.end_idx {
            Some(AtomInfo {
                index: abs_idx,
                coords: &self.parent.get_coord(abs_idx),
                element: &self.parent.get_element(abs_idx),
                atom_name: &self.parent.get_atom_name(abs_idx),
                is_hetero: self.parent.get_is_hetero(abs_idx),
            })
        } else {
            None
        }
    }

    // Find atom by name within this residue
    pub fn find_atom_by_name(&self, name: &str) -> Option<AtomInfo> {
        (self.start_idx..self.end_idx)
            .find(|&i| self.parent.get_atom_name(i) == name)
            .map(|i| AtomInfo {
                index: i,
                coords: &self.parent.get_coord(i),
                element: &self.parent.get_element(i),
                atom_name: &self.parent.get_atom_name(i),
                is_hetero: self.parent.get_is_hetero(i),
            })
    }
    pub fn is_amino_acid(&self) -> bool {
        is_amino_acid(&self.res_name)
    }
    pub fn is_carbohydrate(&self) -> bool {
        is_carbohydrate(&self.res_name)
    }
    pub fn is_nucleotide(&self) -> bool {
        is_nucleotide(&self.res_name)
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::AtomCollection;
    use ferritin_test_data::TestFile;

    #[test]
    fn test_iteration() {
        let (prot_file, _temp) = TestFile::protein_01().create_temp().unwrap();
        let (pdb, _) = pdbtbx::open(prot_file).unwrap();
        let ac = AtomCollection::from(&pdb);

        assert_eq!(ac.iter_residues_aminoacid().count(), 154);

        let first_residue: ResidueAtoms = ac
            .iter_residues_aminoacid()
            .take(1)
            .next()
            .expect("Should have at least one amino acid residue");

        assert_eq!(first_residue.start_idx, 0);
        assert_eq!(first_residue.end_idx, 8); // Met has 9 Atoms
        assert_eq!(first_residue.res_id, 0);
        assert_eq!(first_residue.res_name, "MET");
        assert_eq!(first_residue.chain_id, "A");
        // assert_eq!(first_residue.atoms, "A"); // <- should test Selection
        assert_eq!(first_residue.parent.get_size(), 1413);
    }
}