tranz 0.5.1

//! # tranz
//!
//! Point-embedding knowledge graph completion models.
//!
//! Entities are points in vector space. Relations are transformations
//! (translation, rotation, diagonal scaling). Lower distance between
//! `transform(head, relation)` and `tail` indicates a more likely triple.
//!
//! ## Models
//!
//! - [`TransE`]: `head + relation ~ tail` (Bordes et al., 2013)
//! - [`RotatE`]: `head * relation ~ tail` in complex space (Sun et al., 2019)
//! - [`ComplEx`]: Hermitian dot product in complex space (Trouillon et al., 2016)
//! - [`DistMult`]: diagonal bilinear in real space (Yang et al., 2015)
//!
//! ## Feature flags
//!
//! - **`rand`** (default): enables random initialization via `Model::new()`.
//! - **`candle`**: enables GPU training via the `train` module.
//! - **`cuda`**: implies `candle`, enables CUDA acceleration.

#![warn(missing_docs)]

pub mod dataset;
pub mod eval;
pub mod io;
pub mod query;
#[cfg(feature = "candle")]
pub mod train;

/// Errors from tranz operations.
#[derive(Debug, thiserror::Error)]
#[non_exhaustive]
pub enum Error {
    /// Dimension mismatch.
    #[error("Dimension mismatch: expected {expected}, got {actual}")]
    DimensionMismatch {
        /// Expected dimension.
        expected: usize,
        /// Actual dimension.
        actual: usize,
    },
    /// IO error.
    #[error(transparent)]
    Io(#[from] std::io::Error),
}

/// Trait for scoring knowledge graph triples.
///
/// Scores are distances or negative similarities: **lower values indicate
/// more likely triples**.
pub trait Scorer: Sync {
    /// Score a triple `(head, relation, tail)`. Lower = more likely.
    fn score(&self, head: usize, relation: usize, tail: usize) -> f32;

    /// Number of entities in the model.
    fn num_entities(&self) -> usize;

    /// Score all entities as tail replacements for `(head, relation, ?)`.
    ///
    /// Returns a vec of length `num_entities()` where index `t` holds
    /// `score(head, relation, t)`.
    fn score_all_tails(&self, head: usize, relation: usize) -> Vec<f32> {
        (0..self.num_entities())
            .map(|t| self.score(head, relation, t))
            .collect()
    }

    /// Score all entities as head replacements for `(?, relation, tail)`.
    fn score_all_heads(&self, relation: usize, tail: usize) -> Vec<f32> {
        (0..self.num_entities())
            .map(|h| self.score(h, relation, tail))
            .collect()
    }

    /// Return the top-k entities by score for `(head, relation, ?)`.
    ///
    /// Returns `(entity_id, score)` pairs sorted by score ascending
    /// (best first, since lower = more likely).
    fn top_k_tails(&self, head: usize, relation: usize, k: usize) -> Vec<(usize, f32)> {
        let mut scored: Vec<(usize, f32)> = self
            .score_all_tails(head, relation)
            .into_iter()
            .enumerate()
            .collect();
        scored.sort_by(|a, b| a.1.partial_cmp(&b.1).unwrap_or(std::cmp::Ordering::Equal));
        scored.truncate(k);
        scored
    }

    /// Return the top-k entities by score for `(?, relation, tail)`.
    fn top_k_heads(&self, relation: usize, tail: usize, k: usize) -> Vec<(usize, f32)> {
        let mut scored: Vec<(usize, f32)> = self
            .score_all_heads(relation, tail)
            .into_iter()
            .enumerate()
            .collect();
        scored.sort_by(|a, b| a.1.partial_cmp(&b.1).unwrap_or(std::cmp::Ordering::Equal));
        scored.truncate(k);
        scored
    }

    /// Score all relations for `(head, ?, tail)`.
    ///
    /// Returns a vec where index `r` holds `score(head, r, tail)`.
    /// Requires knowing the number of relations (passed as parameter
    /// since the Scorer trait doesn't expose it).
    fn score_all_relations(&self, head: usize, tail: usize, num_relations: usize) -> Vec<f32> {
        (0..num_relations)
            .map(|r| self.score(head, r, tail))
            .collect()
    }

    /// Return the top-k relations by score for `(head, ?, tail)`.
    fn top_k_relations(
        &self,
        head: usize,
        tail: usize,
        num_relations: usize,
        k: usize,
    ) -> Vec<(usize, f32)> {
        let mut scored: Vec<(usize, f32)> = self
            .score_all_relations(head, tail, num_relations)
            .into_iter()
            .enumerate()
            .collect();
        scored.sort_by(|a, b| a.1.partial_cmp(&b.1).unwrap_or(std::cmp::Ordering::Equal));
        scored.truncate(k);
        scored
    }
}

// ---------------------------------------------------------------------------
// TransE
// ---------------------------------------------------------------------------

/// TransE: translational distance model.
///
/// Scores triples by `||head + relation - tail||_2`. Lower = more likely.
///
/// Initialization uses uniform `[-6/sqrt(dim), 6/sqrt(dim)]` (Xavier-like).
///
/// Reference: Bordes et al. (2013), "Translating Embeddings for Modeling
/// Multi-relational Data."
pub struct TransE {
    /// Flat entity embeddings: `[num_entities * dim]` row-major.
    entities: Vec<f32>,
    /// Flat relation embeddings: `[num_relations * dim]` row-major.
    relations: Vec<f32>,
    dim: usize,
    /// Distance norm: 1 = L1, 2 = L2 (default).
    norm: u32,
}

impl TransE {
    /// Create a new TransE model with random initialization.
    ///
    /// Entity and relation embeddings are drawn from
    /// `Uniform(-6/sqrt(dim), 6/sqrt(dim))`.
    #[cfg(feature = "rand")]
    pub fn new(num_entities: usize, num_relations: usize, dim: usize) -> Self {
        let mut rng = rand::rng();
        let scale = 6.0_f32 / (dim as f32).sqrt();
        Self::from_vecs(
            init_vecs(&mut rng, num_entities, dim, scale),
            init_vecs(&mut rng, num_relations, dim, scale),
            dim,
        )
    }

