rustytorch/demos/rtx-alphafold-demo/src/pdb_export.rs

//! PDB and mmCIF export functionality.

use alphafold_shared::{AminoAcid, AtomName, ProteinStructure};

/// Export structure to PDB format.
#[must_use]
pub fn to_pdb(structure: &ProteinStructure) -> String {
    let mut pdb = String::new();

    // Header
    pdb.push_str(&format!(
        "HEADER    PREDICTED STRUCTURE                        {}\n",
        "01-JAN-26"
    ));
    pdb.push_str(&format!("TITLE     {}\n", structure.name.to_uppercase()));
    pdb.push_str("REMARK   1 PREDICTED BY ALPHAFOLD-LITE (RUSTYTORCH++)\n");
    pdb.push_str(&format!(
        "REMARK   2 AVERAGE PLDDT: {:.1}\n",
        structure.model_confidence.avg_plddt
    ));
    pdb.push_str(&format!(
        "REMARK   2 PTM SCORE: {:.3}\n",
        structure.model_confidence.ptm_score
    ));

    // Sequence
    let seq_lines = structure.sequence.as_bytes().chunks(60);
    for (i, chunk) in seq_lines.enumerate() {
        let seq_str: String = chunk.iter().map(|&b| b as char).collect();
        pdb.push_str(&format!(
            "SEQRES {:3} A {:4}  {}\n",
            i + 1,
            structure.num_residues,
            seq_str
        ));
    }

    // Atoms
    let mut atom_serial = 1;
    let mut current_residue = usize::MAX;
    let mut residue_name = "UNK";

    for atom in &structure.atom_coords {
        // Update residue info if changed
        if atom.residue_idx != current_residue {
            current_residue = atom.residue_idx;
            if let Some(aa_char) = structure.sequence.chars().nth(atom.residue_idx)
                && let Some(aa) = AminoAcid::from_code(aa_char)
            {
                residue_name = aa.code3();
            }
        }

        let atom_name = match atom.atom_name {
            AtomName::N => " N  ",
            AtomName::Ca => " CA ",
            AtomName::C => " C  ",
            AtomName::O => " O  ",
            AtomName::Cb => " CB ",
        };

        let element = match atom.atom_name {
            AtomName::N => "N",
            AtomName::Ca | AtomName::C | AtomName::Cb => "C",
            AtomName::O => "O",
        };

        pdb.push_str(&format!(
            "ATOM  {:5} {}{:>3} A{:4}    {:8.3}{:8.3}{:8.3}{:6.2}{:6.2}          {:>2}\n",
            atom_serial,
            atom_name,
            residue_name,
            atom.residue_idx + 1,
            atom.x,
            atom.y,
            atom.z,
            1.00,          // Occupancy
            atom.b_factor, // B-factor (pLDDT)
            element
        ));

        atom_serial += 1;
    }

    // Terminal
    pdb.push_str("TER\n");
    pdb.push_str("END\n");

    pdb
}

/// Export structure to mmCIF format.
#[must_use]
pub fn to_mmcif(structure: &ProteinStructure) -> String {
    let mut cif = String::new();

    // Data block
    let safe_name = structure.name.replace(' ', "_").to_lowercase();
    cif.push_str(&format!("data_{safe_name}\n#\n"));

    // Entry
    cif.push_str(&format!("_entry.id   {safe_name}\n#\n"));

    // Entity
    cif.push_str("_entity.id               1\n");
    cif.push_str("_entity.type             polymer\n");
    cif.push_str("_entity.src_method       man\n");
    cif.push_str("_entity.pdbx_description 'Predicted protein structure'\n");
    cif.push_str("#\n");

    // Entity poly
    cif.push_str("_entity_poly.entity_id      1\n");
    cif.push_str("_entity_poly.type           polypeptide(L)\n");
    cif.push_str(&format!(
        "_entity_poly.pdbx_seq_one_letter_code\n;{}\n;\n",
        structure.sequence
    ));
    cif.push_str("#\n");

    // Software
    cif.push_str("_software.name                    'AlphaFold-Lite'\n");
    cif.push_str("_software.version                 '1.0'\n");
    cif.push_str("_software.classification          'structure prediction'\n");
    cif.push_str("#\n");

    // Model quality
    cif.push_str("loop_\n");
    cif.push_str("_ma_qa_metric_global.id\n");
    cif.push_str("_ma_qa_metric_global.model_id\n");
    cif.push_str("_ma_qa_metric_global.metric_id\n");
    cif.push_str("_ma_qa_metric_global.metric_value\n");
    cif.push_str(&format!(
        "1 1 1 {:.2}\n",
        structure.model_confidence.avg_plddt
    ));
    cif.push_str(&format!(
        "2 1 2 {:.4}\n",
        structure.model_confidence.ptm_score
    ));
    cif.push_str("#\n");

    // Atom sites
    cif.push_str("loop_\n");
    cif.push_str("_atom_site.group_PDB\n");
    cif.push_str("_atom_site.id\n");
    cif.push_str("_atom_site.type_symbol\n");
    cif.push_str("_atom_site.label_atom_id\n");
    cif.push_str("_atom_site.label_alt_id\n");
    cif.push_str("_atom_site.label_comp_id\n");
    cif.push_str("_atom_site.label_asym_id\n");
    cif.push_str("_atom_site.label_entity_id\n");
    cif.push_str("_atom_site.label_seq_id\n");
    cif.push_str("_atom_site.pdbx_PDB_ins_code\n");
    cif.push_str("_atom_site.Cartn_x\n");
    cif.push_str("_atom_site.Cartn_y\n");
    cif.push_str("_atom_site.Cartn_z\n");
    cif.push_str("_atom_site.occupancy\n");
    cif.push_str("_atom_site.B_iso_or_equiv\n");
    cif.push_str("_atom_site.auth_seq_id\n");
    cif.push_str("_atom_site.auth_asym_id\n");
    cif.push_str("_atom_site.pdbx_PDB_model_num\n");

    let mut current_residue = usize::MAX;
    let mut residue_name = "UNK";

    for (i, atom) in structure.atom_coords.iter().enumerate() {
        if atom.residue_idx != current_residue {
            current_residue = atom.residue_idx;
            if let Some(aa_char) = structure.sequence.chars().nth(atom.residue_idx)
                && let Some(aa) = AminoAcid::from_code(aa_char)
            {
                residue_name = aa.code3();
            }
        }

        let atom_name = atom.atom_name.as_pdb_str();
        let element = match atom.atom_name {
            AtomName::N => "N",
            AtomName::Ca | AtomName::C | AtomName::Cb => "C",
            AtomName::O => "O",
        };

        cif.push_str(&format!(
            "ATOM {} {} {} . {} A 1 {} ? {:.3} {:.3} {:.3} 1.00 {:.2} {} A 1\n",
            i + 1,
            element,
            atom_name,
            residue_name,
            atom.residue_idx + 1,
            atom.x,
            atom.y,
            atom.z,
            atom.b_factor,
            atom.residue_idx + 1
        ));
    }

    cif.push_str("#\n");

    cif
}

