permoot 0.2.1

use core::num::NonZeroUsize;

use arrayvec::ArrayVec;

use crate::error::{CycleElementLocation, InvalidArrFnRepr, InvalidCycleRepr, InvalidSwapRepr};
use crate::util::{copy_to_array, cycle_iter, range_array};
use crate::Permutation;

/// Represents a permutation by its [cycles](https://en.wikipedia.org/wiki/Permutation#Cycle_notation)
#[derive(Debug, Clone)]
pub struct Cycles<const N: usize> {
	// invariant: contains all values 0 to N-1 (thus each exactly once)
	data: [usize; N],
	// invariant: is always sorted (increasing)
	cycle_starts: ArrayVec<usize, N>,
}

impl<const N: usize> Cycles<N> {
	/// The number of cycles in this representation
	///
	/// This returns `0` if and only if `N == 0`
	pub fn cycle_count(&self) -> usize {
		self.cycle_starts.len()
	}

	/// Returns the `i`th cycle in this representation
	pub fn get(&self, i: usize) -> Option<&[usize]> {
		let i0 = *self.cycle_starts.get(i)?;
		let i1 = i
			.checked_add(1)
			.and_then(|j| self.cycle_starts.get(j))
			.copied();

		match i1 {
			None => self.data.get(i0..),
			Some(i1) => self.data.get(i0..i1),
		}
	}

	/// Returns an iterator over the cycles in this representation
	pub fn iter(&self) -> impl Iterator<Item = &[usize]> + '_ {
		self.cycle_starts
			.windows(2)
			.map(|w| {
				let [i0, i1] = copy_to_array(w).unwrap_or_else(|| unreachable!());

				self.data.get(i0..i1).unwrap_or_else(|| unreachable!())
			})
			.chain(
				self.cycle_starts
					.last()
					.copied()
					.map(|i0| self.data.get(i0..).unwrap_or_else(|| unreachable!())),
			)
	}

	fn convert_index(&self, i: usize) -> Option<CycleElementLocation> {
		let (cycle_i, &cycle_start) = self
			.cycle_starts
			.iter()
			.enumerate()
			.rfind(|&(_, &s)| s < i)?;
		let sub_i = i - cycle_start;
		Some(CycleElementLocation {
			cycle_index: cycle_i,
			index: sub_i,
		})
	}

	fn deconvert_index(&self, loc: CycleElementLocation) -> Option<usize> {
		let res = *self.cycle_starts.get(loc.cycle_index)? + loc.index;
		(res < N).then_some(res)
	}
}

#[derive(Default)]
struct CyclesBuilder<const N: usize> {
	data: ArrayVec<usize, N>,
	cycle_starts: ArrayVec<usize, N>,
	curr_cycle_outer_values: Option<(usize, usize)>,
}

#[derive(Copy, Clone, Eq, PartialEq)]
#[repr(transparent)]
struct PushResult {
	pub cycle_closed: bool,
}

impl<const N: usize> CyclesBuilder<N> {
	fn new() -> Self {
		Self::default()
	}

	fn curr_cycle_starts_with(&self, i: usize) -> bool {
		matches!(self.curr_cycle_outer_values, Some(t) if t.0 == i)
	}

	fn push(&mut self, i: usize) -> PushResult {
		if self.curr_cycle_starts_with(i) {
			// finish cycle
			self.curr_cycle_outer_values = None;
			PushResult { cycle_closed: true }
		} else if let Some((_, last)) = &mut self.curr_cycle_outer_values {
			// extend cycle
			self.data.push(i);
			*last = i;
			PushResult {
				cycle_closed: false,
			}
		} else {
			// start new cycle
			self.cycle_starts.push(self.data.len());
			self.data.push(i);
			self.curr_cycle_outer_values = Some((i, i));
			PushResult {
				cycle_closed: false,
			}
		}
	}

