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
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
//! This module provides a bitmask representation for various types of 3D molecular visualizations.
//!
//! Note: this should probably be moved to its own crate.
//!
//! It defines a set of flags that can be combined to represent different visual elements
//! in a molecular structure, such as cylinders, spheres, surfaces, and more.
//!
//! The `RepBitmask` struct allows for efficient storage and manipulation of these flags,
//! while the `RepType` enum provides a way to refer to individual representation types.
//!
//! # Examples
//!
//! ```
//! use ferritin_pymol::pymolparsing::representation::RepBitmask;
//!
//! let mut reps = RepBitmask::new();
//! reps.insert(RepBitmask::CYL | RepBitmask::SPHERE);
//!
//! assert!(reps.contains(RepBitmask::CYL));
//! assert!(reps.contains(RepBitmask::SPHERE));
//! ```
use bitflags::bitflags;
use serde::{Deserialize, Serialize};

/// First, define an enum for all representation types
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum RepType {
    Cyl,
    Sphere,
    Surface,
    Label,
    NonbondedSphere,
    Cartoon,
    Ribbon,
    Line,
    Mesh,
    Dot,
    Dash,
    Nonbonded,
    Cell,
    CGO,
    Callback,
    Extent,
    Slice,
    Angle,
    Dihedral,
    Ellipsoid,
    Volume,
}

bitflags! {
     #[derive(Debug)]
    pub struct RepBitmask: u32 {
        const CYL             = 1 << 0;
        const SPHERE          = 1 << 1;
        const SURFACE         = 1 << 2;
        const LABEL           = 1 << 3;
        const NONBONDED_SPHERE = 1 << 4;
        const CARTOON         = 1 << 5;
        const RIBBON          = 1 << 6;
        const LINE            = 1 << 7;
        const MESH            = 1 << 8;
        const DOT             = 1 << 9;
        const DASH            = 1 << 10;
        const NONBONDED       = 1 << 11;
        const CELL            = 1 << 12;
        const CGO             = 1 << 13;
        const CALLBACK        = 1 << 14;
        const EXTENT          = 1 << 15;
        const SLICE           = 1 << 16;
        const ANGLE           = 1 << 17;
        const DIHEDRAL        = 1 << 18;
        const ELLIPSOID       = 1 << 19;
        const VOLUME          = 1 << 20;

        const REPS_ATOM_MASK = Self::CYL.bits() | Self::SPHERE.bits() | Self::SURFACE.bits() |
            Self::LABEL.bits() | Self::NONBONDED_SPHERE.bits() | Self::CARTOON.bits() | Self::RIBBON.bits() |
            Self::LINE.bits() | Self::MESH.bits() | Self::DOT.bits() | Self::NONBONDED.bits() | Self::ELLIPSOID.bits();

        const REPS_OBJECT_MASK = Self::SURFACE.bits() | Self::MESH.bits() | Self::DOT.bits() |
            Self::CELL.bits() | Self::CGO.bits() | Self::CALLBACK.bits() | Self::EXTENT.bits() | Self::SLICE.bits() |
            Self::ANGLE.bits() | Self::DIHEDRAL.bits() | Self::VOLUME.bits() | Self::DASH.bits();
    }
}

// You can also add methods to RepBitmask if needed
impl RepBitmask {
    pub fn new() -> Self {
        RepBitmask::empty()
    }
    // Function to convert an integer to RepBitmask
    pub fn from_int(value: u32) -> RepBitmask {
        RepBitmask::from_bits_truncate(value)
    }
}

// Custom Serde deserialization
impl<'de> Deserialize<'de> for RepBitmask {
    fn deserialize<D>(deserializer: D) -> Result<Self, D::Error>
    where
        D: serde::Deserializer<'de>,
    {
        let bits = u32::deserialize(deserializer)?;
        Ok(RepBitmask::from_bits_truncate(bits))
    }
}

// Custom Serde serialization
impl Serialize for RepBitmask {
    fn serialize<S>(&self, serializer: S) -> Result<S::Ok, S::Error>
    where
        S: serde::Serializer,
    {
        serializer.serialize_u32(self.bits())
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn test_rep_bitmask() {
        let mut reps = RepBitmask::new();
        assert!(reps.is_empty());

        reps.insert(RepBitmask::CYL | RepBitmask::SPHERE);
        assert!(reps.contains(RepBitmask::CYL));
        assert!(reps.contains(RepBitmask::SPHERE));
        assert!(!reps.contains(RepBitmask::SURFACE));

        reps.remove(RepBitmask::CYL);
        assert!(!reps.contains(RepBitmask::CYL));
        assert!(reps.contains(RepBitmask::SPHERE));

        assert!(RepBitmask::REPS_ATOM_MASK.contains(RepBitmask::SPHERE));
        assert!(RepBitmask::REPS_OBJECT_MASK.contains(RepBitmask::SURFACE));
    }

    #[test]
    fn test_bitmask_from_int() {
        let int_value = 5; // This represents CYL | SURFACE (1 | 4)
        let bitmask = RepBitmask::from_int(int_value);

        // Check that the correct flags are set
        assert!(bitmask.contains(RepBitmask::CYL));
        assert!(bitmask.contains(RepBitmask::SURFACE));
    }

    #[test]
    fn test_bitmask_from_int_02() {
        // I see this value in the example.pse
        // Binary: 111110110001111111
        let int_value = 2060287;
        let bitmask = RepBitmask::from_int(int_value);

        // Check that the correct flags are set
        assert!(bitmask.contains(RepBitmask::CYL));
        assert!(bitmask.contains(RepBitmask::SURFACE));

        assert!(bitmask.contains(RepBitmask::SPHERE));
        assert!(bitmask.contains(RepBitmask::SURFACE));
        assert!(bitmask.contains(RepBitmask::LABEL));
        assert!(bitmask.contains(RepBitmask::NONBONDED_SPHERE));
        assert!(bitmask.contains(RepBitmask::CARTOON));
        assert!(bitmask.contains(RepBitmask::RIBBON));
        assert!(bitmask.contains(RepBitmask::LINE));
        assert!(bitmask.contains(RepBitmask::MESH));
        assert!(bitmask.contains(RepBitmask::DOT));
        assert!(bitmask.contains(RepBitmask::DASH));
        assert!(bitmask.contains(RepBitmask::NONBONDED));
        assert!(bitmask.contains(RepBitmask::CGO));
        assert!(bitmask.contains(RepBitmask::CALLBACK));
        assert!(bitmask.contains(RepBitmask::SLICE));
        assert!(bitmask.contains(RepBitmask::ANGLE));
        assert!(bitmask.contains(RepBitmask::DIHEDRAL));
        assert!(bitmask.contains(RepBitmask::ELLIPSOID));
        assert!(bitmask.contains(RepBitmask::VOLUME));

        assert!(!bitmask.contains(RepBitmask::EXTENT));
        assert_eq!(bitmask.bits(), 2060287);
    }
}