    /// Create from pre-built embedding vectors.
    ///
    /// Flattens the input into contiguous row-major storage.
    ///
    /// # Panics
    ///
    /// Panics if any inner vector length differs from `dim`.
    pub fn from_vecs(entities: Vec<Vec<f32>>, relations: Vec<Vec<f32>>, dim: usize) -> Self {
        assert_dims(&entities, dim, "entity");
        assert_dims(&relations, dim, "relation");
        Self {
            entities: flatten(&entities),
            relations: flatten(&relations),
            dim,
            norm: 2,
        }
    }

    /// Create with a specific distance norm (1 = L1, 2 = L2).
    pub fn from_vecs_with_norm(
        entities: Vec<Vec<f32>>,
        relations: Vec<Vec<f32>>,
        dim: usize,
        norm: u32,
    ) -> Self {
        assert_dims(&entities, dim, "entity");
        assert_dims(&relations, dim, "relation");
        Self {
            entities: flatten(&entities),
            relations: flatten(&relations),
            dim,
            norm,
        }
    }

    /// Flat entity embeddings (row-major, stride = `dim()`).
    pub fn entities_flat(&self) -> &[f32] {
        &self.entities
    }

    /// Flat relation embeddings (row-major, stride = `dim()`).
    pub fn relations_flat(&self) -> &[f32] {
        &self.relations
    }

    /// Reconstruct entity embeddings as `Vec<Vec<f32>>`.
    pub fn entity_vecs(&self) -> Vec<Vec<f32>> {
        unflatten(&self.entities, self.dim)
    }

    /// Reconstruct relation embeddings as `Vec<Vec<f32>>`.
    pub fn relation_vecs(&self) -> Vec<Vec<f32>> {
        unflatten(&self.relations, self.dim)
    }

    /// Embedding dimension.
    pub fn dim(&self) -> usize {
        self.dim
    }

    /// Distance norm (1 = L1, 2 = L2).
    pub fn norm(&self) -> u32 {
        self.norm
    }

    /// Score a triple (h, r, t) via distance `||h + r - t||`.
    pub fn score_triple(&self, head: usize, relation: usize, tail: usize) -> f32 {
        let h = row(&self.entities, head, self.dim);
        let r = row(&self.relations, relation, self.dim);
        let t = row(&self.entities, tail, self.dim);

        match self.norm {
            1 => {
                let mut dist = 0.0_f64;
                for i in 0..self.dim {
                    dist += (h[i] as f64 + r[i] as f64 - t[i] as f64).abs();
                }
                dist as f32
            }
            _ => {
                let mut dist_sq = 0.0_f64;
                for i in 0..self.dim {
                    let d = h[i] as f64 + r[i] as f64 - t[i] as f64;
                    dist_sq += d * d;
                }
                dist_sq.sqrt() as f32
            }
        }
    }
}

impl Scorer for TransE {
    fn score(&self, head: usize, relation: usize, tail: usize) -> f32 {
        self.score_triple(head, relation, tail)
    }

    fn num_entities(&self) -> usize {
        self.entities.len() / self.dim
    }

    fn score_all_tails(&self, head: usize, relation: usize) -> Vec<f32> {
        let h = row(&self.entities, head, self.dim);
        let r = row(&self.relations, relation, self.dim);
        let dim = self.dim;
        let norm = self.norm;
        let n = self.num_entities();
        let mut hr = vec![0.0_f64; dim];
        for i in 0..dim {
            hr[i] = h[i] as f64 + r[i] as f64;
        }
        (0..n)
            .map(|ti| {
                let t = row(&self.entities, ti, dim);
                if norm == 1 {
                    let mut dist = 0.0_f64;
                    for i in 0..dim {
                        dist += (hr[i] - t[i] as f64).abs();
                    }
                    dist as f32
                } else {
                    let mut dist_sq = 0.0_f64;
                    for i in 0..dim {
                        let d = hr[i] - t[i] as f64;
                        dist_sq += d * d;
                    }
                    dist_sq.sqrt() as f32
                }
            })
            .collect()
    }

