use std::fs::File;
use std::io;
use std::path::{Path, PathBuf};

use memmap2::Mmap;

use crate::format::{HEADER_SIZE, INDEX_ENTRY_SIZE, MAGIC, OVERFLOW_ENTRY_SIZE};

pub struct PersistentCompactIntVec {
    mmap: Mmap,
    n: usize,
    n_overflow: usize,
    pub step: usize,
    index: Vec<(usize, usize)>, // (slot, pos) — L1-resident sparse index
    primary_offset: usize,      // = HEADER_SIZE
    data_offset: usize,         // = HEADER_SIZE + n
    path: PathBuf,
}

impl PersistentCompactIntVec {
    /// Opens a persistent compact int vector from the given path.
    pub fn open(path: &Path) -> io::Result<Self> {
        let mmap = unsafe { Mmap::map(&File::open(path)?)? };

        if mmap.len() < HEADER_SIZE {
            return Err(io::Error::new(
                io::ErrorKind::InvalidData,
                "PCIV file too short",
            ));
        }
        if &mmap[0..4] != &MAGIC {
            return Err(io::Error::new(io::ErrorKind::InvalidData, "bad PCIV magic"));
        }

        let n = u64::from_le_bytes(mmap[8..16].try_into().unwrap()) as usize;
        let n_overflow = u64::from_le_bytes(mmap[16..24].try_into().unwrap()) as usize;
        let n_index = u64::from_le_bytes(mmap[24..32].try_into().unwrap()) as usize;
        let step = u64::from_le_bytes(mmap[32..40].try_into().unwrap()) as usize;

        let primary_offset = HEADER_SIZE;
        let data_offset = primary_offset + n;
        let index_offset = data_offset + n_overflow * OVERFLOW_ENTRY_SIZE;

        let mut index = Vec::with_capacity(n_index);
        for i in 0..n_index {
            let off = index_offset + i * INDEX_ENTRY_SIZE;
            let slot = u64::from_le_bytes(mmap[off..off + 8].try_into().unwrap()) as usize;
            let pos = u64::from_le_bytes(mmap[off + 8..off + 16].try_into().unwrap()) as usize;
            index.push((slot, pos));
        }

        Ok(Self {
            mmap,
            n,
            n_overflow,
            step,
            index,
            primary_offset,
            data_offset,
            path: path.to_path_buf(),
        })
    }

    /// Returns the path of the compact int vector file.
    pub fn path(&self) -> &Path {
        &self.path
    }

    /// Returns the length of the compact int vector.
    pub fn len(&self) -> usize {
        self.n
    }

    /// Returns whether the compact int vector is empty.
    pub fn is_empty(&self) -> bool {
        self.n == 0
    }

    /// Returns the value at the given slot.
    pub fn get(&self, slot: usize) -> u32 {
        match self.mmap[self.primary_offset + slot] {
            255 => self.overflow_get(slot),
            v => v as u32,
        }
    }

    /// Returns the value at the given slot from the overflow region.
    fn overflow_get(&self, slot: usize) -> u32 {
        let pos_start;
        let pos_end;

        if self.step == 0 {
            pos_start = 0;
            pos_end = self.n_overflow;
        } else {
            let i = self
                .index
                .partition_point(|&(s, _)| s <= slot)
                .saturating_sub(1);
            pos_start = self.index[i].1;
            pos_end = if i + 1 < self.index.len() {
                self.index[i + 1].1
            } else {
                self.n_overflow
            };
        }

        let mut lo = pos_start;
        let mut hi = pos_end;
        while lo < hi {
            let mid = lo + (hi - lo) / 2;
            match self.data_slot(mid).cmp(&slot) {
                std::cmp::Ordering::Equal => return self.data_value(mid),
                std::cmp::Ordering::Less => lo = mid + 1,
                std::cmp::Ordering::Greater => hi = mid,
            }
        }
        panic!("slot {slot} marked as overflow but not found")
    }

    #[inline]
    /// Returns the slot at the given index in the overflow region.
    fn data_slot(&self, i: usize) -> usize {
        let off = self.data_offset + i * OVERFLOW_ENTRY_SIZE;
        u64::from_le_bytes(self.mmap[off..off + 8].try_into().unwrap()) as usize
    }

    #[inline]
    /// Returns the value at the given index in the overflow region.
    fn data_value(&self, i: usize) -> u32 {
        let off = self.data_offset + i * OVERFLOW_ENTRY_SIZE + 8;
        u32::from_le_bytes(self.mmap[off..off + 4].try_into().unwrap())
    }

