Poker hands kata

Poker.hs

{-# LANGUAGE BangPatterns #-}
{-# LANGUAGE OverloadedStrings #-}

module Main where

import GHC.Exts (sortWith)

import Data.Function (on)
import Data.List (maximumBy, sortBy, nub, foldl')
import qualified Data.Map.Strict as M
import Data.Maybe (isJust, fromJust)
import Data.Text (Text, unpack)

import Test.Hspec
import Test.Hspec.Runner

import Control.Applicative ((<$>))
import Control.Monad (void, forM_)
import System.Random.Shuffle (shuffleM)
import Text.Printf (printf)

data Suit = Spades
          | Hearts
          | Diamonds
          | Clubs
          deriving (Eq, Ord)

instance Show Suit where
  show Spades   = "♠"
  show Hearts   = "♥"
  show Diamonds = "♦"
  show Clubs    = "♣"

data Rank = Digit Int
          | Jack
          | Queen
          | King
          | Ace
          deriving (Eq, Ord)

instance Show Rank where
  show Ace        = "A"
  show King       = "K"
  show Queen      = "Q"
  show Jack       = "J"
  show (Digit 10) = "T"
  show (Digit d)  = show d

instance Enum Rank where
  toEnum = (!!) $ map Digit [2..10] ++ [Jack, Queen, King, Ace]
  fromEnum (Digit 10) = 8
  fromEnum (Digit n)  = n-2
  fromEnum Jack       = 9
  fromEnum Queen      = 10
  fromEnum King       = 11
  fromEnum Ace        = 12

data Card = Card { suit :: Suit, rank :: Rank }
          deriving (Eq, Ord)

instance Show Card where
  show (Card s r) = show r ++ show s

defaultDeck :: [Card]
defaultDeck = [ Card s r | s <- [Spades, Hearts, Diamonds, Clubs], r <- [Digit 2 .. Ace] ]
                    
newtype Hand = Hand { unHand :: [Card] }
            deriving Eq

instance Show Hand where
  show (Hand h) = show h

createHand :: [(Suit, Rank)] -> Hand
createHand hand
  | length hand == 5 = Hand $ map (uncurry Card) hand
  | otherwise = error "a hand should consist of exactly 5 cards"

cardRanks :: Hand -> [Rank]
cardRanks = sortBy (flip compare) . map rank . unHand

-- | the ranks given to this should be sorted (for example using cardRanks)
straight :: [Rank] -> Maybe Rank
straight []    = Nothing
straight ranks
  | min' == Digit 2 =
    if ranks' == [Digit 2 .. Digit 5] ++ [Ace]
    then Just (Digit 5)
    else matches
  | otherwise = matches
  where min' = minimum ranks
        ranks' = reverse ranks
        matches = if ranks' == take 5 [min'..Ace]
                  then Just (maximum ranks)
                  else Nothing


flush :: Hand -> Bool
flush = (== 1) . length . nub . map suit . unHand

-- | the ranks given to this should be sorted 8for example using cardRanks)
twoPair :: [Rank] -> Maybe (Rank, Rank)
twoPair ranks = case (first, second) of
  (Just f, Just s) | f /= s -> Just (f,s)
  (_,_) -> Nothing
  where first = kindOf 2 ranks
        second = kindOf 2 (reverse ranks)

type Signature = [(Int, Rank)]

-- | the ranks given to this should be sorted 8for example using cardRanks)
kindOf :: Int -> [Rank] -> Maybe Rank
kindOf n rs = kind (signature rs) n

kind :: Signature -> Int -> Maybe Rank
kind s n = lookup n s

