{-- Examples from Monad lecture.
    Thanks to Andrew Tolmach for the examples.
    The outline is drawn from "Monad for Functional Programming" by Phil Wadler.
--}

import IO
import Control.Monad.ST.Lazy
import Data.STRef.Lazy
import Data.Array.ST
import Data.Array
import Control.Monad

echo :: IO ()
echo = do { b <- isEOF;
            if not b then do 
              { x <- getChar; putChar x; echo }
            else return () }

data Exp = Plus  Exp Exp
         | Minus Exp Exp
         | Times Exp Exp
         | Div   Exp Exp
         | Const Int
      deriving (Show)

expA = (Div (Const 3)
            (Plus (Const 4) (Const 2)))
buggyExpA = Plus (Const 4) 
                 (Div (Const 2) (Const 0))

buggyExpB = (Div (Plus (Const 4) (Const 2)) 
                 (Times (Const 0) (Const 10)))

{--
Original, simple interpreter
--}
eval :: Exp -> Int
eval (Plus  e1 e2) = (eval e1) + (eval e2)
eval (Minus e1 e2) = (eval e1) - (eval e2)
eval (Times e1 e2) = (eval e1) * (eval e2)
eval (Div   e1 e2) = (eval e1) `div` (eval e2)
eval (Const i)     = i

answerA = eval expA
answerBugA = eval buggyExpA
answerBugB = eval buggyExpB

-- Suppose we want to improve this code by trapping attempts to divide by zero.

data Hopefully a = Ok a | Error String
     deriving (Show)

{-- 
 Original interpreter extended to check for division by zero.
--}
eval1 :: Exp -> Hopefully Int
eval1 (Plus  e1 e2) = 
   case eval1 e1 of
      Ok v1 -> 
        case eval1 e2 of
          Ok v2 -> Ok (v1 + v2)
          Error s -> Error s
      Error s -> Error s
eval1 (Minus e1 e2) = 
   case eval1 e1 of
      Ok v1 -> 
        case eval1 e2 of
          Ok v2 -> Ok (v1 - v2)
          Error s -> Error s
      Error s -> Error s
eval1 (Times e1 e2) = 
   case eval1 e1 of
      Ok v1 -> 
        case eval1 e2 of
          Ok v2 -> Ok (v1 * v2)
          Error s -> Error s
      Error s -> Error s
eval1 (Div   e1 e2) = 
   case eval1 e1 of
      Ok v1 -> 
        case eval1 e2 of
          Ok v2 -> if v2 == 0 then Error "divby0"
                   else Ok (v1 `div` v2)
          Error s -> Error s
      Error s -> Error s
eval1 (Const i)     = Ok i

answer1A = eval1 expA
answerBug1A = eval1 buggyExpA
answerBug1B = eval1 buggyExpB


{-- This solution exposes a lot of ugly plumbing.
    Every time an expression evalutes to Error, the Error 
    propagates to the final result.

    We can abstract this behavior to a higher-order function.
--}

ifOKthen :: Hopefully a -> (a -> Hopefully b) -> Hopefully b
e `ifOKthen` k = 
   case e of 
     Ok x -> k x
     Error s -> Error s
   
{-- Interpreter that uses the `ifOKthen` combinator --}
eval2 :: Exp -> Hopefully Int
eval2 (Plus  e1 e2) = 
    eval2 e1 `ifOKthen` (\v1 ->
    eval2 e2 `ifOKthen` (\v2 ->
    Ok(v1 + v2)))
eval2 (Minus e1 e2) = 
    eval2 e1 `ifOKthen` (\v1 ->
    eval2 e2 `ifOKthen` (\v2 ->
    Ok(v1 - v2)))
eval2 (Times e1 e2) = 
    eval2 e1 `ifOKthen` (\v1 ->
    eval2 e2 `ifOKthen` (\v2 ->
    Ok(v1 * v2)))