    fn score_all_heads(&self, relation: usize, tail: usize) -> Vec<f32> {
        let r = row(&self.relations, relation, self.dim);
        let t = row(&self.entities, tail, self.dim);
        let dim = self.dim;
        let norm = self.norm;
        let n = self.num_entities();
        let mut neg_rt = vec![0.0_f64; dim];
        for i in 0..dim {
            neg_rt[i] = r[i] as f64 - t[i] as f64;
        }
        (0..n)
            .map(|hi| {
                let h = row(&self.entities, hi, dim);
                if norm == 1 {
                    let mut dist = 0.0_f64;
                    for i in 0..dim {
                        dist += (h[i] as f64 + neg_rt[i]).abs();
                    }
                    dist as f32
                } else {
                    let mut dist_sq = 0.0_f64;
                    for i in 0..dim {
                        let d = h[i] as f64 + neg_rt[i];
                        dist_sq += d * d;
                    }
                    dist_sq.sqrt() as f32
                }
            })
            .collect()
    }
}

// ---------------------------------------------------------------------------
// RotatE
// ---------------------------------------------------------------------------

/// RotatE: rotation in complex space.
///
/// Entities are complex vectors. Relations are element-wise rotations
/// (unit-modulus complex numbers parameterized by angles).
///
/// Score: `||head * relation - tail||_2` where `*` is element-wise complex
/// multiplication and `|relation_i| = 1`.
///
/// Initialization: entities from `Uniform(-gamma/sqrt(dim), gamma/sqrt(dim))`,
/// relation angles from `Uniform(-pi, pi)`.
///
/// Reference: Sun et al. (2019), "RotatE: Knowledge Graph Embedding by
/// Relational Rotation in Complex Space."
pub struct RotatE {
    /// Flat entity embeddings: `[num_entities * dim * 2]` row-major, re then im per entity.
    entities: Vec<f32>,
    /// Flat relation angles: `[num_relations * dim]` row-major.
    relation_angles: Vec<f32>,
    /// Complex dimension (half the real storage per entity).
    dim: usize,
    /// Margin parameter used for initialization scaling.
    gamma: f32,
}

impl RotatE {
    /// Create a new RotatE model with random initialization.
    ///
    /// `dim` is the complex dimension (each entity stores `2 * dim` floats).
    /// `gamma` is the margin, used to scale entity initialization to
    /// `Uniform(-gamma/sqrt(dim), gamma/sqrt(dim))`.
    #[cfg(feature = "rand")]
    pub fn new(num_entities: usize, num_relations: usize, dim: usize, gamma: f32) -> Self {
        use rand::Rng;
        let mut rng = rand::rng();
        let entity_scale = gamma / (dim as f32).sqrt();
        let entities = init_vecs(&mut rng, num_entities, dim * 2, entity_scale);
        let relation_angles: Vec<Vec<f32>> = (0..num_relations)
            .map(|_| {
                (0..dim)
                    .map(|_| rng.random_range(-std::f32::consts::PI..std::f32::consts::PI))
                    .collect()
            })
            .collect();
        Self {
            entities: flatten(&entities),
            relation_angles: flatten(&relation_angles),
            dim,
            gamma,
        }
    }

    /// Create from pre-built embedding vectors.
    ///
    /// `entities` must have inner length `dim * 2` (interleaved re/im).
    /// `relation_angles` must have inner length `dim`.
    ///
    /// # Panics
    ///
    /// Panics if any dimension is wrong.
    pub fn from_vecs(
        entities: Vec<Vec<f32>>,
        relation_angles: Vec<Vec<f32>>,
        dim: usize,
        gamma: f32,
    ) -> Self {
        assert_dims(&entities, dim * 2, "entity (re+im)");
        assert_dims(&relation_angles, dim, "relation angle");
        Self {
            entities: flatten(&entities),
            relation_angles: flatten(&relation_angles),
            dim,
            gamma,
        }
    }

    /// Flat entity embeddings (row-major, stride = `dim() * 2`).
    pub fn entities_flat(&self) -> &[f32] {
        &self.entities
    }

    /// Flat relation angle embeddings (row-major, stride = `dim()`).
    pub fn relation_angles_flat(&self) -> &[f32] {
        &self.relation_angles
    }

    /// Reconstruct entity embeddings as `Vec<Vec<f32>>`.
    pub fn entity_vecs(&self) -> Vec<Vec<f32>> {
        unflatten(&self.entities, self.dim * 2)
    }

    /// Reconstruct relation angle embeddings as `Vec<Vec<f32>>`.
    pub fn relation_angle_vecs(&self) -> Vec<Vec<f32>> {
        unflatten(&self.relation_angles, self.dim)
    }

    /// Complex dimension.
    pub fn dim(&self) -> usize {
        self.dim
    }

    /// Margin parameter.
    pub fn gamma(&self) -> f32 {
        self.gamma
    }

    /// Score a triple by `||h * r - t||_2` in complex space.
    ///
    /// Entity layout: first `dim` floats are real parts, next `dim` are imaginary.
    pub fn score_triple(&self, head: usize, relation: usize, tail: usize) -> f32 {
        let h = row(&self.entities, head, self.dim * 2);
        let r = row(&self.relation_angles, relation, self.dim);
        let t = row(&self.entities, tail, self.dim * 2);
        let dim = self.dim;

        let mut dist_sq = 0.0_f64;
        for i in 0..dim {
            let h_re = h[i] as f64;
            let h_im = h[dim + i] as f64;
            let (r_sin, r_cos) = (r[i] as f64).sin_cos();
            let t_re = t[i] as f64;
            let t_im = t[dim + i] as f64;

            let hr_re = h_re * r_cos - h_im * r_sin;
            let hr_im = h_re * r_sin + h_im * r_cos;

            let d_re = hr_re - t_re;
            let d_im = hr_im - t_im;
            dist_sq += d_re * d_re + d_im * d_im;
        }
        dist_sq.sqrt() as f32
    }
}

impl Scorer for RotatE {
    fn score(&self, head: usize, relation: usize, tail: usize) -> f32 {
        self.score_triple(head, relation, tail)
    }

    fn num_entities(&self) -> usize {
        self.entities.len() / (self.dim * 2)
    }

    fn score_all_tails(&self, head: usize, relation: usize) -> Vec<f32> {
        let h = row(&self.entities, head, self.dim * 2);
        let r = row(&self.relation_angles, relation, self.dim);
        let dim = self.dim;
        let n = self.num_entities();
        // Precompute h * r (complex rotation) once.
        let mut hr_re = vec![0.0_f64; dim];
        let mut hr_im = vec![0.0_f64; dim];
        for i in 0..dim {
            let h_re = h[i] as f64;
            let h_im = h[dim + i] as f64;
            let (r_sin, r_cos) = (r[i] as f64).sin_cos();
            hr_re[i] = h_re * r_cos - h_im * r_sin;
            hr_im[i] = h_re * r_sin + h_im * r_cos;
        }
        (0..n)
            .map(|ti| {
                let t = row(&self.entities, ti, dim * 2);
                let mut dist_sq = 0.0_f64;
                for i in 0..dim {
                    let d_re = hr_re[i] - t[i] as f64;
                    let d_im = hr_im[i] - t[dim + i] as f64;
                    dist_sq += d_re * d_re + d_im * d_im;
                }
                dist_sq.sqrt() as f32
            })
            .collect()
    }

