//! Interface and implementation to Flat index type.

use super::*;

use error::{Error, Result};
use std::mem;
use std::ptr;

/// Alias for the native implementation of a flat index.
pub type FlatIndex = FlatIndexImpl;

/// Native implementation of a flat index.
#[derive(Debug)]
pub struct FlatIndexImpl {
    inner: *mut FaissIndexFlat,
}

unsafe impl Send for FlatIndexImpl {}
unsafe impl Sync for FlatIndexImpl {}

impl CpuIndex for FlatIndexImpl {}

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

impl FlatIndexImpl {
    /// Create a new flat index.
    pub fn new(d: u32, metric: MetricType) -> Result<Self> {
        unsafe {
            let metric = metric as c_uint;
            let mut inner = ptr::null_mut();
            faiss_try!(faiss_IndexFlat_new_with(
                &mut inner,
                (d & 0x7FFF_FFFF) as idx_t,
                metric
            ));
            Ok(FlatIndexImpl { inner })
        }
    }

    /// Create a new flat index with L2 as the metric type.
    pub fn new_l2(d: u32) -> Result<Self> {
        FlatIndexImpl::new(d, MetricType::L2)
    }

    /// Create a new flat index with IP (inner product) as the metric type.
    pub fn new_ip(d: u32) -> Result<Self> {
        FlatIndexImpl::new(d, MetricType::InnerProduct)
    }

    /// Obtain a reference to the indexed data.
    pub fn xb(&self) -> &[f32] {
        unsafe {
            let mut xb = ptr::null_mut();
            let mut len = 0;
            faiss_IndexFlat_xb(self.inner, &mut xb, &mut len);
            ::std::slice::from_raw_parts(xb, len)
        }
    }

    /// Compute distance with a subset of vectors. `x` is a sequence of query
    /// vectors, size `n * d`, where `n` is inferred from the length of `x`.
    /// `labels` is a sequence of indexed vector ID's that should be compared
    /// for each query vector, size `n * k`, where `k` is inferred from the
    /// length of `labels`. Returns the corresponding output distances, size
    /// `n * k`.
    pub fn compute_distance_subset(&mut self, x: &[f32], labels: &[Idx]) -> Result<Vec<f32>> {
        unsafe {
            let n = x.len() / self.d() as usize;
            let k = labels.len() / n;
            let mut distances = vec![0.; n * k];
            faiss_try!(faiss_IndexFlat_compute_distance_subset(
                self.inner,
                n as idx_t,
                x.as_ptr(),
                k as idx_t,
                distances.as_mut_ptr(),
                labels.as_ptr()
            ));
            Ok(distances)
        }
    }
}

impl IndexImpl {
    /// Attempt a dynamic cast of an index to the flat index type.
    pub fn as_flat(self) -> Result<FlatIndexImpl> {
        unsafe {
            let new_inner = faiss_IndexFlat_cast(self.inner_ptr());
            if new_inner.is_null() {
                Err(Error::BadCast)
            } else {
                mem::forget(self);
                Ok(FlatIndexImpl { inner: new_inner })
            }
        }
    }
}

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

impl FromInnerPtr for FlatIndexImpl {
    unsafe fn from_inner_ptr(inner_ptr: *mut FaissIndex) -> Self {
        FlatIndexImpl {
            inner: inner_ptr as *mut FaissIndexFlat,
        }
    }
}

impl_native_index!(FlatIndex);

impl_native_index_clone!(FlatIndex);

