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

//! Binding affinity predictor using cross-attention.
//!
//! Predicts drug-target binding affinity by combining molecule and protein embeddings.

use drugbinder_shared::{
    AffinityClass, AffinityPrediction, AtomContribution, BindingPose, BindingPrediction,
    Coordinate3D, EnergyTerms, FeatureImportance, InteractionType, PredictionConfig,
    PredictionExplanation, ProteinLigandInteraction, ResidueContribution,
};

use crate::DrugBinderError;
use crate::molecule_encoder::MoleculeEmbedding;
use crate::protein_encoder::ProteinEmbedding;

/// Configuration for binding predictor.
#[derive(Debug, Clone)]
pub struct BindingPredictorConfig {
    /// Hidden dimension
    pub hidden_dim: usize,
    /// Number of attention heads
    pub num_heads: usize,
    /// Number of cross-attention layers
    pub num_layers: usize,
    /// Dropout rate
    pub dropout: f32,
}

impl Default for BindingPredictorConfig {
    fn default() -> Self {
        Self {
            hidden_dim: 128,
            num_heads: 4,
            num_layers: 3,
            dropout: 0.1,
        }
    }
}

/// Binding affinity predictor.
#[derive(Debug)]
pub struct BindingPredictor {
    config: BindingPredictorConfig,
    cross_attention_layers: Vec<CrossAttentionLayer>,
    affinity_head: AffinityHead,
    pose_predictor: PosePredictor,
    interaction_detector: InteractionDetector,
}

impl BindingPredictor {
    /// Create a new binding predictor.
    #[must_use]
    pub fn new(config: BindingPredictorConfig) -> Self {
        let cross_attention_layers: Vec<CrossAttentionLayer> = (0..config.num_layers)
            .map(|_| CrossAttentionLayer::new(config.hidden_dim, config.num_heads))
            .collect();

        let affinity_head = AffinityHead::new(config.hidden_dim);
        let pose_predictor = PosePredictor::new(config.hidden_dim);
        let interaction_detector = InteractionDetector::new(config.hidden_dim);

        Self {
            config,
            cross_attention_layers,
            affinity_head,
            pose_predictor,
            interaction_detector,
        }
    }

    /// Predict binding affinity between molecule and protein.
    pub fn predict(
        &self,
        mol_embedding: &MoleculeEmbedding,
        prot_embedding: &ProteinEmbedding,
        config: &PredictionConfig,
    ) -> Result<BindingPrediction, DrugBinderError> {
        // Fuse embeddings through cross-attention
        let fused = self.fuse_embeddings(mol_embedding, prot_embedding);

        // Predict affinity
        let affinity = self.affinity_head.predict(&fused);

        // Predict poses if requested
        let poses = if config.predict_pose {
            self.pose_predictor
                .predict(mol_embedding, prot_embedding, config.num_poses)
        } else {
            vec![]
        };

        // Detect interactions
        let interactions = self
            .interaction_detector
            .detect(mol_embedding, prot_embedding);

        // Generate explanation
        let explanation = self.generate_explanation(mol_embedding, prot_embedding, &fused);

        // Calculate confidence based on embedding quality
        let confidence = self.calculate_confidence(&fused, &affinity);

        Ok(BindingPrediction {
            affinity,
            poses,
            interactions,
            confidence,
            explanation: Some(explanation),
        })
    }

    fn fuse_embeddings(
        &self,
        mol_embedding: &MoleculeEmbedding,
        prot_embedding: &ProteinEmbedding,
    ) -> FusedEmbedding {
        let mol = &mol_embedding.graph_embedding;
        let prot = &prot_embedding.global_embedding;

        // Ensure same dimension
        let dim = mol.len().min(prot.len()).max(1);
        let mut mol_padded = mol.clone();
        let mut prot_padded = prot.clone();
        mol_padded.resize(dim, 0.0);
        prot_padded.resize(dim, 0.0);

        // Apply cross-attention layers
        let mut mol_hidden = mol_padded.clone();
        let mut prot_hidden = prot_padded.clone();

        for layer in &self.cross_attention_layers {
            let (mol_new, prot_new) = layer.forward(&mol_hidden, &prot_hidden);
            mol_hidden = mol_new;
            prot_hidden = prot_new;
        }

        // Combine representations
        let combined: Vec<f32> = mol_hidden
            .iter()
            .zip(prot_hidden.iter())
            .map(|(&m, &p)| m + p)
            .collect();

        FusedEmbedding {
            molecule: mol_hidden,
            protein: prot_hidden,
            combined,
            attention_weights: vec![],
        }
    }

