import Control.Monad


type AssocList k v = [(k,v)]
testList :: AssocList Int Int
testList = [ (1,2), (2,4), (3,5), (4,1), (5,2), (6,4), (7,7) ]

lookup1 :: Eq k => AssocList k v -> k -> Maybe v
lookup1 [] key = Nothing
lookup1 ((key1,value):list) key = if (key1 == key)
  then Just value
  else lookup1 list key 

--lookup2plus :: Num k => AssocList k v -> k -> k -> Maybe v
--lookup2plus list key1 key2 = let
                              --val1 = lookup1 list key1
                              --cont Nothing = Nothing
                              --cont (Just v) = lookup1 list $ v+key2
                              --in cont val1

andThen :: Maybe a -> (a -> Maybe b) -> Maybe b
andThen Nothing _ =  Nothing
andThen (Just a) f = f a
-- v `andThen` f bud bude pokracovat v uspesnem vypoctu predanim vysledku
-- vypoctu v do funkce f, nebo selze (Nothing)

--lookup3plus :: (Num a, Eq a) => AssocList a a -> a -> a -> a -> Maybe a
lookup3plus list k1 k2 k3 =
  lookup1 list k1 `andThen`
  \r1 -> lookup1 list (r1+k2) `andThen`
  \r2 -> lookup1 list (r2+k3) `andThen`
  \r3 -> theResultIs $ r1 + r2 + r3

  --tosame pro 3 pomoci andThen



theResultIs = Just --hezky obalit finalni Just, at o Maybe funkce nemaji ani tuseni

lookup3plus' list k1 k2 k3 =
  lookup1 list k1 >>=
  \r1 -> lookup1 list (r1+k2) >>=
  \r2 -> lookup1 list (r2+k3) >>=
  \r3 -> return $ r1 + r2 + r3

lookup3plus'' list k1 k2 k3 = do
                  r1 <- lookup1 list k1
                  r2 <- lookup1 list (r1+k2)
                  r3 <- lookup1 list (r2+k3)
                  return $r1 + r2 + r3
                    --tosame pomoci do syntaxe


pythTriples = [ (a,b,c) | c<-[1..], b<-[1..c], a<-[1..b], a*a+b*b==c*c]
pythTriplesM = do
  c <- [1..]
  b <- [1..c]
  a <- [1..b]
  guard $ a*a+b*b==c*c
  return (a, b, c)


-- reprezentuje (slozenou) funkci, ktere nejak zmeni dodany stav
data StateM s val = StateChange (s -> (val,s))

instance Functor (StateM s) where
  fmap func (StateChange f) = StateChange (\s -> let
                                          (val,s2) = f s
                                          in (func val, s2))

instance Applicative (StateM s) where
  pure val = StateChange (\s -> (val, s))
  StateChange f <*> StateChange v = StateChange (\s -> let
    (val1,s1) = f s
    (val2,s2) = v s1
    in (val1 val2, s2))
  
instance Monad (StateM s) where
  return = pure 
  StateChange f1 >>= f2 = StateChange (\s -> let
                                        (a,s') = f1 s
                                        StateChange f3 = f2 a
                                        in f3 s')

getState :: StateM s s
getState = StateChange (\s -> (s,s))

setState :: s -> StateM s ()
setState s1 = StateChange (\s -> ((),s1))

updateState :: (s -> s) -> StateM s ()
updateState func = StateChange (\s -> ((),func s))

runStateM :: s -> StateM s a -> (a,s)
runStateM s (StateChange f) = f s 

testStateM = runStateM 0 $ do
    setState 5
    updateState (*3)
    x <- getState
    return $ x+1

main = do
  putStrLn "Ahoj"
  str <- getLine
  putStrLn str