rustytorch/demos/rtx-drugbinder-demo/src/molecule_encoder.rs

//! Molecular graph encoder using message passing neural networks.
//!
//! Encodes molecules as graphs where atoms are nodes and bonds are edges.

use drugbinder_shared::{Element, Hybridization, Molecule};

/// Configuration for molecule encoder.
#[derive(Debug, Clone)]
pub struct MoleculeEncoderConfig {
    /// Atom embedding dimension
    pub atom_dim: usize,
    /// Bond embedding dimension
    pub bond_dim: usize,
    /// Hidden dimension
    pub hidden_dim: usize,
    /// Number of message passing layers
    pub num_layers: usize,
    /// Dropout rate
    pub dropout: f32,
}

impl Default for MoleculeEncoderConfig {
    fn default() -> Self {
        Self {
            atom_dim: 64,
            bond_dim: 32,
            hidden_dim: 128,
            num_layers: 4,
            dropout: 0.1,
        }
    }
}

/// Molecule encoder using graph neural networks.
#[derive(Debug)]
pub struct MoleculeEncoder {
    config: MoleculeEncoderConfig,
    atom_embeddings: AtomEmbeddings,
    bond_embeddings: BondEmbeddings,
    mp_layers: Vec<MessagePassingLayer>,
    readout: GraphReadout,
}

impl MoleculeEncoder {
    /// Create a new molecule encoder.
    #[must_use]
    pub fn new(config: MoleculeEncoderConfig) -> Self {
        let atom_embeddings = AtomEmbeddings::new(config.atom_dim);
        let bond_embeddings = BondEmbeddings::new(config.bond_dim);

        let mp_layers: Vec<MessagePassingLayer> = (0..config.num_layers)
            .map(|_| MessagePassingLayer::new(config.hidden_dim, config.bond_dim))
            .collect();

        let readout = GraphReadout::new(config.hidden_dim);

        Self {
            config,
            atom_embeddings,
            bond_embeddings,
            mp_layers,
            readout,
        }
    }

    /// Encode a molecule into a fixed-size embedding.
    #[must_use]
    pub fn encode(&self, molecule: &Molecule) -> MoleculeEmbedding {
        // Get atom features
        let atom_features = self.atom_embeddings.embed(molecule);

        // Get bond features
        let bond_features = self.bond_embeddings.embed(molecule);

        // Build adjacency from bonds
        let adjacency = build_adjacency(molecule);

        // Message passing
        let mut hidden = atom_features;
        for layer in &self.mp_layers {
            hidden = layer.forward(&hidden, &bond_features, &adjacency);
        }

        // Graph-level readout
        self.readout.forward(&hidden)
    }

    /// Get hidden dimension.
    #[must_use]
    pub fn hidden_dim(&self) -> usize {
        self.config.hidden_dim
    }
}

/// Molecule embedding.
#[derive(Debug, Clone)]
pub struct MoleculeEmbedding {
    /// Graph-level embedding
    pub graph_embedding: Vec<f32>,
    /// Atom-level embeddings (optional)
    pub atom_embeddings: Option<Vec<Vec<f32>>>,
}

/// Atom embeddings.
#[derive(Debug)]
struct AtomEmbeddings {
    dim: usize,
    element_embeddings: Vec<Vec<f32>>,
    hybridization_embeddings: Vec<Vec<f32>>,
}

impl AtomEmbeddings {
    fn new(dim: usize) -> Self {
        use rand::SeedableRng;
        use rand_distr::{Distribution, Normal};

        let mut rng = rand::rngs::StdRng::seed_from_u64(42);
        let std = (2.0 / dim as f32).sqrt();
        let normal = Normal::new(0.0_f32, std).unwrap();

        // Embeddings for different elements (up to 20)
        let element_embeddings: Vec<Vec<f32>> = (0..20)
            .map(|_| (0..dim).map(|_| normal.sample(&mut rng)).collect())
            .collect();

        // Embeddings for hybridization states (6 types)
        let hybridization_embeddings: Vec<Vec<f32>> = (0..6)
            .map(|_| (0..dim / 4).map(|_| normal.sample(&mut rng)).collect())
            .collect();

        Self {
            dim,
            element_embeddings,
            hybridization_embeddings,
        }
    }