	fn unwrap_finished(self) -> Cycles<N> {
		debug_assert!(
			self.curr_cycle_outer_values.is_none(),
			"tried to unwrap unfinished `CyclesBuilder`"
		);
		Cycles {
			data: self.data.into_inner().unwrap(),
			cycle_starts: self.cycle_starts,
		}
	}
}

/// Marker trait for the order that swaps are stored in [`Swaps`]
#[allow(missing_docs)] // the trait items are undocumented because this is a type-level enum and the variants are documented themselves
pub trait SwapsReprOrder: private::Sealed {
	fn new() -> Self;

	const IS_FCO: bool;

	type Reverse: SwapsReprOrder;
}

mod private {
	pub trait Sealed {}

	impl Sealed for super::FunctionCompositionOrder {}
	impl Sealed for super::SequentialApplicationOrder {}
}

/// *Marker type*: Swaps are stored in the order of function composition.
///
/// Example:
/// `[f1, f2, f3]` would represent the permutation `f1 . f2 . f3` where `.` is function composition.
///
/// It is the reverse order to [`SequentialApplicationOrder`].
///
/// This order plays well with `slice::swap` beacuse when `sl` represents the permutation `f`,
/// then after `sl.swap(i, j)` it will represent the permuattion `f . τ_ij` where `τ_ij` is a swap of `i` and `j`.
pub struct FunctionCompositionOrder;
/// *Marker type*: Swaps are stored in the order of sequential application.
///
/// Example:
/// `[f1, f2, f3]` represents the permutation given by `i -> { j1 := f1(i); j2 := f2(j1); j3 := f3(j2); j3 }`.
///
/// It is the reverse order to [`FunctionCompositionOrder`].
pub struct SequentialApplicationOrder;

impl SwapsReprOrder for FunctionCompositionOrder {
	fn new() -> Self {
		Self
	}

	const IS_FCO: bool = true;

	type Reverse = SequentialApplicationOrder;
}
impl SwapsReprOrder for SequentialApplicationOrder {
	fn new() -> Self {
		Self
	}

	const IS_FCO: bool = false;

	type Reverse = FunctionCompositionOrder;
}

/// Represents a permuation as a composition of `N` or less swaps
pub struct Swaps<const N: usize, O: SwapsReprOrder> {
	left: [usize; N],
	right: [usize; N],
	len: usize,
	_order: O,
}

impl<const N: usize, O: SwapsReprOrder> Swaps<N, O> {
	fn new() -> Self {
		Self {
			left: [0; N],
			right: [0; N],
			len: 0,
			_order: O::new(),
		}
	}

	/// Returns the number of swaps in this representation
	pub fn len(&self) -> usize {
		self.len
	}

	/// Returns whether there are no swaps in this representation (i.e. whether this is the identity)
	pub fn is_empty(&self) -> bool {
		self.len == 0
	}

	/// Returns the `i`th swap in this representation
	pub fn get(&self, i: usize) -> Option<(&usize, &usize)> {
		self.left
			.get(i)
			.and_then(|l| self.right.get(i).map(|r| (l, r)))
	}

	/// Returns an iterator over the swaps in this representation
	pub fn iter(&self) -> impl Iterator<Item = (&usize, &usize)> + '_ {
		self.left.iter().zip(self.right.iter()).take(self.len)
	}

	/// Get the representation that is stored as the same list of swaps as `self`, but represents a permutation by the specified order
	///
	/// Either `O2 == O`, then this does nothing,
	/// or `O2` is the opposite order to `O`,
	/// in which case the permutation represented by the output is the inverse
	/// of the permutation represented by the input.
	pub fn interpret_as_order<O2: SwapsReprOrder>(self) -> Swaps<N, O2> {
		Swaps {
			left: self.left,
			right: self.right,
			len: self.len,
			_order: O2::new(),
		}
	}

	/// Returns the swaps representation for the inverse permutation
	///
	/// This is equivalent to `self.interpret_as_order::<<O as SwapsReprOrder>::Reverse>()`.
	#[inline]
	pub fn inverse(self) -> Swaps<N, <O as SwapsReprOrder>::Reverse> {
		self.interpret_as_order()
	}

	/// Get the representation with order `O2` that represents the same permutation as `self` does with order `O`
	pub fn into_equivalent_with_order<O2: SwapsReprOrder>(mut self) -> Swaps<N, O2> {
		if O::IS_FCO != O2::IS_FCO {
			self.left[..self.len].reverse();
			self.right[..self.len].reverse();
		}
		self.interpret_as_order::<O2>()
	}
}

