module Main where

import qualified Data.HashMap.Strict as Map
import           Data.Tree
import           System.Environment

main :: IO ()
main = do args <- getArgs
          case args of
            ["count", m, n] -> print $ count (read m) (read n)
            ["full" , m, n] -> putStrLn $ drawTree $ fmap (\x -> show x) $ countFull (read m) (read n)
            ["test"]        -> test

test :: IO ()
test = do print $ count 2 2
          print $ count 2 3
          print $ count 3 2
          print $ count 3 3
          print $ count 2 4
          print $ count 4 2
          print $ count 4 4

{- Memoizing approach; no sub-tree will ever be calculated twice -}
count :: Int -> Int -> Integer
count m' n' = fst $ cnt m' n' Map.empty
    where cnt 1 _ s = (1, s)
          cnt _ 1 s = (1, s)
          cnt m n s | m < n     = lkp m n s
                    | otherwise = lkp n m s

          lkp m n s = case Map.lookup (m, n) s of
                        Just  v -> (v, s)
                        Nothing -> let (v1, s1) = cnt m (n - 1) s
                                       (v2, s2) = cnt (m - 1) n s1
                                       v = v1 + v2
                                    in (v, Map.insert (m, n) v s2)

{- Simple mathematical approach with not optimization -}
countSlow :: Int -> Int -> Int
countSlow 1 _ = 1
countSlow _ 1 = 1
countSlow m n = (countSlow m (n - 1)) + (countSlow (m - 1) n)

{- Show full tree representation -}
countFull :: Int -> Int -> Tree (Int, Int)
countFull 1 n = Node (1, n) []
countFull m 1 = Node (m, 1) []
countFull m n = Node (m, n) [countFull m (n - 1), countFull (m - 1) n]