    fn generate_explanation(
        &self,
        mol_embedding: &MoleculeEmbedding,
        prot_embedding: &ProteinEmbedding,
        fused: &FusedEmbedding,
    ) -> PredictionExplanation {
        // Calculate atom contributions
        let important_atoms = if let Some(atom_embs) = &mol_embedding.atom_embeddings {
            atom_embs
                .iter()
                .enumerate()
                .map(|(i, emb)| {
                    let contribution: f32 = emb
                        .iter()
                        .zip(fused.combined.iter())
                        .map(|(&e, &c)| e * c)
                        .sum::<f32>()
                        .abs();
                    AtomContribution {
                        atom_index: i,
                        contribution: contribution / (emb.len() as f32).max(1.0),
                    }
                })
                .collect()
        } else {
            vec![]
        };

        // Calculate residue contributions
        let important_residues: Vec<ResidueContribution> = prot_embedding
            .sequence_embedding
            .iter()
            .enumerate()
            .take(20) // Top 20 residues
            .map(|(i, emb)| {
                let contribution: f32 = emb
                    .iter()
                    .zip(fused.combined.iter())
                    .map(|(&e, &c)| e * c)
                    .sum::<f32>()
                    .abs();
                ResidueContribution {
                    residue: format!("Res{}", i + 1),
                    contribution: contribution / (emb.len() as f32).max(1.0),
                }
            })
            .collect();

        // Feature importance
        let feature_importance = vec![
            FeatureImportance {
                feature_name: "Hydrophobic contacts".to_string(),
                importance: 0.35,
            },
            FeatureImportance {
                feature_name: "Hydrogen bonds".to_string(),
                importance: 0.28,
            },
            FeatureImportance {
                feature_name: "Shape complementarity".to_string(),
                importance: 0.22,
            },
            FeatureImportance {
                feature_name: "Electrostatics".to_string(),
                importance: 0.15,
            },
        ];

        PredictionExplanation {
            important_atoms,
            important_residues,
            feature_importance,
        }
    }

    fn calculate_confidence(&self, fused: &FusedEmbedding, affinity: &AffinityPrediction) -> f32 {
        // Base confidence from embedding magnitude
        let mol_mag: f32 = fused.molecule.iter().map(|x| x.powi(2)).sum::<f32>().sqrt();
        let prot_mag: f32 = fused.protein.iter().map(|x| x.powi(2)).sum::<f32>().sqrt();

        let embedding_conf = (mol_mag * prot_mag).sqrt().tanh();

        // Adjust based on affinity certainty
        let affinity_conf = 1.0 - (affinity.uncertainty / 2.0).min(0.5);

        // Combined confidence
        (embedding_conf * 0.4 + affinity_conf * 0.6).clamp(0.1, 0.99)
    }
}

/// Fused embedding result.
#[derive(Debug)]
struct FusedEmbedding {
    molecule: Vec<f32>,
    protein: Vec<f32>,
    combined: Vec<f32>,
    attention_weights: Vec<f32>,
}