impl<const N: usize> Swaps<N, FunctionCompositionOrder> {
	/// Same as [`Permutation::apply_in_place`] but avoids having to recalculate the swaps
	/// when wanting to apply the same permutation multiple times.
	///
	/// # Panics
	/// - if there is a swap in `self` that has an index out of bounds for `slice`
	pub fn apply_in_place<T>(&self, slice: &mut [T]) {
		for (&i, &j) in self.iter() {
			slice.swap(i, j);
		}
	}
}

impl<const N: usize> Swaps<N, SequentialApplicationOrder> {
	fn swap(&mut self, i: usize, j: usize) {
		if self.len >= N {
			panic!("ran out of capacity while constructing `Swaps`");
		}
		self.left[self.len] = i;
		self.right[self.len] = j;
		self.len += 1;
	}
}

/// Representations
impl<const N: usize> Permutation<N> {
	/// Construct a permutation from its array representation
	///
	/// - The array is interpreted as the [one-line notation](https://en.wikipedia.org/wiki/Permutation#One-line_notation).
	/// - Elements start at zero, unlike in mathematics where they usually start at 1
	pub fn from_array_repr(arr: [usize; N]) -> Result<Self, InvalidArrFnRepr> {
		for (i, x) in arr.iter().enumerate() {
			if *x >= N {
				return Err(InvalidArrFnRepr::InvalidValue {
					input: i,
					value: *x,
					// this can't fail because the loop won't run if `N == 0`
					max: N - 1,
				});
			}
			if let Some(j) = arr[..i].iter().position(|y| y == x) {
				return Err(InvalidArrFnRepr::RepeatedValue {
					first_input: j,
					second_input: i,
					value: *x,
				});
			}
		}

		Ok(Self { arr_repr: arr })
	}

	/// Same as [`from_array_repr`](Self::from_array_repr) but uses the more familliar one-based values (used in mathematics)
	pub fn from_one_based_array_repr(arr: [NonZeroUsize; N]) -> Result<Self, InvalidArrFnRepr> {
		Self::from_array_repr(arr.map(|x| x.get() - 1))
	}

	/// Obtain the array representation for a permutation (see [`from_array_repr`](Self::from_array_repr))
	pub fn to_array_repr(self) -> [usize; N] {
		self.arr_repr
	}

	/// Construct a permutation from its [cycles representation](https://en.wikipedia.org/wiki/Permutation#Cycle_notation)
	pub fn from_cycles_repr(cycles: Cycles<N>) -> Result<Self, InvalidCycleRepr> {
		let mut arr_repr = [0; N];

		for (i, c) in cycles.iter().enumerate() {
			for (j, (curr, next)) in cycle_iter(c).enumerate() {
				let loc = CycleElementLocation {
					cycle_index: i,
					index: j,
				};

				let data_i = cycles
					.deconvert_index(loc)
					.unwrap_or_else(|| unreachable!());
				if let Some(prev_data_i) = cycles.data[..data_i].iter().position(|&x| x == curr) {
					let prev_loc = cycles
						.convert_index(prev_data_i)
						.unwrap_or_else(|| unreachable!());

					return Err(InvalidCycleRepr::RepeatedValue {
						first_occurrence: prev_loc,
						second_occurrence: loc,
						value: curr,
					});
				}

				match arr_repr.get_mut(curr) {
					// curr >= N
					None => {
						return Err(InvalidCycleRepr::InvalidValue {
							location: loc,
							value: curr,
							// this can't fail because the loop won't run if `N == 0`
							max: N - 1,
						});
					}
					Some(x) => *x = next,
				}
			}
		}

		Ok(Self { arr_repr })
	}

