{-# OPTIONS_GHC -fplugin-opt GHC.TypeLits.Normalise:allow-negated-numbers #-}

{-# LANGUAGE AllowAmbiguousTypes    #-}
{-# LANGUAGE FlexibleInstances      #-}
{-# LANGUAGE FunctionalDependencies #-}
{-# LANGUAGE GADTs                  #-}
{-# LANGUAGE MultiParamTypeClasses  #-}
{-# LANGUAGE PolyKinds              #-}
{-# LANGUAGE ScopedTypeVariables    #-}
{-# LANGUAGE TypeApplications       #-}
{-# LANGUAGE TypeFamilies           #-}
{-# LANGUAGE TypeInType             #-}
{-# LANGUAGE TypeOperators          #-}
{-# LANGUAGE UndecidableInstances   #-}
module HTree.Types where

import Data.Kind    (Type)
import GHC.TypeLits
import Utils        (CmpType, If, Max, Min)

data Tree (count :: Nat) (element :: Type) where
  Leaf :: Tree 0 element

  Node :: Tree m element
       -> element
       -> Tree n element
       -> Tree (1 + m + n) element

data HTree (structure :: Tree (count :: Nat) Type) where
  HLeaf :: HTree 'Leaf

  HNode :: HTree left -> centre -> HTree right
        -> HTree ('Node left centre right)

data Weight
  = LeftHeavy
  | Balanced
  | RightHeavy

type family Factor (left :: Nat) (right :: Nat) :: Weight where
  Factor 0    0     = Balanced
  Factor 1    0     = Balanced
  Factor 0    1     = Balanced
  Factor _    0     = LeftHeavy
  Factor 0    _     = RightHeavy
  Factor left right
    = Factor (left  - Min left right)
             (right - Min left right)

-- Balance by count, not height - maybe not the best idea for
-- performance?
class Balance
    (input  :: Tree (count :: Nat) Type)
    (output :: Tree  count         Type) where
  balance :: HTree input -> HTree output

instance {-# OVERLAPPING #-} Balance 'Leaf 'Leaf where
  balance = id

instance
    ( input  ~ 'Node left  centre  right
    , output ~ 'Node left' centre' right'
    , weight ~ Factor leftcount rightcount
    , BalanceLoop weight input output
    )
    => Balance
          ( 'Node (left   :: Tree leftcount  Type)
                  centre
                  (right  :: Tree rightcount Type)
          )
          ( 'Node (left'  :: Tree leftcount  Type)
                  centre'
                  (right' :: Tree rightcount Type)
          ) where
  balance = balanceLoop @_ @weight @input @output

class BalanceLoop
    (weight :: Weight)
    (input  :: Tree (count :: Nat) Type)
    (output :: Tree  count         Type) where
  balanceLoop :: HTree input -> HTree output

instance BalanceLoop 'Balanced noop noop where
  balanceLoop = id

instance
    ( Insert centre right right'
    , PopMax left centre' left'
    )
    => BalanceLoop LeftHeavy ('Node left  centre  right )
                             ('Node left' centre' right') where
  balanceLoop (HNode left centre right)
    = HNode left' centre' (insert centre right)
    where (centre', left') = popMax left

instance
    ( Insert centre left left'
    , PopMin right centre' right'
    )
    => BalanceLoop RightHeavy ('Node left  centre  right )
                              ('Node left' centre' right') where
  balanceLoop (HNode left centre right)
    = HNode (insert centre left) centre' right'
    where (centre', right') = popMin right

-- Still writing all these bits, don't look too closely...

class PopMax
    (input  :: Tree  count      Type)
    (max    ::                  Type)
    (output :: Tree (count - 1) Type) where
  popMax :: HTree input -> (max, HTree output)

instance {-# OVERLAPPING #-} PopMax ('Node left centre 'Leaf) centre left where
  popMax (HNode left centre HLeaf) = (centre, left)

instance PopMax right biggest right'
    => PopMax ('Node left centre right) biggest ('Node left centre right') where
  popMax (HNode left centre HLeaf)
    = (centre, left)

class PopMin
    (input  :: Tree  count      Type)
    (max    ::                  Type)
    (output :: Tree (count - 1) Type) where
  popMin :: HTree input -> (max, HTree output)

instance {-# OVERLAPPING #-} PopMin ('Node 'Leaf centre right) centre right where
  popMin (HNode HLeaf centre right) = (centre, right)

instance PopMin left smallest left'
    => PopMin ('Node left centre right) smallest ('Node left' centre right) where
  popMin (HNode left centre HLeaf)
    = (centre, left)

class Insert
    (element ::                  Type)
    (input   :: Tree  count      Type)
    (output  :: Tree (count + 1) Type) where
  insert :: element -> HTree input -> HTree output