    fn score_all_heads(&self, relation: usize, tail: usize) -> Vec<f32> {
        let r = row(&self.relation_angles, relation, self.dim);
        let t = row(&self.entities, tail, self.dim * 2);
        let dim = self.dim;
        let n = self.num_entities();
        // Precompute t * conj(r) once: inverse rotation applied to tail.
        // conj(r) has angles -theta, so cos(-t)=cos(t), sin(-t)=-sin(t).
        let mut tr_re = vec![0.0_f64; dim];
        let mut tr_im = vec![0.0_f64; dim];
        for i in 0..dim {
            let t_re = t[i] as f64;
            let t_im = t[dim + i] as f64;
            let (r_sin, r_cos) = (r[i] as f64).sin_cos();
            // t * conj(r) = t * (cos, -sin)
            tr_re[i] = t_re * r_cos + t_im * r_sin;
            tr_im[i] = t_im * r_cos - t_re * r_sin;
        }
        (0..n)
            .map(|hi| {
                let h = row(&self.entities, hi, dim * 2);
                let mut dist_sq = 0.0_f64;
                for i in 0..dim {
                    let d_re = h[i] as f64 - tr_re[i];
                    let d_im = h[dim + i] as f64 - tr_im[i];
                    dist_sq += d_re * d_re + d_im * d_im;
                }
                dist_sq.sqrt() as f32
            })
            .collect()
    }
}

// ---------------------------------------------------------------------------
// ComplEx
// ---------------------------------------------------------------------------

/// ComplEx: complex bilinear model.
///
/// Entities and relations are complex vectors. Score is the real part of the
/// Hermitian dot product: `Re(sum_i h_i * r_i * conj(t_i)))`.
///
/// **Higher scores = more likely**, so the `Scorer` implementation returns
/// the negated score (lower = more likely) for compatibility with
/// distance-based evaluation.
///
/// Initialization: Xavier uniform `Uniform(-sqrt(6/dim), sqrt(6/dim))`.
///
/// Reference: Trouillon et al. (2016), "Complex Embeddings for Simple Link
/// Prediction."
pub struct ComplEx {
    /// Flat entity embeddings: `[num_entities * dim * 2]` row-major.
    entities: Vec<f32>,
    /// Flat relation embeddings: `[num_relations * dim * 2]` row-major.
    relations: Vec<f32>,
    /// Complex dimension.
    dim: usize,
}

impl ComplEx {
    /// Create a new ComplEx model with random initialization.
    ///
    /// `dim` is the complex dimension (each embedding stores `2 * dim` floats).
    /// Xavier uniform initialization: `Uniform(-sqrt(6/dim), sqrt(6/dim))`.
    #[cfg(feature = "rand")]
    pub fn new(num_entities: usize, num_relations: usize, dim: usize) -> Self {
        let mut rng = rand::rng();
        let scale = (6.0_f32 / dim as f32).sqrt();
        Self::from_vecs(
            init_vecs(&mut rng, num_entities, dim * 2, scale),
            init_vecs(&mut rng, num_relations, dim * 2, scale),
            dim,
        )
    }

    /// Create from pre-built embedding vectors.
    ///
    /// Both `entities` and `relations` must have inner length `dim * 2`.
    ///
    /// # Panics
    ///
    /// Panics if any dimension is wrong.
    pub fn from_vecs(entities: Vec<Vec<f32>>, relations: Vec<Vec<f32>>, dim: usize) -> Self {
        assert_dims(&entities, dim * 2, "entity (re+im)");
        assert_dims(&relations, dim * 2, "relation (re+im)");
        Self {
            entities: flatten(&entities),
            relations: flatten(&relations),
            dim,
        }
    }

    /// Flat entity embeddings (row-major, stride = `dim() * 2`).
    pub fn entities_flat(&self) -> &[f32] {
        &self.entities
    }

    /// Flat relation embeddings (row-major, stride = `dim() * 2`).
    pub fn relations_flat(&self) -> &[f32] {
        &self.relations
    }

    /// Reconstruct entity embeddings as `Vec<Vec<f32>>`.
    pub fn entity_vecs(&self) -> Vec<Vec<f32>> {
        unflatten(&self.entities, self.dim * 2)
    }

    /// Reconstruct relation embeddings as `Vec<Vec<f32>>`.
    pub fn relation_vecs(&self) -> Vec<Vec<f32>> {
        unflatten(&self.relations, self.dim * 2)
    }

    /// Complex dimension.
    pub fn dim(&self) -> usize {
        self.dim
    }

    /// Raw score: `Re(sum_i h_i * r_i * conj(t_i))`.
    ///
    /// Higher = more likely. Use [`Scorer::score`] for the negated
    /// distance-compatible version.
    ///
    /// Entity/relation layout: first `dim` floats are real, next `dim` are imaginary.
    pub fn score_triple(&self, head: usize, relation: usize, tail: usize) -> f32 {
        let h = row(&self.entities, head, self.dim * 2);
        let r = row(&self.relations, relation, self.dim * 2);
        let t = row(&self.entities, tail, self.dim * 2);
        let dim = self.dim;

        let mut dot = 0.0_f64;
        for i in 0..dim {
            let h_re = h[i] as f64;
            let h_im = h[dim + i] as f64;
            let r_re = r[i] as f64;
            let r_im = r[dim + i] as f64;
            let t_re = t[i] as f64;
            let t_im = t[dim + i] as f64;

            let hr_re = h_re * r_re - h_im * r_im;
            let hr_im = h_re * r_im + h_im * r_re;

            dot += hr_re * t_re + hr_im * t_im;
        }
        dot as f32
    }
}

impl Scorer for ComplEx {
    /// Returns `-score_triple()` so that lower = more likely.
    fn score(&self, head: usize, relation: usize, tail: usize) -> f32 {
        -self.score_triple(head, relation, tail)
    }

