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
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
//! # PSEData
//!
//! PSEData is a struct for loading and serializing pymol PSE data.
//!
//! Currently, the parsers are working for small test cases of molecules and selections. Additional parser structs would be required for
//! other PSE data types which include the following:
//!

use crate::pymolparsing::parsing::{
    PyObjectMolecule, PymolSessionObjectData, SceneView, SessionName, SessionSelectorList,
    Settings, SettingsEnum,
};
use ferritin_molviewspec::molviewspec::nodes::{self as mvsnodes, ColorNamesT, State};
use pdbtbx::PDB;
use serde::{Deserialize, Serialize};
use serde_bytes;
use serde_pickle::de::{from_reader, DeOptions};
use std::collections::HashMap;
use std::fs;
use std::fs::File;
use std::io::Read;
use std::path::Path;

/// PSEData represents the structure of a PyMOL Session File (PSE).
///
/// We lean heavily on `serde_pickle` to deserialize the PSE binary
/// file into named structs, of which `PSEData` is the highest level object containing
/// most of the required methods for operating on PSE files.
///
///  This struct contains various components of a PyMOL session, including:
/// - Version information
/// - Color settings
/// - View settings
/// - Movie scenes
/// - Custom settings
/// - Cached data
/// - Session names and associated objects
///
/// To be implemented in time:
///
/// PyObject :
/// -[pyobject](https://github.com/schrodinger/pymol-open-source/blob/03d7a7fcf0bd95cd93d710a1268dbace2ed77765/layer1/PyMOLObject.cpp#L681)
///
/// Python Obects:
/// - Object
/// - Gadget
/// - Molecule    ---> WIP.
/// - Dist
/// - Map
/// - Mesh
/// - Slice
/// - Surface
/// - CGO
/// - Alignment
/// - Group
/// - Volume
/// - Callback
/// - Curve
/// - Selection ---> WIP.
///
#[derive(Debug, Deserialize, Serialize)]
pub struct PSEData {
    pub version: i32,
    pub main: Vec<i64>,
    pub colors: Vec<i32>,
    pub color_ext: Vec<i32>,
    pub unique_settings: Vec<i32>,
    pub selector_secrets: Vec<i32>,
    pub editor: Vec<i32>,
    pub view: SceneView,
    pub view_dict: HashMap<String, String>,
    #[serde(with = "serde_bytes")]
    pub wizard: Vec<u8>,
    pub moviescenes: Vec<Vec<i32>>,
    // High level state settings; we need to prpogate these.
    pub settings: Vec<Settings>,
    pub movie: (
        i32,
        i32,
        Vec<f32>,
        i32,
        Option<bool>, // this will probably need to be modified
        Option<bool>,
        Option<bool>,
    ),
    // not needed?
    // session: HashMap<String, Value>,
    pub cache: Vec<usize>,
    // name is the trickiest bit
    pub names: Vec<Option<SessionName>>,
}

impl PSEData {
    /// Load pymole `.pse` binary file to struct.
    pub fn load(file_path: &str) -> Result<Self, Box<dyn std::error::Error>> {
        let mut file = File::open(file_path)?;
        let mut buffer = Vec::new();
        file.read_to_end(&mut buffer)?;
        let options = DeOptions::new()
            .replace_unresolved_globals()
            .decode_strings();
        let pse_data_vals: serde_pickle::Value = from_reader(&buffer[..], options).unwrap();
        let pse_data: PSEData = serde_pickle::from_value(pse_data_vals).unwrap();
        Ok(pse_data)
    }
    /// Serialize to JSON
    pub fn to_json(&self, file_path: &str) -> Result<(), Box<dyn std::error::Error>> {
        let json = serde_json::to_string_pretty(self)?;
        fs::write(file_path, json)?;
        Ok(())
    }
    // adds custom colors to auto colors to get the index of colors
    pub fn get_full_colorlist() {
        // https://github.com/schrodinger/pymol-open-source/blob/master/layer1/Color.cpp#L415
        unimplemented!()
    }

    /// session is where all the action happens
    pub fn get_session_names(&self) -> Vec<String> {
        self.names
            .iter()
            .filter_map(|session_name| {
                session_name
                    .as_ref()
                    .map(|session| session.name.to_string())
            })
            .collect()
    }
    /// Global Pymol Settings
    pub fn get_setting(&self, setting: SettingsEnum) -> Option<Settings> {
        self.settings.iter().find(|s| s.setting == setting).cloned()
    }
    // pub fn get_view(&self) -> &Vec<f32> {
    //     let view = self.view; // f32: 25
    // }
    pub fn get_molecule_data(&self) -> Vec<&PyObjectMolecule> {
        self.names
            .iter()
            .filter_map(|session_name| session_name.as_ref())
            .filter_map(|session| match &session.data {
                PymolSessionObjectData::PyObjectMolecule(a) => Some(a),
                _ => None,
            })
            .collect()
    }

    pub fn get_selection_data(&self) -> Vec<&SessionSelectorList> {
        self.names
            .iter()
            .filter_map(|session_name| session_name.as_ref())
            .filter_map(|session| match &session.data {
                PymolSessionObjectData::SessionSelectorList(a) => Some(a),
                _ => None,
            })
            .collect()
    }

