{-# LANGUAGE MultiParamTypeClasses, FlexibleInstances, Safe #-}

{-
  This module is part of Chatty.
  Copyleft (c) 2014 Marvin Cohrs

  All wrongs reversed. Sharing is an act of love, not crime.
  Please share Chatty with everyone you like.

  Chatty is free software: you can redistribute it and/or modify
  it under the terms of the GNU Affero General Public License as published by
  the Free Software Foundation, either version 3 of the License, or
  (at your option) any later version.

  Chatty is distributed in the hope that it will be useful,
  but WITHOUT ANY WARRANTY; without even the implied warranty of
  MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
  GNU Affero General Public License for more details.

  You should have received a copy of the GNU Affero General Public License
  along with Chatty. If not, see <http://www.gnu.org/licenses/>.
-}

-- | Provides an AVL tree.
module Data.Chatty.AVL (avlMax,avlMin,avlLookup,avlHeight,avlSize,avlInsert,avlRemove,AVL (EmptyAVL,AVL),avlRoot,avlPreorder,avlPostorder,avlInorder) where

import Data.Maybe
import Data.Chatty.BST
import Data.Chatty.None

instance Indexable i o v => AnyBST AVL i o v where
  anyBstMax = avlMax
  anyBstMin = avlMin
  anyBstLookup = avlLookup
  anyBstEmpty = EmptyAVL
  anyBstInsert = avlInsert
  anyBstRemove = avlRemove
  anyBstHead = avlHead
  anyBstInorder = avlInorder

instance None (AVL a) where
  none = EmptyAVL

-- | An AVL tree.
data AVL a = EmptyAVL | AVL a Int Int !(AVL a) !(AVL a)

-- | Get the greatest element.
avlMax :: AVL i -> Maybe i
avlMax EmptyAVL = Nothing
avlMax (AVL a _ _ _ EmptyAVL) = Just a
avlMax (AVL _ _ _ _ r) = avlMax r

-- | Get the least element.
avlMin :: AVL i -> Maybe i
avlMin EmptyAVL = Nothing
avlMin (AVL a _ _ EmptyAVL _) = Just a
avlMin (AVL _ _ _ l _) = avlMin l

-- | Lookup a given key.
avlLookup :: Indexable i o v => o -> AVL i -> Maybe v
avlLookup _ EmptyAVL = Nothing
avlLookup o (AVL a _ _ l r)
  | o == indexOf a = Just $ valueOf a
  | o < indexOf a = avlLookup o l
  | o > indexOf a = avlLookup o r

-- | Lookup if a given key is contained
avlContains :: Indexable i o v => o -> AVL i -> Bool
avlContains o = isJust . avlLookup o

-- | Get the height of the tree.
avlHeight :: AVL i -> Int
avlHeight EmptyAVL = 0
avlHeight (AVL _ _ h _ _) = h

-- | Get the size of the tree.
avlSize :: AVL i -> Int
avlSize EmptyAVL = 0
avlSize (AVL _ s _ _ _) = s

avlBalance :: AVL i -> AVL i
avlBalance EmptyAVL = EmptyAVL
avlBalance t@(AVL a _ _ l r)
  | abs (avlHeight l - avlHeight r) < 2 = t
  | avlHeight l < avlHeight r = case r of
    AVL a1 _ _ l1 r1 ->
      let child = AVL a (findSize l l1) (findHeight l l1) l l1
      in AVL a1 (findSize child r1) (findHeight child r1) child r1
  | otherwise = case l of
    AVL a1 _ _ l1 r1 ->
      let child = AVL a (findSize r1 r) (findHeight r1 r) r1 r
      in AVL a1 (findSize l1 child) (findHeight l1 child) l1 child

findSize :: AVL i -> AVL i -> Int
findSize a b = 1 + avlSize a + avlSize b

findHeight :: AVL i -> AVL i -> Int
findHeight a b = 1 + max (avlHeight a) (avlHeight b)

-- | Insert into the tree.
avlInsert :: Indexable i o v => i -> AVL i -> AVL i
avlInsert a EmptyAVL = AVL a 1 1 EmptyAVL EmptyAVL
avlInsert a (AVL a1 s h l r)
  | indexOf a == indexOf a1 = AVL a s h l r
  | indexOf a < indexOf a1 =
    let l' = avlInsert a l
    in avlBalance $ AVL a1 (s+1) (findHeight l' r) l' r
  | otherwise =
    let r' = avlInsert a r
    in avlBalance $ AVL a1 (s+1) (findHeight l r') l r' 

-- | Remove from the tree.
avlRemove :: Indexable i o v => o -> AVL i -> AVL i
avlRemove _ EmptyAVL = EmptyAVL
avlRemove o t@(AVL a _ _ EmptyAVL EmptyAVL)
  | indexOf a == o = EmptyAVL
  | otherwise = t
avlRemove o t@(AVL a _ _ l r)
  | indexOf a == o =
    case t of
      AVL _ _ _ EmptyAVL _ -> case getLeft r of
        (Just a',r') -> avlBalance $ AVL a' (findSize EmptyAVL r') (findHeight EmptyAVL r') EmptyAVL r'
      _ -> case getRight l of
        (Just a',l') -> avlBalance $ AVL a' (findSize l' r) (findHeight l' r) l' r
  | o < indexOf a =
    let l' = avlRemove o l
    in avlBalance $ AVL a (findSize l' r) (findHeight l' r) l' r
  | otherwise =
    let r' = avlRemove o r
    in avlBalance $ AVL a (findSize l r') (findHeight l r') l r'

getLeft :: AVL i -> (Maybe i,AVL i)
getLeft EmptyAVL = (Nothing,EmptyAVL)
getLeft (AVL a _ _ EmptyAVL EmptyAVL) = (Just a,EmptyAVL)
getLeft (AVL a _ _ EmptyAVL r) = (Just a,r)
getLeft (AVL a _ _ l r) =
  case getLeft l of
    (p, t2) -> (p, AVL a (findSize r t2) (findHeight r t2) t2 r)

getRight :: AVL i -> (Maybe i,AVL i)
getRight EmptyAVL = (Nothing,EmptyAVL)
getRight (AVL a _ _ EmptyAVL EmptyAVL) = (Just a,EmptyAVL)
getRight (AVL a _ _ l EmptyAVL) = (Just a,l)
getRight (AVL a _ _ l r) =
  case getRight r of
    (p, t2) -> (p, AVL a (findSize l t2) (findHeight l t2) l t2)

instance Functor AVL where
  fmap _ EmptyAVL = EmptyAVL
  fmap f (AVL a s h l r) = AVL (f a) s h (fmap f l) (fmap f r)

-- | Get the root of the tree.
avlRoot :: AVL i -> i
avlRoot EmptyAVL = error "Trying to get the root of an empty AVL tree."
avlRoot (AVL a _ _ _ _) = a

-- | Get the root of the tree (safely)
avlHead :: AVL i -> Maybe i
avlHead EmptyAVL = Nothing
avlHead t = Just $ avlRoot t

-- | Traverse the tree, order (head, left, right)
avlPreorder :: AVL i -> [i]
avlPreorder EmptyAVL = []
avlPreorder (AVL a _ _ l r) = a : avlPreorder l ++ avlPreorder r

-- | Traverse the tree, order (left, right, head)
avlPostorder :: AVL i -> [i]
avlPostorder EmptyAVL = []
avlPostorder (AVL a _ _ l r) = avlPostorder l ++ avlPostorder r ++ [a]

-- | Traverse the tree, order (left, head, right)
avlInorder :: AVL i -> [i]
avlInorder EmptyAVL = []
avlInorder (AVL a _ _ l r) = avlInorder l ++ [a] ++ avlInorder r