    fn num_entities(&self) -> usize {
        self.entities.len() / (self.dim * 2)
    }

    fn score_all_tails(&self, head: usize, relation: usize) -> Vec<f32> {
        let h = row(&self.entities, head, self.dim * 2);
        let r = row(&self.relations, relation, self.dim * 2);
        let dim = self.dim;
        let n = self.num_entities();
        // Precompute h * r (complex) once.
        let mut hr_re = vec![0.0_f64; dim];
        let mut hr_im = vec![0.0_f64; dim];
        for i in 0..dim {
            let h_re = h[i] as f64;
            let h_im = h[dim + i] as f64;
            let r_re = r[i] as f64;
            let r_im = r[dim + i] as f64;
            hr_re[i] = h_re * r_re - h_im * r_im;
            hr_im[i] = h_re * r_im + h_im * r_re;
        }
        (0..n)
            .map(|ti| {
                let t = row(&self.entities, ti, dim * 2);
                let mut dot = 0.0_f64;
                for i in 0..dim {
                    // Re(hr * conj(t)) = hr_re * t_re + hr_im * t_im
                    dot += hr_re[i] * t[i] as f64 + hr_im[i] * t[dim + i] as f64;
                }
                -(dot as f32)
            })
            .collect()
    }

    fn score_all_heads(&self, relation: usize, tail: usize) -> Vec<f32> {
        let r = row(&self.relations, relation, self.dim * 2);
        let t = row(&self.entities, tail, self.dim * 2);
        let dim = self.dim;
        let n = self.num_entities();
        // Precompute r * conj(t) once.
        // conj(t) = (t_re, -t_im)
        // r * conj(t) = (r_re*t_re + r_im*t_im, r_im*t_re - r_re*t_im)
        let mut rc_re = vec![0.0_f64; dim];
        let mut rc_im = vec![0.0_f64; dim];
        for i in 0..dim {
            let r_re = r[i] as f64;
            let r_im = r[dim + i] as f64;
            let t_re = t[i] as f64;
            let t_im = t[dim + i] as f64;
            rc_re[i] = r_re * t_re + r_im * t_im;
            rc_im[i] = r_im * t_re - r_re * t_im;
        }
        (0..n)
            .map(|hi| {
                let h = row(&self.entities, hi, dim * 2);
                let mut dot = 0.0_f64;
                for i in 0..dim {
                    // Re(h * rc) = h_re * rc_re - h_im * rc_im
                    dot += h[i] as f64 * rc_re[i] - h[dim + i] as f64 * rc_im[i];
                }
                -(dot as f32)
            })
            .collect()
    }
}

// ---------------------------------------------------------------------------
// DistMult
// ---------------------------------------------------------------------------

/// DistMult: diagonal bilinear model in real space.
///
/// Score: `sum_i h_i * r_i * t_i`. Higher = more likely, so the `Scorer`
/// implementation returns the negated score.
///
/// DistMult can only model symmetric relations (`score(h,r,t) == score(t,r,h)`).
///
/// Initialization: Xavier uniform `Uniform(-sqrt(6/dim), sqrt(6/dim))`.
///
/// Reference: Yang et al. (2015), "Embedding Entities and Relations for
/// Learning and Inference in Knowledge Bases."
pub struct DistMult {
    /// Flat entity embeddings: `[num_entities * dim]` row-major.
    entities: Vec<f32>,
    /// Flat relation embeddings: `[num_relations * dim]` row-major.
    relations: Vec<f32>,
    dim: usize,
}

impl DistMult {
    /// Create a new DistMult model with random initialization.
    #[cfg(feature = "rand")]
    pub fn new(num_entities: usize, num_relations: usize, dim: usize) -> Self {
        let mut rng = rand::rng();
        let scale = (6.0_f32 / dim as f32).sqrt();
        Self::from_vecs(
            init_vecs(&mut rng, num_entities, dim, scale),
            init_vecs(&mut rng, num_relations, dim, scale),
            dim,
        )
    }

    /// Create from pre-built embedding vectors.
    ///
    /// # Panics
    ///
    /// Panics if any inner vector length differs from `dim`.
    pub fn from_vecs(entities: Vec<Vec<f32>>, relations: Vec<Vec<f32>>, dim: usize) -> Self {
        assert_dims(&entities, dim, "entity");
        assert_dims(&relations, dim, "relation");
        Self {
            entities: flatten(&entities),
            relations: flatten(&relations),
            dim,
        }
    }

    /// Flat entity embeddings (row-major, stride = `dim()`).
    pub fn entities_flat(&self) -> &[f32] {
        &self.entities
    }

    /// Flat relation embeddings (row-major, stride = `dim()`).
    pub fn relations_flat(&self) -> &[f32] {
        &self.relations
    }

    /// Reconstruct entity embeddings as `Vec<Vec<f32>>`.
    pub fn entity_vecs(&self) -> Vec<Vec<f32>> {
        unflatten(&self.entities, self.dim)
    }

    /// Reconstruct relation embeddings as `Vec<Vec<f32>>`.
    pub fn relation_vecs(&self) -> Vec<Vec<f32>> {
        unflatten(&self.relations, self.dim)
    }

    /// Embedding dimension.
    pub fn dim(&self) -> usize {
        self.dim
    }

