{-# LANGUAGE ScopedTypeVariables #-}
module Data.Enumerable where

import Data.Int
import Data.Word
import Data.Ratio

import Unsafe.Coerce

import Data.List
import Data.Maybe

import Control.Applicative
import Control.Monad
import Control.Monad.Omega

import Data.Tagged

class Enumerable a where
  enumerate :: [a]
  enumerateTotal :: [a]
  enumeratePartial :: [a]

  enumeratePartial = enumerate ++ [error "bottom"]
  enumerateTotal = enumerate

  isPartial :: Tagged a Bool
  isPartial = Tagged False

data Cardinal = Finite Integer | Aleph Integer -- I can dream

class (Enumerable a) => FinitelyEnumerable a where
  cardinality :: Tagged a Integer
  cardinality = Tagged $ genericLength (enumerate :: [a]) -- if you're too lazy to figure it out

instance (FinitelyEnumerable a, FinitelyEnumerable b, Eq a) => Enumerable (a -> b) where
  enumerate = let xs = enumerate in map (\ys z -> fromJust . lookup z $ zip xs ys) (sequence $ map (const enumerate) (enumerate :: [a])) -- probably not very practical :P
  
instance (FinitelyEnumerable a, FinitelyEnumerable b, Eq a) => FinitelyEnumerable (a -> b) where
  cardinality = Tagged $ unTagged (cardinality :: Tagged b Integer) ^ unTagged (cardinality :: Tagged a Integer)

instance (Enumerable a, Enumerable b) => Enumerable (a, b) 
  where enumerate = runOmega $ (,) <$> each enumerate <*> each enumerate

instance (FinitelyEnumerable a, FinitelyEnumerable b) => FinitelyEnumerable (a, b)
  where cardinality = Tagged $ unTagged (cardinality :: Tagged a Integer) * unTagged (cardinality :: Tagged b Integer)

instance (Enumerable a, Enumerable b, Enumerable c) => Enumerable (a, b, c) 
  where enumerate = runOmega $ (,,) <$> each enumerate <*> each enumerate <*> each enumerate

instance (FinitelyEnumerable a, FinitelyEnumerable b, FinitelyEnumerable c) => FinitelyEnumerable (a, b, c)
  where cardinality = Tagged $ unTagged (cardinality :: Tagged a Integer) * unTagged (cardinality :: Tagged b Integer) * unTagged (cardinality :: Tagged c Integer)

instance (Enumerable a, Enumerable b, Enumerable c, Enumerable d) => Enumerable (a, b, c, d) 
  where enumerate = runOmega $ (,,,) <$> each enumerate <*> each enumerate <*> each enumerate <*> each enumerate

instance (FinitelyEnumerable a, FinitelyEnumerable b, FinitelyEnumerable c, FinitelyEnumerable d) => FinitelyEnumerable (a, b, c, d)
  where cardinality = Tagged $ unTagged (cardinality :: Tagged a Integer) * unTagged (cardinality :: Tagged b Integer) * unTagged (cardinality :: Tagged c Integer) * unTagged (cardinality :: Tagged d Integer)
  
instance (Enumerable a, Enumerable b, Enumerable c, Enumerable d, Enumerable e) => Enumerable (a, b, c, d, e) 
  where enumerate = runOmega $ (,,,,) <$> each enumerate <*> each enumerate <*> each enumerate <*> each enumerate <*> each enumerate

instance (FinitelyEnumerable a, FinitelyEnumerable b, FinitelyEnumerable c, FinitelyEnumerable d, FinitelyEnumerable e) => FinitelyEnumerable (a, b, c, d, e)
  where cardinality = Tagged $ unTagged (cardinality :: Tagged a Integer) * unTagged (cardinality :: Tagged b Integer) * unTagged (cardinality :: Tagged c Integer) * unTagged (cardinality :: Tagged d Integer) * unTagged (cardinality :: Tagged d Integer)

instance (Enumerable a) => Enumerable (Maybe a) where
  enumerate = Nothing : map Just enumerate
  
instance (FinitelyEnumerable a) => FinitelyEnumerable (Maybe a) where
  cardinality = Tagged $ 1 + unTagged (cardinality :: Tagged a Integer)

instance (Enumerable a, Enumerable b) => Enumerable (Either a b) where
  enumerate = concat . transpose $ [map Left enumerate, map Right enumerate]

instance (FinitelyEnumerable a, FinitelyEnumerable b) => FinitelyEnumerable (Either a b) where
  cardinality = Tagged $ unTagged (cardinality :: Tagged a Integer) + unTagged (cardinality :: Tagged b Integer)

instance (Enumerable a) => Enumerable [a] where
  enumerate = concatMap (flip replicateM enumerate) [0..]

instance Enumerable         () where enumerate     = [()]
instance FinitelyEnumerable () where cardinality = Tagged 1

instance Enumerable         Bool where enumerate     = [minBound..maxBound]
instance FinitelyEnumerable Bool where cardinality = Tagged 2

instance Enumerable         Ordering where enumerate     = [minBound..maxBound]
instance FinitelyEnumerable Ordering where cardinality = Tagged 3

instance Enumerable         Char where enumerate     = [minBound..maxBound]
instance FinitelyEnumerable Char where cardinality = Tagged 1114112

instance Enumerable         Word where enumerate     = [minBound..maxBound]
instance FinitelyEnumerable Word where cardinality = Tagged $ fromIntegral (maxBound :: Word) - fromIntegral (minBound :: Word) + 1

instance Enumerable         Word8  where enumerate     = [minBound..maxBound]
instance FinitelyEnumerable Word8  where cardinality = Tagged $ 2 ^ (8 :: Int)

instance Enumerable         Word16 where enumerate     = [minBound..maxBound]
instance FinitelyEnumerable Word16 where cardinality = Tagged $ 2 ^ (16 :: Int)

instance Enumerable         Word32 where enumerate     = [minBound..maxBound]
instance FinitelyEnumerable Word32 where cardinality = Tagged $ 2 ^ (32 :: Int)

instance Enumerable         Word64 where enumerate   = [minBound..maxBound]
instance FinitelyEnumerable Word64 where cardinality = Tagged $ 2 ^ (64 :: Int)

enumerateInterleaved :: (Enum a, Num a) => [a]
enumerateInterleaved = 0 : init (concat [[-x, x] | x <- [-1,-2..]])

instance Enumerable         Int where enumerate     = [minBound..maxBound]
instance FinitelyEnumerable Int where cardinality = Tagged $ fromIntegral (maxBound :: Int) - fromIntegral (minBound :: Int) + 1

instance Enumerable         Int8  where enumerate     = [minBound..maxBound]
instance FinitelyEnumerable Int8  where cardinality = Tagged $ 2 ^ (8 :: Int)

instance Enumerable         Int16 where enumerate     = [minBound..maxBound]
instance FinitelyEnumerable Int16 where cardinality = Tagged $ 2 ^ (16 :: Int)

instance Enumerable         Int32 where enumerate     = [minBound..maxBound]
instance FinitelyEnumerable Int32 where cardinality = Tagged $ 2 ^ (32 :: Int)

instance Enumerable         Int64 where enumerate     = [minBound..maxBound]
instance FinitelyEnumerable Int64 where cardinality = Tagged $ 2 ^ (64 :: Int)

instance Enumerable         Float where enumerate     = map unsafeCoerce [0..maxBound :: Word32]
instance FinitelyEnumerable Float where cardinality = Tagged $ 2 ^ (32 :: Int)

instance Enumerable         Double where enumerate     = map unsafeCoerce [0..maxBound :: Word64]
instance FinitelyEnumerable Double where cardinality = Tagged $ 2 ^ (64 :: Int)

instance Enumerable Integer where enumerate = enumerateInterleaved

instance (Enumerable a, Integral a) => Enumerable (Ratio a) where
  enumerate = nub . map (uncurry (%)) . filter ((/=0) . snd) $ enumerate -- does this cover all of them? there's probably a better way of generating them, regardless
  
instance (FinitelyEnumerable a, Integral a) => FinitelyEnumerable (Ratio a)

newtype Partial a = Partial a

instance Enumerable a => Enumerable (Partial a) where
    enumerate = map Partial $ enumeratePartial 
    enumerateTotal = map Partial $ enumerateTotal 
    enumeratePartial = enumerate
    isPartial = Tagged True
    
instance FinitelyEnumerable a => FinitelyEnumerable (Partial a) where
  cardinality = Tagged $ 1 + unTagged (cardinality :: Tagged a Integer)