h3o 0.3.0

use super::CellIndex;
use crate::{
    coord::{CoordCube, CoordIJK, LocalIJK},
    error::{CompactionError, LocalIjError},
    index::bits,
    Direction, Resolution, DIRECTION_BITSIZE,
};
use std::cmp::max;

/// Iterator over a children cell index at a given resolution.
pub struct Children {
    /// Starting cell resolution.
    parent_resolution: Resolution,
    /// Targeted cell resolution.
    target_resolution: Resolution,

    /// Iterator scratch space, used to build cell index iteratively
    scratchpad: u64,
    /// Number of cell index to skip for pentagonal index, there is one per
    /// resolution.
    skip_count: i16,
    /// Remaining children at the targeted resolution.
    count: u64,
}

impl Children {
    /// Returns an iterator over the children cell index at the given
    /// resolution.
    pub fn new(index: CellIndex, resolution: Resolution) -> Self {
        Self {
            parent_resolution: index.resolution(),
            target_resolution: resolution,
            scratchpad: get_starting_state(index, resolution),
            skip_count: index
                .is_pentagon()
                .then(|| i16::from(u8::from(resolution)))
                .unwrap_or(-1),
            count: index.children_count(resolution),
        }
    }

    /// Increment the direction at `resolution` and return it.
    fn next_direction(&mut self, resolution: Resolution) -> u8 {
        // Shift the 1 to apply it on the right direction.
        let one = 1 << resolution.direction_offset();

        // Add one to the direction.
        //
        // Note that if the direction was 7 (unused) this automatically reset
        // the direction to [`Direction::CENTER`] (wraparound) AND increment the
        // direction of lower resolution (thanks to carry propagation).
        self.scratchpad += one;

        bits::get_direction(self.scratchpad, resolution)
    }
}

impl Iterator for Children {
    type Item = CellIndex;

    fn next(&mut self) -> Option<CellIndex> {
        // No more children, we're done.
        if self.count == 0 {
            return None;
        }

        // Extract the current index, to return it.
        let index = CellIndex::new_unchecked(self.scratchpad);
        self.count -= 1;

        // Prepare the next iteration, if any, by incrementing the scratchpad.
        if self.count != 0 {
            for resolution in Resolution::range(
                self.parent_resolution,
                self.target_resolution,
            )
            .rev()
            {
                // Move to the next direction value.
                let direction = self.next_direction(resolution);

                // First K axe of each resolution is skipped for pentagonal
                // index.
                if self.skip_count == i16::from(resolution)
                    && direction == u8::from(Direction::K)
                {
                    self.next_direction(resolution);
                    self.skip_count -= 1;
                }

                // If we have exhausted this resolution, move to the lower one.
                if Direction::try_from(direction).is_err() {
                    self.scratchpad =
                        bits::clr_direction(self.scratchpad, resolution);
                    continue;
                }

                break;
            }
        }
        Some(index)
    }

    fn size_hint(&self) -> (usize, Option<usize>) {
        let count = usize::try_from(self.count).unwrap_or(usize::MAX);
        (count, Some(count))
    }
}

impl ExactSizeIterator for Children {}

// -----------------------------------------------------------------------------

/// Return the starting state for the listing process.
fn get_starting_state(index: CellIndex, resolution: Resolution) -> u64 {
    let parent_resolution = index.resolution();

    // Compute the range of resolution to iterate over.
    //
    // e.g. if we list children for cell index at resolution 2 to resolution 6
    // we need to iterate of 4 resolution (resolutions 3, 4, 5 and 6).
    let range =
        usize::from(resolution).saturating_sub(parent_resolution.into());

    let mut scratchpad = u64::from(index);
    // If we have resolution between current and targeted one we clear their
    // directions.
    if range != 0 {
        // Mask with the right number of bit to cover the directions.
        let mask = (1 << (range * DIRECTION_BITSIZE)) - 1;
        // Mask offset required to clear the directions.
        let offset = resolution.direction_offset();

        // Clear directions.
        scratchpad &= !(mask << offset);
        // Set resolution.
        scratchpad = bits::set_resolution(scratchpad, resolution);
    }

    scratchpad
}

// -----------------------------------------------------------------------------

/// Iterator over a children cell index at a given resolution.
pub struct GridPathCells {
    /// Starting cell .
    anchor: CellIndex,
    /// Starting coordinate.
    start: CoordCube,
    // Path length.
    distance: i32,
    // Current position in the path.
    n: i32,

    /// Translation offset for the i component.
    i_step: f64,
    /// Translation offset for the j component.
    j_step: f64,
    /// Translation offset for the k component.
    k_step: f64,
}

