wordle.hs

import qualified Data.Set as S
import Data.Set (Set)
import Data.Char (isLower)
import Data.Ord (comparing, Down (Down))
import Data.List (sortBy, subsequences, minimumBy, maximumBy)
import Control.Monad.Trans.Writer.CPS
import Data.Monoid
import Data.Foldable (traverse_)

wordFilter :: String -> Bool
wordFilter w = length w == 5 && all (flip elem letters) w


type Dict = Set String

data Pos = P1 | P2 | P3 | P4 | P5
  deriving (Eq, Ord, Show, Enum, Bounded)

data Result = Exact Char Pos | Has Char | Hasnt Char
  deriving (Eq, Ord, Show)

data Hit = Yup | Hit | Miss
  deriving (Eq, Ord, Show)

parseHit :: Char -> Maybe Hit
parseHit 'x' = Just Yup
parseHit '.' = Just Hit
parseHit ' ' = Just Miss
parseHit _  = Nothing

parseHits :: String -> Maybe [Hit]
parseHits = traverse parseHit . take 5

makeResult :: [Char] -> [Hit] -> [Result]
makeResult = go P1
  where
    go :: Pos -> [Char] -> [Hit] -> [Result]
    go n (s : ss) (Yup : hs) = Exact s n : go (succ n) ss hs
    go n (s : ss) (Hit : hs) = Has s : go (succ n) ss hs
    go n (s : ss) (Miss : hs) = Hasnt s : go (succ n) ss hs
    go _ [] [] = []
    go _ _ _ = error "bad bad man"

refineDict :: Result -> Dict -> Dict
refineDict (Exact c pos) ws = S.filter ((== c) . posToChar pos) ws
refineDict (Has c) ws = S.filter (elem c) ws
refineDict (Hasnt c) ws = S.filter (not . elem c) ws

posToChar :: Pos -> String -> Char
posToChar p s = s !! fromEnum p


entropy :: Dict -> Char -> Int
entropy d c =
  let without = refineDict (Hasnt c) d
      with = d S.\\ without
   in abs $ S.size without - S.size with


check :: String -> String -> [Result]
check word' = go P1 (S.fromList word') word'
  where
    go n bag (w : word) (g : guess)
      | w == g               = Exact g n : go (succ n) bag word guess
      | (S.member g bag) =       Has g   : go (succ n) bag word guess
      | not (S.member g bag) = Hasnt g   : go (succ n) bag word guess
      | otherwise            =             go (succ n) bag word guess
    go _ _ [] [] = []
    go _ _ _ _ = error "broken invariant"

letters :: [Char]
letters = ['a' .. 'z']

best :: Dict -> [Char]
best d = sortBy (comparing $ entropy d) letters

counts :: Dict -> [(Char, Int)]
counts d = fmap (\x -> (x, entropy d x)) letters

wordScore :: Dict -> String -> (Int, Down Int)
wordScore d s =
  let s' = S.toList $ S.fromList s
      num_dups = 5 - length s'
      k = S.size d
   in (length s', Down $ sum (fmap (entropy d) s'))

nextGuess :: Dict -> Dict -> String
nextGuess all_words dict = maximumBy (comparing $ wordScore dict) $ S.elems all_words

search :: String -> Dict -> Dict -> IO ()
search word d0 d | S.size d == 1 = putStrLn $ head $ S.elems d
search word d0 d | S.null d = error "NO MORE WORDS"
search word d0 d = do
  let g = nextGuess d0 d
  putStrLn g
  let res = check word g
  let d' = appEndo (foldMap (Endo . refineDict) res) d
  print $ S.toList d'
  print $ log (fromIntegral (S.size d) / fromIntegral (S.size d')) / log 2
  search word d0 d'


seek :: Dict -> Dict -> IO ()
seek d0 d | S.size d == 1 = putStrLn $ head $ S.elems d
seek d0 d | S.null d = error "NO MORE WORDS"
seek d0 d = do
  let g = nextGuess d0 d
  putStrLn g
  putStr "> "
  Just x <- fmap parseHits getLine
  let res = makeResult g x
  seek d0 $ appEndo (foldMap (Endo . refineDict) res) d


main :: IO ()
main = do
  dict <- fmap (S.fromList . filter wordFilter . lines) $ readFile "words"
  -- dict <- fmap (S.fromList . filter wordFilter . lines) $ readFile "/usr/share/dict/words"
  let word = "pilot"
  search word dict dict