use crate::prelude::*;
use polars_arrow::trusted_len::TrustedLen;
use polars_arrow::utils::CustomIterTools;
use polars_core::prelude::*;
use polars_core::POOL;
use polars_core::{export::rayon::prelude::*, utils::split_offsets};
use polars_utils::flatten;
#[cfg(feature = "serde")]
use serde::{Deserialize, Serialize};
use std::cmp::Ordering;

#[derive(Clone, Copy, Debug)]
#[cfg_attr(feature = "serde", derive(Serialize, Deserialize))]
pub enum ClosedWindow {
    Left,
    Right,
    Both,
    None,
}

/// Based on the given `Window`, which has an
/// - every
/// - period
/// - offset
/// window boundaries are created. And every window boundary we search for the values
/// that fit that window by the given `ClosedWindow`. The groups are return as `GroupTuples`
/// together with the lower bound and upper bound timestamps. These timestamps indicate the start (lower)
/// and end (upper) of the window of that group.
///
/// If `include_boundaries` is `false` those `lower` and `upper` vectors will be empty.
pub fn groupby_windows(
    window: Window,
    time: &[i64],
    include_boundaries: bool,
    closed_window: ClosedWindow,
    tu: TimeUnit,
) -> (GroupsSlice, Vec<i64>, Vec<i64>) {
    let start = time[0];
    let boundary = if time.len() > 1 {
        // +1 because left or closed boundary could match the next window if it is on the boundary
        let stop = time[time.len() - 1] + 1;
        Bounds::new_checked(start, stop)
    } else {
        let stop = start + 1;
        Bounds::new_checked(start, stop)
    };

    let size = if include_boundaries {
        match tu {
            TimeUnit::Nanoseconds => window.estimate_overlapping_bounds_ns(boundary),
            TimeUnit::Microseconds => window.estimate_overlapping_bounds_us(boundary),
            TimeUnit::Milliseconds => window.estimate_overlapping_bounds_ms(boundary),
        }
    } else {
        0
    };
    let mut lower_bound = Vec::with_capacity(size);
    let mut upper_bound = Vec::with_capacity(size);

    let mut groups = match tu {
        TimeUnit::Nanoseconds => {
            Vec::with_capacity(window.estimate_overlapping_bounds_ns(boundary))
        }
        TimeUnit::Microseconds => {
            Vec::with_capacity(window.estimate_overlapping_bounds_us(boundary))
        }
        TimeUnit::Milliseconds => {
            Vec::with_capacity(window.estimate_overlapping_bounds_ms(boundary))
        }
    };
    let mut start_offset = 0;

    for bi in window.get_overlapping_bounds_iter(boundary, tu) {
        let mut skip_window = false;
        // find starting point of window
        while start_offset < time.len() {
            let t = time[start_offset];
            if bi.is_future(t, closed_window) {
                // the window is behind the time values.
                skip_window = true;
                break;
            }
            if bi.is_member(t, closed_window) {
                break;
            }
            start_offset += 1;
        }
        if skip_window {
            start_offset = start_offset.saturating_sub(1);
            continue;
        }
        if start_offset == time.len() {
            start_offset = start_offset.saturating_sub(1);
        }

        // find members of this window
        let mut i = start_offset;
        // start next iteration 1 index back because of boundary conditions.
        // e.g. "closed left" could match the next iteration, but did not this one.
        start_offset = start_offset.saturating_sub(1);

        // last value
        if i == time.len() - 1 {
            let t = time[i];
            if bi.is_member(t, closed_window) {
                if include_boundaries {
                    lower_bound.push(bi.start);
                    upper_bound.push(bi.stop);
                }
                groups.push([i as IdxSize, 1])
            }
            continue;
        }

        let first = i as IdxSize;

        while i < time.len() {
            let t = time[i];
            if !bi.is_member(t, closed_window) {
                break;
            }
            i += 1;
        }
        let len = (i as IdxSize) - first;

        if include_boundaries {
            lower_bound.push(bi.start);
            upper_bound.push(bi.stop);
        }
        groups.push([first, len])
    }
    (groups, lower_bound, upper_bound)
}

