moc 0.19.2

//! Code specific to HEALPix ranges
//! Assuming the NSIDE is a power of 2, the code for NESTED indices or RING indices is the same.

use std::ops::Range;

use num::{CheckedAdd, One, Zero};

use crate::elemset::range::{HpxRanges, SNORanges};
use crate::idx::Idx;
use crate::qty::{Hpx, MocQty};

pub struct HpxToUniqIter<T>
where
  T: Idx + CheckedAdd,
{
  ranges: HpxRanges<T>,
  id: usize,
  buffer: Vec<Range<T>>,
  depth: u8,
  shift: u32,
  off: T,
  depth_off: T,
}

impl<T> HpxToUniqIter<T>
where
  T: Idx + CheckedAdd,
{
  pub fn new(ranges: HpxRanges<T>) -> HpxToUniqIter<T> {
    let id = 0;
    let buffer = Vec::<Range<T>>::new();
    let depth = 0;
    // ((T::MAXDEPTH - depth) << 1) as u32;
    let shift = Hpx::<T>::MAX_SHIFT; // ok since depth init == 0

    let mut off: T = One::one();
    off = off.unsigned_shl(shift) - One::one();

    let mut depth_off: T = One::one();
    depth_off = depth_off.unsigned_shl(((depth + 1) << 1) as u32);

    HpxToUniqIter {
      ranges,
      id,
      buffer,

      depth,
      shift,
      off,
      depth_off,
    }
  }
}

impl<T> Iterator for HpxToUniqIter<T>
where
  T: Idx + CheckedAdd,
{
  type Item = Range<T>;

  fn next(&mut self) -> Option<Self::Item> {
    while !self.ranges.is_empty() {
      let start_id = self.id;
      let end_id = self.ranges.0 .0.len();
      for i in start_id..end_id {
        let range = &self.ranges.0 .0[i];
        self.id += 1;

        let t1 = range.start + self.off;
        let t2 = range.end;

        let pix1 = t1.unsigned_shr(self.shift);
        let pix2 = t2.unsigned_shr(self.shift);

        let c1 = pix1.unsigned_shl(self.shift);
        let c2 = pix2.unsigned_shl(self.shift);

        if c2 > c1 {
          self.buffer.push(c1..c2);

          let e1 = self.depth_off.checked_add(&pix1).unwrap();
          let e2 = self.depth_off.checked_add(&pix2).unwrap();

          return Some(e1..e2);
        }
      }

      // new_from_sorted?
      let buffer_ranges = HpxRanges::<T>::new_from(self.buffer.clone());
      self.ranges = self.ranges.difference(&buffer_ranges);
      self.id = 0;
      self.buffer.clear();

      self.depth += 1;
      assert!(self.depth <= Hpx::<T>::MAX_DEPTH || self.ranges.is_empty());
      if self.depth > Hpx::<T>::MAX_DEPTH && self.ranges.is_empty() {
        break;
      }

      // Recompute the constants for the new depth
      self.shift = Hpx::<T>::shift_from_depth_max(self.depth) as u32;
      self.off = One::one();
      self.off = self.off.unsigned_shl(self.shift) - One::one();

      self.depth_off = One::one();
      self.depth_off = self.depth_off.unsigned_shl(((self.depth + 1) << 1) as u32);
    }
    None
  }
}

// Iterator responsible for converting
// ranges of uniq numbers to ranges of
// nested numbers
pub struct UniqToHpxIter<T>
where
  T: Idx + CheckedAdd,
{
  ranges: HpxRanges<T>,
  cur: T,
  id: usize,
}

impl<T> UniqToHpxIter<T>
where
  T: Idx + CheckedAdd,
{
  pub fn new(ranges: HpxRanges<T>) -> UniqToHpxIter<T> {
    let id = 0;

    let cur = if ranges.is_empty() {
      Zero::zero()
    } else {
      ranges[id].start
    };
    UniqToHpxIter { ranges, cur, id }
  }
}

