use std::collections::BTreeMap;
use crate::{Atom, AtomId, BondOrder, Conformer3D, Molecule};
fn get_default_atom_number(atomic_number: u8, elem: &mut BTreeMap<u8, u32>) -> String {
let mut ret = [' ', ' '];
let entry = elem.entry(atomic_number).or_insert(0);
*entry += 1;
let tmp = *entry;
if tmp < 10 {
ret[0] = char::from_u32(u32::from(b'0') + tmp).unwrap();
} else if tmp < 100 {
ret[0] = char::from_u32(u32::from(b'0') + tmp / 10).unwrap();
ret[1] = char::from_u32(u32::from(b'0') + tmp % 10).unwrap();
} else if tmp < 360 {
ret[0] = char::from_u32(u32::from(b'A') + (tmp - 100) / 10).unwrap();
ret[1] = char::from_u32(u32::from(b'0') + (tmp - 100) % 10).unwrap();
} else if tmp < 1036 {
ret[0] = char::from_u32(u32::from(b'A') + (tmp - 360) / 26).unwrap();
ret[1] = char::from_u32(u32::from(b'A') + (tmp - 360) % 26).unwrap();
}
ret.iter().collect()
}
fn element_symbol(atomic_num: u8) -> &'static str {
match atomic_num {
0 => "*",
1 => "H",
2 => "He",
3 => "Li",
4 => "Be",
5 => "B",
6 => "C",
7 => "N",
8 => "O",
9 => "F",
10 => "Ne",
11 => "Na",
12 => "Mg",
13 => "Al",
14 => "Si",
15 => "P",
16 => "S",
17 => "Cl",
18 => "Ar",
19 => "K",
20 => "Ca",
21 => "Sc",
22 => "Ti",
23 => "V",
24 => "Cr",
25 => "Mn",
26 => "Fe",
27 => "Co",
28 => "Ni",
29 => "Cu",
30 => "Zn",
31 => "Ga",
32 => "Ge",
33 => "As",
34 => "Se",
35 => "Br",
36 => "Kr",
37 => "Rb",
38 => "Sr",
39 => "Y",
40 => "Zr",
41 => "Nb",
42 => "Mo",
43 => "Tc",
44 => "Ru",
45 => "Rh",
46 => "Pd",
47 => "Ag",
48 => "Cd",
49 => "In",
50 => "Sn",
51 => "Sb",
52 => "Te",
53 => "I",
54 => "Xe",
55 => "Cs",
56 => "Ba",
57 => "La",
58 => "Ce",
59 => "Pr",
60 => "Nd",
61 => "Pm",
62 => "Sm",
63 => "Eu",
64 => "Gd",
65 => "Tb",
66 => "Dy",
67 => "Ho",
68 => "Er",
69 => "Tm",
70 => "Yb",
71 => "Lu",
72 => "Hf",
73 => "Ta",
74 => "W",
75 => "Re",
76 => "Os",
77 => "Ir",
78 => "Pt",
79 => "Au",
80 => "Hg",
81 => "Tl",
82 => "Pb",
83 => "Bi",
84 => "Po",
85 => "At",
86 => "Rn",
87 => "Fr",
88 => "Ra",
89 => "Ac",
90 => "Th",
91 => "Pa",
92 => "U",
93 => "Np",
94 => "Pu",
95 => "Am",
96 => "Cm",
97 => "Bk",
98 => "Cf",
99 => "Es",
100 => "Fm",
101 => "Md",
102 => "No",
103 => "Lr",
104 => "Rf",
105 => "Db",
106 => "Sg",
107 => "Bh",
108 => "Hs",
109 => "Mt",
110 => "Ds",
111 => "Rg",
112 => "Cn",
113 => "Nh",
114 => "Fl",
115 => "Mc",
116 => "Lv",
117 => "Ts",
118 => "Og",
_ => "?",
}
}
fn get_pdb_atom_line(
atom: &Atom,
conf: Option<&Conformer3D>,
elem: &mut BTreeMap<u8, u32>,
) -> String {
let symb = element_symbol(atom.atomic_number());
let (at1, at2): (char, char) = match symb.len() {
0 => (' ', 'X'),
1 => (' ', symb.chars().next().unwrap()),
_ => {
let chars: Vec<char> = symb.chars().collect();
let first = chars[0];
let mut second = chars[1];
if second.is_ascii_lowercase() {
second = second.to_ascii_uppercase();
}
(first, second)
}
};
let info = atom.pdb_residue_info();
let mut line = String::with_capacity(80);
if let Some(pdb_info) = info {
if pdb_info.is_hetero_atom() {
line.push_str("HETATM");
} else {
line.push_str("ATOM ");
}
line.push_str(&format!("{:>5}", atom.id().index() + 1));
line.push(' ');
let name = pdb_info.atom_name();
if name.is_empty() {
let atnum = get_default_atom_number(atom.atomic_number(), elem);
line.push(at1);
line.push(at2);
line.push_str(&atnum);
} else {
let truncated: String = name.chars().take(4).collect();
line.push_str(&format!("{:>4}", truncated));
}
let alt = pdb_info.alt_loc();
line.push(alt.chars().next().unwrap_or(' '));
let resname: String = pdb_info.residue_name().chars().take(3).collect();
line.push_str(&format!("{:>3}", resname));
line.push(' ');
let chain = pdb_info.chain_id();
line.push(chain.chars().next().unwrap_or(' '));
line.push_str(&format!("{:>4}", pdb_info.residue_number()));
let icode = pdb_info.insertion_code();
line.push(icode.chars().next().unwrap_or(' '));
line.push_str(" ");
} else {
let atnum = get_default_atom_number(atom.atomic_number(), elem);
line.push_str("HETATM");
line.push_str(&format!("{:>5}", atom.id().index() + 1));
line.push(' ');
line.push(at1);
line.push(at2);
line.push_str(&atnum);
line.push_str(" UNL 1 ");
}
if let Some(conf3d) = conf {
let coords = conf3d.coords();
if atom.id().index() < coords.len() {
let pos = coords[atom.id().index()];
line.push_str(&format!("{:+8.3}{:+8.3}{:+8.3}", pos[0], pos[1], pos[2]));
} else {
line.push_str(" 0.000 0.000 0.000");
}
} else {
line.push_str(" 0.000 0.000 0.000");
}
if let Some(pdb_info) = info {
line.push_str(&format!(
"{:>6.2}{:>6.2}",
pdb_info.occupancy(),
pdb_info.temp_factor()
));
line.push_str(" ");
} else {
line.push_str(" 1.00 0.00 ");
}
line.push(at1);
line.push(at2);
let charge = atom.formal_charge();
if charge > 0 && charge < 10 {
line.push(char::from_u32(u32::from(b'0') + charge as u32).unwrap());
line.push('+');
} else if charge < 0 && charge > -10 {
line.push(char::from_u32(u32::from(b'0') + (-charge) as u32).unwrap());
line.push('-');
} else {
line.push(' ');
line.push(' ');
}
line
}
fn get_pdb_bond_lines(
atom_idx: AtomId,
bonds: &[crate::Bond],
all: bool,
both: bool,
mult: bool,
conect_count: &mut u32,
) -> String {
let src = atom_idx.index() + 1;
let mut v: Vec<usize> = Vec::new();
for bond in bonds {
let begin = bond.begin();
let end = bond.end();
let other_idx = if begin == atom_idx {
end
} else if end == atom_idx {
begin
} else {
continue;
};
let dst = other_idx.index() + 1;
if dst < src && !both {
continue;
}
let btype = if mult {
bond.order()
} else {
BondOrder::Single
};
match btype {
BondOrder::Single | BondOrder::Aromatic => {
if all {
v.push(dst);
}
}
BondOrder::Quadruple => {
v.push(dst);
v.push(dst);
v.push(dst);
v.push(dst);
}
BondOrder::Triple => {
v.push(dst);
v.push(dst);
v.push(dst);
}
BondOrder::Double => {
v.push(dst);
v.push(dst);
}
_ => {}
}
}
if v.is_empty() {
return String::new();
}
v.sort_unstable();
let mut result = String::new();
for (i, dst) in v.iter().enumerate() {
if (i & 3) == 0 {
if i != 0 {
result.push('\n');
}
result.push_str(&format!("CONECT{:>5}", src));
*conect_count += 1;
}
result.push_str(&format!("{:>5}", dst));
}
result.push('\n');
result
}
fn mol_to_pdb_body(
mol: &Molecule,
conf: Option<&Conformer3D>,
flavor: u32,
atm_count: &mut u32,
ter_count: &mut u32,
conect_count: &mut u32,
) -> String {
let mut res = String::new();
let mut last = String::new();
let mut elem: BTreeMap<u8, u32> = BTreeMap::new();
for atom in mol.atoms() {
last = get_pdb_atom_line(atom, conf, &mut elem);
res.push_str(&last);
res.push('\n');
*atm_count += 1;
}
if *ter_count == 0 && *atm_count > 0 && (flavor & 32) != 0 {
let mut ter_line = String::with_capacity(80);
ter_line.push_str("TER ");
ter_line.push_str(&format!("{:>5}", *atm_count + 1));
if last.len() >= 27 {
let end = (17 + 10).min(last.len());
ter_line.push_str(" ");
ter_line.push_str(&last[17..end]);
}
ter_line.push('\n');
res.push_str(&ter_line);
*ter_count = 1;
}
let all = (flavor & 2) == 0;
let both = (flavor & 4) != 0;
let mult = (flavor & 8) == 0;
if all || mult {
for atom in mol.atoms() {
res.push_str(&get_pdb_bond_lines(
atom.id(),
mol.bonds(),
all,
both,
mult,
conect_count,
));
}
}
res
}
pub fn mol_to_pdb_block(mol: &Molecule, conf_id: i32, flavor: u32) -> String {
let mut res = String::new();
if let Some(name) = mol.properties().name() {
if !name.is_empty() {
res.push_str("COMPND ");
res.push_str(name);
res.push('\n');
}
}
let mut atm_count: u32 = 0;
let mut ter_count: u32 = 0;
let mut conect_count: u32 = 0;
let conformers = mol.conformers_3d();
if conf_id < 0 && conformers.len() > 1 {
for (idx, conf) in conformers.iter().enumerate() {
res.push_str(&format!("MODEL {:>4}\n", idx + 1));
res.push_str(&mol_to_pdb_body(
mol,
Some(conf),
flavor,
&mut atm_count,
&mut ter_count,
&mut conect_count,
));
res.push_str("ENDMDL\n");
}
} else {
let conf = if conf_id < 0 {
conformers.first()
} else {
conformers.get(conf_id as usize)
};
res.push_str(&mol_to_pdb_body(
mol,
conf,
flavor,
&mut atm_count,
&mut ter_count,
&mut conect_count,
));
}
if (flavor & 16) != 0 {
res.push_str(&format!(
"MASTER 0 0 0 0 0 0 0 0{:>5}{:>5}{:>5} 0\n",
atm_count, ter_count, conect_count
));
}
res.push_str("END\n");
res
}
#[cfg(test)]
mod tests {
use super::*;
use crate::Molecule;
#[test]
fn test_element_symbol() {
assert_eq!(super::element_symbol(0), "*");
assert_eq!(super::element_symbol(1), "H");
assert_eq!(super::element_symbol(6), "C");
assert_eq!(super::element_symbol(8), "O");
assert_eq!(super::element_symbol(118), "Og");
assert_eq!(super::element_symbol(200), "?");
}
#[test]
fn test_get_default_atom_number() {
let mut elem: BTreeMap<u8, u32> = BTreeMap::new();
let n1 = get_default_atom_number(6, &mut elem);
assert_eq!(n1, "1 ");
let n2 = get_default_atom_number(6, &mut elem);
assert_eq!(n2, "2 ");
}
#[test]
fn test_mol_to_pdb_block_empty() {
let mol = Molecule::new();
let result = mol_to_pdb_block(&mol, -1, 0);
assert!(result.ends_with("END\n"));
assert_eq!(result, "END\n");
}
#[test]
fn test_mol_to_pdb_block_simple() {
let mol = Molecule::from_smiles("CO").unwrap();
let result = mol_to_pdb_block(&mol, -1, 0);
assert!(result.contains("HETATM"));
assert!(result.contains("END\n"));
assert!(result.contains("CONECT"));
}
#[test]
fn test_mol_to_pdb_block_with_name() {
let mol = Molecule::from_smiles("CCO").unwrap().with_name("Ethanol");
let result = mol_to_pdb_block(&mol, -1, 0);
assert!(result.contains("COMPND Ethanol"));
}
#[test]
fn test_mol_to_pdb_block_no_conect() {
let mol = Molecule::from_smiles("CO").unwrap();
let result = mol_to_pdb_block(&mol, -1, 2);
assert!(!result.contains("CONECT"));
}
#[test]
fn test_mol_to_pdb_block_master() {
let mol = Molecule::from_smiles("O").unwrap();
let result = mol_to_pdb_block(&mol, -1, 16);
assert!(result.contains("MASTER"));
}
}