impl GridPathCells {
    /// Returns an iterator over the children cell index at the given
    /// resolution.
    pub fn new(start: CellIndex, end: CellIndex) -> Result<Self, LocalIjError> {
        let anchor = start;

        // Get IJK coords for the start and end.
        let src = start.to_local_ijk(start)?;
        let dst = end.to_local_ijk(start)?;
        let distance = src.coord().distance(dst.coord());

        // Convert IJK to cube coordinates suitable for linear interpolation
        let start = CoordCube::from(*src.coord());
        let end = CoordCube::from(*dst.coord());

        let (i_step, j_step, k_step) = if distance == 0 {
            (0., 0., 0.)
        } else {
            (
                f64::from(end.i - start.i) / f64::from(distance),
                f64::from(end.j - start.j) / f64::from(distance),
                f64::from(end.k - start.k) / f64::from(distance),
            )
        };

        Ok(Self {
            anchor,
            start,
            distance,
            n: 0,
            i_step,
            j_step,
            k_step,
        })
    }
}

impl Iterator for GridPathCells {
    type Item = Result<CellIndex, LocalIjError>;

    fn next(&mut self) -> Option<Self::Item> {
        (self.n <= self.distance).then(|| {
            let coord = self.start.translate((
                self.i_step * f64::from(self.n),
                self.j_step * f64::from(self.n),
                self.k_step * f64::from(self.n),
            ));
            self.n += 1;

            // Convert cube -> ijk -> h3 index
            let local_ijk = LocalIJK {
                anchor: self.anchor,
                coord: CoordIJK::from(coord),
            };
            CellIndex::try_from(local_ijk)
        })
    }

    fn size_hint(&self) -> (usize, Option<usize>) {
        let count = usize::try_from(max(self.distance - self.n, 0))
            .unwrap_or(usize::MAX);
        (count, Some(count))
    }
}

impl ExactSizeIterator for GridPathCells {}

// -----------------------------------------------------------------------------

/// Iterator over a compacted set of cells.
pub struct Compact {
    /// Sorted list of unique uncompacted cells.
    cells: Vec<CellIndex>,
    /// Current position in the vector.
    index: usize,
    /// Resolution of the uncompacted cells.
    resolution: Resolution,
}

impl Compact {
    /// Returns an iterator over the compacted version of the input iterator.
    ///
    /// # Errors
    ///
    /// All cell indexes must be unique and have the same resolution, otherwise
    /// [`CompactionError`] is returned.
    pub fn new(
        cells: impl IntoIterator<Item = CellIndex>,
    ) -> Result<Self, CompactionError> {
        let mut cells = cells.into_iter();
        let size = cells.size_hint();
        let mut vec = Vec::with_capacity(size.1.unwrap_or(size.0));

        let resolution = if let Some(cell) = cells.next() {
            let resolution = cell.resolution();
            vec.push(cell);
            for cell in cells {
                if cell.resolution() != resolution {
                    return Err(CompactionError::HeterogeneousResolution);
                }
                vec.push(cell);
            }

            let old_len = vec.len();
            vec.sort_unstable();
            vec.dedup();
            let new_len = vec.len();

            // Dups were removed, not good.
            if new_len < old_len {
                return Err(CompactionError::DuplicateInput);
            }

            resolution
        } else {
            // If we're here the input stream is empty, so the value is not
            // important.
            Resolution::Zero
        };

        Ok(Self {
            cells: vec,
            index: 0,
            resolution,
        })
    }
}

impl Iterator for Compact {
    type Item = CellIndex;

    fn next(&mut self) -> Option<Self::Item> {
        let cell = *self.cells.get(self.index)?;

        // Base cell cannot be compacted, return as-is.
        if u8::from(self.resolution) == 0 {
            self.index += 1;
            return Some(cell);
        }

        // If this is the first cell, we may be able to compact it with the next
        // ones.
        if bits::get_direction(cell.into(), self.resolution) == 0 {
            for res in Resolution::range(Resolution::Zero, self.resolution) {
                let parent = cell.parent(res).expect("parent exists");
                let count =
                    usize::try_from(parent.children_count(self.resolution))
                        .expect("too many children");
                let end_index = self.index + count - 1;

                // Compute the expected last cell index at the given targeted
                // resolution.
                //
                // First compute a mask wide enough to cover the bit range to
                // update. Also compute the offset of the bit range.
                //
                // Next, compute the bits of the direction to update by masking
                // a constant (repeated `0b110`, i.e `6` on 3-bit) and shifting
                // accordingly.
                //
                // Finally, clear the targeted directions using the mask
                // (shifted and negated) and applies the new values.
                let diff =
                    usize::from(u8::from(self.resolution) - u8::from(res));
                let mask = (1_u64 << (diff * DIRECTION_BITSIZE)) - 1;
                let offset = self.resolution.direction_offset();
                let new_dirs = (0x0000_1b6d_b6db_6db6 & mask) << offset;
                // SAFETY: this bit twiddling produces a valid cell index.
                let expected = CellIndex::new_unchecked(
                    (u64::from(cell) & !(mask << offset)) | new_dirs,
                );
                if self.cells.get(end_index) == Some(&expected) {
                    self.index += count;
                    return Some(parent);
                }
            }
        }
        self.index += 1;
        Some(cell)
    }
}