    #[inline]
    /// Returns the sum of all values in the compact int vector.
    pub fn sum(&self) -> u64 {
        self.iter().map(|v| v as u64).sum()
    }

    #[inline]
    /// Returns the Bray-Curtis distance between two compact int vectors.
    pub fn bray_dist(&self, other: &PersistentCompactIntVec) -> f64 {
        let sum_min = self.partial_bray_dist(other);
        let denom = self.sum() + other.sum();
        if denom == 0 {
            return 0.0;
        }
        1.0 - 2.0 * sum_min as f64 / denom as f64
    }

    /// Returns `Σ_slot min(self[slot], other[slot])` — the additive numerator of Bray-Curtis.
    /// The denominator `sum_a + sum_b` is obtained from `self.sum() + other.sum()`.
    pub fn partial_bray_dist(&self, other: &PersistentCompactIntVec) -> u64 {
        assert_eq!(self.n, other.len(), "length mismatch");
        self.iter()
            .zip(other.iter())
            .map(|(a, b)| a.min(b) as u64)
            .sum()
    }

    /// Returns the relative frequency Bray-Curtis distance between two compact int vectors.
    ///
    /// This is a variant of [`bray_dist`] that uses relative frequencies instead of raw counts.
    pub fn relfreq_bray_dist(&self, other: &PersistentCompactIntVec) -> f64 {
        assert_eq!(self.n, other.len(), "length mismatch");
        let sum_a = self.sum() as f64;
        let sum_b = other.sum() as f64;
        if sum_a == 0.0 && sum_b == 0.0 {
            return 0.0;
        }
        let sum_min = self.partial_relfreq_bray_dist(other, sum_a, sum_b);
        1.0 - sum_min
    }

    /// Returns the partial relative frequency Bray-Curtis distance between two compact int vectors.
    ///
    /// This is used internally by [`relfreq_bray_dist`] and to easily compute the relative frequency
    /// Bray-Curtis distance over a set of vector pairs.
    ///
    /// Arguments:
    /// - `other`: the other compact int vector to compare with
    /// - `sum_a`: the sum of the first vector's counts
    /// - `sum_b`: the sum of the second vector's counts
    ///
    /// Returns the sum of the minimum relative frequencies at each index.
    pub fn partial_relfreq_bray_dist(
        &self,
        other: &PersistentCompactIntVec,
        sum_a: f64,
        sum_b: f64,
    ) -> f64 {
        assert_eq!(self.n, other.len(), "length mismatch");
        let sum_min: f64 = self
            .iter()
            .zip(other.iter())
            .map(|(a, b)| {
                let pa = if sum_a > 0.0 { a as f64 / sum_a } else { 0.0 };
                let pb = if sum_b > 0.0 { b as f64 / sum_b } else { 0.0 };
                pa.min(pb)
            })
            .sum();
        sum_min
    }

    /// Returns the euclidean distance between two compact int vectors.
    pub fn euclidean_dist(&self, other: &PersistentCompactIntVec) -> f64 {
        self.partial_euclidean_dist(other).sqrt()
    }

    /// Returns the partial euclidean distance between two compact int vectors.
    ///
    /// This is used internally by [`euclidean_dist`] and to easily compute the euclidean distance
    /// over a set of vector pairs.
    ///
    /// The result is the sum of the squared differences between corresponding elements of the two
    /// vectors.
    pub fn partial_euclidean_dist(&self, other: &PersistentCompactIntVec) -> f64 {
        assert_eq!(self.n, other.len(), "length mismatch");
        self.iter()
            .zip(other.iter())
            .map(|(a, b)| {
                let d = a as f64 - b as f64;
                d * d
            })
            .sum()
    }

    /// Returns the relative frequency euclidean distance between two compact int vectors.
    ///
    /// This is a variant of [`euclidean_dist`] that uses relative frequencies instead of raw counts.
    pub fn relfreq_euclidean_dist(&self, other: &PersistentCompactIntVec) -> f64 {
        assert_eq!(self.n, other.len(), "length mismatch");
        let sum_a = self.sum() as f64;
        let sum_b = other.sum() as f64;
        if sum_a == 0.0 && sum_b == 0.0 {
            return 0.0;
        }
        self.partial_relfreq_euclidean_dist(other, sum_a, sum_b)
            .sqrt()
    }