    pub fn get_location_as_transform(&self) -> mvsnodes::TransformParams {
        let location = self.view.get_location();
        mvsnodes::TransformParams {
            rotation: Some(location.to_vec()),
            ..Default::default()
        }
    }
    /// Convert Pymol to pdbtbx::PDB
    pub fn create_pdb(&self) -> PDB {
        // todo: extend this to more than one molecule and/or to modify the global scene
        let moldata = &self.get_molecule_data();
        let first_mol = moldata[0];
        first_mol.to_pdb()
    }
    /// Pymol --> PDBTBX --> PDB --> file
    pub fn save_pdbs(&self, file_path: &str) -> std::io::Result<()> {
        let path = Path::new(file_path);
        let pdb_folder = path.join("pdb");
        fs::create_dir_all(&pdb_folder)?;
        let mut file_list = Vec::new();
        for molecule in self.get_molecule_data().iter() {
            let pdb = molecule.to_pdb();
            let filename = format!("{}.pdb", molecule.get_name());
            let file_path = pdb_folder.join(&filename);
            let _ = pdbtbx::save_pdb(
                &pdb,
                file_path.to_str().expect("Invalid UTF-8 in file path"),
                pdbtbx::StrictnessLevel::Strict,
            );
            file_list.push(filename);
        }
        let contents = file_list.join("\n");
        fs::write(path.join("pdb_contents.txt"), contents)?;
        Ok(())
    }
    /// Convert to MSVJ Formay.
    /// see also: [mol-view-spec](https://github.com/molstar/mol-view-spec)
    pub fn create_molviewspec(&self) -> State {
        // write state for loading the PDB files
        let mut state = State::new();

        // Add Global Camera Data
        // let [o1, o2, o3] = self.view.origin;
        // let [p1, p2, p3] = self.view.position;
        //
        // let [o1, o2, o3] = self.view.get_translated_origin();
        // let [p1, p2, p3] = self.view.get_translated_position();
        //
        // let camparam = CameraParams {
        //     // https://molstar.org/mol-view-spec-docs/camera-settings/
        //     target: (o1, o2, o3), // <--- Todo
        //     position: (p1, p2, p3),
        //     ..Default::default() // <--- Todo
        // };
        // state.camera(camparam);

        // It will be easier to set the focus based on all the components in the PDB then trying to match pymol exactly
        // let focus = FocusInlineParams {};

        // Add Molecule Data
        for molecule in self.get_molecule_data() {
            let molname = molecule.get_name();

            let structure = state
                .download(&format!("pdb/{}.pdb", molname))
                .expect("Create a Download node with a URL")
                .parse(mvsnodes::ParseParams {
                    format: mvsnodes::ParseFormatT::Pdb,
                })
                .expect("Parseable option")
                .assembly_structure(mvsnodes::StructureParams {
                    structure_type: mvsnodes::StructureTypeT::Model,
                    ..Default::default()
                })
                .expect("a set of Structure options");

            // add base structure component then any selections that may be relevant
            structure
                .component(mvsnodes::ComponentSelector::default())
                .expect("defined a valid component")
                .representation(mvsnodes::RepresentationTypeT::Cartoon);

            // this works great
            let transform = self.get_location_as_transform();
            structure.transform(transform);

            // selections return MVD ComponentExpression
            let selection_data = self.get_selection_data()[0];
            for selector in selection_data.get_selectors() {
                if selector.id == molname {
                    println!("Found a selction for Model {}!!!!", molname);
                    let component = selector.to_component();
                    structure
                        .component(component)
                        .expect("defined a valid component")
                        .representation(mvsnodes::RepresentationTypeT::BallAndStick)
                        .expect("a representation")
                        // to do add colors and other settings....
                        .color(
                            mvsnodes::ColorT::Named(ColorNamesT::Magenta),
                            mvsnodes::ComponentSelector::default(),
                        )
                        .expect("a color");
                }
            }
        }

        state
    }
    ///  Pymol --> MSVJ --? disk
    pub fn to_disk(&self, file_path: &str) -> std::io::Result<()> {
        let path = Path::new(file_path);
        let msvj_file = path.join("state.mvsj");
        let state = self.create_molviewspec();
        let pretty_json = serde_json::to_string_pretty(&state)?;
        self.save_pdbs(file_path)?;
        fs::write(msvj_file, pretty_json)?;
        Ok(())
    }
    /// this one will write  a ready-to-go folder with pdbs, a msvj file, and the
    /// html/css/js needed to load them
    pub fn to_disk_full(&self, file_path: &str) -> std::io::Result<()> {
        // Create the directory if it doesn't exist
        fs::create_dir_all(file_path)?;

        // Copy our standard files
        let resources_dir = Path::new("resources");
        let files_to_copy = ["index.html", "molstar.css", "molstar.js"];
        for file_name in files_to_copy.iter() {
            let src_path = resources_dir.join(file_name);
            let dest_path = Path::new(file_path).join(file_name);
            fs::copy(&src_path, &dest_path)?;
        }
        // copy our custom files
        let _ = self.to_disk(file_path);
        Ok(())
    }
    /// To loadable MSVJ URL
    /// See also [msvj URL encoding](https://molstar.org/mol-view-spec-docs/mvs-molstar-extension/)
    pub fn to_mvsj_url(&self) -> String {
        let state = self.create_molviewspec();
        state.to_url()
    }
}