module FAM where

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

fmap_List :: (a -> b) -> [] a -> [] b
-- "[] a" means "[a]" in types.
fmap_List f [] = []
fmap_List f (x:xs) = f x : fmap_List f xs

fmap_Maybe :: (a -> b) -> Maybe a -> Maybe b
fmap_Maybe f Nothing = Nothing
fmap_Maybe f (Just a) = Just (f a)

fmap_Either :: (a -> b) -> (Either e) a -> (Either e) b
fmap_Either f (Left e) = Left e
fmap_Either f (Right a) = Right (f a)

data BinTree a = BTNil | BTNode a (BinTree a) (BinTree a) deriving Show

instance Functor BinTree where
    -- fmap :: (a -> b) -> BinTree a -> BinTree b
    fmap f BTNil = BTNil
    fmap f (BTNode a lt rt) = BTNode (f a) (fmap f lt) (fmap f rt)


liftA2_List :: (a -> b -> c) -> [a] -> [b] -> [c]
liftA2_List f [] _ = []
liftA2_List f (a:as) bs = map (f a) bs ++ liftA2_List f as bs

liftA2_Maybe :: (a -> b -> c) -> Maybe a -> Maybe b -> Maybe c
liftA2_Maybe f Nothing _ = Nothing
liftA2_Maybe f (Just a) Nothing = Nothing
liftA2_Maybe f (Just a) (Just b) = Just (f a b)
-- The last two lines could be merged into:
-- liftA2_Maybe f (Just a) mb = fmap (f a) mb

liftA2_Either :: (a -> b -> c) -> Either e a -> Either e b -> Either e c
liftA2_Either f (Left e) _ = Left e
liftA2_Either f (Right a) (Left e) = Left e
liftA2_Either f (Right a) (Right b) = Right (f a b)
-- The last two lines could be merged into:
-- liftA2_Either f (Right a) eb = fmap (f a) eb

ap_List :: [a -> b] -> [a] -> [b]
-- Example:
--   ap_List [f,g] [1,2,3]
-- = [f 1, f 2, f 3, g 1, g 2, g 3]
-- = map f [1,2,3] ++ map g [1,2,3] ++ []
ap_List [] as = []
ap_List (f:fs) as = map f as ++ ap_List fs as

ap_ListV2 fs as = liftA2_List (\f a -> f a) fs as

liftA2_ListV2 f as bs = ap_List (fmap f as) bs

liftA3_List :: (a -> b -> c -> d) -> [a] -> [b] -> [c] -> [d]
liftA3_List f as bs cs = ap_List (ap_List (fmap f as) bs) cs


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

addRecipV1 x y =
    case recipMay x of
      Nothing -> Nothing
      Just x_recip -> case recipMay y of
        Nothing -> Nothing
        Just y_recip -> Just (1/x_recip + 1/y_recip)

addRecipV2 :: Double -> Double -> Maybe Double
addRecipV2 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]

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)

-- 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.

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 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 nothing 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 be consistent with >>='s order.)
    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 } )


-- Loop through a list of numbers, add the numbers to the state variable.
-- When done, return the latest value of the state variable.
sumProg :: [Integer] -> State Integer Integer
sumProg [] = get
sumProg (x:xs) =
    get
    >>= \s -> put (s+x)
    >> sumProg xs

statefulSum :: [Integer] -> Integer
statefulSum xs = functionize (sumProg xs) 0


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