impl ConcurrentIndex for FlatIndexImpl {
    fn assign(&self, query: &[f32], k: usize) -> Result<AssignSearchResult> {
        unsafe {
            let nq = query.len() / self.d() as usize;
            let mut out_labels = vec![0 as Idx; k * nq];
            faiss_try!(faiss_Index_assign(
                self.inner,
                nq as idx_t,
                query.as_ptr(),
                out_labels.as_mut_ptr(),
                k as i64
            ));
            Ok(AssignSearchResult { labels: out_labels })
        }
    }
    fn search(&self, query: &[f32], k: usize) -> Result<SearchResult> {
        unsafe {
            let nq = query.len() / self.d() as usize;
            let mut distances = vec![0_f32; k * nq];
            let mut labels = vec![0 as Idx; k * nq];
            faiss_try!(faiss_Index_search(
                self.inner,
                nq as idx_t,
                query.as_ptr(),
                k as idx_t,
                distances.as_mut_ptr(),
                labels.as_mut_ptr()
            ));
            Ok(SearchResult { distances, labels })
        }
    }
    fn range_search(&self, query: &[f32], radius: f32) -> Result<RangeSearchResult> {
        unsafe {
            let nq = (query.len() / self.d() as usize) as idx_t;
            let mut p_res: *mut FaissRangeSearchResult = ptr::null_mut();
            faiss_try!(faiss_RangeSearchResult_new(&mut p_res, nq));
            faiss_try!(faiss_Index_range_search(
                self.inner,
                nq,
                query.as_ptr(),
                radius,
                p_res
            ));
            Ok(RangeSearchResult { inner: p_res })
        }
    }
}

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

    const D: u32 = 8;

    #[test]
    fn flat_index_from_cast() {
        let mut index = index_factory(8, "Flat", MetricType::L2).unwrap();
        assert_eq!(index.is_trained(), true); // Flat index does not need training
        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 index: FlatIndexImpl = index.as_flat().unwrap();
        assert_eq!(index.is_trained(), true);
        assert_eq!(index.ntotal(), 5);
        let xb = index.xb();
        assert_eq!(xb.len(), 8 * 5);
    }

    #[test]
    fn index_clone() {
        let mut index = FlatIndexImpl::new_l2(D).unwrap();
        assert_eq!(index.is_trained(), true); // Flat index does not need training
        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 mut index: FlatIndexImpl = index.try_clone().unwrap();
            assert_eq!(index.is_trained(), true);
            assert_eq!(index.ntotal(), 5);
            {
                let xb = index.xb();
                assert_eq!(xb.len(), 8 * 5);
            }
            index.reset().unwrap();
            assert_eq!(index.ntotal(), 0);
        }
        assert_eq!(index.ntotal(), 5);
    }

    #[test]
    fn flat_index_search() {
        let mut index = FlatIndexImpl::new_l2(D).unwrap();
        assert_eq!(index.d(), D);
        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.; D as usize];
        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.; D as usize];
        // flat index can be used behind an immutable ref
        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.));

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

    #[test]
    fn flat_index_assign() {
        let mut index = FlatIndexImpl::new(D, MetricType::L2).unwrap();
        assert_eq!(index.d(), D);
        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.; D as usize];
        let result = index.assign(&my_query, 5).unwrap();
        assert_eq!(result.labels, vec![2, 1, 0, 3, 4]);

        let my_query = [100.; D as usize];
        // flat index can be used behind an immutable ref
        let result = (&index).assign(&my_query, 5).unwrap();
        assert_eq!(result.labels, vec![3, 4, 0, 1, 2]);

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

    #[test]
    fn flat_index_range_search() {
        let mut index = FlatIndexImpl::new(D, MetricType::L2).unwrap();
        assert_eq!(index.d(), D);
        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.; D as usize];
        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.));
    }

    #[test]
    fn index_transition() {
        let index = {
            let mut index = FlatIndexImpl::new_l2(D).unwrap();
            assert_eq!(index.d(), D);
            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., 4.,
                -4., -8., 1., 1., 2., 4., -1., 8., 8., 10., -10., -10., 10., -10., 10., 16., 16.,
                32., 25., 20., 20., 40., 15.,
            ];
            index.add(some_data).unwrap();
            assert_eq!(index.ntotal(), 5);

            unsafe {
                let inner = index.inner_ptr();
                // forget index, rebuild it into another object
                ::std::mem::forget(index);
                FlatIndexImpl::from_inner_ptr(inner)
            }
        };
        assert_eq!(index.ntotal(), 5);
    }
}