//! Index interface and native implementations.
//!
//! This module hosts the vital [`Index`] trait, through which index building
//! and searching is made. It also contains [`index_factory`], a generic
//! function through which the user can retrieve most of the available index
//! implementations. A very typical usage scenario of this crate is to create
//! the index through this function, but some statically verified index types
//! are available as well.
//!
//! [`Index`]: trait.Index.html
//! [`index_factory`]: fn.index_factory.html

use error::{Error, Result};
use metric::MetricType;
use selector::IdSelector;
use std::ffi::CString;
use std::os::raw::c_uint;
use std::ptr;

use faiss_sys::*;

pub mod flat;
pub mod id_map;
pub mod io;
pub mod lsh;

#[cfg(feature = "gpu")]
pub mod gpu;

/// Primitive data type for identifying a vector in an index.
pub type Idx = idx_t;

/// Interface for a Faiss index. Most methods in this trait match the ones in
/// the native library, whereas some others serve as getters to the index'
/// parameters.
///
/// Although all methods appear to be available for all index implementations,
/// some methods may not be supported. For instance, a [`FlatIndex`] stores
/// vectors sequentially, and so does not support `add_with_ids` nor
/// `remove_ids`. Users are advised to read the Faiss wiki pages in order
/// to understand which index algorithms support which operations.
///
/// [`FlatIndex`]: flat/struct.FlatIndex.html
pub trait Index {
    /// Whether the Index does not require training, or if training is done already
    fn is_trained(&self) -> bool;

    /// The total number of vectors indexed
    fn ntotal(&self) -> u64;

    /// The dimensionality of the indexed vectors
    fn d(&self) -> u32;

    /// The metric type assumed by the index
    fn metric_type(&self) -> MetricType;

    /// Add new data vectors to the index.
    /// This assumes a C-contiguous memory slice of vectors, where the total
    /// number of vectors is `x.len() / d`.
    fn add(&mut self, x: &[f32]) -> Result<()>;

    /// Add new data vectors to the index with IDs.
    /// This assumes a C-contiguous memory slice of vectors, where the total
    /// number of vectors is `x.len() / d`.
    /// Not all index types may support this operation.
    fn add_with_ids(&mut self, x: &[f32], xids: &[Idx]) -> Result<()>;

    /// Train the underlying index with the given data.
    fn train(&mut self, x: &[f32]) -> Result<()>;

    /// Similar to `search`, but only provides the labels.
    fn assign(&mut self, q: &[f32], k: usize) -> Result<AssignSearchResult>;

    /// Perform a search for the `k` closest vectors to the given query vectors.
    fn search(&mut self, q: &[f32], k: usize) -> Result<SearchResult>;

    /// Perform a ranged search for the vectors closest to the given query vectors
    /// by the given radius.
    fn range_search(&mut self, q: &[f32], radius: f32) -> Result<RangeSearchResult>;

    /// Clear the entire index.
    fn reset(&mut self) -> Result<()>;

    /// Remove data vectors represented by IDs.
    fn remove_ids(&mut self, sel: &IdSelector) -> Result<i64>;
}

/// Sub-trait for native implementations of a Faiss index.
pub trait NativeIndex: Index {
    /// Retrieve a pointer to the native index object.
    fn inner_ptr(&self) -> *mut FaissIndex;
}

/// Trait for a Faiss index that can be safely searched over multiple threads.
/// Operations which do not modify the index are given a method taking an
/// immutable reference. This is not the default for every index type because
/// some implementations (such as the ones running on the GPU) do not allow
/// concurrent searches.
///
/// Users of these methods should still note that batched querying is
/// considerably faster than running queries one by one, even in parallel.
pub trait ConcurrentIndex: Index {
    /// Similar to `search`, but only provides the labels.
    fn assign(&self, q: &[f32], k: usize) -> Result<AssignSearchResult>;

    /// Perform a search for the `k` closest vectors to the given query vectors.
    fn search(&self, q: &[f32], k: usize) -> Result<SearchResult>;

    /// Perform a ranged search for the vectors closest to the given query vectors
    /// by the given radius.
    fn range_search(&self, q: &[f32], radius: f32) -> Result<RangeSearchResult>;
}

/// Trait for Faiss index types known to be running on the CPU.
pub trait CpuIndex: Index {}