    /// Raw score: `sum_i h_i * r_i * t_i`. Higher = more likely.
    pub fn score_triple(&self, head: usize, relation: usize, tail: usize) -> f32 {
        let h = row(&self.entities, head, self.dim);
        let r = row(&self.relations, relation, self.dim);
        let t = row(&self.entities, tail, self.dim);

        let mut dot = 0.0_f64;
        for i in 0..self.dim {
            dot += h[i] as f64 * r[i] as f64 * t[i] as f64;
        }
        dot as f32
    }
}

impl Scorer for DistMult {
    /// Returns `-score_triple()` so that lower = more likely.
    fn score(&self, head: usize, relation: usize, tail: usize) -> f32 {
        -self.score_triple(head, relation, tail)
    }

    fn num_entities(&self) -> usize {
        self.entities.len() / self.dim
    }

    fn score_all_tails(&self, head: usize, relation: usize) -> Vec<f32> {
        let h = row(&self.entities, head, self.dim);
        let r = row(&self.relations, relation, self.dim);
        let dim = self.dim;
        let n = self.num_entities();
        let mut hr = vec![0.0_f64; dim];
        for i in 0..dim {
            hr[i] = h[i] as f64 * r[i] as f64;
        }
        (0..n)
            .map(|ti| {
                let t = row(&self.entities, ti, dim);
                let mut dot = 0.0_f64;
                for i in 0..dim {
                    dot += hr[i] * t[i] as f64;
                }
                -(dot as f32)
            })
            .collect()
    }

    fn score_all_heads(&self, relation: usize, tail: usize) -> Vec<f32> {
        // DistMult is symmetric: score(h,r,t) == score(t,r,h).
        // Precompute r * t once, then dot with each head.
        let r = row(&self.relations, relation, self.dim);
        let t = row(&self.entities, tail, self.dim);
        let dim = self.dim;
        let n = self.num_entities();
        let mut rt = vec![0.0_f64; dim];
        for i in 0..dim {
            rt[i] = r[i] as f64 * t[i] as f64;
        }
        (0..n)
            .map(|hi| {
                let h = row(&self.entities, hi, dim);
                let mut dot = 0.0_f64;
                for i in 0..dim {
                    dot += rt[i] * h[i] as f64;
                }
                -(dot as f32)
            })
            .collect()
    }
}

// ---------------------------------------------------------------------------
// Helpers
// ---------------------------------------------------------------------------

/// Return a row slice from a flat embedding buffer.
#[inline]
fn row(data: &[f32], i: usize, stride: usize) -> &[f32] {
    &data[i * stride..(i + 1) * stride]
}

/// Flatten `Vec<Vec<f32>>` into contiguous row-major `Vec<f32>`.
fn flatten(vecs: &[Vec<f32>]) -> Vec<f32> {
    let total: usize = vecs.iter().map(|v| v.len()).sum();
    let mut flat = Vec::with_capacity(total);
    for v in vecs {
        flat.extend_from_slice(v);
    }
    flat
}

/// Reconstruct `Vec<Vec<f32>>` from a flat buffer.
fn unflatten(flat: &[f32], stride: usize) -> Vec<Vec<f32>> {
    flat.chunks_exact(stride).map(|c| c.to_vec()).collect()
}

fn assert_dims(vecs: &[Vec<f32>], expected: usize, label: &str) {
    for (i, v) in vecs.iter().enumerate() {
        assert_eq!(
            v.len(),
            expected,
            "{label} embedding {i} has length {}, expected {expected}",
            v.len()
        );
    }
}

#[cfg(feature = "rand")]
fn init_vecs(rng: &mut impl rand::Rng, count: usize, len: usize, scale: f32) -> Vec<Vec<f32>> {
    (0..count)
        .map(|_| (0..len).map(|_| rng.random_range(-scale..scale)).collect())
        .collect()
}

// ---------------------------------------------------------------------------
// Ensemble
// ---------------------------------------------------------------------------

/// Ensemble scorer that averages scores from multiple models.
///
/// Useful for combining snapshot models from cosine annealing (SnapE) or
/// models trained with different random seeds. Score averaging reduces
/// predictive multiplicity (Zhu et al., 2024).
///
/// ```
/// # use tranz::{DistMult, Scorer, EnsembledScorer};
/// let m1 = DistMult::new(100, 10, 50);
/// let m2 = DistMult::new(100, 10, 50);
/// let ensemble = EnsembledScorer::new(vec![Box::new(m1), Box::new(m2)]);
/// assert_eq!(ensemble.num_entities(), 100);
/// let score = ensemble.score(0, 0, 1);
/// ```
pub struct EnsembledScorer {
    models: Vec<Box<dyn Scorer>>,
}

impl EnsembledScorer {
    /// Create an ensemble from multiple scorers.
    ///
    /// # Panics
    ///
    /// Panics if `models` is empty or if models disagree on `num_entities()`.
    pub fn new(models: Vec<Box<dyn Scorer>>) -> Self {
        assert!(!models.is_empty(), "ensemble requires at least one model");
        let n = models[0].num_entities();
        for (i, m) in models.iter().enumerate().skip(1) {
            assert_eq!(
                m.num_entities(),
                n,
                "model {i} has {} entities, expected {n}",
                m.num_entities()
            );
        }
        Self { models }
    }
}

impl Scorer for EnsembledScorer {
    fn score(&self, head: usize, relation: usize, tail: usize) -> f32 {
        let sum: f32 = self
            .models
            .iter()
            .map(|m| m.score(head, relation, tail))
            .sum();
        sum / self.models.len() as f32
    }

    fn num_entities(&self) -> usize {
        self.models[0].num_entities()
    }

