{-# OPTIONS_GHC -Wall #-}
{-# OPTIONS_GHC -fno-warn-incomplete-patterns #-}
{-# OPTIONS_GHC -fno-warn-unused-binds #-}
{-# OPTIONS_GHC -fno-warn-unused-imports #-}

{-# LANGUAGE CPP #-}
{-# LANGUAGE StandaloneDeriving #-}
{-# LANGUAGE MultiParamTypeClasses #-}
{-# LANGUAGE TypeFamilies #-}
{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE FlexibleInstances #-}
{-# LANGUAGE TypeSynonymInstances #-}
{-# LANGUAGE UndecidableInstances #-} -- Instances for Fix f
{-# LANGUAGE GADTs #-}
{-# LANGUAGE KindSignatures #-}
{-# LANGUAGE Rank2Types #-}

module Main (main) where

--------------------------------------------------------------------------------

-- base-4.0.0.0
import Prelude hiding (zipWith, mod, init, cycle, words)
import qualified Prelude as P (zipWith, words)
import Data.List (genericLength, foldl')
import Data.Maybe (fromJust)
import Control.Arrow (second)
import Control.Monad (liftM3)
import System.CPUTime (getCPUTime)
import System.IO (openFile, IOMode(..), hGetContents)
import Text.Printf (printf)
import System.Random (mkStdGen)

-- deepseq-1.1.0.0
import Control.DeepSeq (NFData(..))

-- criterion-0.4
import Criterion.Main

-- QuickCheck 2.1
import Test.QuickCheck.Arbitrary
import Test.QuickCheck.Gen

--------------------------------------------------------------------------------

data family F t :: * -> *

class (Functor (F t)) => InOut t where
  inF  :: F t t -> t
  outF :: t -> F t t

--------------------------------------------------------------------------------

type Alg f s = f s -> s
type AlgF t s = F t s -> s

fold :: (InOut t) => AlgF t s -> t -> s
fold alg = alg . fmap (fold alg) . outF

--------------------------------------------------------------------------------

data Tree a
  = Tip
  | Bin a (Tree a) (Tree a)
  deriving (Eq, Show)

sampleTree :: (Num a) => Tree a
sampleTree = Bin 2 (Bin 1 Tip Tip) (Bin 3 Tip (Bin 4 Tip Tip))

data instance F (Tree a) r
  = TipF
  | BinF a r r
  deriving (Eq, Show)

instance Functor (F (Tree a)) where
  fmap _ TipF            = TipF
  fmap f (BinF x rL rR)  = BinF x (f rL) (f rR)

instance InOut (Tree a) where
  inF TipF            = Tip
  inF (BinF x tL tR)  = Bin x tL tR
  outF Tip            = TipF
  outF (Bin x tL tR)  = BinF x tL tR

sizeAlg :: AlgF (Tree a) Int
sizeAlg  TipF            = 0
sizeAlg  (BinF _ sL sR)  = 1 + sL + sR

--------------------------------------------------------------------------------

type Set = Tree

empty :: Set a
empty = Tip

singleton :: a -> Set a
singleton x = Bin x empty empty

size :: Set a -> Int
size = foldSet 0 (\ _ sL sR -> 1 + sL + sR)

foldSet :: s -> (a -> s -> s -> s) -> Set a -> s
foldSet t b = fold alg
  where  alg  TipF            = t
         alg  (BinF x tL tR)  = b x tL tR

insert :: (Ord a) => a -> Set a -> Set a
insert x Tip            =  singleton x
insert x (Bin y tL tR)  =  case compare x y of
                             LT  -> balance  y  (insert x tL)  tR
                             GT  -> balance  y  tL             (insert x tR)
                             EQ  -> Bin      x  tL             tR

balance :: a -> Set a -> Set a -> Set a
balance x tL tR
  | sL + sR <= 1  = Bin x tL tR
  | sR >= 4 * sL  = rotateL x tL tR sRL sRR
  | sL >= 4 * sR  = rotateR x tL tR sLL sLR
  | otherwise     = Bin x tL tR
  where  (sL, sR)  = (size tL, size tR)
         Bin _ tRL tRR = tR
         (sRL, sRR)  = (size tRL, size tRR)
         Bin _ tLL tLR = tL
         (sLL, sLR)  = (size tLL, size tLR)

rotateL :: a -> Tree a -> Tree a -> Int -> Int -> Tree a
rotateL x tL tR sRL sRR
  | sRL < 2 * sRR  = singleL x tL tR
  | otherwise      = doubleL x tL tR
  where
    singleL x' tA (Bin y tB            tC) = Bin y (Bin x' tA tB) tC
    doubleL x' tA (Bin y (Bin z tB tC) tD) = Bin z (Bin x' tA tB) (Bin y tC tD)

rotateR :: a -> Tree a -> Tree a -> Int -> Int -> Tree a
rotateR x tL tR sLL sLR
  | sLR < 2 * sLL  = singleR x tL tR
  | otherwise      = doubleR x tL tR
  where
    singleR x' (Bin y tA tB           ) tC = Bin y tA            (Bin x' tB tC)
    doubleR x' (Bin y tA (Bin z tB tC)) tD = Bin z (Bin y tA tB) (Bin x' tC tD)

--------------------------------------------------------------------------------

data TreeI a
  = TipI Int
  | BinI Int a (TreeI a) (TreeI a)
  deriving (Eq, Show)

instance Functor (F (TreeI a)) where
  fmap _ (TipFI i)          = TipFI i
  fmap f (BinFI i x tL tR)  = BinFI i x (f tL) (f tR)

data instance F (TreeI a) r
  = TipFI Int
  | BinFI Int a r r

instance InOut (TreeI a) where
  inF (TipFI i)          = TipI i
  inF (BinFI i x tL tR)  = BinI i x tL tR
  outF (TipI i)          = TipFI i
  outF (BinI i x tL tR)  = BinFI i x tL tR

sizeI :: TreeI a -> Int
sizeI (TipI  i)        = i
sizeI (BinI  i _ _ _)  = i

--------------------------------------------------------------------------------

class TreeS t where
  type Elem t
  tip       :: t
  bin       :: Elem t -> t -> t -> t
  caseTree  :: r -> (Elem t -> t -> t -> r) -> t -> r

instance TreeS (Tree a) where
  type Elem (Tree a) = a
  tip             = Tip
  bin x tL tR     = Bin  x tL tR
  caseTree t b n  = case n of { Tip -> t ; Bin x tL tR -> b x tL tR }

instance TreeS (TreeI a) where
  type Elem (TreeI a) = a
  tip             = TipI  0
  bin x tL tR     = BinI  (1 + sizeI tL + sizeI tR) x tL tR
  caseTree t b n  = case n of { TipI _ -> t ; BinI _ x tL tR -> b x tL tR }

--------------------------------------------------------------------------------

emptyI :: TreeI a
emptyI = tip

singletonI :: a -> TreeI a
singletonI x = bin x emptyI emptyI

insertI :: (Ord a) => a -> TreeI a -> TreeI a
insertI x = caseTree  (singletonI x)
                      $ \ y tL tR ->  case compare x y of
                                            LT  -> balanceI    y  (insertI x tL)  tR
                                            GT  -> balanceI    y  tL              (insertI x tR)
                                            EQ  -> bin         x  tL              tR

balanceI :: a -> TreeI a -> TreeI a -> TreeI a
balanceI x tL tR
  | sL + sR <= 1  = bin x tL tR
  | sR >= 4 * sL  = rotateLI x tL tR
  | sL >= 4 * sR  = rotateRI x tL tR
  | otherwise     = bin x tL tR
  where
    sL   = sizeI tL
    sR   = sizeI tR
    rotateLI x' tL' tR'@(BinI _ _ tRL tRR)
      | sizeI tRL < 2 * sizeI tRR  = singleLI x' tL' tR'
      | otherwise                  = doubleLI x' tL' tR'
      where
        singleLI :: a -> TreeI a -> TreeI a -> TreeI a
        singleLI x'' tA (BinI _ y tB tC)               = bin y (bin x'' tA tB) tC
        doubleLI :: a -> TreeI a -> TreeI a -> TreeI a
        doubleLI x'' tA (BinI _ y (BinI _ z tB tC) tD) = bin z (bin x'' tA tB) (bin y tC tD)
    rotateRI x' tL'@(BinI _ _ tLL tLR) tR'
      | sizeI tLR < 2 * sizeI tLL    = singleRI x' tL' tR'
      | otherwise                    = doubleRI x' tL' tR'
      where
        singleRI :: a -> TreeI a -> TreeI a -> TreeI a
        singleRI x'' (BinI _ y tA tB) tC               = bin y tA (bin x'' tB tC)
        doubleRI :: a -> TreeI a -> TreeI a -> TreeI a
        doubleRI x'' (BinI _ y tA (BinI _ z tB tC)) tD = bin z (bin y tA tB) (bin x'' tC tD)

--------------------------------------------------------------------------------

class SizeX t where
  sizeX :: t -> Int

emptyX :: (TreeS t) => t
emptyX = tip

singletonX :: (TreeS t) => Elem t -> t
singletonX x = bin x emptyX emptyX

insertX :: (TreeS t, Ord (Elem t), SizeX t) => Elem t -> t -> t
insertX x = caseTree  (singletonX x)
                      $ \ y tL tR ->  case compare x y of
                                            LT  -> balanceX    y  (insertX x tL)  tR
                                            GT  -> balanceX    y  tL              (insertX x tR)
                                            EQ  -> bin         x  tL              tR

caseBin :: (TreeS t) => (Elem t -> t -> t -> b) -> t -> b
caseBin = caseTree undefined

balanceX :: (TreeS t, SizeX t) => Elem t -> t -> t -> t
balanceX x tL tR
  | sL + sR <= 1  = bin x tL tR
  | sR >= 4 * sL  = rotateLX x tL tR
  | sL >= 4 * sR  = rotateRX x tL tR
  | otherwise     = bin x tL tR
  where
    sL   = sizeX tL
    sR   = sizeX tR
    rotateLX x' tL' tR'
      | sizeX tRL < 2 * sizeX tRR  = singleLX x' tL' tR'
      | otherwise                  = doubleLX x' tL' tR'
      where
        (tRL, tRR) = caseBin (const (,)) tR'
        singleLX x'' tA = caseBin (\ y tB -> bin y (bin x'' tA tB))
        doubleLX x'' tA = caseBin $
          \ y tZ tD -> caseBin (\z tB tC -> bin z (bin x'' tA tB) (bin y tC tD)) tZ
    rotateRX x' tL' tR'
      | sizeX tLR < 2 * sizeX tLL  = singleRX x' tL' tR'
      | otherwise                  = doubleRX x' tL' tR'
      where
        (tLL, tLR) = caseBin (const (,)) tL'
        singleRX x'' tZ tC = caseBin (\ y tA tB -> bin y tA (bin x'' tB tC)) tZ
        doubleRX x'' tZ tD = caseBin (\ y tA tY -> caseBin (\ z tB tC -> bin z (bin y tA tB) (bin x'' tC tD)) tY ) tZ

--------------------------------------------------------------------------------

newtype Fix f = In { out :: f (Fix f) }

deriving instance (Eq (f (Fix f))) => Eq (Fix f)

-- Don't print the record syntax: it makes it difficult to read
instance (Show (f (Fix f))) => Show (Fix f) where
  showsPrec p (In x) = showParen (p >= 11) $ showString "In " . showsPrec 11 x

--------------------------------------------------------------------------------

data Ann s f r = Ann s (f r)
  deriving (Eq, Show)

--------------------------------------------------------------------------------

type FixA s f = Fix (Ann s f)

data instance F (FixA s f) r = InAF { annAF :: s, outAF :: (f r) }

instance (Functor f) => Functor (F (FixA s f)) where
  fmap f (InAF s x) = InAF s (fmap f x)

instance (Functor f) => InOut (FixA s f) where
  inF   (InAF s x)      = inA s x
  outF  (In (Ann s x))  = InAF s x

inA :: s -> f (FixA s f) -> FixA s f
inA s x = In (Ann s x)

outA :: FixA s f -> f (FixA s f)
outA (In (Ann _ x)) = x

ann :: FixA s f -> s
ann (In (Ann s _)) = s

foldMapA :: (Functor f) => (r -> s) -> FixA r f -> FixA s f
foldMapA f = fold (\ (InAF s x) -> inA (f s) x)

type TypeA s t = FixA s (F t)

--------------------------------------------------------------------------------

upwards :: (InOut t) => AlgF t s -> t -> TypeA s t
upwards = fold . pullUp

pullUp :: (Functor f) => Alg f s -> Alg f (FixA s f)
pullUp alg fs = inA (alg (fmap ann fs)) fs

ex_upwards :: TypeA Int (Tree Int)
ex_upwards = upwards sizeAlg sampleTree

--------------------------------------------------------------------------------

newtype Tree2 a = Tree2 { unTree2 :: TypeA Int (Tree a) }
  deriving Show

instance TreeS (Tree2 a) where
  type Elem (Tree2 a) = a
  tip             = Tree2 $ pullUp sizeAlg TipF
  bin x tL tR     = Tree2 $ pullUp sizeAlg (BinF x (unTree2 tL) (unTree2 tR))
  caseTree t b n  = case outA (unTree2 n) of { TipF -> t ; BinF x tL tR -> b x (Tree2 tL) (Tree2 tR) }

instance SizeX (Tree2 a) where
  sizeX = ann . unTree2

--------------------------------------------------------------------------------

data family Path t

data instance Path (Tree a)
  = PRoot
  | PBinL a (Path (Tree a))
  | PBinR a (Path (Tree a))
  deriving (Eq, Show)

data instance F (Path (Tree a)) r
  = PRootF
  | PBinLF a r
  | PBinRF a r

instance Functor (F (Path (Tree a))) where
  fmap _ PRootF        = PRootF
  fmap f (PBinLF x r)  = PBinLF x (f r)
  fmap f (PBinRF x r)  = PBinRF x (f r)

instance InOut (Path (Tree a)) where
  inF   PRootF        = PRoot
  inF   (PBinLF x t)  = PBinL x t
  inF   (PBinRF x t)  = PBinR x t
  outF  PRoot         = PRootF
  outF  (PBinL x t)   = PBinLF x t
  outF  (PBinR x t)   = PBinRF x t

class (InOut t, InOut (Path t), ZipWith (F t)) => Paths t where
  proot  :: Path t
  pnode  :: F t r -> F t (Path t -> Path t)

instance Paths (Tree t) where
  proot                =  PRoot
  pnode  TipF          =  TipF
  pnode  (BinF x _ _)  =  BinF x (PBinL x) (PBinR x)

paths0 :: Tree a -> TypeA (Path (Tree a)) (Tree a)
paths0 = pfrom PRoot
  where  pfrom p t = inA p $
           case t of
             Tip          -> TipF
             Bin x tL tR  -> BinF x (pfrom (PBinL x p) tL) (pfrom (PBinR x p) tR)

paths :: (Paths t) => t -> TypeA (Path t) t
paths = appA proot . fold (inA id . zipApp compA pnode)

appA :: (Functor f) => a -> FixA (a -> b) f -> FixA b f
appA x = foldMapA ($ x)

compA :: (Functor f) => (a -> b) -> FixA (b -> c) f -> FixA (a -> c) f
compA f = foldMapA (. f)

zipApp :: (ZipWith f) => (a -> b -> c) -> (f b -> f a) -> f b -> f c
zipApp f g x = zipWith f (g x) x

ex_paths :: TypeA (Path (Tree Int)) (Tree Int)
ex_paths = paths (Bin 2 Tip Tip)

--------------------------------------------------------------------------------

downwards :: (Paths t) => AlgF (Path t) s -> t -> TypeA s t
downwards alg = foldMapA (fold alg) . paths

depthAlgD :: AlgF (Path (Tree a)) Int
depthAlgD p = case p of { PRootF -> 1 ; PBinLF _ i -> succ i ; PBinRF _ i -> succ i }

ex_downwards :: TypeA Int (Tree Int)
ex_downwards = downwards depthAlgD sampleTree

--------------------------------------------------------------------------------

type AlgD f i = forall s . i -> f s -> f i
type AlgDF t i = AlgD (F t) i

pushDown :: (ZipWith f) => i -> AlgD f i -> Alg f (FixA i f)
pushDown init alg = inA init . zipApp push (alg init)
  where push i = pushDown i alg . outA

pushDownM :: (ZipWith f, Eq i) => i -> AlgD f i -> Alg f (FixA i f)
pushDownM init alg = inA init . zipApp push (alg init)
  where  push i x  | i == ann x  = x
                   | otherwise   = pushDownM i alg (outA x)

depthInit :: Int
depthInit = 1

depthAlg :: AlgDF (Tree a) Int
depthAlg i t = case t of { TipF -> TipF ; BinF x _ _ -> BinF x (succ i) (succ i) }

--------------------------------------------------------------------------------

depthTree :: Tree a -> Tree Int
depthTree = go 1
  where  go _ Tip            = Tip
         go d (Bin _ tL tR)  = let go' = go (succ d) in Bin d (go' tL) (go' tR)

--------------------------------------------------------------------------------

newtype Tree3 a = Tree3 { unTree3 :: TypeA Int (Tree a) }
  deriving Show

instance TreeS (Tree3 a) where
  type Elem (Tree3 a) = a
  tip             = Tree3 $ pushDown depthInit depthAlg TipF
  bin x tL tR     = Tree3 $ pushDown depthInit depthAlg (BinF x (unTree3 tL) (unTree3 tR))
  caseTree t b n  = case outA (unTree3 n) of { TipF -> t ; BinF x tL tR -> b x (Tree3 tL) (Tree3 tR) }

--------------------------------------------------------------------------------

newtype Tree3M a = Tree3M { unTree3M :: TypeA Int (Tree a) }
  deriving Show

instance TreeS (Tree3M a) where
  type Elem (Tree3M a) = a
  tip             = Tree3M $ pushDownM depthInit depthAlg TipF
  bin x tL tR     = Tree3M $ pushDownM depthInit depthAlg (BinF x (unTree3M tL) (unTree3M tR))
  caseTree t b n  = case outA (unTree3M n) of { TipF -> t ; BinF x tL tR -> b x (Tree3M tL) (Tree3M tR) }

--------------------------------------------------------------------------------

class ZipWith f where
  zipWith :: (a -> b -> c) -> f a -> f b -> f c

instance ZipWith (F (Tree a)) where
  zipWith _  TipF               TipF              = TipF
  zipWith f  (BinF x1 sL1 sR1)  (BinF _ sL2 sR2)  = BinF x1 (f sL1 sL2) (f sR1 sR2)
  zipWith _  _                  _                 = error "zipWith: mismatch!"

--------------------------------------------------------------------------------

inh :: FixA (i, s) f -> i
inh = fst . ann

syn :: FixA (i, s) f -> s
syn = snd . ann

--------------------------------------------------------------------------------

data family Context t

data instance Context (Tree a)
  = CRoot
  | CBinL a (Context (Tree a)) (Tree a)
  | CBinR a (Tree a) (Context (Tree a))
  deriving (Eq, Show)

data instance F (Context (Tree a)) r
  = CRootF
  | CBinLF a r (Tree a)
  | CBinRF a (Tree a) r

instance Functor (F (Context (Tree a))) where
  fmap _ CRootF            = CRootF
  fmap f (CBinLF x cL tR)  = CBinLF x (f cL) tR
  fmap f (CBinRF x tL cR)  = CBinRF x tL (f cR)

instance InOut (Context (Tree a)) where
  inF   CRootF            = CRoot
  inF   (CBinLF x cL tR)  = CBinL x cL tR
  inF   (CBinRF x tL cR)  = CBinR x tL cR
  outF  CRoot             = CRootF
  outF  (CBinL x cL tR)   = CBinLF x cL tR
  outF  (CBinR x tL cR)   = CBinRF x tL cR

class (InOut t, InOut (Context t), ZipWith (F t)) => Contexts t where
  croot :: Context t
  cnode :: F t t -> F t (Context t -> Context t)

instance Contexts (Tree a) where
  croot = CRoot
  cnode TipF = TipF
  cnode (BinF x tL tR) = BinF x (\c -> CBinL x c tR) (CBinR x tL)

contexts0 :: (Contexts t) => t -> TypeA (Context t) t
contexts0 = contextsFrom croot
    where  contextsFrom ctx = inA ctx . zipApp (\c -> contextsFrom (c ctx)) cnode . outF

contexts :: (Contexts t) => t -> TypeA (Context t) t
contexts = appA croot . fold (inA id . zipApp compA (cnode . fmap rmA))

rmA :: (InOut t) => TypeA s t -> t
rmA = fold (inF . outAF)

ex_contexts :: TypeA (Context (Tree Int)) (Tree Int)
ex_contexts = contexts sampleTree

--------------------------------------------------------------------------------

circular :: (Contexts t) => AlgF (Context t) (s -> i) -> AlgF t (i -> s) -> t -> TypeA (i, s) t
circular algD algU = foldMapA cycle . subtrees . contexts
  where  cycle (ct, st) = let (i, s) = (fold algD ct s, fold algU st i) in (i, s)

subtrees :: (InOut t) => TypeA c t -> TypeA (c, t) t
subtrees = fold (\ (InAF c x) -> inA (c, inF (fmap rmA x)) x)

--------------------------------------------------------------------------------

diff :: [Float] -> [Float]
diff xs = let avg = sum xs / genericLength xs in map (subtract avg) xs

newtype DiffI  = DI  { avgI :: Float }
  deriving (Eq, Show)

data    DiffS  = DS  { sumS :: Float, sizeS :: Float, diffS :: Float }
  deriving (Eq, Show)

diffAlgD :: AlgF (Context (Tree Float)) (DiffS -> DiffI)
diffAlgD CRootF           s   =  DI { avgI = sumS s / sizeS s }
diffAlgD (CBinLF x i tR)  sL  =  let  j = i $ dbinS x sL (fold diffAlgU tR j) j in j
diffAlgD (CBinRF x tL i)  sR  =  let  j = i $ dbinS x (fold diffAlgU tL j) sR j in j

diffAlgU :: AlgF (Tree Float) (DiffI -> DiffS)
diffAlgU TipF            _  =  DS { sumS = 0, sizeS = 0, diffS = 0 }
diffAlgU (BinF x sL sR)  i  =  dbinS x (sL i) (sR i) i

dbinS :: Float -> DiffS -> DiffS -> DiffI -> DiffS
dbinS x sL sR i = DS {  sumS   = x + sumS sL + sumS sR
                     ,  sizeS  = 1 + sizeS sL + sizeS sR
                     ,  diffS  = x - avgI i }

ex_circular :: TypeA (DiffI, DiffS) (Tree Float)
ex_circular = circular diffAlgD diffAlgU sampleTree

--------------------------------------------------------------------------------

type AlgC f i s = i -> f s -> (s, f i)
type AlgCF t i s = AlgC (F t) i s

passAround :: (Functor f, ZipWith f) => (s -> i) -> AlgC f i s -> Alg f (FixA (i, s) f)
passAround fun alg fis = inA (i, s) (zipWith pass fi fis)
  where  i        = fun s
         (s, fi)  = alg i (fmap syn fis)
         pass j   = passAround (const j) alg . outA

passAroundM :: (Functor f, ZipWith f, Eq i) => (s -> i) -> AlgC f i s -> Alg f (FixA (i, s) f)
passAroundM fun alg fx = inA (init, s) (zipWith pass fi fx)
  where  init     = fun s
         (s, fi)  = alg init (fmap syn fx)
         pass i x  | i == inh x  = x
                   | otherwise   = passAroundM (const i) alg (outA x)

--------------------------------------------------------------------------------

diffFunC :: DiffS -> DiffI
diffFunC s = DI { avgI = sumS s / sizeS s }

diffAlgC :: AlgCF (Tree Float) DiffI DiffS
diffAlgC _ TipF            = (DS { sumS = 0, sizeS = 0, diffS = 0 }, TipF)
diffAlgC i (BinF x sL sR)  = (dbinS x sL sR i, BinF x i i)

tipC :: TypeA (DiffI, DiffS) (Tree Float)
tipC = passAround diffFunC diffAlgC TipF
binC :: Float -> TypeA (DiffI, DiffS) (Tree Float) -> TypeA (DiffI, DiffS) (Tree Float) -> TypeA (DiffI, DiffS) (Tree Float)
binC x tL tR = passAround diffFunC diffAlgC (BinF x tL tR)

--------------------------------------------------------------------------------

diffList :: [Float] -> [Float]
diffList xs =
  let  nil _ = (0.0, 0.0, [])
       cons x fs avg =
         let  (s,l,ds) = fs avg
         in   (s+x,l+1.0,x-avg : ds)
       (sum',length',ds') = foldr cons nil xs (sum' / length')
  in   ds'

diffTree :: Tree Float -> Tree Float
diffTree t =
  let  alg TipF _ = (0, 0, Tip)
       alg (BinF x tL tR) avg =
         let  (sL, lL, tL') = tL avg
              (sR, lR, tR') = tR avg
         in   (x + sL + sR, 1 + lL + lR, Bin (x - avg) tL' tR')
       (sum', length', t') = fold alg t (sum' / length')
  in   t'

--------------------------------------------------------------------------------

newtype Tree4 = Tree4 { unTree4 :: TypeA (DiffI, DiffS) (Tree Float) }
  deriving Show

instance TreeS Tree4 where
  type Elem Tree4 = Float
  tip             = Tree4 $ passAround diffFunC diffAlgC TipF
  bin x tL tR     = Tree4 $ passAround diffFunC diffAlgC (BinF x (unTree4 tL) (unTree4 tR))
  caseTree t b n  = case outA (unTree4 n) of { TipF -> t ; BinF x tL tR -> b x (Tree4 tL) (Tree4 tR) }

--------------------------------------------------------------------------------

newtype Tree4M = Tree4M { unTree4M :: TypeA (DiffI, DiffS) (Tree Float) }
  deriving Show

instance TreeS Tree4M where
  type Elem Tree4M = Float
  tip             = Tree4M $ passAroundM diffFunC diffAlgC TipF
  bin x tL tR     = Tree4M $ passAroundM diffFunC diffAlgC (BinF x (unTree4M tL) (unTree4M tR))
  caseTree t b n  = case outA (unTree4M n) of { TipF -> t ; BinF x tL tR -> b x (Tree4M tL) (Tree4M tR) }

--------------------------------------------------------------------------------

class (Functor f) => Annot s f where
  data Ann_ s (f :: * -> *) :: * -> *
  attach :: s -> f r -> Ann_ s f r
  detach :: Ann_ s f r -> f r
  annot :: Ann_ s f r -> s

instance Annot Int (F (Tree a)) where
  data Ann_ Int (F (Tree a)) r = AnnInt Int (F (Tree a) r)
  attach = AnnInt
  detach (AnnInt _ x) = x
  annot (AnnInt s _) = s

type FixA' s f = Fix (Ann_ s f)

inA' :: (Annot s f) => s -> f (FixA' s f) -> FixA' s f
inA' s = In . attach s
outA' :: (Annot s f) => FixA' s f -> f (FixA' s f)
outA' = detach . out
ann'  :: (Annot s f) => FixA' s f -> s
ann' = annot . out

pullUp' :: (Annot s f) => Alg f s -> Alg f (FixA' s f)
pullUp' alg fx = inA' (alg (fmap ann' fx)) fx

--------------------------------------------------------------------------------

type Ctx t = F (Context t)

class (InOut t) => Zipper t where
  fill    :: Ctx t r -> r -> F t r
  first   :: (r -> Ctx t r -> a) -> F t r -> Maybe a
  next    :: (r -> Ctx t r -> a) -> Ctx t r -> r -> Maybe a

  fillS   :: (Zipper u) => Ctx t (TypeA (i, s) u) -> s -> F t s
  seekI   :: Ctx t s -> F t i -> i

data Locc :: * -> * -> * -> * where
  Loc ::  (Zipper t, ZipWith (F t), Eq i)
          =>  AlgCF t i s ->  (s -> i) ->  TypeA (i, s) t ->  [Ctx t (TypeA (i, s) t)] ->  Locc i s t

enter   ::  (Zipper t, ZipWith (F t), Eq i)
            => AlgCF t i s ->  (s -> i) ->  TypeA (i, s) t -> Locc i s t
enter alg fun foc = Loc alg fun foc []

leave   :: Locc i s t -> TypeA (i, s) t
leave (Loc _ _ foc []) = foc
leave loc              = leave (fromJust (up loc))

leaveM :: (Eq i, ZipWith (F t), Monad m) => Locc i s t -> m (TypeA (i, s) t)
leaveM = return . leave

up      :: Locc i s t -> Maybe (Locc i s t)
up (Loc _   _   _   []    )  = Nothing
up (Loc alg fun foc (c:cs))  = Just (Loc alg fun foc' cs)
  where
    fun' s | null cs    = fun s    --  top-level
           | otherwise  = inh foc  --  every other level
    foc' = passAround fun' alg (fill c foc)

down    :: Locc i s t -> Maybe (Locc i s t)
down (Loc alg fun foc cs) =
  first (\ foc' c -> Loc alg fun foc' (c:cs)) (outA foc)

right   :: Locc i s t -> Maybe (Locc i s t)
right (Loc _   _   _   []    ) = Nothing
right (Loc alg fun foc (c:cs)) = next (\ foc' c' -> Loc alg fun foc' (c':cs)) c foc

--  Works for given Diff algebra, but may not work with algebras that pass
--  values between siblings. I fmap over each context, so there is only an 'i'
--  coming from the top of the context, not the siblings. This seems to work for
--  Count, so I'm not sure how to show that it doesn't work.
update  :: (Functor (Ctx t)) => (TypeA (i, s) t -> TypeA (i, s) t) -> Locc i s t -> Locc i s t
update mod (Loc alg fun foc cs) = Loc alg fun foc2 cs1
  where
    foc1 = mod foc
    (s, i, cs1) = pushCtxs alg (syn foc1) (fun s) cs
    foc2 = passAround (const i) alg (outA foc1)

updateM :: (Monad m, Functor (Ctx t)) => (TypeA (i, s) t -> TypeA (i, s) t) -> Locc i s t -> m (Locc i s t)
updateM f = return . update f

--  See note for 'update'.
pushCtxs ::  (Functor (Ctx t), ZipWith (F t), Eq i, Zipper t)
             => AlgCF t i s -> s -> i -> [Ctx t (TypeA (i, s) t)] -> (s, i, [Ctx t (TypeA (i, s) t)])
pushCtxs _   s     i     []     = (s, i, [])
pushCtxs alg s_bot i_top (c:cs) = (s_top, i_bot, c':cs')
  where
    (s, i_bot) = digest alg c s_bot i
    c' = fmap (passAround (const i) alg . outA) c
    (s_top, i, cs') = pushCtxs alg s i_top cs

digest ::  (Zipper t, Zipper u) => AlgCF t i s -> Ctx t (TypeA (i, s) u) -> s -> i -> (s, i)
digest alg c s = second (seekI c) . flip alg (fillS c s)

on :: (TypeA (i, s) t -> TypeA (i, s) t) -> Locc i s t -> TypeA (i, s) t
on f (Loc _ _ foc _) = f foc

--------------------------------------------------------------------------------

fromList1 :: (Ord a) => [a] -> Tree a
fromList1 = foldl' (flip insert) empty

fromList2 :: (Ord a) => [a] -> TreeI a
fromList2 = foldl' (flip insertI) emptyI

fromList3 :: (Ord a) => [a] -> Tree2 a
fromList3 = foldl' (flip insertX) emptyX

--------------------------------------------------------------------------------

eqTrees :: (Eq a) => TreeI a -> Tree2 a -> Bool
eqTrees m n = go m (unTree2 n)
  where
    go (TipI  a)         (In (Ann b TipF))           = a == b
    go (BinI  a x tL tR) (In (Ann b (BinF y uL uR))) =
      a == b && x == y && go tL uL && go tR uR
    go _                 _                           = False

cmpTrees :: IO (Bool, Bool, Bool)
cmpTrees = do
  smallInput  <- toWords "inputs/dicte.txt"
  mediumInput <- toWords "inputs/acronym-dictionary.txt"
  largeInput  <- toWords "inputs/German.txt"
  let c x = fromList2 x `eqTrees` fromList3 x
  return (c smallInput, c mediumInput, c largeInput)

--------------------------------------------------------------------------------

instance (NFData a) => NFData (Tree a) where
  rnf Tip            = ()
  rnf (Bin x tL tR)  = rnf x `seq` rnf tL `seq` rnf tR

instance (NFData a) => NFData (TreeI a) where
  rnf (TipI i)          = rnf i
  rnf (BinI i x tL tR)  = rnf i `seq` rnf x `seq` rnf tL `seq` rnf tR

instance (NFData (f (Fix f))) => NFData (Fix f) where
  rnf (In x) = rnf x

instance (NFData s, NFData (f r)) => NFData (Ann s f r) where
  rnf (Ann s x) = rnf s `seq` rnf x

instance (NFData a, NFData r) => NFData (F (Tree a) r) where
  rnf TipF            = ()
  rnf (BinF x tL tR)  = rnf x `seq` rnf tL `seq` rnf tR

instance (NFData a) => NFData (Tree2 a) where
  rnf (Tree2 t) = rnf t

instance (NFData a) => NFData (Tree3 a) where
  rnf (Tree3 t) = rnf t

instance (NFData a) => NFData (Tree3M a) where
  rnf (Tree3M t) = rnf t

instance NFData DiffI where
  rnf (DI x) = rnf x

instance NFData DiffS where
  rnf (DS x y z) = rnf x `seq` rnf y `seq` rnf z

instance NFData Tree4 where
  rnf (Tree4 t) = rnf t

instance NFData Tree4M where
  rnf (Tree4M t) = rnf t

data NF a where
  NF1 :: (NFData b) => (a -> b) -> NF a
  NF2 :: (NFData b) => (a -> a -> b) -> NF (a -> a)

--------------------------------------------------------------------------------

arbitree :: (Arbitrary (Elem a), TreeS a) => Gen a
arbitree = sized f
    where f 0 = return tip
          f n = let g = f (n `div` 2) in
                frequency [(1, return tip), (3, liftM3 bin arbitrary g g)]

instance (Arbitrary a) => Arbitrary (Tree a) where arbitrary = arbitree
instance (Arbitrary a) => Arbitrary (TreeI a) where arbitrary = arbitree
instance (Arbitrary a) => Arbitrary (Tree2 a) where arbitrary = arbitree
instance (Arbitrary a) => Arbitrary (Tree3 a) where arbitrary = arbitree
instance (Arbitrary a) => Arbitrary (Tree3M a) where arbitrary = arbitree
instance Arbitrary Tree4 where arbitrary = arbitree
instance Arbitrary Tree4M where arbitrary = arbitree

-- Args: seed then size
gentree :: (Arbitrary a) => Int -> Int -> a
gentree = unGen arbitrary . mkStdGen

-- Args: seed then size
gentreefloat :: Int -> Int -> Tree Float
gentreefloat = gentree

genseeds :: [Int]
genseeds = [11, -34, 833]

gensizes :: [Int]
gensizes = [1000, 10000, 100000]

--------------------------------------------------------------------------------

toWords :: FilePath -> IO [String]
toWords file = do
  h <- openFile file ReadMode
  text <- hGetContents h
  let wds = P.words text
  writeFile "/dev/null" (show wds)
  return wds

-- Derived from Lennart Augustsson's timeit package
time :: (Show a) => String -> a -> IO ()
time name val = do
  putStr ("Time: " ++ name)
  t1 <- getCPUTime
  writeFile "/dev/null" (show val)
  t2 <- getCPUTime
  let t = fromIntegral (t2-t1) * 1e-12 :: Double
  printf " %6.3fs\n" t

runTimes :: (Show a) => (b, b, b) -> [Char] -> (b -> a) -> IO ()
runTimes (s, m, l) nm f = do
  putStrLn $ "\n" ++ nm ++ ":"
  time "small:   " $ f s
  time "medium:  " $ f m
  time "large:   " $ f l

--------------------------------------------------------------------------------

namedGroup :: NF a -> [(String, a)] -> [Benchmark]
namedGroup (NF1 f) = map (\(nm, arg) -> bench nm (nf f arg))

orderedGroup :: NF a -> [a] -> [Benchmark]
orderedGroup fut = namedGroup fut . P.zipWith (\n x -> ('n':show n, x)) [1 :: Int ..]

manyOrderedGroup :: [a] -> [(String, NF a)] -> [Benchmark]
manyOrderedGroup orders = map (\(nm, fut) -> bgroup nm (orderedGroup fut orders))

seededGroup :: String -> (NF (Int -> Int)) -> [Int] -> Benchmark
seededGroup nm (NF2 f) orders =
  let sg s = bgroup ("seed(" ++ show s ++ ")") (orderedGroup (NF1 (f s)) orders)
  in  bgroup nm $ P.zipWith ($) (repeat sg) genseeds

manySeededGroup :: [Int] -> [(String, NF (Int -> Int))] -> [Benchmark]
manySeededGroup sizes = map (\(nm, fut) -> seededGroup nm fut sizes)

fromListGroup :: (Ord a, NFData a) => [[a]] -> Benchmark
fromListGroup lists = bgroup "fromList" $
  manyOrderedGroup lists [("1", NF1 fromList1) , ("2", NF1 fromList2) , ("3", NF1 fromList3)]

specGroup :: [Int] -> Benchmark
specGroup sizes = bgroup "spec" $
  manySeededGroup sizes
    [ ("N", NF2 gentreefloat)
    , ("U", NF2 $ \i -> upwards sizeAlg . gentreefloat i)
    , ("D", NF2 $ \i -> downwards depthAlgD . gentreefloat i)
    , ("C", NF2 $ \i -> circular diffAlgD diffAlgU . gentreefloat i)
    ]

accGroup :: [Int] -> Benchmark
accGroup sizes = bgroup "acc" $
  manySeededGroup sizes
    [  ("depthTree", NF2 $ \i -> depthTree . gentreefloat i)
    , ("diffTree", NF2 $ \i -> diffTree . gentreefloat i)
    ]

incrGroup :: [Int] -> Benchmark
incrGroup sizes = bgroup "incr" $
  manySeededGroup sizes
    [ ("pullUp"      , NF2 $ \sd sz -> gentree sd sz :: Tree2 Float)
    , ("pushDown"    , NF2 $ \sd sz -> gentree sd sz :: Tree3 Float)
    , ("pushDownM"   , NF2 $ \sd sz -> gentree sd sz :: Tree3M Float)
    , ("passAround"  , NF2 $ \sd sz -> gentree sd sz :: Tree4)
    , ("passAroundM" , NF2 $ \sd sz -> gentree sd sz :: Tree4M)
    ]

perf :: IO ()
perf = do
  smallInput  <- toWords "inputs/dicte.txt"
  mediumInput <- toWords "inputs/acronym-dictionary.txt"
  largeInput  <- toWords "inputs/German.txt"
  let words = [smallInput, mediumInput, largeInput]

  defaultMain
    [ fromListGroup words
    , specGroup gensizes
    , incrGroup gensizes
    , accGroup gensizes
    ]

main :: IO ()
main = perf