/// Trait for Faiss index types which can be built from a pointer
/// to a native implementation.
pub trait FromInnerPtr: NativeIndex {
    /// Create an index using the given pointer to a native object.
    ///
    /// # Safety
    ///
    /// `inner_ptr` must point to a valid, non-freed CPU index, and cannot be
    /// shared across multiple instances. The inner index must also be
    /// compatible with the target `NativeIndex` type according to the native
    /// class hierarchy. For example, creating an `IndexImpl` out of a pointer
    /// to `FaissIndexFlatL2` is valid, but creating a `FlatIndex` out of a
    /// plain `FaissIndex` can cause undefined behaviour.
    unsafe fn from_inner_ptr(inner_ptr: *mut FaissIndex) -> Self;
}

/// The outcome of an index assign operation.
#[derive(Debug, Clone, PartialEq)]
pub struct AssignSearchResult {
    pub labels: Vec<Idx>,
}

/// The outcome of an index search operation.
#[derive(Debug, Clone, PartialEq)]
pub struct SearchResult {
    pub distances: Vec<f32>,
    pub labels: Vec<Idx>,
}

/// The outcome of an index range search operation.
#[derive(Debug, Clone, PartialEq)]
pub struct RangeSearchResult {
    inner: *mut FaissRangeSearchResult,
}

impl RangeSearchResult {
    pub fn nq(&self) -> usize {
        unsafe { faiss_RangeSearchResult_nq(self.inner) }
    }

    pub fn lims(&self) -> &[usize] {
        unsafe {
            let mut lims_ptr = ptr::null_mut();
            faiss_RangeSearchResult_lims(self.inner, &mut lims_ptr);
            ::std::slice::from_raw_parts(lims_ptr, self.nq() + 1)
        }
    }

    /// getter for labels and respective distances (not sorted):
    /// result for query `i` is `labels[lims[i] .. lims[i+1]]`
    pub fn distance_and_labels(&self) -> (&[f32], &[Idx]) {
        let lims = self.lims();
        let full_len = lims.last().cloned().unwrap_or(0);
        unsafe {
            let mut distances_ptr = ptr::null_mut();
            let mut labels_ptr = ptr::null_mut();
            faiss_RangeSearchResult_labels(self.inner, &mut labels_ptr, &mut distances_ptr);
            let distances = ::std::slice::from_raw_parts(distances_ptr, full_len);
            let labels = ::std::slice::from_raw_parts(labels_ptr, full_len);
            (distances, labels)
        }
    }

    /// getter for labels and respective distances (not sorted):
    /// result for query `i` is `labels[lims[i] .. lims[i+1]]`
    pub fn distance_and_labels_mut(&self) -> (&mut [f32], &mut [Idx]) {
        unsafe {
            let buf_size = faiss_RangeSearchResult_buffer_size(self.inner);
            let mut distances_ptr = ptr::null_mut();
            let mut labels_ptr = ptr::null_mut();
            faiss_RangeSearchResult_labels(self.inner, &mut labels_ptr, &mut distances_ptr);
            let distances = ::std::slice::from_raw_parts_mut(distances_ptr, buf_size);
            let labels = ::std::slice::from_raw_parts_mut(labels_ptr, buf_size);
            (distances, labels)
        }
    }

    /// getter for distances (not sorted):
    /// result for query `i` is `distances[lims[i] .. lims[i+1]]`
    pub fn distances(&self) -> &[f32] {
        self.distance_and_labels().0
    }

    /// getter for distances (not sorted):
    /// result for query `i` is `distances[lims[i] .. lims[i+1]]`
    pub fn distances_mut(&mut self) -> &mut [f32] {
        self.distance_and_labels_mut().0
    }

    /// getter for labels (not sorted):
    /// result for query `i` is `labels[lims[i] .. lims[i+1]]`
    pub fn labels(&self) -> &[Idx] {
        self.distance_and_labels().1
    }

    /// getter for labels (not sorted):
    /// result for query `i` is `labels[lims[i] .. lims[i+1]]`
    pub fn labels_mut(&mut self) -> &mut [Idx] {
        self.distance_and_labels_mut().1
    }
}

impl Drop for RangeSearchResult {
    fn drop(&mut self) {
        unsafe {
            faiss_RangeSearchResult_free(self.inner);
        }
    }
}

/// Native implementation of a Faiss Index
/// running on the CPU.
#[derive(Debug)]
pub struct IndexImpl {
    inner: *mut FaissIndex,
}

unsafe impl Send for IndexImpl {}
unsafe impl Sync for IndexImpl {}

impl CpuIndex for IndexImpl {}