eval2 (Div   e1 e2) = 
    eval2 e1 `ifOKthen` (\v1 ->
    eval2 e2 `ifOKthen` (\v2 ->
    if v2 == 0 then Error "divby0" else Ok(v1 `div` v2)))
eval2 (Const i)     = Ok i

answer2A = eval2 expA
answerBug2A = eval2 buggyExpA
answerBug2B = eval2 buggyExpB


{-- 
Compare the type of `isOKthen` with the type of >>=
ifOKthen :: Hopefully a -> (a -> Hopefully b) -> Hopefully b
(>>=)    :: IO a        -> (a -> IO b)        -> IO b

and the type of Ok with the type of return
Ok     :: a -> Hopefully a
return :: a -> IO a

This similarity isn't accidental.  

Like IO, Hopefully with ifOKthen and Ok forms a monad.
--}

{-- 
The standard prelude defines a type class for monads.
We can use it to tell the system that Hopefully is a monad.
This allows us to use a library of operations that work over any moand.
And, it enables us to use the "do" notation for Hopefully just
as we did for IO
--}

instance Monad Hopefully where
   return = Ok
   (>>=) = ifOKthen

{-- Monadic error-checking interpreter --}   
eval3 :: Exp -> Hopefully Int
eval3 (Plus  e1 e2) = do {
    v1 <- eval3 e1;
    v2 <- eval3 e2;
    return (v1 + v2)     }
eval3 (Minus e1 e2) = do {
    v1 <- eval3 e1;
    v2 <- eval3 e2;
    return (v1 - v2)    }
eval3 (Times e1 e2) = do {
    v1 <- eval3 e1;
    v2 <- eval3 e2;
    return (v1 * v2)     }
eval3 (Div   e1 e2) =  do {
    v1 <- eval3 e1;
    v2 <- eval3 e2;
    if v2 == 0 then Error "divby0" else return (v1 `div` v2)}
eval3 (Const i)     = return i

answer3A = eval3 expA
answerBug3A = eval3 buggyExpA
answerBug3B = eval3 buggyExpB

{-- 
  Suppose we wanted to modify our original evaluator to record a trace of its ops.
--}

evalT :: Exp -> [String] -> ([String], Int)
evalT (Plus  e1 e2) s = 
  let (s1,v1) = evalT e1 s
      (s2,v2) = evalT e2 s1
  in (s2++["+"], v1 + v2)
evalT (Minus e1 e2) s = 
  let (s1,v1) = evalT e1 s
      (s2,v2) = evalT e2 s1
  in (s2++["-"], v1 - v2)
evalT (Times e1 e2) s = 
  let (s1,v1) = evalT e1 s
      (s2,v2) = evalT e2 s1
  in (s2++["*"], v1 * v2)
evalT (Div e1 e2) s = 
  let (s1,v1) = evalT e1 s
      (s2,v2) = evalT e2 s1
  in (s2++["/"], v1 `div` v2)
evalT (Const i) s = (s++[show i], i)

(traceTA,anwerTA) = evalT expA []
-- (["3","4","2","+","/"],0)

{-
  In an imperative language, it would be more convenient to store trace
  in a global variable.  In Haskell, we can use a tracing moad.
-}

