module Monads where

import Control.Applicative (liftA2)
import Data.List (delete)

recipMay :: Double -> Maybe Double
recipMay a | a == 0 = Nothing
           | otherwise = Just (1 / a)    -- or: pure (1 / a)

twoPlusRecip :: Double -> Maybe Double
twoPlusRecip x = fmap (\r -> 2 + r) (recipMay x)

addRecips :: Double -> Double -> Maybe Double
addRecips x y = liftA2 (+) (recipMay x) (recipMay y)


sqrts :: Double -> [Double]
sqrts a | a < 0 = []
        | a == 0 = [0]    -- or: pure 0
        | otherwise = [- sqrt a, sqrt a]

twoPlusSqrts :: Double -> [Double]
twoPlusSqrts x = fmap (\r -> 2 + r) (sqrts x)

addSqrts :: Double -> Double -> [Double]
addSqrts x y = liftA2 (+) (sqrts x) (sqrts y)


bind_Maybe :: Maybe a -> (a -> Maybe b) -> Maybe b
bind_Maybe Nothing k = Nothing
bind_Maybe (Just a) k = k a

bind_Either :: Either e a -> (a -> Either e b) -> Either e b
bind_Either (Left e) _ = Left e
bind_Either (Right a) k = k a

bind_List :: [a] -> (a -> [b]) -> [b]
bind_List xs k = concat (map k xs)
-- Explanation: map k xs :: [[b]] is almost there, use concat to flatten

sqrtsTwoPlusSqrts :: Double -> [Double]
sqrtsTwoPlusSqrts x = bind_List (sqrts x) (\r -> sqrts (2 + r))


-- All permutations of the input.
permsV1 :: Eq a => [a] -> [[a]]
permsV1 xs = permsAddV1 xs []

-- Helper: permsAdd xs ys = all permutations of xs prepended to ys
-- E.g.,
-- permsAdd [1,2] [6,4,7] = [2:1:[6,4,7], 1:2:[6,4,7]]
-- permsAdd [] [6,4,7] = [[6,4,7]]
permsAddV1 :: Eq a => [a] -> [a] -> [[a]]
permsAddV1 [] ys = [ys]
permsAddV1 xs ys = bind_List xs (\x -> permsAddV1 (delete x xs) (x : ys))
-- For each x in xs, I want to delete x from xs, add x to ys, and recurse.
-- OR:
-- Non-deterministically choose x from xs, ... ditto.


sqrtsTwoPlusSqrtsV2 :: Double -> [Double]
sqrtsTwoPlusSqrtsV2 x = sqrts x >>= \r -> sqrts (2 + r)

permsV2 :: Eq a => [a] -> [[a]]
permsV2 xs = permsAddV2 xs []

permsAddV2 :: Eq a => [a] -> [a] -> [[a]]
permsAddV2 [] ys = return ys
permsAddV2 xs ys = xs >>= \x -> permsAddV2 (delete x xs) (x : ys)


data CExpr = Lit Double
           | Add CExpr CExpr
           | Sub CExpr CExpr
           | Mul CExpr CExpr
           | Div CExpr CExpr
           | C   -- read value of c
           | CPP -- c++
           | PPC -- ++c
    deriving Show
-- e.g., (c++ + c) + (14.4 / ++c): Add (Add CPP C) (Div (Lit 14.4) PPC)
-- c is an int variable, but the whole expression is intended to be double

sampleCExpr :: CExpr
sampleCExpr = Add (Add CPP C) (Div (Lit 14.4) PPC)

evalCExpr :: CExpr -> State Int Double
evalCExpr (Lit x) = pure x
evalCExpr (Add e1 e2) = liftA2 (+) (evalCExpr e1) (evalCExpr e2)
evalCExpr (Sub e1 e2) = liftA2 (-) (evalCExpr e1) (evalCExpr e2)
evalCExpr (Mul e1 e2) = liftA2 (*) (evalCExpr e1) (evalCExpr e2)
evalCExpr (Div e1 e2) = liftA2 (/) (evalCExpr e1) (evalCExpr e2)
evalCExpr C = fmap fromIntegral get
evalCExpr CPP = get
                >>= \i -> put (i+1)
                >> pure (fromIntegral i)
evalCExpr PPC = get
                >>= \i -> put (i+1)
                >> pure (fromIntegral (i+1))


data State s a = MkState (s -> (s, a))

-- Unwrap MkState.
deState :: State s a -> s -> (s, a)
deState (MkState stf) = stf

-- Bridge from stateful fantasy to mathematical reality! "functionize prog s0"
-- runs prog starting with initial state value s0 and gives you the final
-- answer.  Or, turns prog into a math pure function.
functionize :: State s a -> s -> a
functionize prog s0 = snd (deState prog s0)

-- "get" reads and returns the current value of the state variable.
get :: State s s
get = MkState (\s0 -> (s0, s0))
-- old state = s0, new state = old state = s0, answer s0 too.