    fn embed(&self, molecule: &Molecule) -> Vec<Vec<f32>> {
        if molecule.atoms.is_empty() {
            // Generate from SMILES if no atoms provided
            return self.embed_from_smiles(&molecule.smiles);
        }

        molecule
            .atoms
            .iter()
            .map(|atom| {
                let mut embedding = vec![0.0; self.dim];

                // Element embedding
                let elem_idx = element_to_index(&atom.element);
                let elem_emb = &self.element_embeddings[elem_idx];
                for (i, &v) in elem_emb.iter().enumerate() {
                    embedding[i] = v;
                }

                // Add hybridization
                let hyb_idx = hybridization_to_index(&atom.hybridization);
                let hyb_emb = &self.hybridization_embeddings[hyb_idx];
                for (i, &v) in hyb_emb.iter().enumerate() {
                    embedding[self.dim * 3 / 4 + i] += v;
                }

                // Add charge and aromaticity features
                embedding[self.dim - 2] = f32::from(atom.formal_charge) * 0.5;
                embedding[self.dim - 1] = if atom.is_aromatic { 1.0 } else { 0.0 };

                embedding
            })
            .collect()
    }

    fn embed_from_smiles(&self, smiles: &str) -> Vec<Vec<f32>> {
        use rand::SeedableRng;
        use rand_distr::{Distribution, Normal};

        let mut rng = rand::rngs::StdRng::seed_from_u64(smiles.len() as u64);
        let normal = Normal::new(0.0_f32, 0.1).unwrap();

        // Parse SMILES to extract atom types
        let atoms = parse_smiles_atoms(smiles);

        atoms
            .iter()
            .map(|elem| {
                let elem_idx = element_to_index(elem);
                let mut embedding = self.element_embeddings[elem_idx].clone();

                // Add some noise
                for v in &mut embedding {
                    *v += normal.sample(&mut rng);
                }

                embedding
            })
            .collect()
    }
}

/// Bond embeddings.
#[derive(Debug)]
struct BondEmbeddings {
    dim: usize,
    bond_type_embeddings: Vec<Vec<f32>>,
}

impl BondEmbeddings {
    fn new(dim: usize) -> Self {
        use rand::SeedableRng;
        use rand_distr::{Distribution, Normal};

        let mut rng = rand::rngs::StdRng::seed_from_u64(123);
        let std = (2.0 / dim as f32).sqrt();
        let normal = Normal::new(0.0_f32, std).unwrap();

        // 5 bond types
        let bond_type_embeddings: Vec<Vec<f32>> = (0..5)
            .map(|_| (0..dim).map(|_| normal.sample(&mut rng)).collect())
            .collect();

        Self {
            dim,
            bond_type_embeddings,
        }
    }

    fn embed(&self, molecule: &Molecule) -> Vec<Vec<f32>> {
        if molecule.bonds.is_empty() {
            // Generate from SMILES
            return self.embed_from_smiles(&molecule.smiles);
        }

        molecule
            .bonds
            .iter()
            .map(|bond| {
                let bond_idx = bond_type_to_index(&bond.bond_type);
                let mut embedding = self.bond_type_embeddings[bond_idx].clone();

                // Add ring and conjugation features
                embedding[self.dim - 2] = if bond.is_in_ring { 1.0 } else { 0.0 };
                embedding[self.dim - 1] = if bond.is_conjugated { 1.0 } else { 0.0 };

                embedding
            })
            .collect()
    }

    fn embed_from_smiles(&self, smiles: &str) -> Vec<Vec<f32>> {
        // Simplified: generate bond embeddings based on SMILES patterns
        let num_bonds = smiles.chars().filter(|&c| c == '=' || c == '#').count()
            + smiles
                .chars()
                .filter(|c| c.is_uppercase())
                .count()
                .saturating_sub(1);

        (0..num_bonds.max(1))
            .map(|i| {
                let bond_idx = i % 5;
                self.bond_type_embeddings[bond_idx].clone()
            })
            .collect()
    }
}

/// Message passing layer.
#[derive(Debug)]
struct MessagePassingLayer {
    hidden_dim: usize,
    edge_dim: usize,
    weights_msg: Vec<Vec<f32>>,
    weights_update: Vec<Vec<f32>>,
}

impl MessagePassingLayer {
    fn new(hidden_dim: usize, edge_dim: usize) -> Self {
        use rand::SeedableRng;
        use rand_distr::{Distribution, Normal};

        let mut rng = rand::rngs::StdRng::seed_from_u64(456);
        let std = (2.0 / hidden_dim as f32).sqrt();
        let normal = Normal::new(0.0_f32, std).unwrap();

        let msg_input_dim = hidden_dim * 2 + edge_dim;
        let weights_msg: Vec<Vec<f32>> = (0..msg_input_dim)
            .map(|_| (0..hidden_dim).map(|_| normal.sample(&mut rng)).collect())
            .collect();

        let weights_update: Vec<Vec<f32>> = (0..hidden_dim * 2)
            .map(|_| (0..hidden_dim).map(|_| normal.sample(&mut rng)).collect())
            .collect();

        Self {
            hidden_dim,
            edge_dim,
            weights_msg,
            weights_update,
        }
    }