impl Drop for IndexImpl {
    fn drop(&mut self) {
        unsafe {
            faiss_Index_free(self.inner);
        }
    }
}

impl IndexImpl {
    pub fn inner_ptr(&self) -> *mut FaissIndex {
        self.inner
    }
}

impl NativeIndex for IndexImpl {
    fn inner_ptr(&self) -> *mut FaissIndex {
        self.inner
    }
}

impl FromInnerPtr for IndexImpl {
    unsafe fn from_inner_ptr(inner_ptr: *mut FaissIndex) -> Self {
        IndexImpl { inner: inner_ptr }
    }
}

impl_native_index!(IndexImpl);

impl_native_index_clone!(IndexImpl);

/// Use the index factory to create a native instance of a Faiss index, for `d`-dimensional
/// vectors. `description` should follow the exact guidelines as the native Faiss interface
/// (see the [Faiss wiki](https://github.com/facebookresearch/faiss/wiki/Faiss-indexes) for examples).
///
/// # Error
///
/// This function returns an error if the description contains any byte with the value `\0` (since
/// it cannot be converted to a C string), or if the internal index factory operation fails.
pub fn index_factory<D>(d: u32, description: D, metric: MetricType) -> Result<IndexImpl>
where
    D: AsRef<str>,
{
    unsafe {
        let metric = metric as c_uint;
        let description =
            CString::new(description.as_ref()).map_err(|_| Error::IndexDescription)?;
        let mut index_ptr = ::std::ptr::null_mut();
        faiss_try!(faiss_index_factory(
            &mut index_ptr,
            (d & 0x7FFF_FFFF) as i32,
            description.as_ptr(),
            metric
        ));
        Ok(IndexImpl { inner: index_ptr })
    }
}

#[cfg(test)]
mod tests {
    use super::{index_factory, Index};
    use metric::MetricType;

    #[test]
    fn index_factory_flat() {
        let index = index_factory(64, "Flat", MetricType::L2).unwrap();
        assert_eq!(index.is_trained(), true); // Flat index does not need training
        assert_eq!(index.ntotal(), 0);
    }

    #[test]
    fn index_factory_ivf_flat() {
        let index = index_factory(64, "IVF8,Flat", MetricType::L2).unwrap();
        assert_eq!(index.is_trained(), false);
        assert_eq!(index.ntotal(), 0);
    }

    #[test]
    fn index_factory_sq() {
        let index = index_factory(64, "SQ8", MetricType::L2).unwrap();
        assert_eq!(index.is_trained(), false);
        assert_eq!(index.ntotal(), 0);
    }

    #[test]
    fn index_factory_pq() {
        let index = index_factory(64, "PQ8", MetricType::L2).unwrap();
        assert_eq!(index.is_trained(), false);
        assert_eq!(index.ntotal(), 0);
    }

    #[test]
    fn index_factory_ivf_sq() {
        let index = index_factory(64, "IVF8,SQ4", MetricType::L2).unwrap();
        assert_eq!(index.is_trained(), false);
        assert_eq!(index.ntotal(), 0);

        let index = index_factory(64, "IVF8,SQ8", MetricType::L2).unwrap();
        assert_eq!(index.is_trained(), false);
        assert_eq!(index.ntotal(), 0);
    }

    #[test]
    fn index_factory_hnsw() {
        let index = index_factory(64, "HNSW8", MetricType::L2).unwrap();
        assert_eq!(index.is_trained(), true); // training is not required
        assert_eq!(index.ntotal(), 0);
    }

    #[test]
    fn bad_index_factory_description() {
        let r = index_factory(64, "fdnoyq", MetricType::L2);
        assert!(r.is_err());
        let r = index_factory(64, "Flat\0Flat", MetricType::L2);
        assert!(r.is_err());
    }

    #[test]
    fn index_clone() {
        let mut index = index_factory(4, "Flat", MetricType::L2).unwrap();
        let some_data = &[
            7.5_f32, -7.5, 7.5, -7.5, 7.5, 7.5, 7.5, 7.5, -1., 1., 1., 1., 1., 1., 1., -1., 0., 0.,
            0., 1., 1., 0., 0., -1.,
        ];

        index.add(some_data).unwrap();
        assert_eq!(index.ntotal(), 6);

        let mut index2 = index.try_clone().unwrap();
        assert_eq!(index2.ntotal(), 6);

        let some_more_data = &[
            100., 100., 100., 100., -100., 100., 100., 100., 120., 100., 100., 105., -100., 100.,
            100., 105.,
        ];

        index2.add(some_more_data).unwrap();
        assert_eq!(index.ntotal(), 6);
        assert_eq!(index2.ntotal(), 10);
    }