-- "put s1" sets the state variable to s1.  It returns the 0-tuple because there
-- is no information to return.
put :: s -> State s ()
put s = MkState (\s0 -> (s , ()))
-- ignore old state, new state = s, answer the 0-tuple ().

instance Functor (State s) where
    -- fmap :: (a -> b) -> State s a -> State s b
    fmap f (MkState stf) = MkState
        (\s0 ->
           -- Goal: Like stf but use f to convert a to b
           -- old state = s0, give to stf for new state s1 and answer a
           case stf s0 of (s1, a) ->
           -- overall new state is also s1, but change answer to f a
                              (s1, f a))

testStateFunctor = deState (fmap length program) 10
  where
    program :: State Integer String
    program = MkState (\s0 -> (s0+2, "hello"))
-- should give (12, 5)

instance Applicative (State s) where
    -- pure :: a -> State s a
    -- Goal: Give the answer a and try not to have an effect.
    -- "effect" for State means state change.
    pure a = MkState (\s0 -> (s0, a))
    -- so new state = old state

    -- liftA2 :: (a -> b -> c) -> State s a -> State s b -> State s c
    --
    -- State transition goal:
    --                  overall old state
    -- --1st-program--> intermediate state
    -- --2nd-program--> overall new state
    --
    -- (Why not the other order? Actually would be legitimate, but we usually
    -- desire liftA2's order to follow >>='s order below.)
    liftA2 op (MkState stf1) (MkState stf2) =
      MkState
        (\s0 ->
           -- overall old state = s0, give to stf1
           case stf1 s0 of { (s1, a) ->
           -- intermediate state = s1, give to stf2                  
           case stf2 s1 of { (s2, b) ->
           -- overall new state = s2
           -- overall answer = op a b
           (s2, op a b) }} )

testStateApplicative = deState (liftA2 (:) prog1 prog2) 10
  where
    prog1 :: State Integer Char
    prog1 = MkState (\s0 -> (s0+2, 'h'))
    prog2 :: State Integer String
    prog2 = MkState (\s0 -> (s0*2, "ello"))
-- should give (24, "hello").  24 = (10+2)*2.

instance Monad (State s) where
    return = pure

    -- (>>=) :: State s a -> (a -> State s b) -> State s b
    -- Goal:
    -- 1. overall old state --1st-program--> (intermediate state, a)
    -- 2. give a and intermedate state to 2nd program.
    MkState stf1 >>= k = MkState
        (\s0 ->
           -- overall old state = s0, give to stf1
           case stf1 s0 of { (s1, a) ->
           -- k is waiting for the answer a
           --      and also the intermediate state s1
           -- technicality: "(k a) s1" is conceptually right but nominally a
           -- type error because (k a) :: State s b, not s -> (s, b)
           -- Ah but deState can unwrap! (Or use pattern matching.)
           deState (k a) s1 } )


toyCheckV1 :: IO Bool
toyCheckV1 =
    getChar
    >>= \c1 -> getChar
    >>= \c2 -> getChar
    >>= \c3 -> return ([c1, c2, c3] == "AL\n")

class Monad f => MonadToyCheck f where
    toyGetChar :: f Char
-- Simplifying assumptions: Enough characters, no failure.  A practical version
-- should add methods for raising and catching EOF exceptions.

toyCheckV2 :: MonadToyCheck f => f Bool
toyCheckV2 =
    toyGetChar
    >>= \c1 -> toyGetChar
    >>= \c2 -> toyGetChar
    >>= \c3 -> return ([c1, c2, c3] == "AL\n")

-- For production code:

instance MonadToyCheck IO where
    toyGetChar = getChar

realProgram :: IO Bool
realProgram = toyCheckV2

-- For mock testing:

data Feeder a = MkFeeder (String -> (String, a))
-- Again, simplifying assumptions etc.  But basically like the state monad, with
-- the state being what's not yet consumed in the string.

-- Unwrap MkFeeder.
unFeeder :: Feeder a -> String -> (String, a)
unFeeder (MkFeeder sf) = sf

instance Monad Feeder where
    return a = MkFeeder (\s -> (s, a))
    prog1 >>= k = MkFeeder (\s0 -> case unFeeder prog1 s0 of
                                     (s1, a) -> unFeeder (k a) s1)

instance MonadToyCheck Feeder where
    -- toyGetChar :: Feeder Char
    toyGetChar = MkFeeder (\(c:cs) -> (cs, c))

instance Functor Feeder where
    fmap f p = p >>= \a -> return (f a)

instance Applicative Feeder where
    pure a = MkFeeder (\s -> (s, a))
    pf <*> pa = pf >>= \f -> pa >>= \a -> return (f a)

testToyChecker2 :: String -> Bool
testToyChecker2 str = snd (unFeeder toyCheckV2 str)

toyTest1 = testToyChecker2 "ALhello"   -- should be False
toyTest2 = testToyChecker2 "AL\nhello" -- should be True