use crate::stats::{SingletonSplitStat, Stat};
use derive_more::{Deref, DerefMut};
use fxhash::FxHashMap;
use std::convert::Infallible;
use std::fmt::{Debug, Formatter};
use std::hash::Hash;

#[derive(Deref, DerefMut)]
pub struct Map<T, A: MapKey<T>> {
	inner: FxHashMap<A::Key, T>,
}

impl<T: Debug, A: MapKey<T>> Debug for Map<T, A> {
	fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
		self.inner.fmt(f)
	}
}

impl<T: PartialEq, A: MapKey<T>> PartialEq for Map<T, A> {
	fn eq(&self, other: &Self) -> bool {
		self.inner == other.inner
	}
}

impl<T: Clone, A: MapKey<T>> Clone for Map<T, A> {
	fn clone(&self) -> Self {
		Self {
			inner: self.inner.clone(),
		}
	}
}

impl<T, A: MapKey<T>> Default for Map<T, A> {
	fn default() -> Self {
		Self { inner: FxHashMap::default() }
	}
}

impl<T: SingletonSplitStat, A: MapKey<T>> Stat for Map<T, A> {
	type Split = T;
	type TryFromSplitError = Infallible;

	fn from_splits(splits: impl Iterator<Item=T>) -> Result<Self, Self::TryFromSplitError>
	where
		Self: Sized
	{
		let mut this = Self::default();
		for split in splits {
			let id: A::Key = A::get_key(&split);
			let id = id.clone();
			this.inner.entry(id.clone()).insert_entry(split);
		}
		Ok(this)
	}
}

#[derive(Deref, DerefMut)]
pub struct Map2D<T, A: MapKey<T>, B: MapKey<T>> {
	inner: FxHashMap<A::Key, FxHashMap<B::Key, T>>,
}

impl<T: Clone, A: MapKey<T>, B: MapKey<T>> Clone for Map2D<T, A, B> {
	fn clone(&self) -> Self {
		Self {
			inner: self.inner.clone(),
		}
	}
}

impl<T: Debug, A: MapKey<T>, B: MapKey<T>> Debug for Map2D<T, A, B> {
	fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
		self.inner.fmt(f)
	}
}

impl<T: PartialEq, A: MapKey<T>, B: MapKey<T>> PartialEq for Map2D<T, A, B> {
	fn eq(&self, other: &Self) -> bool {
		self.inner == other.inner
	}
}

impl<T, A: MapKey<T>, B: MapKey<T>> Default for Map2D<T, A, B> {
	fn default() -> Self {
		Self { inner: FxHashMap::default() }
	}
}

impl<T: SingletonSplitStat, A: MapKey<T>, B: MapKey<T>> Stat for Map2D<T, A, B> {
	type Split = T;
	type TryFromSplitError = Infallible;

	fn from_splits(splits: impl Iterator<Item=Self::Split>) -> Result<Self, Self::TryFromSplitError>
	where
		Self: Sized
	{
		let mut this = Self::default();
		for split in splits {
			let id: A::Key = A::get_key(&split);
			let inner_id: B::Key = B::get_key(&split);
			this.inner.entry(id).or_insert_with(FxHashMap::default).entry(inner_id.clone()).insert_entry(split);
		}
		Ok(this)
	}
}

pub trait MapKey<T> {
	type Key: Hash + Clone + Debug + Eq;

	fn get_key(this: &T) -> Self::Key;
}