impl<T> Iterator for UniqToHpxIter<T>
where
  T: Idx + CheckedAdd,
{
  type Item = Range<T>;

  fn next(&mut self) -> Option<Self::Item> {
    // Iteration through the ranges
    if self.id < self.ranges.0 .0.len() {
      // We get the depth/ipix values of the
      // current uniq number
      let (depth, ipix) = Hpx::<T>::from_uniq_hpx(self.cur);

      // We compute the number of bit to shift
      let shift = ((Hpx::<T>::MAX_DEPTH - depth) << 1) as u32;

      let one: T = One::one();
      // Compute the final nested range
      // for the depth given
      let e1 = ipix.unsigned_shl(shift);
      let e2 = ipix.checked_add(&one).unwrap().unsigned_shl(shift);

      self.cur = self.cur.checked_add(&one).unwrap();

      let end = self.ranges[self.id].end;
      if self.cur == end {
        self.id += 1;

        if self.id < self.ranges.0 .0.len() {
          self.cur = self.ranges[self.id].start;
        }
      }

      Some(e1..e2)
    } else {
      None
    }
  }
}

/// IMPORTANT: the iterator is ordered according first to the cell depth
/// and then to the cell index.
/// See `ranges2cells` bench, using `CellMOCIteratorFromRanges` and then sorting the result
/// may be more efficient (x3 on the only bench done so far).
pub struct HpxUniq2DepthIdxIter<T>
where
  T: Idx + CheckedAdd,
{
  ranges: HpxRanges<T>,
  current: Option<Range<T>>,

  last: Option<T>,
  depth: u8,
  shift: u32,

  offset: T,
  depth_offset: T,
}

impl<T> HpxUniq2DepthIdxIter<T>
where
  T: Idx + CheckedAdd,
{
  pub fn new(ranges: HpxRanges<T>) -> HpxUniq2DepthIdxIter<T> {
    let depth = 0;
    let shift = ((Hpx::<T>::MAX_DEPTH - depth) << 1) as u32;
    // More generic (but slower): let shift = Hpx::<T>::shift_from_depth_max(depth)

    let mut offset: T = One::one();
    offset = offset.unsigned_shl(shift) - One::one();

    let mut depth_offset: T = One::one();
    depth_offset = depth_offset.unsigned_shl(((depth + 1) << 1) as u32);

    let current = None;
    let last = None;
    HpxUniq2DepthIdxIter {
      ranges,
      current,
      last,
      depth,
      shift,
      offset,
      depth_offset,
    }
  }

  fn next_item_range(&mut self) -> Option<(i8, T)> {
    if let Some(current) = self.current.clone() {
      let last = self.last.unwrap();
      if last < current.end {
        let (depth, pix) = Hpx::<T>::from_uniq_hpx(last);
        self.last = last.checked_add(&One::one());

        Some((depth as i8, pix))
      } else {
        self.current = None;
        self.last = None;
        None
      }
    } else {
      None
    }
  }
}

impl<T> Iterator for HpxUniq2DepthIdxIter<T>
where
  T: Idx + CheckedAdd,
{
  type Item = (i8, T);

  fn next(&mut self) -> Option<Self::Item> {
    let next_depth_pix = self.next_item_range();
    if next_depth_pix.is_some() {
      next_depth_pix
    } else {
      while !self.ranges.is_empty() {
        for range in self.ranges.iter() {
          let t1 = range.start + self.offset;
          let t2 = range.end;

          let pix1 = t1.unsigned_shr(self.shift);
          let pix2 = t2.unsigned_shr(self.shift);

          let c1 = pix1.unsigned_shl(self.shift);
          let c2 = pix2.unsigned_shl(self.shift);

          if c2 > c1 {
            let range_to_remove = HpxRanges::<T>::new_unchecked(vec![c1..c2]);
            self.ranges = self.ranges.difference(&range_to_remove);

            let e1 = self.depth_offset.checked_add(&pix1).unwrap();
            let e2 = self.depth_offset.checked_add(&pix2).unwrap();

            self.last = Some(e1);
            self.current = Some(e1..e2);

            return self.next_item_range();
          }
        }
        self.depth += 1;

        // Recompute the constants for the new depth
        self.shift = ((Hpx::<T>::MAX_DEPTH - self.depth) << 1) as u32;
        // <=> let shift = Hpx::<T>::shift_from_depth_max(depth)

        self.offset = One::one();
        self.offset = self.offset.unsigned_shl(self.shift) - One::one();

        self.depth_offset = One::one();
        self.depth_offset = self.depth_offset.unsigned_shl((2 * self.depth + 2) as u32);
      }
      None
    }
  }
}