    fn score_all_tails(&self, head: usize, relation: usize) -> Vec<f32> {
        let n = self.num_entities();
        let k = self.models.len() as f32;
        let mut avg = vec![0.0_f32; n];
        for m in &self.models {
            let scores = m.score_all_tails(head, relation);
            for (i, &s) in scores.iter().enumerate() {
                avg[i] += s;
            }
        }
        for v in &mut avg {
            *v /= k;
        }
        avg
    }

    fn score_all_heads(&self, relation: usize, tail: usize) -> Vec<f32> {
        let n = self.num_entities();
        let k = self.models.len() as f32;
        let mut avg = vec![0.0_f32; n];
        for m in &self.models {
            let scores = m.score_all_heads(relation, tail);
            for (i, &s) in scores.iter().enumerate() {
                avg[i] += s;
            }
        }
        for v in &mut avg {
            *v /= k;
        }
        avg
    }
}

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

    // -- TransE ---------------------------------------------------------------

    #[test]
    fn transe_manual_score() {
        // h = [1, 0], r = [0, 1], t = [1, 1]
        // h + r - t = [0, 0], ||.||_2 = 0
        let model = TransE::from_vecs(
            vec![vec![1.0, 0.0], vec![1.0, 1.0]],
            vec![vec![0.0, 1.0]],
            2,
        );
        let score = model.score_triple(0, 0, 1);
        assert!((score - 0.0).abs() < 1e-6, "expected 0, got {score}");
    }

    #[test]
    fn transe_manual_nonzero() {
        // h = [3, 0], r = [0, 0], t = [0, 4]
        // h + r - t = [3, -4], ||.||_2 = 5
        let model = TransE::from_vecs(
            vec![vec![3.0, 0.0], vec![0.0, 4.0]],
            vec![vec![0.0, 0.0]],
            2,
        );
        let score = model.score_triple(0, 0, 1);
        assert!((score - 5.0).abs() < 1e-5, "expected 5, got {score}");
    }

    #[test]
    fn transe_scorer_trait() {
        let model = TransE::new(10, 3, 8);
        let scorer: &dyn Scorer = &model;
        assert_eq!(scorer.num_entities(), 10);
        let s = scorer.score(0, 0, 1);
        assert!(s.is_finite());
        assert!(s >= 0.0);
    }

    #[test]
    #[should_panic(expected = "entity embedding 0 has length 3, expected 2")]
    fn transe_rejects_bad_dims() {
        TransE::from_vecs(vec![vec![1.0, 2.0, 3.0]], vec![vec![1.0, 2.0]], 2);
    }

    // -- RotatE ---------------------------------------------------------------

    #[test]
    fn rotate_identity_rotation() {
        // Angle = 0 means rotation is identity (cos=1, sin=0).
        // h * r = h, so score = ||h - t||.
        // h = [1,0] (re=1, im=0), t = [1,0], score should be 0.
        let model = RotatE::from_vecs(
            vec![vec![1.0, 0.0], vec![1.0, 0.0]],
            vec![vec![0.0]], // angle = 0
            1,
            12.0,
        );
        let score = model.score_triple(0, 0, 1);
        assert!((score - 0.0).abs() < 1e-6, "expected 0, got {score}");
    }

    #[test]
    fn rotate_90_degrees() {
        use std::f32::consts::FRAC_PI_2;
        // h = (1, 0), rotate by pi/2 => (0, 1)
        // t = (0, 1), score should be 0.
        let model = RotatE::from_vecs(
            vec![
                vec![1.0, 0.0], // entity 0: re=1, im=0
                vec![0.0, 1.0], // entity 1: re=0, im=1
            ],
            vec![vec![FRAC_PI_2]], // 90 degrees
            1,
            12.0,
        );
        let score = model.score_triple(0, 0, 1);
        assert!(score < 1e-5, "expected ~0, got {score}");
    }

    #[test]
    fn rotate_scorer_trait() {
        let model = RotatE::new(10, 3, 8, 12.0);
        let scorer: &dyn Scorer = &model;
        assert_eq!(scorer.num_entities(), 10);
        let s = scorer.score(0, 0, 1);
        assert!(s.is_finite());
        assert!(s >= 0.0);
    }

    #[test]
    fn rotate_contiguous_layout_dim2() {
        use std::f32::consts::FRAC_PI_2;
        // dim=2: entity layout is [re0, re1, im0, im1]
        // Entity 0: (1+0i, 0+0i) -> [1, 0, 0, 0]
        // Entity 1: (0+1i, 0+0i) -> [0, 0, 1, 0]
        // Rotation by [pi/2, 0]: first component rotates 90 deg, second is identity.
        // (1+0i) * (cos(pi/2)+i*sin(pi/2)) = 0+1i
        // (0+0i) * (cos(0)+i*sin(0)) = 0+0i
        // Result: (0+1i, 0+0i) -> matches entity 1.
        let model = RotatE::from_vecs(
            vec![
                vec![1.0, 0.0, 0.0, 0.0], // entity 0: [re0=1, re1=0, im0=0, im1=0]
                vec![0.0, 0.0, 1.0, 0.0], // entity 1: [re0=0, re1=0, im0=1, im1=0]
            ],
            vec![vec![FRAC_PI_2, 0.0]], // rotate first dim by 90 deg
            2,
            12.0,
        );
        let score = model.score_triple(0, 0, 1);
        assert!(score < 1e-5, "expected ~0, got {score}");
    }

    // -- ComplEx ---------------------------------------------------------------

    #[test]
    fn complex_manual_score() {
        // h = (1+0i), r = (1+0i), t = (1+0i), dim=1
        // h*r*conj(t) = 1*1*1 = 1. Re = 1.
        let model = ComplEx::from_vecs(vec![vec![1.0, 0.0]], vec![vec![1.0, 0.0]], 1);
        let score = model.score_triple(0, 0, 0);
        assert!((score - 1.0).abs() < 1e-6, "expected 1.0, got {score}");
    }