pub(crate) fn groupby_values_iter_full_lookbehind(
    period: Duration,
    offset: Duration,
    time: &[i64],
    closed_window: ClosedWindow,
    tu: TimeUnit,
    start_offset: usize,
) -> impl Iterator<Item = (IdxSize, IdxSize)> + TrustedLen + '_ {
    debug_assert_eq!(offset.nanoseconds(), period.nanoseconds());
    debug_assert!(offset.negative);

    let add = match tu {
        TimeUnit::Nanoseconds => Duration::add_ns,
        TimeUnit::Microseconds => Duration::add_us,
        TimeUnit::Milliseconds => Duration::add_ms,
    };

    let mut last_lookbehind_i = 0;
    time[start_offset..]
        .iter()
        .enumerate()
        .map(move |(mut i, lower)| {
            i += start_offset;
            let lower = add(&offset, *lower);
            let upper = add(&period, lower);

            let b = Bounds::new(lower, upper);

            // we have a complete lookbehind so we know that `i` is the upper bound.
            // Safety
            // we are in bounds
            let slice = unsafe { time.get_unchecked(last_lookbehind_i..i) };
            let offset = slice.partition_point(|v| !b.is_member(*v, closed_window));

            let lookbehind_i = offset + last_lookbehind_i;
            // -1 for window boundary effects
            last_lookbehind_i = lookbehind_i.saturating_sub(1);

            let mut len = i - lookbehind_i;
            if matches!(closed_window, ClosedWindow::Right | ClosedWindow::Both) {
                len += 1;
            }

            (lookbehind_i as IdxSize, len as IdxSize)
        })
}

// this one is correct for all lookbehind/lookaheads, but is slower
pub(crate) fn groupby_values_iter_partial_lookbehind(
    period: Duration,
    offset: Duration,
    time: &[i64],
    closed_window: ClosedWindow,
    tu: TimeUnit,
) -> impl Iterator<Item = (IdxSize, IdxSize)> + TrustedLen + '_ {
    let add = match tu {
        TimeUnit::Nanoseconds => Duration::add_ns,
        TimeUnit::Microseconds => Duration::add_us,
        TimeUnit::Milliseconds => Duration::add_ms,
    };

    let mut lagging_offset = 0;
    time.iter().enumerate().map(move |(i, lower)| {
        let lower = add(&offset, *lower);
        let upper = add(&period, lower);

        let b = Bounds::new(lower, upper);

        for &t in &time[lagging_offset..] {
            if b.is_member(t, closed_window) || lagging_offset == i {
                break;
            }
            lagging_offset += 1;
        }

        // Safety
        // we just iterated over value i.
        let slice = unsafe { time.get_unchecked(lagging_offset..) };
        let len = slice.partition_point(|v| b.is_member(*v, closed_window));

        (lagging_offset as IdxSize, len as IdxSize)
    })
}

pub(crate) fn groupby_values_iter_full_lookahead(
    period: Duration,
    offset: Duration,
    time: &[i64],
    closed_window: ClosedWindow,
    tu: TimeUnit,
    start_offset: usize,
    upper_bound: Option<usize>,
) -> impl Iterator<Item = (IdxSize, IdxSize)> + TrustedLen + '_ {
    let upper_bound = upper_bound.unwrap_or(time.len());
    debug_assert!(!offset.negative);

    let add = match tu {
        TimeUnit::Nanoseconds => Duration::add_ns,
        TimeUnit::Microseconds => Duration::add_us,
        TimeUnit::Milliseconds => Duration::add_ms,
    };

    time[start_offset..upper_bound]
        .iter()
        .enumerate()
        .map(move |(mut i, lower)| {
            i += start_offset;
            let lower = add(&offset, *lower);
            let upper = add(&period, lower);

            let b = Bounds::new(lower, upper);

            let slice = unsafe { time.get_unchecked(i..) };
            let len = slice.partition_point(|v| b.is_member(*v, closed_window));

            (i as IdxSize, len as IdxSize)
        })
}

