module AssocTree where

import Data.Tree
import Data.List
import qualified Data.Map as M
import Data.Maybe
import Data.Either


data Assoc = Assoc { label :: String
                   , assocId :: Int
                   , assocPid :: Int
                   }
                   deriving ( Show )

r = Assoc "root" 1 0
a = Assoc "A" 10 1
b = Assoc "B" 20 1
c = Assoc "C" 30 1
d = Assoc "D" 40 30
e = Assoc "e" 50 30
f = Assoc "f" 60 30
g = Assoc "g" 70 40
h = Assoc "h" 71 40
i = Assoc "i" 72 70
j = Assoc "j" 73 70
k = Assoc "k" 74 73
l = Assoc "l" 80 88
m = Assoc "m" 90 50


assocs :: [ Assoc ]
assocs = [ f, e, d, c, b, a, g, h, i, j, k, l, m ]

nullAssoc = Assoc "" (-1) (-1)

emptyTree = Node nullAssoc []

isEmptyTree ( Node ( Assoc "" ( -1 ) ( -1 ) ) [] ) = True
isEmptyTree _ = False

listToMap :: [ Assoc ] -> M.Map Int Assoc
listToMap xs =
  foldl insertAssocIntoMap M.empty xs
  where
    insertAssocIntoMap :: M.Map Int Assoc -> Assoc -> M.Map Int Assoc
    insertAssocIntoMap m a = M.insert ( assocId a ) a m

isDescendedFrom :: M.Map Int Assoc -> Assoc -> Assoc -> Bool
isDescendedFrom m parent@( Assoc _ pnid _ ) node@( Assoc _ _ npid ) =
  if npid == pnid 
    then True
    else
      case M.lookup npid m of
        Just a -> isDescendedFrom m parent a
        Nothing -> False

filterUnrelated :: [ Assoc ] -> [ Assoc ]
filterUnrelated xs =
  filter ( isDescendedFrom m r ) xs
  where
    m = listToMap xs

growTree :: Tree Assoc -> [ Assoc ] -> Tree Assoc
growTree t [] = t
growTree t ( x:xs ) =
  case insertAssoc t x of
    Left t' -> growTree t' xs
    Right _ -> growTree t $ xs ++ [ x ]


printTree :: Tree Assoc -> IO ()
printTree t = putStrLn $ drawTree $ fmap ( \x -> label x ) t


newTree :: a -> Tree a
newTree x = Node x []

insertNode :: Tree a -> ( a -> a -> Bool ) -> a -> Maybe ( Tree a )
insertNode t@( Node n [] ) p x 
  | p n x == True = Just $ Node n $ [ newTree x ]
  | otherwise = Nothing
  
insertNode t@( Node n kids@( k:ks ) ) p x
  | p n x == True = Just $ Node n $ kids ++ [ newTree x ]
  | otherwise =
      if inserted
        then Just $ Node n nkotb
        else Nothing
    where
      f = ( `insertNode` p )
      f' = ( `f` x )
      ms = map f' kids
      pairs = zip kids ms
      nkotb = map ( \(t', mt) -> if isNothing mt then t' else fromJust mt ) pairs
      inserted = any ( \m -> isJust m ) ms
 

         

insertAssoc :: Tree Assoc -> Assoc -> Either ( Tree Assoc ) Assoc
insertAssoc t a =
  case insertNode t p a of
    Just t' -> Left t'
    Nothing -> Right a
  where
    p :: Assoc -> Assoc -> Bool
    p n@(Assoc _ nid _) x@(Assoc _ _ pid) =
      if pid == 0
        then True
        else nid == pid