    #[test]
    fn complex_imaginary_parts() {
        // h = (0+1i), r = (0+1i), t = (1+0i) [entity 1]
        // h*r = (0*0 - 1*1) + (0*1 + 1*0)i = -1 + 0i
        // conj(t) = (1 - 0i)
        // Re((-1+0i)*(1-0i)) = -1
        let model = ComplEx::from_vecs(
            vec![
                vec![0.0, 1.0], // entity 0: 0+1i
                vec![1.0, 0.0], // entity 1: 1+0i
            ],
            vec![vec![0.0, 1.0]],
            1,
        );
        let score = model.score_triple(0, 0, 1);
        assert!((score - (-1.0)).abs() < 1e-6, "expected -1.0, got {score}");
    }

    #[test]
    fn complex_scorer_negates() {
        // Scorer returns -score_triple, so a positive raw score becomes negative.
        let model = ComplEx::from_vecs(vec![vec![1.0, 0.0]], vec![vec![1.0, 0.0]], 1);
        let raw = model.score_triple(0, 0, 0);
        let via_scorer = model.score(0, 0, 0);
        assert!((via_scorer - (-raw)).abs() < 1e-6);
    }

    // -- DistMult -------------------------------------------------------------

    #[test]
    fn distmult_manual_score() {
        // h = [2, 3], r = [1, -1], t = [4, 5]
        // sum = 2*1*4 + 3*(-1)*5 = 8 - 15 = -7
        let model = DistMult::from_vecs(
            vec![vec![2.0, 3.0], vec![4.0, 5.0]],
            vec![vec![1.0, -1.0]],
            2,
        );
        let score = model.score_triple(0, 0, 1);
        assert!((score - (-7.0)).abs() < 1e-5, "expected -7.0, got {score}");
    }

    #[test]
    fn distmult_symmetric() {
        // DistMult is symmetric: score(h, r, t) == score(t, r, h).
        let model = DistMult::new(10, 3, 16);
        let s1 = model.score_triple(0, 0, 1);
        let s2 = model.score_triple(1, 0, 0);
        assert!(
            (s1 - s2).abs() < 1e-5,
            "DistMult should be symmetric: {s1} vs {s2}"
        );
    }

    #[test]
    fn distmult_scorer_negates() {
        let model = DistMult::from_vecs(vec![vec![1.0], vec![2.0]], vec![vec![3.0]], 1);
        let raw = model.score_triple(0, 0, 1);
        let via_scorer = model.score(0, 0, 1);
        assert!((via_scorer - (-raw)).abs() < 1e-6);
    }

    // -- Batch scoring --------------------------------------------------------

    #[test]
    fn score_all_tails_length() {
        let model = TransE::new(10, 3, 8);
        let scores = model.score_all_tails(0, 0);
        assert_eq!(scores.len(), 10);
        assert!(scores.iter().all(|s| s.is_finite()));
    }

    #[test]
    fn top_k_tails_sorted() {
        let model = TransE::new(20, 3, 8);
        let top = model.top_k_tails(0, 0, 5);
        assert_eq!(top.len(), 5);
        for w in top.windows(2) {
            assert!(w[0].1 <= w[1].1, "top_k should be sorted ascending");
        }
    }

    #[test]
    fn top_k_heads_sorted() {
        let model = TransE::new(20, 3, 8);
        let top = model.top_k_heads(0, 0, 5);
        assert_eq!(top.len(), 5);
        for w in top.windows(2) {
            assert!(w[0].1 <= w[1].1);
        }
    }

    // -- TransE norm --------------------------------------------------------

    #[test]
    fn transe_l1_vs_l2_differ() {
        let entities = vec![vec![3.0, 0.0], vec![0.0, 4.0]];
        let relations = vec![vec![0.0, 0.0]];
        let l1 = TransE::from_vecs_with_norm(entities.clone(), relations.clone(), 2, 1);
        let l2 = TransE::from_vecs_with_norm(entities, relations, 2, 2);
        let s1 = l1.score_triple(0, 0, 1); // |3|+|4| = 7
        let s2 = l2.score_triple(0, 0, 1); // sqrt(9+16) = 5
        assert!((s1 - 7.0).abs() < 1e-4, "L1 score should be 7, got {s1}");
        assert!((s2 - 5.0).abs() < 1e-4, "L2 score should be 5, got {s2}");
    }

    #[test]
    fn transe_l1_score_all_tails_consistent() {
        let model = TransE::from_vecs_with_norm(
            vec![vec![1.0, 0.0], vec![0.0, 1.0], vec![1.0, 1.0]],
            vec![vec![0.0, 0.0]],
            2,
            1,
        );
        let all = model.score_all_tails(0, 0);
        for (t, &score) in all.iter().enumerate() {
            let individual = model.score(0, 0, t);
            assert!(
                (score - individual).abs() < 1e-5,
                "L1 score_all_tails[{t}]={score} vs score()={individual}"
            );
        }
    }

    // -- Relation prediction ------------------------------------------------

    #[test]
    fn relation_prediction_returns_correct_count() {
        let model = DistMult::new(10, 5, 8);
        let scores = model.score_all_relations(0, 1, 5);
        assert_eq!(scores.len(), 5);
        assert!(scores.iter().all(|s| s.is_finite()));
    }

    #[test]
    fn top_k_relations_sorted() {
        let model = ComplEx::new(10, 5, 8);
        let top = model.top_k_relations(0, 1, 5, 3);
        assert_eq!(top.len(), 3);
        for w in top.windows(2) {
            assert!(w[0].1 <= w[1].1);
        }
    }
}