    fn forward(
        &self,
        node_features: &[Vec<f32>],
        edge_features: &[Vec<f32>],
        adjacency: &[(usize, usize, usize)], // (src, dst, edge_idx)
    ) -> Vec<Vec<f32>> {
        let n = node_features.len();
        let mut messages = vec![vec![0.0; self.hidden_dim]; n];

        // Aggregate messages
        for &(src, dst, edge_idx) in adjacency {
            if src >= n || dst >= n {
                continue;
            }

            // Concatenate src node, edge, dst node features
            let src_feat = &node_features[src];
            let dst_feat = &node_features[dst];
            let edge_feat = edge_features
                .get(edge_idx)
                .cloned()
                .unwrap_or_else(|| vec![0.0; self.edge_dim]);

            let mut concat = Vec::with_capacity(self.hidden_dim * 2 + self.edge_dim);
            concat.extend_from_slice(src_feat);
            for (i, v) in edge_feat.iter().enumerate() {
                if i < self.edge_dim {
                    concat.push(*v);
                }
            }
            // Pad if needed
            while concat.len() < self.hidden_dim + self.edge_dim {
                concat.push(0.0);
            }
            concat.extend_from_slice(dst_feat);

            // Compute message
            let msg = self.linear(&concat, &self.weights_msg);

            // Add to destination
            for (i, &v) in msg.iter().enumerate() {
                messages[dst][i] += v;
            }
        }

        // Update node features
        node_features
            .iter()
            .zip(messages.iter())
            .map(|(node, msg)| {
                let mut concat = node.clone();
                concat.extend_from_slice(msg);

                let updated = self.linear(&concat, &self.weights_update);

                // ReLU + residual
                node.iter()
                    .zip(updated.iter())
                    .map(|(&n, &u)| n + u.max(0.0))
                    .collect()
            })
            .collect()
    }

    fn linear(&self, input: &[f32], weights: &[Vec<f32>]) -> Vec<f32> {
        let output_dim = weights.first().map_or(0, std::vec::Vec::len);
        let mut output = vec![0.0; output_dim];

        for (i, &x) in input.iter().enumerate() {
            if i < weights.len() {
                for (j, &w) in weights[i].iter().enumerate() {
                    output[j] += x * w;
                }
            }
        }

        output
    }
}

/// Graph readout layer.
#[derive(Debug)]
struct GraphReadout {
    hidden_dim: usize,
    attention_weights: Vec<f32>,
    output_weights: Vec<Vec<f32>>,
}

impl GraphReadout {
    fn new(hidden_dim: usize) -> Self {
        use rand::SeedableRng;
        use rand_distr::{Distribution, Normal};

        let mut rng = rand::rngs::StdRng::seed_from_u64(789);
        let std = (2.0 / hidden_dim as f32).sqrt();
        let normal = Normal::new(0.0_f32, std).unwrap();

        let attention_weights: Vec<f32> =
            (0..hidden_dim).map(|_| normal.sample(&mut rng)).collect();

        let output_weights: Vec<Vec<f32>> = (0..hidden_dim)
            .map(|_| (0..hidden_dim).map(|_| normal.sample(&mut rng)).collect())
            .collect();

        Self {
            hidden_dim,
            attention_weights,
            output_weights,
        }
    }

    fn forward(&self, node_features: &[Vec<f32>]) -> MoleculeEmbedding {
        if node_features.is_empty() {
            return MoleculeEmbedding {
                graph_embedding: vec![0.0; self.hidden_dim],
                atom_embeddings: None,
            };
        }

        // Compute attention scores
        let scores: Vec<f32> = node_features
            .iter()
            .map(|node| {
                let score: f32 = node
                    .iter()
                    .zip(self.attention_weights.iter())
                    .map(|(&n, &w)| n * w)
                    .sum();
                score.exp()
            })
            .collect();

        let sum_scores: f32 = scores.iter().sum::<f32>() + 1e-8;
        let attention: Vec<f32> = scores.iter().map(|s| s / sum_scores).collect();

        // Weighted sum
        let mut graph_embedding = vec![0.0; self.hidden_dim];
        for (node, &attn) in node_features.iter().zip(attention.iter()) {
            for (i, &v) in node.iter().enumerate() {
                if i < self.hidden_dim {
                    graph_embedding[i] += attn * v;
                }
            }
        }

        // Final projection
        let projected = self.linear(&graph_embedding, &self.output_weights);

        MoleculeEmbedding {
            graph_embedding: projected,
            atom_embeddings: Some(node_features.to_vec()),
        }
    }