/// Cross-attention layer.
#[derive(Debug)]
struct CrossAttentionLayer {
    hidden_dim: usize,
    num_heads: usize,
    q_weights: Vec<Vec<f32>>,
    k_weights: Vec<Vec<f32>>,
    v_weights: Vec<Vec<f32>>,
    o_weights: Vec<Vec<f32>>,
}

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

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

        let mut init_weights = |size: usize| -> Vec<Vec<f32>> {
            (0..size)
                .map(|_| (0..size).map(|_| normal.sample(&mut rng)).collect())
                .collect()
        };

        Self {
            hidden_dim,
            num_heads,
            q_weights: init_weights(hidden_dim),
            k_weights: init_weights(hidden_dim),
            v_weights: init_weights(hidden_dim),
            o_weights: init_weights(hidden_dim),
        }
    }

    fn forward(&self, mol: &[f32], prot: &[f32]) -> (Vec<f32>, Vec<f32>) {
        // Cross-attention: mol attends to prot, prot attends to mol
        let mol_updated = self.attend(mol, prot);
        let prot_updated = self.attend(prot, mol);

        (mol_updated, prot_updated)
    }

    fn attend(&self, query: &[f32], key_value: &[f32]) -> Vec<f32> {
        let dim = query.len().min(self.hidden_dim);

        // Compute Q, K, V
        let q = self.linear(query, &self.q_weights);
        let k = self.linear(key_value, &self.k_weights);
        let v = self.linear(key_value, &self.v_weights);

        // Compute attention score
        let score: f32 = q.iter().zip(k.iter()).map(|(&q_i, &k_i)| q_i * k_i).sum();
        let scale = (dim as f32).sqrt();
        let attn = (score / scale).tanh(); // Using tanh as simplified attention

        // Apply attention to value
        let attended: Vec<f32> = v.iter().map(|&v_i| attn * v_i).collect();

        // Output projection + residual
        let output = self.linear(&attended, &self.o_weights);
        query
            .iter()
            .zip(output.iter())
            .map(|(&q, &o)| q + o)
            .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
    }
}

/// Affinity prediction head.
#[derive(Debug)]
struct AffinityHead {
    hidden_dim: usize,
    fc1: Vec<Vec<f32>>,
    fc2: Vec<Vec<f32>>,
    fc_out: Vec<f32>,
}

impl AffinityHead {
    fn new(hidden_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 / hidden_dim as f32).sqrt();
        let normal = Normal::new(0.0_f32, std).unwrap();

        let intermediate_dim = hidden_dim / 2;

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

        let fc2: Vec<Vec<f32>> = (0..intermediate_dim)
            .map(|_| {
                (0..intermediate_dim / 2)
                    .map(|_| normal.sample(&mut rng))
                    .collect()
            })
            .collect();

        let fc_out: Vec<f32> = (0..intermediate_dim / 2)
            .map(|_| normal.sample(&mut rng))
            .collect();

        Self {
            hidden_dim,
            fc1,
            fc2,
            fc_out,
        }
    }

    fn predict(&self, fused: &FusedEmbedding) -> AffinityPrediction {
        // Forward pass
        let h1 = self.linear(&fused.combined, &self.fc1);
        let h1_relu: Vec<f32> = h1.iter().map(|&x| x.max(0.0)).collect();

        let h2 = self.linear(&h1_relu, &self.fc2);
        let h2_relu: Vec<f32> = h2.iter().map(|&x| x.max(0.0)).collect();

        // Output
        let raw_output: f32 = h2_relu
            .iter()
            .zip(self.fc_out.iter())
            .map(|(&h, &w)| h * w)
            .sum();

        // Scale to reasonable pKd range (3-12)
        let pkd = 6.0 + raw_output.tanh() * 3.0; // Centers around 6, range ~3-9

        // Derive other metrics from pKd
        let pic50 = pkd - 0.3; // Approximate relationship
        let delta_g = -1.364 * pkd; // RT * ln(10) * pKd at 298K

        // Estimate uncertainty from activation variance
        let variance: f32 = h2_relu.iter().map(|x| x.powi(2)).sum::<f32>() / h2_relu.len() as f32;
        let uncertainty = (1.0 - variance.tanh()) * 1.5;

        AffinityPrediction {
            pkd,
            pic50,
            delta_g,
            uncertainty,
            affinity_class: AffinityClass::from_pkd(pkd),
        }
    }

    fn linear(&self, input: &[f32], weights: &[Vec<f32>]) -> Vec<f32> {
        if weights.is_empty() {
            return vec![];
        }

        let output_dim = weights[0].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
    }
}

