content: {-# LANGUAGE CPP #-}
  {-# LANGUAGE Rank2Types #-}
  {-# LANGUAGE TypeFamilies #-}
  {-# LANGUAGE FlexibleContexts #-}
  {-# LANGUAGE FlexibleInstances #-}
  {-# LANGUAGE MultiParamTypeClasses #-}
  {-# LANGUAGE UndecidableInstances #-}
  {-# LANGUAGE DeriveDataTypeable #-}
  {-# OPTIONS_GHC -fno-warn-name-shadowing #-}
  {-# OPTIONS_HADDOCK not-home #-}
  
  -----------------------------------------------------------------------------
  -- |
  -- Copyright   : (c) Edward Kmett 2010-2021
  -- License     : BSD3
  -- Maintainer  : ekmett@gmail.com
  -- Stability   : experimental
  -- Portability : GHC only
  --
  -----------------------------------------------------------------------------
  
  module Numeric.AD.Internal.Tower
    ( Tower(..)
    , zeroPad
    , zeroPadF
    , transposePadF
    , d
    , d'
    , withD
    , tangents
    , bundle
    , apply
    , getADTower
    , tower
    ) where
  
  import Prelude hiding (all, sum)
  import Control.Monad (join)
  import Data.Foldable
  import Data.Data (Data)
  import Data.Number.Erf
  import Numeric
  import Numeric.AD.Internal.Combinators
  import Numeric.AD.Jacobian
  import Numeric.AD.Mode
  
  -- | @Tower@ is an AD 'Mode' that calculates a tangent tower by forward AD, and provides fast 'diffsUU', 'diffsUF'
  newtype Tower a = Tower { getTower :: [a] } deriving Data
  
  instance Show a => Show (Tower a) where
    showsPrec n (Tower as) = showParen (n > 10) $ showString "Tower " . showList as
  
  -- Local combinators
  
  zeroPad :: Num a => [a] -> [a]
  zeroPad xs = xs ++ repeat 0
  {-# INLINE zeroPad #-}
  
  zeroPadF :: (Functor f, Num a) => [f a] -> [f a]
  zeroPadF fxs@(fx:_) = fxs ++ repeat (0 <$ fx)
  zeroPadF _ = error "zeroPadF :: empty list"
  {-# INLINE zeroPadF #-}
  
  transposePadF :: (Foldable f, Functor f) => a -> f [a] -> [f a]
  transposePadF pad fx
    | all null fx = []
    | otherwise = fmap headPad fx : transposePadF pad (drop1 <$> fx)
    where
      headPad [] = pad
      headPad (x:_) = x
      drop1 (_:xs) = xs
      drop1 xs = xs
  
  d :: Num a => [a] -> a
  d (_:da:_) = da
  d _ = 0
  {-# INLINE d #-}
  
  d' :: Num a => [a] -> (a, a)
  d' (a:da:_) = (a, da)
  d' (a:_)    = (a, 0)
  d' _        = (0, 0)
  {-# INLINE d' #-}
  
  tangents :: Tower a -> Tower a
  tangents (Tower []) = Tower []
  tangents (Tower (_:xs)) = Tower xs
  {-# INLINE tangents #-}
  
  truncated :: Tower a -> Bool
  truncated (Tower []) = True
  truncated _ = False
  {-# INLINE truncated #-}
  
  bundle :: a -> Tower a -> Tower a
  bundle a (Tower as) = Tower (a:as)
  {-# INLINE bundle #-}
  
  withD :: (a, a) -> Tower a
  withD (a, da) = Tower [a,da]
  {-# INLINE withD #-}
  
  apply :: Num a => (Tower a -> b) -> a -> b
  apply f a = f (Tower [a,1])
  {-# INLINE apply #-}
  
  getADTower :: Tower a -> [a]
  getADTower = getTower
  {-# INLINE getADTower #-}
  
  tower :: [a] -> Tower a
  tower = Tower
  
  primal :: Num a => Tower a -> a
  primal (Tower (x:_)) = x
  primal _ = 0
  
  instance Num a => Mode (Tower a) where
    type Scalar (Tower a) = a
    auto a = Tower [a]
    isKnownZero (Tower xs) = null xs
    asKnownConstant (Tower []) = Just 0
    asKnownConstant (Tower [a]) = Just a
    asKnownConstant Tower {} = Nothing
    isKnownConstant (Tower []) = True
    isKnownConstant (Tower [_]) = True
    isKnownConstant Tower {} = False
    zero = Tower []
    a *^ Tower bs = Tower (map (a*) bs)
    Tower as ^* b = Tower (map (*b) as)
    Tower as ^/ b = Tower (map (/b) as)
  
  infixr 6 <+>
  
  (<+>) :: Num a => Tower a -> Tower a -> Tower a
  Tower [] <+> bs = bs
  as <+> Tower [] = as
  Tower (a:as) <+> Tower (b:bs) = Tower (c:cs) where
    c = a + b
    Tower cs = Tower as <+> Tower bs
  
  instance Num a => Jacobian (Tower a) where
    type D (Tower a) = Tower a
    unary f dadb b = bundle (f (primal b)) (tangents b * dadb)
    lift1 f df b   = bundle (f (primal b)) (tangents b * df b)
    lift1_ f df b = a where
      a = bundle (f (primal b)) (tangents b * df a b)
  
    binary f dadb dadc b c = bundle (f (primal b) (primal c)) (tangents b * dadb + tangents c * dadc)
    lift2 f df b c = bundle (f (primal b) (primal c)) tana where
       (dadb, dadc) = df b c
       tanb = tangents b
       tanc = tangents c
       tana = case (truncated tanb, truncated tanc) of
         (False, False) -> tanb * dadb + tanc * dadc
         (True, False) -> tanc * dadc
         (False, True) -> tanb * dadb
         (True, True) -> zero
    lift2_ f df b c = a where
      a0 = f (primal b) (primal c)
      da = tangents b * dadb + tangents c * dadc
      a = bundle a0 da
      (dadb, dadc) = df a b c
  
  -- mul xs ys = [ sum [xs!!j * ys!!(k-j)*bin k j | j <- [0..k]] | k <- [0..] ]
  -- adapted for efficiency and to handle finite lists xs, ys
  mul:: Num a => Tower a -> Tower a -> Tower a
  mul (Tower []) _ = Tower []
  mul (Tower (a:as)) (Tower bs) = Tower (convs' [1] [a] as bs)
    where convs' _ _ _ [] = []
          convs' ps ars as bs = sumProd3 ps ars bs :
                case as of
                   [] -> convs'' (next' ps) ars bs
                   a:as -> convs' (next ps) (a:ars) as bs
          convs'' _ _ [] = undefined -- convs'' never called with last argument empty
          convs'' _ _ [_] = []
          convs'' ps ars (_:bs) = sumProd3 ps ars bs : convs'' (next' ps) ars bs
          next xs = 1 : zipWith (+) xs (tail xs) ++ [1] -- next row in Pascal's triangle
          next' xs = zipWith (+) xs (tail xs) ++ [1] -- end part of next row in Pascal's triangle
          sumProd3 as bs cs = sum (zipWith3 (\x y z -> x*y*z) as bs cs)
  
  #define HEAD (Tower a)
  #include <instances.h>
  