{-# OPTIONS -Wall #-}

import Data.Maybe ( isJust )
import Data.List ( sort )

--------------------------------------------------------------------------------

data Tree = Leaf | Node Tree Tree
 deriving (Eq,Show)

--------------------------------------------------------------------------------

-- True iff a complete tree.
-- Given N nodes, time complexity is O(N).
complete1 :: Tree -> Bool
complete1 t = isJust (check t)
  where
  -- Computes height and if perfect or not.
  check :: Monad m => Tree -> m (Int,Bool)
  check Leaf = return (-1,True)
  check (Node l r) = do
    (lh,lp) <- check l
    (rh,rp) <- check r
    if lh == rh && lp || lh == rh+1 && rp
      then return (max lh rh + 1, lp && rp && lh == rh)
      else fail "check: not a complete tree"

--------------------------------------------------------------------------------

-- True iff a complete tree.
-- Given N nodes, time complexity is O(N).
complete2 :: Tree -> Bool
complete2 t = isJust (check t)
  where
  -- Computes the size of a complete tree. Fails if not a complete tree.
  check :: Monad m => Tree -> m Int
  check Leaf = return 0
  check (Node l r) = do
    ln <- check l
    rn <- check r
    let height n = if n==0
                     then -1
                     else floor (logBase 2 $ fromIntegral n :: Double)
        perfect n = let m = logBase 2 (fromIntegral n + 1) :: Double
                    in floor m == (ceiling m :: Int)
        --let perfect n = 2^(floor $ logBase 2 (fromIntegral n + 1)) == n+1
        lh = height ln :: Int
        rh = height rn
        lp = perfect ln
        rp = perfect rn
    if lh == rh && lp || lh == rh+1 && rp
      then return (ln + rn + 1)
      else fail "check: not a complete tree"

--------------------------------------------------------------------------------

-- True iff a complete tree.
-- Given N nodes, time complexity is O(N).
complete3 :: Tree -> Bool
complete3 Leaf = True
complete3 t@(Node _ _) =
  -- Produces heights (in order) of all nodes with a leaf child.
  -- First height must be same or one greater than last height, and the heights
  -- must be in decreasing order.
  let hs = leafDepths t 0    -- depths of all leaves
      rhs = reverse hs
      lmd = head hs          -- left-most leaf depth
      rmd = last hs          -- right-most leaf depth
  in (lmd == rmd || lmd == rmd+1)
       &&
     sort rhs == rhs         -- leaf depths in reverse sorted order

leafDepths :: Tree -> Int -> [Int]
leafDepths Leaf _ = error "leafDepths: invalid argument"
leafDepths (Node Leaf Leaf) h = [h,h]
leafDepths (Node Leaf r@(Node _ _)) h = h : leafDepths r (h+1)
leafDepths (Node l@(Node _ _) Leaf) h = leafDepths l (h+1) ++ [h]
leafDepths (Node l@(Node _ _) r@(Node _ _)) h =
  leafDepths l (h+1) ++ leafDepths r (h+1)

--------------------------------------------------------------------------------

-- Produces all binary trees of size n.
makeTrees :: Int -> [Tree]
makeTrees 0 = [Leaf]
makeTrees n = [ Node l r |
                i <- [n-1,n-2..0],
                l <- makeTrees i,
                r <- makeTrees (n-1-i)]

testsOK :: Bool
testsOK =
  -- all binary trees of a given size
  makeTrees 0 == [Leaf] &&
  makeTrees 1 == [Node Leaf Leaf] &&
  makeTrees 2 == [Node (Node Leaf Leaf) Leaf, Node Leaf (Node Leaf Leaf)] &&
  map (length . makeTrees) [3..5] == [5,14,42] &&

  -- there is exactly one complete binary tree of size n
  let oneCompleteTree n =
       let c1s = [ t | t <- makeTrees n, complete1 t ]
           c2s = [ t | t <- makeTrees n, complete2 t ]
           c3s = [ t | t <- makeTrees n, complete3 t ]
       in length c1s == 1 && c1s == c2s && c2s == c3s
  in all oneCompleteTree ([1..10])

main :: IO ()
main = print testsOK

--------------------------------------------------------------------------------