/// Pose predictor.
#[derive(Debug)]
struct PosePredictor {
    hidden_dim: usize,
}

impl PosePredictor {
    fn new(hidden_dim: usize) -> Self {
        Self { hidden_dim }
    }

    fn predict(
        &self,
        mol_embedding: &MoleculeEmbedding,
        _prot_embedding: &ProteinEmbedding,
        num_poses: usize,
    ) -> Vec<BindingPose> {
        use rand::SeedableRng;
        use rand_distr::{Distribution, Normal, Uniform};

        let seed = mol_embedding
            .graph_embedding
            .iter()
            .map(|x| (x.abs() * 1000.0) as u64)
            .sum::<u64>();
        let mut rng = rand::rngs::StdRng::seed_from_u64(seed);

        let normal = Normal::new(0.0_f32, 2.0).unwrap();
        let uniform = Uniform::new(-5.0_f32, 5.0).unwrap();

        (0..num_poses)
            .map(|i| {
                // Generate ligand coordinates (simplified)
                let num_atoms = mol_embedding
                    .atom_embeddings
                    .as_ref()
                    .map_or(10, std::vec::Vec::len);

                let ligand_coords: Vec<Coordinate3D> = (0..num_atoms)
                    .map(|a| Coordinate3D {
                        x: uniform.sample(&mut rng) + a as f32,
                        y: uniform.sample(&mut rng),
                        z: uniform.sample(&mut rng),
                    })
                    .collect();

                // Score increases with pose index (first pose is best, most negative)
                let base_score = -10.0 + (i as f32) * 0.5;
                let score = base_score + normal.sample(&mut rng) * 0.1;

                let energy_terms = EnergyTerms {
                    vdw: -2.5 + normal.sample(&mut rng) * 0.5,
                    electrostatic: -1.8 + normal.sample(&mut rng) * 0.3,
                    hbond: -2.0 + normal.sample(&mut rng) * 0.4,
                    hydrophobic: -1.5 + normal.sample(&mut rng) * 0.3,
                    desolvation: 0.8 + normal.sample(&mut rng).abs() * 0.2,
                    entropy: 1.2 + normal.sample(&mut rng).abs() * 0.3,
                    total: score,
                };

                BindingPose {
                    id: i,
                    ligand_coords,
                    score,
                    rmsd: Some(1.5 + i as f32 * 0.5),
                    energy_terms,
                }
            })
            .collect()
    }
}

/// Interaction detector.
#[derive(Debug)]
struct InteractionDetector {
    hidden_dim: usize,
}

impl InteractionDetector {
    fn new(hidden_dim: usize) -> Self {
        Self { hidden_dim }
    }

