import Data.List

main :: IO()
main = do
  print $ totalBalance (1, "Ivan", "Varna") db
  print $ getCriticalBalance db fromVarna         10 -- → [(3,5.0)] 
  print $ getCriticalBalance db (not . fromVarna) 15 -- → [(2,4.0),(4,12.0)]
  print $ studAvg [ (1, 100), (1, 50), (2, 100)
                  , (2, 93), (1, 39), (2, 87)
                  , (1, 89), (1, 87), (1, 90)
                  , (2, 100), (2, 76)] -- → [(1, 83.2), (2, 91.2)]
  print $ studAvg [ (3, 55), (2, 50), (1, 21)
                  , (3, 53), (2, 48), (1, 3)
                  , (3, 4), (2, 28), (1, 10)
                  , (3, 80), (2, 68), (1, 15)
                  , (3, 91), (2, 45), (1, 49)] -- → [(1,19.6),(2,47.8),(3,56.6)]
  print $ contains bt2 0
  print $ contains bt2 10
  print $ pruneTree bt0 1
  print $ pruneTree bt1 1
  print $ pruneTree bt2 1
  print $ pruneTree bt3 1
  -- print $ generateNum t1 1 -- ➝ 3
  print $ generateNum t1 2 -- ➝ 27
  -- print $ generateNum t1 3 -- ➝ 4
  -- print $ generateNum t2 1 -- ➝ 1
  -- print $ generateNum t2 2 -- ➝ 0
  -- print $ generateNum t2 3 -- ➝ 6
  
db = ( [ (1, 1, 10), (2, 1, 11)
        , (3, 1, 12), (4, 2, 3)
        , (5, 2, 1), (6, 3, 2)
        , (7, 3, 3), (8, 4, 12) ]
      , [ (1, "Ivan", "Varna")
        , (2, "Petar", "Burgas")
        , (3, "Georgi", "Varna")
        , (4, "Yordan", "Plovdiv") ] )

fromVarna :: Person -> Bool
fromVarna (_, _, "Varna") = True
fromVarna _ = False

type PersonID = Int
type Name = String
type City = String
type AccountID = Int
type Balance = Double
type Person = (PersonID, Name, City)
type Account = (AccountID, PersonID, Balance)

totalBalance :: Person -> ([Account], [Person]) -> (Person, Balance)
totalBalance p@(id, _, _) (as, _) =
  (p, sum [b | (_, pid, b) <- as, pid == id])

getCriticalBalance :: ([Account], [Person]) -> (Person -> Bool) -> Balance -> [(PersonID, Balance)]
getCriticalBalance db@(as, ps) p criticalBalance =
  [(id, balance) | ((id, _, _), balance) <- personTotalBalance,
                   balance < criticalBalance]
  where
    personTotalBalance = [totalBalance person db | person <- ps, p person]


studAvg :: [(Int, Double)] -> [(Int, Double)]
studAvg grades = [(id, getTop5Average id) | id <- ids]
  where
    ids = sort (nub (map fst grades))
    average xs = sum xs / (fromIntegral (length xs))
    getTop5Average id =
      (average . (take 5) . reverse . sort)
        [grade | (id1, grade) <- grades, id1 == id]

data BTree = Empty | Node Int BTree BTree deriving (Eq, Show)

bt0 = Node 6 Empty Empty
bt1 = Node 1 Empty
             (Node 0 (Node 0 Empty Empty)
                     (Node 1 Empty Empty))
bt2 = Node 1 (Node 0 (Node 0 Empty Empty)
                     (Node 0 Empty Empty))
             (Node 1 (Node 0 Empty Empty)
                     (Node 1 Empty Empty))
bt3 = Node 1 (Node 1 (Node 1 (Node 0 Empty Empty)
                             Empty)
                     (Node 1 Empty Empty))
             (Node 0 (Node 0 Empty Empty)
                     (Node 1 Empty Empty))

contains :: BTree -> Int -> Bool
contains Empty _ = False
contains (Node v lt rt) n =
  v == n || contains lt n || contains rt n

pruneTree :: BTree -> Int -> BTree
pruneTree Empty _ = Empty
pruneTree bt@(Node v lt rt) n =
  if not (contains bt n) then Empty
  else Node v (pruneTree lt n) (pruneTree rt n)

t1 :: BTree
t1 = Node 6 (Node 3 (Node 2 Empty Empty)
                    (Node 5 (Node 4 Empty Empty)
                            Empty))
            (Node 8 (Node 7 Empty Empty)
                    (Node 9 Empty Empty))
t2 :: BTree
t2 = Node 4 (Node 1 Empty
                    (Node 3 Empty Empty))
            (Node 5 Empty
                    (Node 7 (Node 6 Empty Empty)
                            Empty))

generateNum :: BTree -> Int -> Int
generateNum bt k = toNumber (getDigits bt k)
  where
    getDigits Empty _ = []
    getDigits (Node _ (Node d _ _) _) 1 = [d]
    getDigits (Node _ lt rt) k =
      getDigits lt (k - 1) ++ getDigits rt (k - 1)
    toNumber ds = helper 0 ds
      where
        helper res [] = res
        helper res (d:ds) = helper (res * 10 + d) ds