    /// Returns the partial relative frequency euclidean distance between two compact int vectors.
    ///
    /// This is used internally by [`relfreq_euclidean_dist`] and to easily compute the relative frequency
    /// euclidean distance over a set of vector pairs.
    pub fn partial_relfreq_euclidean_dist(
        &self,
        other: &PersistentCompactIntVec,
        sum_a: f64,
        sum_b: f64,
    ) -> f64 {
        assert_eq!(self.n, other.len(), "length mismatch");
        self.iter()
            .zip(other.iter())
            .map(|(a, b)| {
                let pa = if sum_a > 0.0 { a as f64 / sum_a } else { 0.0 };
                let pb = if sum_b > 0.0 { b as f64 / sum_b } else { 0.0 };
                let d = pa - pb;
                d * d
            })
            .sum()
    }

    /// Returns the Euclidean distance between two compact int vectors using the Hellinger transform.
    ///
    /// The Hellinger transform is applied to the raw counts of each vector, and the result is
    /// the Euclidean distance between the transformed vectors. The Hellinger transform is defined
    /// as the square root of the relative frequencies.
    pub fn hellinger_euclidean_dist(&self, other: &PersistentCompactIntVec) -> f64 {
        assert_eq!(self.n, other.len(), "length mismatch");
        let sum_a = self.sum() as f64;
        let sum_b = other.sum() as f64;
        if sum_a == 0.0 && sum_b == 0.0 {
            return 0.0;
        }
        self.partial_hellinger_euclidean_dist(other, sum_a, sum_b)
            .sqrt()
    }

    /// Returns the partial Hellinger Euclidean distance between two compact int vectors.
    ///
    /// This is used internally by [`hellinger_euclidean_dist`] and to easily compute the Hellinger
    /// Euclidean distance over a set of vector pairs.
    pub fn partial_hellinger_euclidean_dist(
        &self,
        other: &PersistentCompactIntVec,
        sum_a: f64,
        sum_b: f64,
    ) -> f64 {
        assert_eq!(self.n, other.len(), "length mismatch");
        self.iter()
            .zip(other.iter())
            .map(|(a, b)| {
                let pa = if sum_a > 0.0 {
                    (a as f64 / sum_a).sqrt()
                } else {
                    0.0
                };
                let pb = if sum_b > 0.0 {
                    (b as f64 / sum_b).sqrt()
                } else {
                    0.0
                };
                let d = pa - pb;
                d * d
            })
            .sum()
    }

    pub fn hellinger_dist(&self, other: &PersistentCompactIntVec) -> f64 {
        self.hellinger_euclidean_dist(other) / std::f64::consts::SQRT_2
    }

    pub fn threshold_jaccard_dist(&self, other: &PersistentCompactIntVec, threshold: u32) -> f64 {
        assert_eq!(self.n, other.len(), "length mismatch");
        let (intersection, union) = self.partial_threshold_jaccard_dist(other, threshold);
        if union == 0 {
            return 0.0;
        }
        1.0 - intersection as f64 / union as f64
    }

    pub fn partial_threshold_jaccard_dist(
        &self,
        other: &PersistentCompactIntVec,
        threshold: u32,
    ) -> (u64, u64) {
        assert_eq!(self.n, other.len(), "length mismatch");
        self.iter()
            .zip(other.iter())
            .fold((0u64, 0u64), |(inter, uni), (a, b)| {
                let ap = a >= threshold;
                let bp = b >= threshold;
                (inter + (ap & bp) as u64, uni + (ap | bp) as u64)
            })
    }

    pub fn jaccard_dist(&self, other: &PersistentCompactIntVec) -> f64 {
        self.threshold_jaccard_dist(other, 1)
    }

    pub fn iter(&self) -> Iter<'_> {
        Iter {
            pciv: self,
            slot: 0,
            overflow_pos: 0,
        }
    }
}

impl<'a> IntoIterator for &'a PersistentCompactIntVec {
    type Item = u32;
    type IntoIter = Iter<'a>;

    fn into_iter(self) -> Iter<'a> {
        self.iter()
    }
}

pub struct Iter<'a> {
    pciv: &'a PersistentCompactIntVec,
    slot: usize,
    overflow_pos: usize,
}

impl ExactSizeIterator for Iter<'_> {}

impl Iterator for Iter<'_> {
    type Item = u32;

    fn next(&mut self) -> Option<u32> {
        if self.slot >= self.pciv.n {
            return None;
        }
        let v = self.pciv.mmap[self.pciv.primary_offset + self.slot];
        self.slot += 1;
        if v < 255 {
            Some(v as u32)
        } else {
            let val = self.pciv.data_value(self.overflow_pos);
            self.overflow_pos += 1;
            Some(val)
        }
    }

    fn size_hint(&self) -> (usize, Option<usize>) {
        let remaining = self.pciv.n - self.slot;
        (remaining, Some(remaining))
    }
}
