module Game.Mastermind.CodeSet.Tree where

import qualified Game.Mastermind.CodeSet as CodeSet

import qualified Data.Map as Map
import qualified Data.Set as Set

import Data.Tuple.HT (mapFst, swap, )

import Control.Monad (liftM2, )

import Prelude hiding (null, )


{- |
@Products [(a,b),(c,d)]@
expresses  a x b  union  c x d,
where @x@ denotes the set product.
-}
data T a =
   End | Products (Map.Map (Set.Set a) (T a))
   deriving (Show)

{-
instance (Ord a, Show a) => Show (T a) where
   showsPrec n cs =
      showParen (n>=10) $
      showString "CodeSet.fromLists " . shows (toLists cs)
-}

instance CodeSet.C T where
   empty = Products Map.empty
   union = union
   intersection = intersection
   unit = End
   leftNonEmptyProduct c xs =
      Products $
      if null xs
        then Map.empty
        else Map.singleton c xs
   flatten = flatten
   symbols = symbols
   null = null
   size = size
   select = select
   representationSize = representationSize
   compress = compress


flatten :: (Ord a) => T a -> [[a]]
flatten End = [[]]
flatten (Products xs) =
   concatMap
      (\(a,b) -> liftM2 (:) (Set.toList a) (flatten b))
      (Map.toList xs)

symbols :: (Ord a) => T a -> Set.Set a
symbols End = Set.empty
symbols (Products xps) =
   Set.unions $
   map (\(x,xs) -> Set.union x (symbols xs)) $
   Map.toList xps


size :: T a -> Integer
size End = 1
size (Products xs) =
   sum (map (\(a,b) -> fromIntegral (Set.size a) * size b) $ Map.toList xs)

-- somehow inefficient, because the sizes of subsets are recomputed several times
select :: T a -> Integer -> [a]
select End n =
   case compare n 0 of
     LT -> error "CodeSet.select.end: index negative"
     EQ -> []
     GT -> error "CodeSet.select.end: index too large"
select (Products xps) n0 =
   if n0<0
     then error "CodeSet.select: negative index"
     else
       case dropWhile (\(_, ((n1,sz), _)) -> n1>=sz) $
            zip (Map.toList xps) $
            uncurry zip $
            mapFst (\sizes -> zip (scanl (-) n0 sizes) sizes) $
            unzip $
            map (\(x,xs) ->
               let sz = size xs
               in  (fromIntegral (Set.size x) * sz, sz)) $
            Map.toList xps of
         [] -> error "CodeSet.select: index too large"
         ((x,xs), ((n1,_), xsSize)) : _ ->
            let (j,k) = divMod n1 xsSize
            in  (Set.toList x !! fromInteger j)
                : select xs k

representationSize :: T a -> Int
representationSize End = 1
representationSize (Products xs) =
   sum (map (\(a,b) -> Set.size a + representationSize b) $ Map.toList xs)


{- |
We could try to merge set products.
I'll first want to see, whether this is needed in a relevant number of cases.
-}
union :: (Ord a) => T a -> T a -> T a
union End End = End
union (Products xs) (Products ys) =
   Products (Map.unionWith union xs ys)
union _ _ = error "CodeSet.union: sets with different tuple size"

intersection :: (Ord a) => T a -> T a -> T a
intersection End End = End
intersection (Products xps) (Products yps) =
   Products $ Map.fromListWith union $ normalizeProducts $
   liftM2
      (\(x,xs) (y,ys) -> (Set.intersection x y, intersection xs ys))
      (Map.toList xps)
      (Map.toList yps)
intersection _ _ =
   error "CodeSet.intersection: sets with different tuple size"

{- |
Remove empty set products.
-}
normalizeProducts ::
   [(Set.Set a, T a)] -> [(Set.Set a, T a)]
normalizeProducts =
   filter (\(x,xs) -> not (Set.null x || null xs))


{- |
This allows (T a) to be a key in a Map.
I do not want an Ord (T a) instance,
since it makes no sense and it requires an Eq (T a) instance
that is either expensive (if it means set equality)
or confusing (if it means structural equality).
-}
newtype Indexable a = Indexable (T a)

instance (Eq a) => Eq (Indexable a) where
   (Indexable x) == (Indexable y) =
      case (x,y) of
         (End,End) -> True
         (Products xs, Products ys) ->
            fmap Indexable xs == fmap Indexable ys
         _ -> False

instance (Ord a) => Ord (Indexable a) where
   compare (Indexable x) (Indexable y) =
      case (x,y) of
         (End,End) -> EQ
         (End,Products _) -> LT
         (Products _,End) -> GT
         (Products xs, Products ys) ->
            compare (fmap Indexable xs) (fmap Indexable ys)


compress :: (Ord a) => T a -> T a
compress End = End
compress (Products xs) =
   Products $
   Map.fromListWith union $ map swap $ map (mapFst (\(Indexable set) -> set)) $ Map.toList $
   Map.fromListWith Set.union $ map (mapFst Indexable) $ map swap $ Map.toList $
   fmap compress xs

member :: (Ord a) => [a] -> T a -> Bool
member [] End = True
member (c:cs) (Products xps) =
   any (\(x,xs) -> Set.member c x && member cs xs) $
   Map.toList xps
member _ _ =
   error "CodeSet.member: mismatch of tuple size and tuple size in set"

null :: T a -> Bool
null End = False
null (Products xs) = Map.null xs