    #[test]
    fn flat_index_search() {
        let mut index = index_factory(8, "Flat", MetricType::L2).unwrap();
        let some_data = &[
            7.5_f32, -7.5, 7.5, -7.5, 7.5, 7.5, 7.5, 7.5, -1., 1., 1., 1., 1., 1., 1., -1., 0., 0.,
            0., 1., 1., 0., 0., -1., 100., 100., 100., 100., -100., 100., 100., 100., 120., 100.,
            100., 105., -100., 100., 100., 105.,
        ];
        index.add(some_data).unwrap();
        assert_eq!(index.ntotal(), 5);

        let my_query = [0.; 8];
        let result = index.search(&my_query, 5).unwrap();
        assert_eq!(result.labels, vec![2, 1, 0, 3, 4]);
        assert!(result.distances.iter().all(|x| *x > 0.));

        let my_query = [100.; 8];
        let result = index.search(&my_query, 5).unwrap();
        assert_eq!(result.labels, vec![3, 4, 0, 1, 2]);
        assert!(result.distances.iter().all(|x| *x > 0.));

        let my_query = vec![
            0., 0., 0., 0., 0., 0., 0., 0., 100., 100., 100., 100., 100., 100., 100., 100.,
        ];
        let result = index.search(&my_query, 5).unwrap();
        assert_eq!(result.labels, vec![2, 1, 0, 3, 4, 3, 4, 0, 1, 2]);
        assert!(result.distances.iter().all(|x| *x > 0.));
    }

    #[test]
    fn flat_index_assign() {
        let mut index = index_factory(8, "Flat", MetricType::L2).unwrap();
        assert_eq!(index.d(), 8);
        assert_eq!(index.ntotal(), 0);
        let some_data = &[
            7.5_f32, -7.5, 7.5, -7.5, 7.5, 7.5, 7.5, 7.5, -1., 1., 1., 1., 1., 1., 1., -1., 0., 0.,
            0., 1., 1., 0., 0., -1., 100., 100., 100., 100., -100., 100., 100., 100., 120., 100.,
            100., 105., -100., 100., 100., 105.,
        ];
        index.add(some_data).unwrap();
        assert_eq!(index.ntotal(), 5);

        let my_query = [0.; 8];
        let result = index.assign(&my_query, 5).unwrap();
        assert_eq!(result.labels, vec![2, 1, 0, 3, 4]);

        let my_query = [0.; 32];
        let result = index.assign(&my_query, 5).unwrap();
        assert_eq!(
            result.labels,
            vec![2, 1, 0, 3, 4, 2, 1, 0, 3, 4, 2, 1, 0, 3, 4, 2, 1, 0, 3, 4]
        );

        let my_query = [100.; 8];
        let result = index.assign(&my_query, 5).unwrap();
        assert_eq!(result.labels, vec![3, 4, 0, 1, 2]);

        let my_query = vec![
            0., 0., 0., 0., 0., 0., 0., 0., 100., 100., 100., 100., 100., 100., 100., 100.,
        ];
        let result = index.assign(&my_query, 5).unwrap();
        assert_eq!(result.labels, vec![2, 1, 0, 3, 4, 3, 4, 0, 1, 2]);

        index.reset().unwrap();
        assert_eq!(index.ntotal(), 0);
    }

    #[test]
    fn flat_index_range_search() {
        let mut index = index_factory(8, "Flat", MetricType::L2).unwrap();
        let some_data = &[
            7.5_f32, -7.5, 7.5, -7.5, 7.5, 7.5, 7.5, 7.5, -1., 1., 1., 1., 1., 1., 1., -1., 0., 0.,
            0., 1., 1., 0., 0., -1., 100., 100., 100., 100., -100., 100., 100., 100., 120., 100.,
            100., 105., -100., 100., 100., 105.,
        ];
        index.add(some_data).unwrap();
        assert_eq!(index.ntotal(), 5);

        let my_query = [0.; 8];
        let result = index.range_search(&my_query, 8.125).unwrap();
        let (distances, labels) = result.distance_and_labels();
        assert!(labels == &[1, 2] || labels == &[2, 1]);
        assert!(distances.iter().all(|x| *x > 0.));
    }
}