set.hs

module Set
    ( findSets
    , Shape (..)
    , Color (..)
    , Shade (..)
    , Card  (..)
    , Count (..)
    ) where

import Data.List (nub, sort, tails)

data Shape = Diamond | Squiggle | Pill deriving (Show, Ord, Enum, Eq)
data Color = Red | Green | Purple deriving (Show, Ord, Enum, Eq)
data Shade = Solid | Lined | Clear deriving (Show, Ord, Enum, Eq)
data Count = One | Two | Three deriving (Show, Ord, Enum, Eq)
data Card = Card { shape :: Shape
                 , color :: Color
                 , shade :: Shade
                 , count :: Count } deriving (Show, Ord, Eq)

allSame :: (Eq a) => a -> a -> a -> Bool
allSame a b c = a == b && b == c

allDiff :: (Eq a) => a -> a -> a -> Bool
allDiff a b c = a /= b && b /= c && a /= c

satisfiesSet :: (Eq a) => a -> a -> a -> Bool
satisfiesSet a b c = allSame a b c || allDiff a b c

partition' :: Int -> Int -> [a] -> [[a]]
partition' size offset
  | size <= 0   = error "partition': size must be positive"
  | offset <= 0 = error "partition': offset must be positive"
  | otherwise   = loop
  where
    loop :: [a] -> [[a]]
    loop xs = case splitAt size xs of
                (ys, []) -> if length ys == size then [ys] else []
                (ys, _ ) -> ys : loop (drop offset xs)

comb :: Int -> [a] -> [[a]]
comb 0 _      = [[]]
comb m l = [x:ys | x:xs <- tails l, ys <- comb (m-1) xs]

isSet :: Card -> Card -> Card -> Bool
isSet c1 c2 c3 = satisfiesSet (shape c1) (shape c2) (shape c3) &&
                 satisfiesSet (color c1) (color c2) (color c3) &&
                 satisfiesSet (shade c1) (shade c2) (shade c3) &&
                 satisfiesSet (count c1) (count c2) (count c3)

setCandidates :: [Card] -> [[Card]]
setCandidates cards = comb 3 cards

findSets cards = nub $ map sort $ filter (\[c1, c2, c3] -> isSet c1 c2 c3) (setCandidates cards)