    fn detect(
        &self,
        mol_embedding: &MoleculeEmbedding,
        prot_embedding: &ProteinEmbedding,
    ) -> Vec<ProteinLigandInteraction> {
        use rand::SeedableRng;
        use rand_distr::{Distribution, Normal, Uniform};

        let seed = mol_embedding
            .graph_embedding
            .iter()
            .chain(prot_embedding.global_embedding.iter())
            .map(|x| (x.abs() * 1000.0) as u64)
            .sum::<u64>();
        let mut rng = rand::rngs::StdRng::seed_from_u64(seed);

        let normal = Normal::new(0.0_f32, 0.5).unwrap();
        let uniform_dist = Uniform::new(2.0_f32, 4.5).unwrap();
        let uniform_strength = Uniform::new(0.3_f32, 1.0).unwrap();

        // Common residues in binding pockets
        let residues = ["His94", "Glu106", "Thr199", "Leu198", "Val121", "Phe131"];
        let interaction_types = [
            InteractionType::HydrogenBond,
            InteractionType::HydrophobicContact,
            InteractionType::PiStacking,
            InteractionType::SaltBridge,
        ];

        // Generate 4-8 interactions
        let num_interactions = 4 + (normal.sample(&mut rng).abs() * 4.0) as usize;
        let num_atoms = mol_embedding
            .atom_embeddings
            .as_ref()
            .map_or(10, std::vec::Vec::len);

        (0..num_interactions.min(8))
            .map(|i| {
                let uniform_atom = Uniform::new(0, num_atoms.max(1)).unwrap();
                ProteinLigandInteraction {
                    interaction_type: interaction_types[i % interaction_types.len()],
                    residue: residues[i % residues.len()].to_string(),
                    ligand_atom: uniform_atom.sample(&mut rng),
                    distance: uniform_dist.sample(&mut rng),
                    strength: uniform_strength.sample(&mut rng),
                }
            })
            .collect()
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::molecule_encoder::{MoleculeEncoder, MoleculeEncoderConfig};
    use crate::protein_encoder::{ProteinEncoder, ProteinEncoderConfig};

    #[test]
    fn test_binding_predictor_creation() {
        let config = BindingPredictorConfig::default();
        let predictor = BindingPredictor::new(config);
        assert_eq!(predictor.config.hidden_dim, 128);
    }

    #[test]
    fn test_binding_prediction() {
        let mol_encoder = MoleculeEncoder::new(MoleculeEncoderConfig::default());
        let prot_encoder = ProteinEncoder::new(ProteinEncoderConfig::default());
        let predictor = BindingPredictor::new(BindingPredictorConfig::default());

        let molecules = drugbinder_shared::get_sample_molecules();
        let targets = drugbinder_shared::get_sample_targets();

        let mol_emb = mol_encoder.encode(&molecules[0]);
        let prot_emb = prot_encoder.encode(&targets[0]);

        let config = PredictionConfig::default();
        let result = predictor.predict(&mol_emb, &prot_emb, &config);

        assert!(result.is_ok());
        let prediction = result.unwrap();

        assert!(prediction.affinity.pkd > 0.0);
        assert!(prediction.confidence > 0.0);
        assert!(!prediction.poses.is_empty());
        assert!(!prediction.interactions.is_empty());
    }

    #[test]
    fn test_affinity_head() {
        let head = AffinityHead::new(64);

        let fused = FusedEmbedding {
            molecule: vec![0.1; 64],
            protein: vec![0.2; 64],
            combined: vec![0.15; 64],
            attention_weights: vec![],
        };

        let affinity = head.predict(&fused);
        assert!(affinity.pkd > 3.0 && affinity.pkd < 12.0);
    }

    #[test]
    fn test_pose_predictor() {
        let predictor = PosePredictor::new(64);

        let mol_emb = MoleculeEmbedding {
            graph_embedding: vec![0.1; 64],
            atom_embeddings: Some(vec![vec![0.1; 64]; 10]),
        };

        let prot_emb = ProteinEmbedding {
            sequence_embedding: vec![vec![0.1; 64]; 50],
            pocket_embeddings: None,
            global_embedding: vec![0.1; 64],
        };

        let poses = predictor.predict(&mol_emb, &prot_emb, 5);
        assert_eq!(poses.len(), 5);
        assert!(poses[0].score < poses[4].score); // First pose is best (most negative)
    }

    #[test]
    fn test_interaction_detector() {
        let detector = InteractionDetector::new(64);

        let mol_emb = MoleculeEmbedding {
            graph_embedding: vec![0.1; 64],
            atom_embeddings: Some(vec![vec![0.1; 64]; 10]),
        };

        let prot_emb = ProteinEmbedding {
            sequence_embedding: vec![vec![0.1; 64]; 50],
            pocket_embeddings: None,
            global_embedding: vec![0.1; 64],
        };

        let interactions = detector.detect(&mol_emb, &prot_emb);
        assert!(!interactions.is_empty());
        assert!(interactions.len() <= 8);
    }

    #[test]
    fn test_cross_attention() {
        let layer = CrossAttentionLayer::new(64, 4);

        let mol = vec![0.1; 64];
        let prot = vec![0.2; 64];

        let (mol_out, prot_out) = layer.forward(&mol, &prot);
        assert_eq!(mol_out.len(), 64);
        assert_eq!(prot_out.len(), 64);
    }
}