Tree.hs

module Main where

import Control.Monad (forM_, mapM_)
import Control.Applicative ((<$>))
import Data.Bifunctor (bimap, first, second)
import Data.Array.MArray (newArray, writeArray)
import Data.Array.ST (runSTArray)
import Data.Array (Array, bounds, (!))
import Text.Read (readMaybe)


type Point = (Int, Int)
type Tree = [Point]
type CordSum = (Int -> Int -> Int)
type Canvas = Array (Int, Int) Bool


makePoint :: Int -> Int -> Point
makePoint = (,)


line :: CordSum -> CordSum -> Point -> Int -> Tree
line fX fY p height
    | height > 0 = [bimap (fX h) (fY h) p | h <- [0..(height -1)]]
    | otherwise = []


-- draw vertical line
verticalLine :: Point -> Int -> Tree
verticalLine = line (flip const) (+) 


-- draw diagonal line to right
diagonalLineR :: Point -> Int -> Tree
diagonalLineR = line (+) (+)


-- draw diagonal line to left
diagonalLineL :: Point -> Int -> Tree
diagonalLineL = line subtract (+)


-- draw subtree
subtree :: Point -> Int -> Tree
subtree p height = verticalLine p height ++
                   diagonalLineL pl height ++
                   diagonalLineR pr height
  where
    pl = bimap (subtract 1) (+height) p
    pr = bimap (+1) (+height) p


-- calc cords of next subtree
subtreeNext :: Point -> Int -> (Point, Point)
subtreeNext p h = (pel, per)
  where
    next_h = 2 * h
    pel = bimap (subtract h) (+ next_h) p
    per = bimap (+ h) (+ next_h) p


tree :: Point -> Int -> Int -> Tree
tree _ _ 0 = []
tree _ 0 _ = []
tree start height splits = subtree start height ++ left_tree ++ right_tree
  where
    height' = height `div` 2
    splits' = splits - 1
    (pl, pr) = subtreeNext start height
    left_tree = tree pl height' splits'
    right_tree = tree pr height' splits'


toCanvas :: Int -> Int -> Tree -> Maybe Canvas
toCanvas width height tree
    | width > 0 && height > 0 = Just canvasArr
    | otherwise = Nothing
  where
    pointToIndex = id
    canvasArr = runSTArray $ do
        arr <- newArray ((0, 0), (width - 1, height - 1)) False
        forM_ tree $ \p -> writeArray arr (pointToIndex p) True
        return arr


canvasToStrings :: Char -> Char -> Canvas -> [String]
canvasToStrings f t can = strLine <$> yCords
  where
    (_, (maxX, maxY)) = bounds can
    xCords = enumFromThenTo maxX (maxX - 1) 0
    yCords = enumFromThenTo maxY (maxY - 1) 0
    toChar False = f
    toChar True = t
    strLine y = (\x -> toChar $ can ! (x, y)) <$> xCords


drawCanvas :: Canvas -> IO ()
drawCanvas can = mapM_ putStrLn $ canvasToStrings '_' '1' can


main :: IO ()
main = do
    mSize <- readMaybe <$> getLine
    case mSize of
        Just sp -> case toCanvas 100 63 (tree (makePoint 50 0) 16 sp) of
            Just canv -> drawCanvas canv
            Nothing   -> putStrLn "N is to small"
        Nothing -> putStrLn "Please type integer"