	/// Obtain the [cycles representation](https://en.wikipedia.org/wiki/Permutation#Cycle_notation) for a permutation
	pub fn to_cycles_repr(self) -> Cycles<N> {
		let mut builder = CyclesBuilder::new();

		#[derive(Copy, Clone, Eq, PartialEq)]
		enum ToVisit {
			NotVisited,
			Visited,
		}

		let mut visited = [ToVisit::NotVisited; N];

		let mut curr = 0;
		loop {
			visited[curr] = ToVisit::Visited;
			if builder.push(curr).cycle_closed {
				// look for the first non-visited value and use that to continue
				match visited.iter().position(|&v| v == ToVisit::NotVisited) {
					None => {
						// no values left, i.e. we are done
						break;
					}
					Some(i) => curr = i,
				}
			} else {
				// use `curr <- f(curr)` to walk along the cycle
				curr = self.arr_repr[curr]
			}
		}

		builder.unwrap_finished()
	}

	/// Construct a permutation from its representation as a composition as `M` or less swaps
	pub fn from_swaps_repr<const M: usize, O: SwapsReprOrder>(
		swaps: Swaps<M, O>,
	) -> Result<Self, InvalidSwapRepr> {
		let mut arr_repr = range_array();
		for (i, (&a, &b)) in swaps
			.into_equivalent_with_order::<FunctionCompositionOrder>()
			.iter()
			.enumerate()
		{
			if a >= N {
				return Err(InvalidSwapRepr::InvalidLeftValue {
					index: i,
					value: a,
					// this can't fail because the loop won't run if `N == 0`
					max: N - 1,
				});
			}
			if b >= N {
				return Err(InvalidSwapRepr::InvalidRightValue {
					index: i,
					value: b,
					// this can't fail because the loop won't run if `N == 0`
					max: N - 1,
				});
			}
			arr_repr.swap(a, b);
		}

		Ok(Self { arr_repr })
	}

	/// Represent a permutation as `N` or less swaps
	pub fn to_swaps_repr(self) -> Swaps<N, SequentialApplicationOrder> {
		// the result is reversed function composition order (so sequential application order)
		// because swaps are appended to this as they are popped off the end of the permutation
		// -
		// also note that every permutation of N elements can be represented as max. N swaps
		let mut swaps = Swaps::<N, SequentialApplicationOrder>::new();
		let mut remaining = ArrayVec::from(self.arr_repr);

		while let Some(i) = remaining.iter().position(|x| *x == remaining.len() - 1) {
			// `i` is the index of the largest remaining value

			if i == remaining.len() - 1 {
				remaining.pop();
				continue;
			}

			swaps.swap(i, remaining.len() - 1);
			// see this as also swapping `i` ans `remaining.len() - 1`
			// the latter element is then removed because
			// a) it is more efficient
			// b) it isn't of interest anymore as the remaining permutation maps it to itself
			remaining.swap_remove(i);
		}

		swaps
	}
}

#[cfg(test)]
mod tests {
	use super::*;
	use crate::util::with_random_perms;
	use crate::Permutation;

	#[test]
	fn roundtrip_cycle() {
		with_random_perms::<100>(100, |p| {
			assert_eq!(Permutation::from_cycles_repr(p.to_cycles_repr()), Ok(p));
		})
	}

	#[test]
	fn roundtrip_swap() {
		with_random_perms::<100>(100, |p| {
			assert_eq!(Permutation::from_swaps_repr(p.to_swaps_repr()), Ok(p));
		})
	}

	#[test]
	fn reverse_order_is_inerse() {
		with_random_perms::<100>(100, |p| {
			let sw: Swaps<100, SequentialApplicationOrder> = p.to_swaps_repr();
			let opp =
				sw.interpret_as_order::<<SequentialApplicationOrder as SwapsReprOrder>::Reverse>();
			let opp_p = Permutation::from_swaps_repr(opp).unwrap();

			assert_eq!(opp_p, p.inverse());
		})
	}
}