polaranges 0.3.2

use std::borrow::Cow;

use polars_core::prelude::{DataFrame, Series};

use crate::range_frame::{take_rows_owned, RangeView};
use crate::{
    ClipRangesOptions, ClusterOptions, ComplementRangesOptions, CountOverlapsOptions,
    ExtendRangesOptions, GroupCumsumOptions, IntersectOverlapsOptions, JoinOverlapsOptions,
    MaxDisjointOptions, MergeOptions, NearestOptions, OuterRangesOptions, OverlapOptions,
    OverlapPairs, Result, SetIntersectOverlapsOptions, SetUnionOverlapsOptions, SortRangesOptions,
    SplitOverlapsOptions, SubtractOptions, TileRangesOptions,
};

/// Crate-internal trait used by the public `.r()` accessor to dispatch range
/// operations onto a raw `DataFrame`. Not part of the public API.
pub(crate) trait DataFrameRanges {
    fn range_overlap_pairs(
        &self,
        other: &DataFrame,
        options: OverlapOptions,
    ) -> Result<OverlapPairs>;
    fn range_overlap(&self, other: &DataFrame, options: OverlapOptions) -> Result<DataFrame>;
    fn merge_overlaps(&self, options: MergeOptions) -> Result<DataFrame>;
    fn cluster_overlaps(&self, options: ClusterOptions) -> Result<DataFrame>;
    fn count_overlaps(&self, other: &DataFrame, options: CountOverlapsOptions) -> Result<Series>;
    fn join_overlaps(&self, other: &DataFrame, options: JoinOverlapsOptions) -> Result<DataFrame>;
    fn intersect_overlaps(
        &self,
        other: &DataFrame,
        options: IntersectOverlapsOptions,
    ) -> Result<DataFrame>;
    fn set_intersect_overlaps(
        &self,
        other: &DataFrame,
        options: SetIntersectOverlapsOptions,
    ) -> Result<DataFrame>;
    fn set_union_overlaps(
        &self,
        other: &DataFrame,
        options: SetUnionOverlapsOptions,
    ) -> Result<DataFrame>;
    fn sort_ranges(&self, options: SortRangesOptions) -> Result<DataFrame>;
    fn extend_ranges(&self, options: ExtendRangesOptions) -> Result<DataFrame>;
    fn tile_ranges(&self, options: TileRangesOptions) -> Result<DataFrame>;
    fn clip_ranges(&self, options: ClipRangesOptions) -> Result<DataFrame>;
    fn group_cumsum(&self, options: GroupCumsumOptions) -> Result<DataFrame>;
    fn max_disjoint_overlaps(&self, options: MaxDisjointOptions) -> Result<DataFrame>;
    fn split_overlaps(&self, options: SplitOverlapsOptions) -> Result<DataFrame>;
    fn outer_ranges(&self, options: OuterRangesOptions) -> Result<DataFrame>;
    fn complement_ranges(&self, options: ComplementRangesOptions) -> Result<DataFrame>;
    fn nearest_ranges(&self, other: &DataFrame, options: NearestOptions) -> Result<DataFrame>;
    fn subtract_overlaps(&self, other: &DataFrame, options: SubtractOptions) -> Result<DataFrame>;
}

impl DataFrameRanges for DataFrame {
    fn range_overlap_pairs(
        &self,
        other: &DataFrame,
        options: OverlapOptions,
    ) -> Result<OverlapPairs> {
        overlap_pairs_fast(self, other, &options)
    }

    fn range_overlap(&self, other: &DataFrame, options: OverlapOptions) -> Result<DataFrame> {
        overlap_fast(self, other, &options)
    }

    fn merge_overlaps(&self, options: MergeOptions) -> Result<DataFrame> {
        RangeView::new(self)?.merge_overlaps(&options)
    }

    fn cluster_overlaps(&self, options: ClusterOptions) -> Result<DataFrame> {
        RangeView::new(self)?.cluster_overlaps(&options)
    }

    fn count_overlaps(&self, other: &DataFrame, options: CountOverlapsOptions) -> Result<Series> {
        let left = RangeView::new(self)?;
        let right = RangeView::new(other)?;
        left.count_overlaps(&right, &options)
    }

    fn join_overlaps(&self, other: &DataFrame, options: JoinOverlapsOptions) -> Result<DataFrame> {
        let left = RangeView::new(self)?;
        let right = RangeView::new(other)?;
        left.join_overlaps(&right, &options)
    }

    fn intersect_overlaps(
        &self,
        other: &DataFrame,
        options: IntersectOverlapsOptions,
    ) -> Result<DataFrame> {
        let left = RangeView::new(self)?;
        let right = RangeView::new(other)?;
        left.intersect_overlaps(&right, &options)
    }

    fn set_intersect_overlaps(
        &self,
        other: &DataFrame,
        options: SetIntersectOverlapsOptions,
    ) -> Result<DataFrame> {
        let left = RangeView::new(self)?;
        let right = RangeView::new(other)?;
        left.set_intersect_overlaps(&right, &options)
    }

    fn set_union_overlaps(
        &self,
        other: &DataFrame,
        options: SetUnionOverlapsOptions,
    ) -> Result<DataFrame> {
        let left = RangeView::new(self)?;
        let right = RangeView::new(other)?;
        left.set_union_overlaps(&right, &options)
    }

    fn sort_ranges(&self, options: SortRangesOptions) -> Result<DataFrame> {
        RangeView::new(self)?.sort_ranges(&options)
    }

