module Partition2 where

import Data.Foldable as F
import Data.List as L
import Data.Set as S
import Data.Tuple (swap)
import Data.Map.Strict as M
import Data.Vector as V

type Class = Int
data Partition a = Partition (Map a Class) (Vector (Set a))
  deriving Show

size :: Partition a -> Int
size (Partition _ v) = V.length v

elems :: Ord a => Partition a -> Set a
elems (Partition _ v) = V.foldr S.union S.empty v

flat :: Vector (Set a) -> [(Int, a)]
flat v = F.concatMap (\(i, s) -> foldMap (\x -> [(i, x)]) s) $ indexed v

-- Ord on o is used to determine classes, Ord on a is used to
-- store the elements
partitionWith :: (Ord o, Ord a) => (a -> o) -> [a] -> Partition a
partitionWith f ls = Partition map groups
  where
    map = M.fromList . fmap swap $ flat groups
    groups = V.fromList . fmap S.fromList . groupOn f . sortOn f $ ls

groupOn :: Eq b => (a -> b) -> [a] -> [[a]]
groupOn f ls = L.groupBy (\a b -> f a == f b) ls

instance Ord a => Eq (Partition a) where
  p1@(Partition m1 v1) == p2@(Partition m2 v2)
    | Partition.size p1 /= Partition.size p2 = False
    | Partition.elems p1 /= Partition.elems p2 = False
    | otherwise = wellDefined v1 m2 && wellDefined v2 m1
    where wellDefined v m = V.all (\x -> S.size x == 1) $ imap (\i s -> S.map (\x -> M.lookup x m) s) v

fromJustToSet Nothing = S.empty
fromJustToSet (Just s) = S.singleton s

dseflatten :: Ord a => Set (Set a) -> Set a
dseflatten = S.unions . S.toList

dseconcatMap :: (Ord a, Ord b) => (a -> Set b) -> Set a -> Set b
dseconcatMap f s = dseflatten (S.map f s)


--{-# LANGUAGE ViewPatterns, TupleSections #-}
--import Data.Foldable
--
--M.lookup
--
--let p = partitionWith (== 1) [1,2,3,4]
--let p2 = partitionWith (/= 1) [1,2,3,4]
--let p3 = partitionWith (== 2) [1,2,3,4]
--print p
--print p2
--p == p2
--p == p3