data Tr a = Tr [String] a
instance Monad Tr where
  return a = Tr [] a
  m >>= k = let (trace,  a) = runTr m 
                (trace', b) = runTr (k a)
            in Tr (trace++trace') b


-- This function lets us "run" the Trace monad
runTr :: Tr a -> ([String], a)
runTr (Tr s a) = (s,a)

-- This function adds x to the trace
trace :: String -> Tr ()
trace x = Tr [x] ()

{-- Monadic tracing interpreter --}   
evalTM :: Exp -> Tr Int
evalTM (Plus  e1 e2) = do {
    v1 <- evalTM e1;
    v2 <- evalTM e2;
    trace "+";
    return (v1 + v2)     }
evalTM (Minus e1 e2) = do {
    v1 <- evalTM e1;
    v2 <- evalTM e2;
    trace "-";
    return (v1 - v2)    }
evalTM (Times e1 e2) = do {
    v1 <- evalTM e1;
    v2 <- evalTM e2;
    trace "*";
    return (v1 * v2)      }
evalTM (Div   e1 e2) = do {
    v1 <- evalTM e1;
    v2 <- evalTM e2;
    trace "/";        
    return (v1 `div` v2)  }
evalTM (Const i)     = do{trace (show i); return i}

answerTM = runTr (evalTM expA)
-- (["3","4","2","+","/"],0)

{-- 
   Count the number of divisions during evaluation.

   Doing this non-monadically requires changes similar to the trace changes, where
   an integer variable is threaded through all the computation.

   To do this monadically, we'll introduce a State monad.  
   The structure of this monad is very similar to the structure of the IO monad,
   in that an action takes a "state" (world) as input and returns a value as well
   as a modified "state" (world).

   The Prelude 

   The type constructor ST is defined in Control.Monad.ST.Lazy, so to show the code
   here, we'll use the name State istead.

--}


data State s a = ST {runST' :: s -> (a,s)}
{-- Newtype is just like a datatype except
  . it can only have one constructor
  . its constructor can have only one argument
  . it describes a strict isomorphism between types
  . it can be implemented more efficiently that the corresponding datatype
 -- The curly braces and the label :: defien a record with a single field.
  . the name of the field serves as a deconstructor for the type, so
  . runST' :: State s a -> s -> (s,a)
  . This function is called "runST'" because it will corresponds to 
    "running" the action of the monad.
--}

instance Monad (State s) where
  return a = ST (\s -> (a,s))     -- think about the pipe diagram for IO return
  m >>= k =  ST (\s -> let (a,s') = runST' m s  
                       in runST' (k a) s')

-- Get the value of the state, leave state value unchanged
get :: State s s
get = ST (\s -> (s,s))

-- make put's argument the new state, return the unit value
put :: s -> State s ()
put s = ST (\_ -> ((),s))

-- before update, the state has value s.
-- return s, replace s with f s.
update :: (s->s) -> State s s
update f = ST (\s -> (s, f s))

{-- Monadic interpreter that counts the number of divisions. --}   
evalCD :: Exp -> State Int Int
evalCD (Plus  e1 e2) = do {
    v1 <- evalCD e1;
    v2 <- evalCD e2;
    return (v1 + v2)     }
evalCD (Minus e1 e2) = do {
    v1 <- evalCD e1;
    v2 <- evalCD e2;
    return (v1 - v2)    }
evalCD (Times e1 e2) = do {
    v1 <- evalCD e1;
    v2 <- evalCD e2;
    return (v1 * v2)      }
evalCD (Div   e1 e2) = do {
    v1 <- evalCD e1;
    v2 <- evalCD e2;
    update (+1);
    return (v1 `div` v2)  }
evalCD (Const i)     = do{return i}

-- The second component is the desired answer as it is the final "state".
answerCD = runST' (evalCD expA) 0
-- (0,1)

{-- End of interpreter example  --}

-- Functions for allocating imperative variables in s
{--  From Data.STRef.Lazy
data STRef s a
newSTRef :: a -> ST s (STRef s a)
readSTRef :: STRef s a -> ST s a
writeSTRef :: STRef s a -> a -> ST s ()
modifySTRef :: STRef s a -> (a -> a) -> ST s ()
--}

swap :: STRef s a -> STRef s a -> ST s ()
swap r1 r2 = do {v1 <- readSTRef r1;
                 v2 <- readSTRef r2;
                 writeSTRef r1 v2;
                 writeSTRef r2 v1}

testSwap :: Int
testSwap = runST (do { r1 <- newSTRef 1;
                       r2 <- newSTRef 2;
                       swap r1 r2;
                       readSTRef r2})
-- 1

type Vertex = Char
type Graph = Array Vertex [Vertex]
data Tree a = Node a [Tree a]

aGraph = listArray ('a','e') [['b','c'], [],['d'],['e'],[]]

dfs :: Graph -> [Vertex] -> [Tree Vertex]
dfs g vs = runST(
              do{ marks <- newArray (bounds g) False;
                  search marks vs})
    where search :: STArray s Vertex Bool -> [Vertex] -> ST s [Tree Vertex]
          search marks [] = return []
          search marks (v:vs) = do {
               visited <- readArray marks v;
               if visited then 
                  search marks vs
               else 
                  do { writeArray marks v True;
                       ts <- search marks (g!v);
                       us <- search marks vs;
                       return ((Node v ts) : us)  } }

toPreOrder :: [Tree Vertex] -> [Vertex]
toPreOrder [] = []
toPreOrder (t:ts) = (toOrder t) ++ (toPreOrder ts)
   where toOrder :: Tree Vertex -> [Vertex]
         toOrder (Node n ts) = n:(join(map toOrder ts))

reachable :: Graph -> Vertex -> Vertex -> Bool
reachable g a b = b `elem` (toPreOrder ( dfs g [a]))

testReachable = reachable aGraph 'a' 'c'

{-- 
    A monad of non-determinism.
    
    Like many other algebraic types, lists form a monad

    The instance declaration for monad [] is part of the prelude.
    return :: a -> [a]
    (>>=) :: [a] -> (a -> [b]) -> [b]

    instance Monad [ ] where
      return x = [x]
      []     >>= f = [] 
      (x:xs) >>= f = f x ++ (xs >>= f)

    
    With lazy evaluation, we can view this monad as a representation of
    non-deterministic computations.  Each item in the list corresponds 
    to one possible outcome. 
--}

orelse = (++)
bad = []

multiplyTo :: Int -> [(Int,Int)]
multiplyTo n = do {
  x <- [1..n];
  y <- [x..n];
  if (x * y == n) then return (x,y) else bad }

fstMult = head (multiplyTo 10)
sndMult = head (tail (multiplyTo 10))


type Row = Int 
type Col = Int 
type QPos = (Row,Col) 
type Board = [QPos] 

safe :: QPos -> QPos -> Bool 
safe (r,c) (r',c') = r /= r' && c /= c' && (abs(r-r') /= abs(c-c')) 

pick :: Int -> [Int] 
pick 0 = bad 
pick n = return n `orelse` pick (n-1) 

add :: QPos -> Board -> [Board] 
add q qs | all (safe q) qs = return (q:qs) 
         | otherwise = bad 

nqueens :: Int -> [Board] 
nqueens n = fill_row 1 [] 
 where 
  fill_row r board | r > n = 
                      return board 
                   | otherwise = 
                      do { c <- pick n;
                           board' <- add (r,c) board; 
                           fill_row (r+1) board';     }

queenResult = head (nqueens 8)
-- [(8,5),(7,7),(6,2),(5,6),(4,3),(3,1),(2,4),(1,8)]

data BTree a = Leaf a | BNode (BTree a) (BTree a)
     deriving (Show)
instance Monad BTree where
  return x = Leaf x
  Leaf x >>= f = f x
  BNode tl tr >>= f = BNode (tl>>=f) (tr>>=f)

aTree = BNode (Leaf 3) (BNode (Leaf 4) (Leaf 5))
anf :: Int -> BTree Int
anf i = Leaf (i+1)

testTreeMonad = aTree >>= anf

-- sequence' because sequence is in Prelude
sequence' [] = return []
sequence' (m:ms) = do{ a <- m; 
                       as<-sequence' ms; 
                       return (a:as) }

testSequence = sequence' [getChar, getChar, getChar]

putGetChar c = putChar c >> getChar

mapM' f as = sequence(map f as)
mapMUse = mapM' putGetChar "abc"