/*
// I rename Nested in Hpx because it works the same for Nested or Ring indexes.
// Since Ranges is now MocRanges and we already have the type HpxRanges
// But this is specific to HEALPix, because of methods to_uniq and complement.
#[derive(Debug, Clone)]
pub struct HpxRanges<T: MocIdx> {
    ranges: ranges::HpxRanges<T, Hpx<T>>,
}

pub type HpxRanges =
// pub struct Hpx<T: MocIdx> (std::marker::PhantomData<T>);
// pub type HpxRange<T> = MocRange<T, Hpx<T>>;
// pub type TimeRange<T> = MocRange<T, Time<T>>;

impl<T: MocIdx> HpxRanges<T> {
    pub fn new(data: Vec<Range<T>>) -> Self {
        let ranges = MocRanges::<T, Hpx<T>>::new(data);

        HpxRanges { ranges }
    }

    /// Make the NestedRanges<T> consistent
    ///
    /// # Info
    ///
    /// By construction, the data are sorted so that it is possible (see the new
    /// method definition above) to merge the overlapping ranges.
    pub fn make_consistent(mut self) -> Self {
        self.ranges = self.ranges.make_consistent();
        self
    }

    /// Divide the nested ranges into ranges of length
    /// 4**(<T>::MAXDEPTH - min_depth) if they are bigger than
    /// this size.
    ///
    /// # Info
    ///
    /// This requires min_depth to be defined between [0, <T>::MAXDEPTH]
    pub fn divide(mut self, min_depth: u8) -> Self {
        self.ranges = self.ranges.divide(min_depth);
        self
    }

    pub fn to_uniq(self) -> HpxUniqRanges<T> {
        let uniq_data = HpxToUniqIter::new(self.ranges).collect::<Vec<_>>();
        HpxUniqRanges::<T>::new(uniq_data).make_consistent()
    }

    pub fn depth(&self) -> u8 {
        self.ranges.depth()
    }

    /// # Input
    /// * `dim` dimension of the quantity
    ///     + 1 for time
    ///     + 2 for equatorial coordinates
    /// * `delta_depth`: difference between the target depth and the current depth
    pub fn degrade(&mut self, depth: u8) {
        self.ranges.degrade(depth)
    }

    pub fn intersects(&self, x: &Range<T>) -> bool {
        self.ranges.intersects(x)
    }

    pub fn contains(&self, x: &Range<T>) -> bool {
        self.ranges.contains(x)
    }

    pub fn is_empty(&self) -> bool {
        self.ranges.is_empty()
    }

    pub fn iter(&self) -> Iter<Range<T>> {
        self.ranges.iter()
    }

    pub fn iter_depth_pix(self) -> HpxUniq2DepthIdxIter<T> {
        HpxUniq2DepthIdxIter::<T>::new(self.ranges)
    }

    pub fn union(&self, other: &Self) -> Self {
        let ranges = self.ranges.union(&other.ranges);
        HpxRanges { ranges }
    }

    pub fn intersection(&self, other: &Self) -> Self {
        let ranges = self.ranges.intersection(&other.ranges);
        HpxRanges { ranges }
    }

    pub fn difference(&self, other: &Self) -> Self {
        let ranges = self.ranges.difference(&other.ranges);
        HpxRanges { ranges }
    }

    pub fn complement(&self) -> Self {
        let ranges = self.ranges.complement(T::HPX_MAXPIX);
        HpxRanges { ranges }
    }
}

impl<T: MocIdx> PartialEq for HpxRanges<T> {
    fn eq(&self, other: &Self) -> bool {
        self.ranges == other.ranges
    }
}

impl<T: MocIdx> Eq for HpxRanges<T> { }

impl<T: MocIdx> From<MocRanges<T, Hpx<T>>> for HpxRanges<T> {
    fn from(ranges: MocRanges<T, Hpx<T>>) -> Self {
        HpxRanges::<T> { ranges }
    }
}

impl<T: MocIdx> From<HpxRanges<T>> for MocRanges<T, Hpx<T>> {
    fn from(hpx_ranges: HpxRanges<T>) -> Self {
        hpx_ranges.ranges
    }
}

use ndarray::Array2;

impl From<HpxRanges<u64>> for Array2<u64> {
    fn from(input: HpxRanges<u64>) -> Self {
        input.ranges.into()
    }
}
impl From<HpxRanges<i64>> for Array2<i64> {
    fn from(input: HpxRanges<i64>) -> Self {
        input.ranges.into()
    }
}
*/