    fn linear(&self, input: &[f32], weights: &[Vec<f32>]) -> Vec<f32> {
        let output_dim = weights.first().map_or(0, std::vec::Vec::len);
        let mut output = vec![0.0; output_dim];

        for (i, &x) in input.iter().enumerate() {
            if i < weights.len() {
                for (j, &w) in weights[i].iter().enumerate() {
                    output[j] += x * w;
                }
            }
        }

        output
    }
}

/// Build adjacency list from molecule bonds.
fn build_adjacency(molecule: &Molecule) -> Vec<(usize, usize, usize)> {
    if molecule.bonds.is_empty() {
        // Generate from SMILES (simplified)
        let num_atoms = molecule.num_atoms.max(1);
        return (0..num_atoms.saturating_sub(1))
            .flat_map(|i| vec![(i, i + 1, i), (i + 1, i, i)])
            .collect();
    }

    molecule
        .bonds
        .iter()
        .enumerate()
        .flat_map(|(idx, bond)| vec![(bond.atom1, bond.atom2, idx), (bond.atom2, bond.atom1, idx)])
        .collect()
}

/// Map element to embedding index.
fn element_to_index(element: &Element) -> usize {
    match element {
        Element::H => 0,
        Element::C => 1,
        Element::N => 2,
        Element::O => 3,
        Element::F => 4,
        Element::P => 5,
        Element::S => 6,
        Element::Cl => 7,
        Element::Br => 8,
        Element::I => 9,
        _ => 10,
    }
}

/// Map hybridization to embedding index.
fn hybridization_to_index(hybridization: &Hybridization) -> usize {
    match hybridization {
        Hybridization::Sp => 0,
        Hybridization::Sp2 => 1,
        Hybridization::Sp3 => 2,
        Hybridization::Sp3d => 3,
        Hybridization::Sp3d2 => 4,
        Hybridization::Other => 5,
    }
}

/// Map bond type to embedding index.
fn bond_type_to_index(bond_type: &drugbinder_shared::BondType) -> usize {
    match bond_type {
        drugbinder_shared::BondType::Single => 0,
        drugbinder_shared::BondType::Double => 1,
        drugbinder_shared::BondType::Triple => 2,
        drugbinder_shared::BondType::Aromatic => 3,
        drugbinder_shared::BondType::Other => 4,
    }
}

/// Parse atoms from SMILES string (simplified).
fn parse_smiles_atoms(smiles: &str) -> Vec<Element> {
    let mut atoms = Vec::new();

    let chars: Vec<char> = smiles.chars().collect();
    let mut i = 0;

    while i < chars.len() {
        let c = chars[i];

        // Check for two-letter elements
        if i + 1 < chars.len() && chars[i + 1].is_lowercase() {
            let elem_str: String = vec![c, chars[i + 1]].into_iter().collect();
            let elem = match elem_str.as_str() {
                "Cl" => Element::Cl,
                "Br" => Element::Br,
                "Na" => Element::Na,
                "Mg" => Element::Mg,
                "Ca" => Element::Ca,
                "Fe" => Element::Fe,
                "Zn" => Element::Zn,
                "Cu" => Element::Cu,
                _ => Element::Other,
            };
            atoms.push(elem);
            i += 2;
            continue;
        }

        // Single letter elements
        let elem = match c {
            'C' => Element::C,
            'N' => Element::N,
            'O' => Element::O,
            'S' => Element::S,
            'P' => Element::P,
            'F' => Element::F,
            'I' => Element::I,
            'H' => Element::H,
            _ => {
                i += 1;
                continue;
            }
        };
        atoms.push(elem);
        i += 1;
    }

    if atoms.is_empty() {
        atoms.push(Element::C); // Default
    }

    atoms
}

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

    #[test]
    fn test_molecule_encoder_creation() {
        let config = MoleculeEncoderConfig::default();
        let encoder = MoleculeEncoder::new(config);
        assert_eq!(encoder.hidden_dim(), 128);
    }

    #[test]
    fn test_encode_molecule() {
        let config = MoleculeEncoderConfig {
            hidden_dim: 64,
            num_layers: 2,
            ..Default::default()
        };
        let encoder = MoleculeEncoder::new(config);

        let molecules = drugbinder_shared::get_sample_molecules();
        let mol = &molecules[0];

        let embedding = encoder.encode(mol);
        assert_eq!(embedding.graph_embedding.len(), 64);
    }

    #[test]
    fn test_parse_smiles_atoms() {
        let atoms = parse_smiles_atoms("CC(=O)O");
        assert!(!atoms.is_empty());
        assert!(atoms.iter().any(|e| *e == Element::C));
        assert!(atoms.iter().any(|e| *e == Element::O));
    }

    #[test]
    fn test_element_index() {
        assert_eq!(element_to_index(&Element::C), 1);
        assert_eq!(element_to_index(&Element::N), 2);
        assert_eq!(element_to_index(&Element::O), 3);
    }
}