    fn extend_ranges(&self, options: ExtendRangesOptions) -> Result<DataFrame> {
        RangeView::new(self)?.extend_ranges(&options)
    }

    fn tile_ranges(&self, options: TileRangesOptions) -> Result<DataFrame> {
        RangeView::new(self)?.tile_ranges(&options)
    }

    fn clip_ranges(&self, options: ClipRangesOptions) -> Result<DataFrame> {
        RangeView::new(self)?.clip_ranges(&options)
    }

    fn group_cumsum(&self, options: GroupCumsumOptions) -> Result<DataFrame> {
        RangeView::new(self)?.group_cumsum(&options)
    }

    fn max_disjoint_overlaps(&self, options: MaxDisjointOptions) -> Result<DataFrame> {
        RangeView::new(self)?.max_disjoint_overlaps(&options)
    }

    fn split_overlaps(&self, options: SplitOverlapsOptions) -> Result<DataFrame> {
        RangeView::new(self)?.split_overlaps(&options)
    }

    fn outer_ranges(&self, options: OuterRangesOptions) -> Result<DataFrame> {
        RangeView::new(self)?.outer_ranges(&options)
    }

    fn complement_ranges(&self, options: ComplementRangesOptions) -> Result<DataFrame> {
        RangeView::new(self)?.complement_ranges(&options)
    }

    fn nearest_ranges(&self, other: &DataFrame, options: NearestOptions) -> Result<DataFrame> {
        nearest_fast(self, other, &options)
    }

    fn subtract_overlaps(&self, other: &DataFrame, options: SubtractOptions) -> Result<DataFrame> {
        let left = RangeView::new(self)?;
        let right = RangeView::new(other)?;
        left.subtract_overlaps(&right, &options)
    }
}

/// Run the overlap-pairs sweep against two frames. Column names and
/// match_by lists are read from `options`.
pub(crate) fn overlap_pairs_fast(
    left: &DataFrame,
    right: &DataFrame,
    options: &OverlapOptions,
) -> Result<OverlapPairs> {
    let left_df = rechunk_for_overlap(left);
    let right_df = rechunk_for_overlap(right);
    let left_view = RangeView::with_columns(
        left_df.as_ref(),
        options.left_start_col.as_str(),
        options.left_end_col.as_str(),
    )?;
    let right_view = RangeView::with_columns(
        right_df.as_ref(),
        options.right_start_col.as_str(),
        options.right_end_col.as_str(),
    )?;
    left_view.overlap_pairs_by_match(&right_view, &options.match_by, &options.match_by, options)
}

/// Return the rows of `left` that have at least one overlap in `right`.
/// Duplicates left rows once per overlap (matching `overlap_pairs` semantics
/// on the left side).
pub(crate) fn overlap_fast(
    left: &DataFrame,
    right: &DataFrame,
    options: &OverlapOptions,
) -> Result<DataFrame> {
    let left_df = rechunk_for_overlap(left);
    let right_df = rechunk_for_overlap(right);
    let left_view = RangeView::with_columns(
        left_df.as_ref(),
        options.left_start_col.as_str(),
        options.left_end_col.as_str(),
    )?;
    let right_view = RangeView::with_columns(
        right_df.as_ref(),
        options.right_start_col.as_str(),
        options.right_end_col.as_str(),
    )?;
    let indices = left_view.overlap_indices_by_match(
        &right_view,
        &options.match_by,
        &options.match_by,
        options,
    )?;
    take_rows_owned(left_df.as_ref(), indices)
}

/// Run the nearest sweep against two frames; returns the joined DataFrame
/// (left rows with right columns appended via `options.suffix` and the
/// optional distance column).
pub(crate) fn nearest_fast(
    left: &DataFrame,
    right: &DataFrame,
    options: &NearestOptions,
) -> Result<DataFrame> {
    let left_df = rechunk_for_overlap(left);
    let right_df = rechunk_for_overlap(right);
    let left_view = RangeView::with_columns(
        left_df.as_ref(),
        options.left_start_col.as_str(),
        options.left_end_col.as_str(),
    )?;
    let right_view = RangeView::with_columns(
        right_df.as_ref(),
        options.right_start_col.as_str(),
        options.right_end_col.as_str(),
    )?;
    left_view.nearest_ranges_by_match(&right_view, &options.match_by, &options.match_by, options)
}

/// Return just the indices of `left` rows that have at least one overlap.
/// Used internally (e.g. for `invert` in the bio accessor).
pub(crate) fn overlap_indices_fast(
    left: &DataFrame,
    right: &DataFrame,
    options: &OverlapOptions,
) -> Result<Vec<u32>> {
    let left_df = rechunk_for_overlap(left);
    let right_df = rechunk_for_overlap(right);
    let left_view = RangeView::with_columns(
        left_df.as_ref(),
        options.left_start_col.as_str(),
        options.left_end_col.as_str(),
    )?;
    let right_view = RangeView::with_columns(
        right_df.as_ref(),
        options.right_start_col.as_str(),
        options.right_end_col.as_str(),
    )?;
    left_view.overlap_indices_by_match(
        &right_view,
        &options.match_by,
        &options.match_by,
        options,
    )
}

fn rechunk_for_overlap(df: &DataFrame) -> Cow<'_, DataFrame> {
    if df.max_n_chunks() <= 1 {
        return Cow::Borrowed(df);
    }
    let mut owned = df.clone();
    owned.rechunk_mut_par();
    Cow::Owned(owned)
}