-- | the ranks given to this should be sorted 8for example using cardRanks)
signature :: [Rank] -> Signature
signature [] = []
signature (r:rs) = collect (1,r) rs []
  where collect (n,l) [] acc = reverse $ (n,l):acc
        collect (n,l) (l':ls') acc
          | l==l' = collect (n+1,l) ls' acc
          | otherwise = collect (1,l') ls' $ (n,l):acc

type Hands = [Hand]

randomHands :: [Card] -> Int -> IO Hands
randomHands deck n = do
  deck' <- shuffleM deck
  return $ fst $ foldl takeHand ([], deck') [1..n]
  where takeHand (acc, d) _ = ((Hand $ take 5 d):acc, drop 5 d)

poker :: Hands -> Hand
poker = maximumBy (compare `on` handRank)

data HandRank = HighCard Rank [Rank]
              | Pair Rank [Rank]
              | TwoPairs Rank Rank Rank
              | ThreeOfAKind Rank [Rank]
              | Straight Rank
              | Flush [Rank]
              | FullHouse Rank Rank
              | FourOfAKind Rank Rank
              | StraightFlush Rank
              deriving (Eq, Ord, Show)

handRank :: Hand -> HandRank
handRank hand
  | isJust str && flush hand     = StraightFlush (fromJust str)
  | isJust fourOfAKind          = FourOfAKind (fromJust fourOfAKind) (fromJust single)
  | isJust threeOfAKind && isJust twoOfAKind = FullHouse (fromJust threeOfAKind) (fromJust twoOfAKind)
  | flush hand                  = Flush ranks
  | isJust str                  = Straight (fromJust str)
  | isJust threeOfAKind         = ThreeOfAKind (fromJust threeOfAKind) ranks
  | isJust pair2                = (\ (Just (h, l)) -> TwoPairs h l (fromJust single)) pair2
  | isJust twoOfAKind           = Pair (fromJust twoOfAKind) ranks
  | otherwise                   = HighCard (maximum ranks) ranks
  where ranks        = cardRanks hand
        sig          = signature ranks
        fourOfAKind  = kind sig 4
        threeOfAKind = kind sig 3
        twoOfAKind   = kind sig 2
        single       = kind sig 1
        pair2        = twoPair ranks
        str          = straight ranks

classifyRank :: Hand -> Text
classifyRank hand =
  case handRank hand of
    (HighCard _ _)     -> "High-Card     "
    (Pair _ _)         -> "Pair          "
    (TwoPairs _ _ _)   -> "Two-Pairs     "
    (ThreeOfAKind _ _) -> "3-of-a-Kind   "
    (Straight _)       -> "Straight      "
    (Flush _)          -> "Flush         "
    (FullHouse _ _)    -> "Full-House    "
    (FourOfAKind _ _)  -> "4-of-a-Kind   "
    (StraightFlush _)  -> "Straight-Flush"
    
rankStats :: Hands -> [(Text, Int)]
rankStats = sortWith (negate . snd) . M.toList . foldl' count M.empty . map classifyRank
  where count acc r = M.insertWith (+) r 1 acc

randomStats :: Int -> IO [(Text, Float)]
randomStats n = percentOf . sortWith (negate . snd) . M.toList <$> drawHands n M.empty
  where drawHands n' m =
          if n' == 0
          then return m
          else do
            !r <- classifyRank . head <$> randomHands defaultDeck 1
            let m' = M.insertWith (+) r 1 m
            drawHands (n'-1) m'
        percentOf :: [(Text, Int)] -> [(Text, Float)]
        percentOf = map (\ (x,c) -> (x, 100 * fromIntegral c / fromIntegral n))

main :: IO ()
main = do
  tests
  putStrLn "Stats for poker-hands"
  stats <- randomStats 1000000
  forM_ stats (\ (r,p) -> putStrLn $ unpack r ++ "=\t" ++ printf "%.4f" p ++ "%")


tests :: IO ()
tests =  void . hspecWith (defaultConfig { configColorMode = ColorAlways }) $ do
  let sf = createHand [ (Clubs, Digit 6), (Clubs, Digit 7), (Clubs, Digit 8), (Clubs, Digit 9), (Clubs, Digit 10) ]
  let fk = createHand [ (Diamonds, Digit 9), (Hearts, Digit 9), (Spades, Digit 9), (Clubs, Digit 9), (Diamonds, Digit 7) ]
  let fh = createHand [ (Diamonds, Digit 10), (Clubs, Digit 10), (Hearts, Digit 10), (Clubs, Digit 7), (Diamonds, Digit 7) ]
  let tP = createHand [ (Hearts, King), (Spades, Digit 10), (Diamonds, Digit 10), (Clubs, King), (Clubs, Ace) ]
  let s1 = createHand [ (Spades, Ace), (Spades, Digit 2), (Spades, Digit 3), (Spades, Digit 4), (Clubs, Digit 5) ]
  let s2 = createHand [ (Clubs, Digit 2), (Clubs, Digit 3), (Clubs, Digit 4), (Clubs, Digit 5), (Spades, Digit 6) ]
  let p1 = createHand [ (Clubs, Digit 9), (Spades, Digit 9), (Clubs, Digit 4), (Clubs, Digit 5), (Spades, Digit 6) ]
  let p2 = createHand [ (Clubs, Digit 9), (Spades, Digit 9), (Clubs, Digit 4), (Clubs, Digit 5), (Spades, Digit 7) ]
  let p3 = createHand [ (Clubs, Digit 9), (Spades, Digit 9), (Clubs, Digit 4), (Clubs, Digit 5), (Spades, Ace) ]

  describe "Testing poker with a couple of hands" $ do
    context "finding the best hand with poker" $ do
      it "returns the pair with the better kicker" $ do
        poker [p1,p2] `shouldBe` p2
      it "returns the pair with the better kicker - even if it's higher than the pair" $ do
        poker [p1,p2,p3] `shouldBe` p3
      it "returns the straight flush out of minor hands" $ do
        poker [sf, fk, fh] `shouldBe` sf
      it "returns four-of-a-kind over full-house" $ do
        poker [fk, fh] `shouldBe` fk
      it "returns full-house of two full-houses" $ do
        poker [fh, fh] `shouldBe` fh
      it "returns the only hand if only one is given" $ do
        poker [fh] `shouldBe` fh
      it "returns one of the hand if 100 equal hands are given" $ do
        poker (sf : [fh | _ <- [1..100 :: Int]]) `shouldBe` sf
      it "returns a low straight over a two pair" $ do
        poker [tP, s1] `shouldBe` s1
      it "picks a straight beginning at 2 over one beginning at Ace" $ do
        poker [s2, s1] `shouldBe` s2
    context "ranking hands with handRank" $ do
      it "should rank the straight-flush correctly" $ do
        handRank sf `shouldBe` StraightFlush (Digit 10)
      it "should rank then four-of-a-kind correctly" $ do
        handRank fk `shouldBe` FourOfAKind (Digit 9) (Digit 7)
      it "should rank the full-house correctly" $ do
        handRank fh `shouldBe` FullHouse (Digit 10) (Digit 7)
    context "using kind/kindOf to extract information" $ do
      it "should get 4 nines out of fk" $ do
        kindOf 4 (cardRanks fk) `shouldBe` Just (Digit 9)
      it "should not get 3 nines out of fk" $ do
        kindOf 3 (cardRanks fk) `shouldBe` Nothing
      it "should get 3 tens out of fh" $ do
        kindOf 3 (cardRanks fh) `shouldBe` Just (Digit 10)
      it "should get 2 sevens out of fh" $ do
        kindOf 2 (cardRanks fh) `shouldBe` Just (Digit 7)
    context "using twoPair to extract information" $ do
      it "should get 2 Kings and 2 tens from `kH, 10S, 10D, kC, aC`" $ do
        twoPair (cardRanks tP) `shouldBe` Just (King, Digit 10)
    context "cardRanks returns the ranks of the cards in descending order" $ do
      it "should get 2 Kings and 2 tens from `kH, 10S, 10D, kC, aC`" $ do
        cardRanks tP `shouldBe` [Ace, King, King, Digit 10, Digit 10]
    context "straight checks for a unbrocken line of numbers" $ do
      it "should be true for [9,8,7,6,5]" $ do
        straight [Digit 9, Digit 8, Digit 7, Digit 6, Digit 5] `shouldBe` Just (Digit 9)
      it "should be true for [A,5,4,3,2]" $ do
        straight [Ace, Digit 5, Digit 4, Digit 3, Digit 2] `shouldBe` Just (Digit 5)
      it "should be true for [A,K,D,J,10]" $ do
        straight [Ace, King, Queen, Jack, Digit 10] `shouldBe` Just Ace
      it "should be false for [9,8,8,6,5]" $ do
        straight [Digit 9, Digit 8, Digit 8, Digit 6, Digit 5] `shouldBe` Nothing
    context "flush checks if all cards are in the same suit" $ do
      it "should be true sf" $ do
        flush sf `shouldBe` True
      it "should be false for fk" $ do
        flush fk `shouldBe` False