fn partially_check_sorted(time: &[i64]) {
    // check sortedness of a small subslice.
    if time.len() > 1 {
        assert!(time[..std::cmp::min(time.len(), 10)].windows(2).filter_map(|w| match w[0].cmp(&w[1]) {
            Ordering::Equal => None,
            t => Some(t)
        }).all_equal(), "subslice check showed that the values in `groupby_rolling` were not sorted. Pleasure ensure the index column is sorted.")
    }
}

#[cfg(feature = "rolling_window")]
pub(crate) fn groupby_values_iter<'a>(
    period: Duration,
    offset: Duration,
    time: &'a [i64],
    closed_window: ClosedWindow,
    tu: TimeUnit,
) -> Box<dyn TrustedLen<Item = (IdxSize, IdxSize)> + 'a> {
    partially_check_sorted(time);
    // we have a (partial) lookbehind window
    if offset.negative {
        // only lookbehind
        if offset.nanoseconds() == period.nanoseconds() {
            let iter =
                groupby_values_iter_full_lookbehind(period, offset, time, closed_window, tu, 0);
            Box::new(iter)
        }
        // partial lookbehind
        else {
            let iter =
                groupby_values_iter_partial_lookbehind(period, offset, time, closed_window, tu);
            Box::new(iter)
        }
    } else {
        let iter =
            groupby_values_iter_full_lookahead(period, offset, time, closed_window, tu, 0, None);
        Box::new(iter)
    }
}

/// Different from `groupby_windows`, where define window buckets and search which values fit that
/// pre-defined bucket, this function defines every window based on the:
///     - timestamp (lower bound)
///     - timestamp + period (upper bound)
/// where timestamps are the individual values in the array `time`
pub fn groupby_values(
    period: Duration,
    offset: Duration,
    time: &[i64],
    closed_window: ClosedWindow,
    tu: TimeUnit,
) -> GroupsSlice {
    partially_check_sorted(time);
    let thread_offsets = split_offsets(time.len(), POOL.current_num_threads());

    // we have a (partial) lookbehind window
    if offset.negative {
        // only lookbehind
        if offset.nanoseconds() == period.nanoseconds() {
            let vals = thread_offsets
                .par_iter()
                .copied()
                .map(|(base_offset, len)| {
                    let upper_bound = base_offset + len;
                    let iter = groupby_values_iter_full_lookbehind(
                        period,
                        offset,
                        &time[..upper_bound],
                        closed_window,
                        tu,
                        base_offset,
                    );
                    iter.map(|(offset, len)| [offset as IdxSize, len])
                        .collect_trusted::<Vec<_>>()
                })
                .collect::<Vec<_>>();
            flatten(&vals, Some(time.len()))
        }
        // partial lookbehind
        // this one is still single threaded
        // can make it parallel later, its a bit more complicated because the boundaries are unknown
        else {
            let iter =
                groupby_values_iter_partial_lookbehind(period, offset, time, closed_window, tu);
            iter.map(|(offset, len)| [offset, len]).collect_trusted()
        }
    } else {
        let vals = thread_offsets
            .par_iter()
            .copied()
            .map(|(base_offset, len)| {
                let lower_bound = base_offset;
                let upper_bound = base_offset + len;
                let iter = groupby_values_iter_full_lookahead(
                    period,
                    offset,
                    time,
                    closed_window,
                    tu,
                    lower_bound,
                    Some(upper_bound),
                );
                iter.map(|(offset, len)| [offset as IdxSize, len])
                    .collect_trusted::<Vec<_>>()
            })
            .collect::<Vec<_>>();
        flatten(&vals, Some(time.len()))
    }
}