/// Generate a `PyMOL` script for visualization.
#[must_use]
pub fn to_pymol_script(structure: &ProteinStructure) -> String {
    let mut script = String::new();

    script.push_str("# PyMOL visualization script for AlphaFold-Lite prediction\n");
    script.push_str(&format!("# Protein: {}\n\n", structure.name));

    // Color by pLDDT
    script.push_str("# Color by pLDDT confidence\n");
    script.push_str("cmd.color('0x0053D6', 'b > 90')  # Very high (blue)\n");
    script.push_str("cmd.color('0x65CBF3', 'b > 70 and b <= 90')  # High (cyan)\n");
    script.push_str("cmd.color('0xFFDB13', 'b > 50 and b <= 70')  # Low (yellow)\n");
    script.push_str("cmd.color('0xFF7D45', 'b <= 50')  # Very low (orange)\n\n");

    // Visualization settings
    script.push_str("# Visualization settings\n");
    script.push_str("cmd.show('cartoon')\n");
    script.push_str("cmd.set('cartoon_fancy_helices', 1)\n");
    script.push_str("cmd.set('cartoon_smooth_loops', 1)\n");
    script.push_str("cmd.bg_color('white')\n");
    script.push_str("cmd.zoom()\n");

    script
}

#[cfg(test)]
mod tests {
    use super::*;
    use alphafold_shared::{
        AtomCoord, ChainInfo, ConfidenceCategory, ModelConfidence, SecondaryStructure,
    };

    fn create_test_structure() -> ProteinStructure {
        ProteinStructure {
            name: "Test Protein".to_string(),
            sequence: "ACD".to_string(),
            num_residues: 3,
            atom_coords: vec![
                AtomCoord {
                    residue_idx: 0,
                    atom_name: AtomName::Ca,
                    x: 0.0,
                    y: 0.0,
                    z: 0.0,
                    b_factor: 90.0,
                },
                AtomCoord {
                    residue_idx: 1,
                    atom_name: AtomName::Ca,
                    x: 3.8,
                    y: 0.0,
                    z: 0.0,
                    b_factor: 85.0,
                },
                AtomCoord {
                    residue_idx: 2,
                    atom_name: AtomName::Ca,
                    x: 7.6,
                    y: 0.0,
                    z: 0.0,
                    b_factor: 80.0,
                },
            ],
            plddt_scores: vec![90.0, 85.0, 80.0],
            pae_matrix: None,
            model_confidence: ModelConfidence {
                avg_plddt: 85.0,
                ptm_score: 0.8,
                iptm_score: None,
                category: ConfidenceCategory::High,
            },
            secondary_structure: vec![
                SecondaryStructure::Helix,
                SecondaryStructure::Helix,
                SecondaryStructure::Coil,
            ],
            chains: vec![ChainInfo {
                chain_id: 'A',
                start_residue: 0,
                end_residue: 3,
                sequence: "ACD".to_string(),
            }],
        }
    }

    #[test]
    fn test_to_pdb() {
        let structure = create_test_structure();
        let pdb = to_pdb(&structure);

        assert!(pdb.contains("HEADER"));
        assert!(pdb.contains("ATOM"));
        assert!(pdb.contains("END"));
        assert!(pdb.contains("CA"));
    }

    #[test]
    fn test_to_mmcif() {
        let structure = create_test_structure();
        let cif = to_mmcif(&structure);

        assert!(cif.contains("data_"));
        assert!(cif.contains("_atom_site"));
        assert!(cif.contains("ATOM"));
    }

    #[test]
    fn test_pymol_script() {
        let structure = create_test_structure();
        let script = to_pymol_script(&structure);

        assert!(script.contains("cmd.color"));
        assert!(script.contains("cartoon"));
    }
}