text:
{-# LANGUAGE AllowAmbiguousTypes, DataKinds, FlexibleInstances,
             FunctionalDependencies, QuantifiedConstraints, RankNTypes,
             TypeFamilies, TypeOperators #-}
{-# OPTIONS_GHC -fconstraint-solver-iterations=16 #-}
{-# OPTIONS_GHC -Wno-missing-methods #-}
{-# OPTIONS_GHC -fplugin GHC.TypeLits.KnownNat.Solver #-}
{-# OPTIONS_GHC -fplugin GHC.TypeLits.Normalise #-}
-- | Dual numbers and various operations on them, arithmetic and related
-- to tensors (vectors, matrices and others). This is the high-level API,
-- defined using the low-level API in "HordeAd.Core.DualClass".
module HordeAd.Core.DualNumber
  ( module HordeAd.Core.DualNumber
  , IsScalar, HasDelta, DMode(..)
  , Domain0, Domain1, Domain2, DomainX, Domains  -- an important re-export
  ) where

import Prelude

import qualified Data.Array.Convert
import qualified Data.Array.Dynamic as OTB
import qualified Data.Array.DynamicS as OT
import           Data.Array.Internal (valueOf)
import           Data.Array.Shape (DivRoundUp)
import qualified Data.Array.Shaped as OSB
import qualified Data.Array.ShapedS as OS
import           Data.List.Index (imap)
import           Data.MonoTraversable (MonoFunctor (omap))
import           Data.Proxy (Proxy (Proxy))
import qualified Data.Strict.Vector as Data.Vector
import qualified Data.Vector.Generic as V
import           GHC.TypeLits (KnownNat, type (+), type (-), type (<=))
import           Numeric.LinearAlgebra (Matrix, Numeric, Vector)
import qualified Numeric.LinearAlgebra as HM

import HordeAd.Core.DualClass
import HordeAd.Internal.Delta
  (CodeOut (..), Domain0, Domain1, Domain2, DomainX, Domains)

-- * The main dual number types

-- | Dual numbers with the second type argument being the primal component.
data DualNumber (d :: DMode) a = D a (Dual d a)

class (IsScalar d r, Monad m, Functor m, Applicative m)
      => DualMonad d r m | m -> d r where
  returnLet :: IsPrimalWithScalar d a r
            => DualNumber d a -> m (DualNumber d a)

addParameters :: (Numeric r, Num (Vector r))
              => Domains r -> Domains r -> Domains r
addParameters (a0, a1, a2, aX) (b0, b1, b2, bX) =
  (a0 + b0, V.zipWith (+) a1 b1, V.zipWith (+) a2 b2, V.zipWith (+) aX bX)

-- Dot product and sum respective ranks and sum it all.
dotParameters :: Numeric r => Domains r -> Domains r -> r
dotParameters (a0, a1, a2, aX) (b0, b1, b2, bX) =
  a0 HM.<.> b0
  + V.sum (V.zipWith (HM.<.>) a1 b1)
  + V.sum (V.zipWith (HM.<.>) (V.map HM.flatten a2) (V.map HM.flatten b2))
  + V.sum (V.zipWith (HM.<.>) (V.map OT.toVector aX) (V.map OT.toVector bX))

-- * General operations, for any tensor rank

-- These instances are required by the @Real@ instance, which is required
-- by @RealFloat@, which gives @atan2@. No idea what properties
-- @Real@ requires here, so let it crash if it's really needed.
instance Eq (DualNumber d a) where

instance Ord (DualNumber d a) where

-- These instances are dangerous due to possible subexpression copies
-- leading to allocation explosion. Expressions should be wrapped in
-- the monadic @returnLet@ whenever there is a possibility they can be
-- used multiple times in a larger expression.
instance (Num a, IsPrimal d a) => Num (DualNumber d a) where
  D u u' + D v v' = D (u + v) (dAdd u' v')
  D u u' - D v v' = D (u - v) (dAdd u' (dScale (-1) v'))
  D u u' * D v v' = D (u * v) (dAdd (dScale v u') (dScale u v'))
  negate (D v v') = D (negate v) (dScale (-1) v')
  abs (D v v') = D (abs v) (dScale (signum v) v')
  signum (D v _) = D (signum v) dZero
  fromInteger = constant . fromInteger

instance (Real a, IsPrimal d a) => Real (DualNumber d a) where
  toRational = undefined  -- TODO?

instance (Fractional a, IsPrimal d a) => Fractional (DualNumber d a) where
  D u u' / D v v' =
    let recipSq = recip (v * v)
    in D (u / v) (dAdd (dScale (v * recipSq) u') (dScale (- u * recipSq) v'))
  recip (D v v') =
    let minusRecipSq = - recip (v * v)
    in D (recip v) (dScale minusRecipSq v')
  fromRational = constant . fromRational

instance (Floating a, IsPrimal d a) => Floating (DualNumber d a) where
  pi = constant pi
  exp (D u u') = let expU = exp u
                 in D expU (dScale expU u')
  log (D u u') = D (log u) (dScale (recip u) u')
  sqrt (D u u') = let sqrtU = sqrt u
                  in D sqrtU (dScale (recip (sqrtU + sqrtU)) u')
  D u u' ** D v v' = D (u ** v) (dAdd (dScale (v * (u ** (v - 1))) u')
                                      (dScale ((u ** v) * log u) v'))
  logBase x y = log y / log x
  sin (D u u') = D (sin u) (dScale (cos u) u')
  cos (D u u') = D (cos u) (dScale (- (sin u)) u')
  tan (D u u') = let cosU = cos u
                 in D (tan u) (dScale (recip (cosU * cosU)) u')
  asin (D u u') = D (asin u) (dScale (recip (sqrt (1 - u*u))) u')
  acos (D u u') = D (acos u) (dScale (- recip (sqrt (1 - u*u))) u')
  atan (D u u') = D (atan u) (dScale (recip (1 + u*u)) u')
  sinh (D u u') = D (sinh u) (dScale (cosh u) u')
  cosh (D u u') = D (cosh u) (dScale (sinh u) u')
  tanh (D u u') = let y = tanh u
                  in D y (dScale (1 - y * y) u')
  asinh (D u u') = D (asinh u) (dScale (recip (sqrt (1 + u*u))) u')
  acosh (D u u') = D (acosh u) (dScale (recip (sqrt (u*u - 1))) u')
  atanh (D u u') = D (atanh u) (dScale (recip (1 - u*u)) u')

instance (RealFrac a, IsPrimal d a) => RealFrac (DualNumber d a) where
  properFraction = undefined
    -- very low priority, since these are all extremely not continuous

instance (RealFloat a, IsPrimal d a) => RealFloat (DualNumber d a) where
  atan2 (D u u') (D v v') =
    let t = 1 / (u * u + v * v)
    in D (atan2 u v) (dAdd (dScale (- u * t) v') (dScale (v * t) u'))
      -- we can be selective here and omit the other methods,
      -- most of which don't even have a differentiable codomain

constant :: IsPrimal d a => a -> DualNumber d a
constant a = D a dZero

scale :: (Num a, IsPrimal d a) => a -> DualNumber d a -> DualNumber d a
scale a (D u u') = D (a * u) (dScale a u')

tanhAct :: (DualMonad d r m, IsPrimalAndHasFeatures d a r)
        => DualNumber d a -> m (DualNumber d a)
tanhAct = returnLet . tanh

logistic :: (Floating a, IsPrimal d a) => DualNumber d a -> DualNumber d a
logistic (D u u') =
  let y = recip (1 + exp (- u))
  in D y (dScale (y * (1 - y)) u')

logisticAct :: (DualMonad d r m, IsPrimalAndHasFeatures d a r)
            => DualNumber d a -> m (DualNumber d a)
logisticAct = returnLet . logistic

-- Optimized and more clearly written @u ** 2@.
square :: (Num a, IsPrimal d a) => DualNumber d a -> DualNumber d a
square (D u u') = D (u * u) (dScale (2 * u) u')

squaredDifference :: (Num a, IsPrimal d a)
                  => a -> DualNumber d a -> DualNumber d a
squaredDifference targ res = square $ res - constant targ

lossSquared :: (DualMonad d r m, IsPrimalAndHasFeatures d a r)
            => a -> DualNumber d a -> m (DualNumber d a)
lossSquared targ res = returnLet $ squaredDifference targ res

reluAct :: (DualMonad d r m, IsPrimalAndHasFeatures d a r)
        => DualNumber d a -> m (DualNumber d a)
reluAct v@(D u _) = do
  let oneIfGtZero = omap (\x -> if x > 0 then 1 else 0) u
  returnLet $ scale oneIfGtZero v

reluLeakyAct :: (DualMonad d r m, IsPrimalAndHasFeatures d a r)
             => DualNumber d a -> m (DualNumber d a)
reluLeakyAct v@(D u _) = do
  let oneIfGtZero = omap (\x -> if x > 0 then 1 else 0.01) u
  returnLet $ scale oneIfGtZero v

-- * Operations resulting in a scalar

sumElements0 :: IsScalar d r => DualNumber d (Vector r) -> DualNumber d r
sumElements0 (D u u') = D (HM.sumElements u) (dSumElements0 u' (V.length u))

index0 :: IsScalar d r => DualNumber d (Vector r) -> Int -> DualNumber d r
index0 (D u u') ix = D (u V.! ix) (dIndex0 u' ix (V.length u))

minimum0 :: IsScalar d r => DualNumber d (Vector r) -> DualNumber d r
minimum0 (D u u') =
  D (HM.minElement u) (dIndex0 u' (HM.minIndex u) (V.length u))

maximum0 :: IsScalar d r => DualNumber d (Vector r) -> DualNumber d r
maximum0 (D u u') =
  D (HM.maxElement u) (dIndex0 u' (HM.maxIndex u) (V.length u))

-- If @v'@ is a @Var1@, this is much faster due to the optimization
-- in @Index0@.
foldl'0 :: IsScalar d r
        => (DualNumber d r -> DualNumber d r -> DualNumber d r)
        -> DualNumber d r -> DualNumber d (Vector r)
        -> DualNumber d r
foldl'0 f uu' (D v v') =
  let k = V.length v
      g !acc ix p = f (D p (dIndex0 v' ix k)) acc
  in V.ifoldl' g uu' v

altSumElements0 :: IsScalar d r => DualNumber d (Vector r) -> DualNumber d r
altSumElements0 = foldl'0 (+) 0

-- | Dot product.
infixr 8 <.>!
(<.>!) :: IsScalar d r
       => DualNumber d (Vector r) -> DualNumber d (Vector r) -> DualNumber d r
(<.>!) (D u u') (D v v') = D (u HM.<.> v) (dAdd (dDot0 v u') (dDot0 u v'))

-- | Dot product with a constant vector.
infixr 8 <.>!!
(<.>!!) :: IsScalar d r
        => DualNumber d (Vector r) -> Vector r -> DualNumber d r
(<.>!!) (D u u') v = D (u HM.<.> v) (dDot0 v u')

infixr 8 <.>$
(<.>$) :: (IsScalar d r, KnownNat n)
       => DualNumber d (OS.Array '[n] r) -> DualNumber d (OS.Array '[n] r)
       -> DualNumber d r
(<.>$) d e = fromS1 d <.>! fromS1 e

fromX0 :: IsScalar d r => DualNumber d (OT.Array r) -> DualNumber d r
fromX0 (D u u') = D (OT.unScalar u) (dFromX0 u')

fromS0 :: IsScalar d r => DualNumber d (OS.Array '[] r) -> DualNumber d r
fromS0 (D u u') = D (OS.unScalar u) (dFromS0 u')

sumElementsVectorOfDual
  :: IsScalar d r => Data.Vector.Vector (DualNumber d r) -> DualNumber d r
sumElementsVectorOfDual = V.foldl' (+) 0

softMaxAct :: DualMonad d r m
           => Data.Vector.Vector (DualNumber d r)
           -> m (Data.Vector.Vector (DualNumber d r))
softMaxAct us = do
  expUs <- V.mapM (returnLet . exp) us
  let sumExpUs = sumElementsVectorOfDual expUs
  -- This has to be let-bound, because it's used many times below.
  recipSum <- returnLet $ recip sumExpUs
  V.mapM (\r -> returnLet $ r * recipSum) expUs

-- In terms of hmatrix: @-(log res <.> targ)@.
lossCrossEntropy :: forall d r m. DualMonad d r m
                 => Vector r
                 -> Data.Vector.Vector (DualNumber d r)
                 -> m (DualNumber d r)
lossCrossEntropy targ res = do
  let f :: DualNumber d r -> Int -> DualNumber d r -> DualNumber d r
      f !acc i d = acc + scale (targ V.! i) (log d)
  returnLet $ negate $ V.ifoldl' f 0 res

-- In terms of hmatrix: @-(log res <.> targ)@.
lossCrossEntropyV :: DualMonad d r m
                  => Vector r
                  -> DualNumber d (Vector r)
                  -> m (DualNumber d r)
lossCrossEntropyV targ res = returnLet $ negate $ log res <.>!! targ

-- Note that this is equivalent to a composition of softMax and cross entropy
-- only when @target@ is one-hot. Otherwise, results vary wildly. In our
-- rendering of the MNIST data all labels are on-hot.
lossSoftMaxCrossEntropyV
  :: DualMonad d r m
  => Vector r -> DualNumber d (Vector r) -> m (DualNumber d r)
lossSoftMaxCrossEntropyV target (D u u') = do
  -- The following protects from underflows, overflows and exploding gradients
  -- and is required by the QuickCheck test in TestMnistCNN.
  -- See https://github.com/tensorflow/tensorflow/blob/5a566a7701381a5cf7f70fce397759483764e482/tensorflow/core/kernels/sparse_softmax_op.cc#L106
  -- and https://github.com/tensorflow/tensorflow/blob/5a566a7701381a5cf7f70fce397759483764e482/tensorflow/core/kernels/xent_op.h
  let expU = exp (u - HM.scalar (HM.maxElement u))
      sumExpU = HM.sumElements expU
      recipSum = recip sumExpU
-- not exposed: softMaxU = HM.scaleRecip sumExpU expU
      softMaxU = HM.scale recipSum expU
  returnLet $ D (negate $ log softMaxU HM.<.> target)  -- TODO: avoid: log . exp
                (dDot0 (softMaxU - target) u')

-- * Operations resulting in a vector

-- @1@ means rank one, so the dual component represents a vector.
seq1 :: IsScalar d r
     => Data.Vector.Vector (DualNumber d r) -> DualNumber d (Vector r)
seq1 v = D (V.convert $ V.map (\(D u _) -> u) v)  -- I hope this fuses
           (dSeq1 $ V.map (\(D _ u') -> u') v)

konst1 :: IsScalar d r => DualNumber d r -> Int -> DualNumber d (Vector r)
konst1 (D u u') n = D (HM.konst u n) (dKonst1 u' n)

append1 :: IsScalar d r
        => DualNumber d (Vector r) -> DualNumber d (Vector r)
        -> DualNumber d (Vector r)
append1 (D u u') (D v v') = D (u V.++ v) (dAppend1 u' (V.length u) v')

slice1 :: IsScalar d r
       => Int -> Int -> DualNumber d (Vector r) -> DualNumber d (Vector r)
slice1 i n (D u u') = D (V.slice i n u) (dSlice1 i n u' (V.length u))

sumRows1 :: IsScalar d r => DualNumber d (Matrix r) -> DualNumber d (Vector r)
sumRows1 (D u u') = D (V.fromList $ map HM.sumElements $ HM.toRows u)
                      (dSumRows1 u' (HM.cols u))

sumColumns1 :: IsScalar d r => DualNumber d (Matrix r) -> DualNumber d (Vector r)
sumColumns1 (D u u') = D (V.fromList $ map HM.sumElements $ HM.toColumns u)
                         (dSumColumns1 u' (HM.rows u))

-- If @v'@ is a @Var1@, this is much faster due to the optimization
-- in @Index0@. The detour through a boxed vector (list probably fuses away)
-- is costly, but only matters if @f@ is cheap.
map1 :: IsScalar d r
     => (DualNumber d r -> DualNumber d r) -> DualNumber d (Vector r)
     -> DualNumber d (Vector r)
map1 f (D v v') =
  let k = V.length v
      g ix p = f $ D p (dIndex0 v' ix k)
      ds = imap g $ V.toList v
  in seq1 $ V.fromList ds

-- | Dense matrix-vector product.
infixr 8 #>!
(#>!) :: IsScalar d r
      => DualNumber d (Matrix r) -> DualNumber d (Vector r)
      -> DualNumber d (Vector r)
(#>!) (D u u') (D v v') = D (u HM.#> v) (dAdd (dMD_V1 u' v) (dM_VD1 u v'))

-- | Dense matrix-vector product with a constant vector.
infixr 8 #>!!
(#>!!) :: IsScalar d r
       => DualNumber d (Matrix r) -> Vector r
       -> DualNumber d (Vector r)
(#>!!) (D u u') v = D (u HM.#> v) (dMD_V1 u' v)

fromX1 :: IsScalar d r => DualNumber d (OT.Array r) -> DualNumber d (Vector r)
fromX1 (D u u') = D (OT.toVector u) (dFromX1 u')

fromS1 :: forall len d r. (KnownNat len, IsScalar d r)
       => DualNumber d (OS.Array '[len] r) -> DualNumber d (Vector r)
fromS1 (D u u') = D (OS.toVector u) (dFromS1 u')

reverse1 :: IsScalar d r => DualNumber d (Vector r) -> DualNumber d (Vector r)
reverse1 (D u u') = D (V.reverse u) (dReverse1 u')

flatten1 :: IsScalar d r => DualNumber d (Matrix r) -> DualNumber d (Vector r)
flatten1 (D u u') = let (rows, cols) = HM.size u
                    in D (HM.flatten u) (dFlatten1 rows cols u')

flattenX1 :: IsScalar d r => DualNumber d (OT.Array r) -> DualNumber d (Vector r)
flattenX1 (D u u') = let sh = OT.shapeL u
                     in D (OT.toVector u) (dFlattenX1 sh u')

flattenS1 :: (IsScalar d r, OS.Shape sh)
          => DualNumber d (OS.Array sh r) -> DualNumber d (Vector r)
flattenS1 (D u u') = D (OS.toVector u) (dFlattenS1 u')

corr1 :: IsScalar d r
      => DualNumber d (Vector r) -> DualNumber d (Vector r)
      -> DualNumber d (Vector r)
corr1 ker@(D u _) vv@(D v _) = case (V.length u, V.length v) of
  (0, lenV) -> konst1 0 lenV
  (lenK, lenV) -> if lenK <= lenV
                  then vectorSlices2 lenK vv #>! ker
                  else error $ "corr1: len kernel " ++ show lenK
                               ++ " > len vector " ++ show lenV

-- This is not optimally implemented: @append1@ is costly compared
-- to a @mconcat@ counterpart and @z@ is used twice without
-- assigning it to a variable.
conv1 :: IsScalar d r
      => DualNumber d (Vector r) -> DualNumber d (Vector r)
      -> DualNumber d (Vector r)
conv1 ker@(D u _) vv@(D v _) =
  let lenK = V.length u
      lenV = V.length v
      kerRev = reverse1 ker
      z = konst1 0 (lenK - 1)
      vvPadded = append1 z $ append1 vv z
  in if lenK == 0
     then konst1 0 lenV
     else corr1 kerRev vvPadded

-- No padding; remaining areas ignored.
maxPool1 :: IsScalar d r
         => Int -> Int -> DualNumber d (Vector r) -> DualNumber d (Vector r)
maxPool1 ksize stride v@(D u _) =
  let slices = [slice1 i ksize v | i <- [0, stride .. V.length u - ksize]]
  in seq1 $ V.fromList $ map maximum0 slices

softMaxActV :: DualMonad d r m
            => DualNumber d (Vector r) -> m (DualNumber d (Vector r))
softMaxActV d@(D u _) = do
  expU <- returnLet $ exp d
  let sumExpU = sumElements0 expU
  -- This has to be let-bound, because it's used many times below.
  recipSum <- returnLet $ recip sumExpU
  returnLet $ konst1 recipSum (V.length u) * expU

-- Note that this is equivalent to a composition of softMax and cross entropy
-- only when @target@ is one-hot. Otherwise, results vary wildly. In our
-- rendering of the MNIST data all labels are one-hot.
lossSoftMaxCrossEntropyL
  :: DualMonad d r m
  => Matrix r
  -> DualNumber d (Matrix r)
  -> m (DualNumber d (Vector r))
lossSoftMaxCrossEntropyL target (D u u') = do
  let expU = exp (u - HM.scalar (HM.maxElement u))  -- vs exploding gradients
      sumExpU = V.fromList $ map HM.sumElements $ HM.toColumns expU
      recipSum = recip sumExpU
      softMaxU = HM.asRow recipSum * expU
                   -- this @asRow@ is safe; multiplied at once
      scaled = D (negate $ log softMaxU * target)
                 (dScale (softMaxU - target) u')
  returnLet $ sumColumns1 scaled

-- * Operations resulting in a matrix

-- @2@ means rank two, so the dual component represents a matrix.
fromRows2 :: IsScalar d r
          => Data.Vector.Vector (DualNumber d (Vector r))
          -> DualNumber d (Matrix r)
fromRows2 v = D (HM.fromRows $ map (\(D u _) -> u) $ V.toList v)
                (dFromRows2 $ V.map (\(D _ u') -> u') v)

fromColumns2 :: IsScalar d r
             => Data.Vector.Vector (DualNumber d (Vector r))
             -> DualNumber d (Matrix r)
fromColumns2 v = D (HM.fromRows $ map (\(D u _) -> u) $ V.toList v)
                   (dFromColumns2 $ V.map (\(D _ u') -> u') v)

konst2 :: IsScalar d r => DualNumber d r -> (Int, Int) -> DualNumber d (Matrix r)
konst2 (D u u') sz = D (HM.konst u sz) (dKonst2 u' sz)

transpose2 :: IsScalar d r => DualNumber d (Matrix r) -> DualNumber d (Matrix r)
transpose2 (D u u') = D (HM.tr' u) (dTranspose2 u')

-- | Dense matrix-matrix product.
--
-- If @u@ is a m x n (number of rows x number of columns) matrix
-- and @v@ is a n x p matrix then the result of @u <>! v@ is a m x p matrix.
infixr 8 <>!
(<>!) :: IsScalar d r
      => DualNumber d (Matrix r) -> DualNumber d (Matrix r)
      -> DualNumber d (Matrix r)
(<>!) (D u u') (D v v') = D (u HM.<> v) (dAdd (dMD_M2 u' v) (dM_MD2 u v'))

-- | Dense matrix-matrix product with a constant matrix.
infixr 8 <>!!
(<>!!) :: IsScalar d r
       => DualNumber d (Matrix r) -> Matrix r
       -> DualNumber d (Matrix r)
(<>!!) (D u u') v = D (u HM.<> v) (dMD_M2 u' v)

rowAppend2 :: IsScalar d r
           => DualNumber d (Matrix r) -> DualNumber d (Matrix r)
           -> DualNumber d (Matrix r)
rowAppend2 (D u u') (D v v') =
  D (u HM.=== v) (dRowAppend2 u' (HM.rows u) v')

columnAppend2 :: IsScalar d r
              => DualNumber d (Matrix r) -> DualNumber d (Matrix r)
              -> DualNumber d (Matrix r)
columnAppend2 (D u u') (D v v') =
  D (u HM.||| v) (dColumnAppend2 u' (HM.cols u) v')

rowSlice2 :: IsScalar d r
          => Int -> Int -> DualNumber d (Matrix r)
          -> DualNumber d (Matrix r)
rowSlice2 i n (D u u') = D (HM.subMatrix (i, 0) (n, HM.cols u) u)
                           (dRowSlice2 i n u' (HM.rows u))

columnSlice2 :: IsScalar d r
             => Int -> Int -> DualNumber d (Matrix r)
             -> DualNumber d (Matrix r)
columnSlice2 i n (D u u') = D (HM.subMatrix (0, i) (HM.rows u, n) u)
                              (dColumnSlice2 i n u' (HM.rows u))

asRow2 :: IsScalar d r
       => DualNumber d (Vector r) -> Int -> DualNumber d (Matrix r)
asRow2 (D u u') n = D (HM.fromRows $ replicate n u) (dAsRow2 u')

asColumn2 :: IsScalar d r
          => DualNumber d (Vector r) -> Int -> DualNumber d (Matrix r)
asColumn2 (D u u') n = D (HM.fromColumns $ replicate n u) (dAsColumn2 u')

fromX2 :: IsScalar d r => DualNumber d (OT.Array r) -> DualNumber d (Matrix r)
fromX2 (D u u') = case OT.shapeL u of
  [_, cols] -> D (HM.reshape cols $ OT.toVector u) (dFromX2 u')
  dims -> error $ "fromX2: the tensor has wrong dimensions " ++ show dims

fromS2 :: forall rows cols d r.
          (KnownNat rows, KnownNat cols, IsScalar d r)
       => DualNumber d (OS.Array '[rows, cols] r) -> DualNumber d (Matrix r)
fromS2 (D u u') = D (HM.reshape (valueOf @cols) $ OS.toVector u) (dFromS2 u')

flipud2 :: IsScalar d r => DualNumber d (Matrix r) -> DualNumber d (Matrix r)
flipud2 (D u u') = D (HM.flipud u) (dFlipud2 u')

fliprl2 :: IsScalar d r => DualNumber d (Matrix r) -> DualNumber d (Matrix r)
fliprl2 (D u u') = D (HM.fliprl u) (dFliprl2 u')

vectorSlices2 :: IsScalar d r
              => Int -> DualNumber d (Vector r) -> DualNumber d (Matrix r)
vectorSlices2 n vv@(D v _) =
  fromRows2 $ V.fromList [slice1 i n vv | i <- [0 .. V.length v - n]]

reshape2 :: IsScalar d r
         => Int -> DualNumber d (Vector r) -> DualNumber d (Matrix r)
reshape2 cols (D u u') = D (HM.reshape cols u) (dReshape2 cols u')

-- TODO: This has list of matrices result instead of a cube tensor.
matrixSlices2 :: DualMonad d r m
              => Int -> DualNumber d (Matrix r) -> m [DualNumber d (Matrix r)]
matrixSlices2 dr m@(D u _) = do
  let (rows, cols) = HM.size u
      n = dr * cols
  v <- returnLet $ flatten1 m  -- used many times below
  let f k = returnLet $ reshape2 cols $ slice1 (k * cols) n v
  mapM f [0 .. rows - dr]

-- Not optimal: matrix is constructed and destructed immediately,
-- which is costly when evaluating delta expressions. The transposes
-- may not be optimal, either. This goes down to individual deltas
-- of scalars, which is horrible for performance. Unlike @corr1@
-- this uses the slow dot product instead of the fast matrix-vector
-- (or matrix-matrix) multiplication.
corr2 :: forall d r m. DualMonad d r m
      => DualNumber d (Matrix r) -> DualNumber d (Matrix r)
      -> m (DualNumber d (Matrix r))
corr2 ker@(D u _) m@(D v _) = do
  let (rowsK, colsK) = HM.size u
      (rowsM, colsM) = HM.size v
      rr = rowsM - rowsK + 1
      rc = colsM - colsK + 1
  if | rowsK <= 0 || colsK <= 0 ->
       error $ "corr2: empty kernel not handled: " ++ show (rowsK, colsK)
     | rr <= 0 || rc <= 0 ->
       error $ "corr2: dim kernel " ++ show (rowsK, colsK)
               ++ " > dim matrix " ++ show (rowsM, colsM)
     | otherwise -> do
       kerTransV <- returnLet $ flatten1 (transpose2 ker)
       let dotColSlices :: DualNumber d (Matrix r) -> m [DualNumber d r]
           dotColSlices tm = do
             ttm <- returnLet $ transpose2 tm
             colSlices <- matrixSlices2 colsK ttm
             let f :: DualNumber d (Matrix r) -> DualNumber d r
                 f sm = kerTransV <.>! flatten1 sm
             return $ map f colSlices
       rowSlices <- matrixSlices2 rowsK m
       dotSlicesOfSlices <- mapM dotColSlices rowSlices
       returnLet $ reshape2 rc $ seq1 $ V.fromList $ concat dotSlicesOfSlices

conv2 :: forall d r m. DualMonad d r m
      => DualNumber d (Matrix r) -> DualNumber d (Matrix r)
      -> m (DualNumber d (Matrix r))
conv2 ker@(D u _) m@(D v _) = do
  let (rowsK, colsK) = HM.size u
      (rowsM, colsM) = HM.size v
  if | rowsK <= 0 || colsK <= 0 ->
       returnLet $ konst2 0 (rowsM + rowsK - 1, colsM + colsK - 1)
     | otherwise -> do
       let zRow = konst2 0 (rowsK - 1, colsM)
           rowPadded = rowAppend2 zRow $ rowAppend2 m zRow
           zCol = konst2 0 (rowsM + 2 * (rowsK - 1), colsK - 1)
           padded = columnAppend2 zCol $ columnAppend2 rowPadded zCol
       corr2 (fliprl2 . flipud2 $ ker) padded

conv2' :: IsScalar d r
       => DualNumber d (Matrix r) -> DualNumber d (Matrix r)
       -> DualNumber d (Matrix r)
conv2' (D u u') (D v v') = D (HM.conv2 u v) (dAdd (dConv2 u v') (dConv2 v u'))

-- A variant with limited padding, corresponding to SAME padding
-- from Tensorflow. Data size does not change with this padding.
-- It also performs convolution wrt flipped kernel (and so saves
-- on flipping it here), which makes no practical difference when
-- the kernel is initialized randomly.
convSame2 :: forall d r m. DualMonad d r m
          => DualNumber d (Matrix r) -> DualNumber d (Matrix r)
          -> m (DualNumber d (Matrix r))
convSame2 ker@(D u _) m@(D v _) = do
  let (rowsK, colsK) = HM.size u
      (rowsM, colsM) = HM.size v
  if | rowsK <= 0 || colsK <= 0 ->
       returnLet $ konst2 0 (rowsM, colsM)
     | otherwise -> do
       let zRow = konst2 0 ((rowsK - 1) `div` 2, colsM)
           rowPadded = rowAppend2 zRow $ rowAppend2 m zRow
           zCol = konst2 0 (rowsM + rowsK - 1, (colsK - 1) `div` 2)
           padded = columnAppend2 zCol $ columnAppend2 rowPadded zCol
       corr2 ker padded

-- No padding; remaining areas ignored.
maxPool2 :: forall d r m. DualMonad d r m
         => Int -> Int -> DualNumber d (Matrix r) -> m (DualNumber d (Matrix r))
maxPool2 ksize stride m@(D u _) = do
  let (rows, cols) = HM.size u
      colsOut = cols `div` stride
      resultRows = [0, stride .. rows - ksize]
      resultCols = [0, stride .. cols - ksize]
      resultCoords = [(r, c) | r <- resultRows, c <- resultCols]
  v <- returnLet $ flatten1 m  -- used many times below
  let getArea :: (Int, Int) -> DualNumber d (Vector r)
      getArea (r0, c0) =
        let getAreaAtRow r1 = append1 (slice1 (r1 * cols + c0) ksize v)
        in foldr getAreaAtRow (seq1 V.empty) [r0 .. r0 + ksize - 1]
      mins = map (maximum0 . getArea) resultCoords
  returnLet $ reshape2 colsOut $ seq1 $ V.fromList mins

-- * Operations resulting in an arbitrary untyped tensor

konstX :: IsScalar d r => DualNumber d r -> OT.ShapeL -> DualNumber d (OT.Array r)
konstX (D u u') sh = D (OT.constant sh u) (dKonstX u' sh)

appendX :: IsScalar d r
        => DualNumber d (OT.Array r) -> DualNumber d (OT.Array r)
        -> DualNumber d (OT.Array r)
appendX (D u u') (D v v') =
  D (u `OT.append` v) (dAppendX u' (head $ OT.shapeL u) v')

sliceX :: IsScalar d r
       => Int -> Int -> DualNumber d (OT.Array r) -> DualNumber d (OT.Array r)
sliceX i n (D u u') = D (OT.slice [(i, n)] u)
                        (dSliceX i n u' (head $ OT.shapeL u))

indexX :: IsScalar d r
       => DualNumber d (OT.Array r) -> Int -> DualNumber d (OT.Array r)
indexX (D u u') ix = D (OT.index u ix)
                       (dIndexX u' ix (head $ OT.shapeL u))

ravelFromListX :: IsScalar d r
               => [DualNumber d (OT.Array r)] -> DualNumber d (OT.Array r)
ravelFromListX ld =
  let (lu, lu') = unzip $ map (\(D u u') -> (u, u')) ld
      sh = case lu of
        u : _ -> length lu : OT.shapeL u
        [] -> []
  in D (OT.ravel $ OTB.fromList sh lu) (dRavelFromListX lu')

unravelToListX :: IsScalar d r
               => DualNumber d (OT.Array r) -> [DualNumber d (OT.Array r)]
unravelToListX (D v v') = case OT.shapeL v of
  k : _ ->
    let g ix p = D p (dIndexX v' ix k)
    in imap g $ OTB.toList $ OT.unravel v
  [] -> error "unravelToListX: wrong tensor dimensions"  -- catch early

mapX :: IsScalar d r
     => (DualNumber d (OT.Array r) -> DualNumber d (OT.Array r))
     -> DualNumber d (OT.Array r)
     -> DualNumber d (OT.Array r)
mapX f = ravelFromListX . map f . unravelToListX

zipWithX :: IsScalar d r
         => (DualNumber d (OT.Array r) -> DualNumber d (OT.Array r)
             -> DualNumber d (OT.Array r))
         -> DualNumber d (OT.Array r) -> DualNumber d (OT.Array r)
         -> DualNumber d (OT.Array r)
zipWithX f d e =
  ravelFromListX $ zipWith f (unravelToListX d) (unravelToListX e)

reshapeX :: IsScalar d r
         => OT.ShapeL -> DualNumber d (OT.Array r) -> DualNumber d (OT.Array r)
reshapeX sh' (D u u') = D (OT.reshape sh' u) (dReshapeX (OT.shapeL u) sh' u')

from0X :: IsScalar d r => DualNumber d r -> DualNumber d (OT.Array r)
from0X (D u u') = D (OT.scalar u) (dFrom0X u')

from1X :: IsScalar d r => DualNumber d (Vector r) -> DualNumber d (OT.Array r)
from1X (D u u') = D (OT.fromVector [V.length u] u) (dFrom1X u')

from2X :: IsScalar d r => DualNumber d (Matrix r) -> DualNumber d (OT.Array r)
from2X (D u u') = D (OT.fromVector [HM.rows u, HM.cols u] $ HM.flatten u)
                    (dFrom2X u' (HM.cols u))

fromSX :: forall sh d r. (IsScalar d r, OS.Shape sh)
       => DualNumber d (OS.Array sh r) -> DualNumber d (OT.Array r)
fromSX (D u u') = D (Data.Array.Convert.convert u) (dFromSX u')

-- * Operations resulting in an arbitrary fully typed Shaped tensor

konstS :: (IsScalar d r, OS.Shape sh)
       => DualNumber d r -> DualNumber d (OS.Array sh r)
konstS (D u u') = D (OS.constant u) (dKonstS u')

appendS :: (KnownNat m, KnownNat n, IsScalar d r, OS.Shape sh)
        => DualNumber d (OS.Array (m ': sh) r)
        -> DualNumber d (OS.Array (n ': sh) r)
        -> DualNumber d (OS.Array ((m + n) ': sh) r)
appendS (D u u') (D v v') = D (u `OS.append` v) (dAppendS u' v')

sliceS :: forall i n k rest d r.
          (KnownNat i, KnownNat n, KnownNat k, IsScalar d r, OS.Shape rest)
       => DualNumber d (OS.Array (i + n + k ': rest) r)
       -> DualNumber d (OS.Array (n ': rest) r)
sliceS (D u u') = D (OS.slice @'[ '(i, n) ] u)
                    (dSliceS (Proxy :: Proxy i) Proxy u')

indexS :: forall ix k rest d r.
          (KnownNat ix, KnownNat k, IsScalar d r, OS.Shape rest)
       => DualNumber d (OS.Array (ix + 1 + k ': rest) r)
       -> DualNumber d (OS.Array rest r)
indexS (D u u') = D (OS.index u (valueOf @ix))
                    (dIndexS u' (Proxy :: Proxy ix))

ravelFromListS :: forall rest k d r.
                  (KnownNat k, IsScalar d r, OS.Shape rest)
               => [DualNumber d (OS.Array rest r)]
               -> DualNumber d (OS.Array (k : rest) r)
ravelFromListS ld =
  let (lu, lu') = unzip $ map (\(D u u') -> (u, u')) ld
  in D (OS.ravel $ OSB.fromList lu) (dRavelFromListS lu')

unravelToListS :: forall k rest d r.
                  (KnownNat k, IsScalar d r, OS.Shape rest)
               => DualNumber d (OS.Array (k : rest) r)
               -> [DualNumber d (OS.Array rest r)]
unravelToListS (D v v') =
  -- @dIndexS@ is rigid, with type-level bound-checking, so we have to switch
  -- to @dIndexX@ for this function.
  let g ix p = D p (dFromXS $ dIndexX (dFromSX v') ix (valueOf @k))
  in imap g $ OSB.toList $ OS.unravel v

mapS :: forall k sh1 sh d r. (KnownNat k, IsScalar d r, OS.Shape sh, OS.Shape sh1)
     => (DualNumber d (OS.Array sh1 r) -> DualNumber d (OS.Array sh r))
     -> DualNumber d (OS.Array (k : sh1) r)
     -> DualNumber d (OS.Array (k : sh) r)
mapS f = ravelFromListS . map f . unravelToListS

mapMS :: forall k sh1 sh d r m.
         (Monad m, KnownNat k, IsScalar d r, OS.Shape sh, OS.Shape sh1)
      => (DualNumber d (OS.Array sh1 r) -> m (DualNumber d (OS.Array sh r)))
      -> DualNumber d (OS.Array (k : sh1) r)
      -> m (DualNumber d (OS.Array (k : sh) r))
mapMS f d = do
  let ld = unravelToListS d
  ld2 <- mapM f ld
  return $! ravelFromListS ld2

zipWithS :: forall k sh1 sh2 sh d r.
            ( KnownNat k, IsScalar d r, OS.Shape sh, OS.Shape sh1, OS.Shape sh2)
         => (DualNumber d (OS.Array sh1 r) -> DualNumber d (OS.Array sh2 r)
             -> DualNumber d (OS.Array sh r))
         -> DualNumber d (OS.Array (k : sh1) r)
         -> DualNumber d (OS.Array (k : sh2) r)
         -> DualNumber d (OS.Array (k : sh) r)
zipWithS f d e =
  ravelFromListS $ zipWith f (unravelToListS d) (unravelToListS e)

reshapeS :: (IsScalar d r, OS.Shape sh, OS.Shape sh', OS.Size sh ~ OS.Size sh')
         => DualNumber d (OS.Array sh r) -> DualNumber d (OS.Array sh' r)
reshapeS (D u u') = D (OS.reshape u) (dReshapeS u')

-- TODO: generalize as broadcast or stretch
asRowS :: forall k n d r. (IsScalar d r, KnownNat k, KnownNat n)
       => DualNumber d (OS.Array '[k] r) -> DualNumber d (OS.Array '[n, k] r)
asRowS d = from2S $ asRow2 (fromS1 d) (valueOf @n)

asColumnS :: forall k n d r. (IsScalar d r, KnownNat k, KnownNat n)
          => DualNumber d (OS.Array '[k] r) -> DualNumber d (OS.Array '[k, n] r)
asColumnS d = from2S $ asColumn2 (fromS1 d) (valueOf @n)

from0S :: IsScalar d r => DualNumber d r -> DualNumber d (OS.Array '[] r)
from0S (D u u') = D (OS.scalar u) (dFrom0S u')

from1S :: (KnownNat n, IsScalar d r)
       => DualNumber d (Vector r) -> DualNumber d (OS.Array '[n] r)
from1S (D u u') = D (OS.fromVector u) (dFrom1S u')

from2S :: (KnownNat rows, KnownNat cols, IsScalar d r)
       => DualNumber d (Matrix r) -> DualNumber d (OS.Array '[rows, cols] r)
from2S (D u u') = D (OS.fromVector $ HM.flatten u) (dFrom2S Proxy u')

fromXS :: (IsScalar d r, OS.Shape sh)
       => DualNumber d (OT.Array r) -> DualNumber d (OS.Array sh r)
fromXS (D u u') = D (Data.Array.Convert.convert u) (dFromXS u')

-- TODO: generalize to arbitrary permutations of arbitrarily many ranks using https://hackage.haskell.org/package/orthotope/docs/Data-Array-ShapedS.html#v:transpose
transpose2S :: (IsScalar d r, KnownNat rows, KnownNat cols)
            => DualNumber d (OS.Array '[rows, cols] r)
            -> DualNumber d (OS.Array '[cols, rows] r)
transpose2S = from2S . transpose2 . fromS2

infixr 8 #>$
(#>$) :: (IsScalar d r, KnownNat rows, KnownNat cols)
      => DualNumber d (OS.Array '[rows, cols] r)
      -> DualNumber d (OS.Array '[cols] r)
      -> DualNumber d (OS.Array '[rows] r)
(#>$) d e = from1S $ fromS2 d #>! fromS1 e

infixr 8 <>$
(<>$) :: (IsScalar d r, KnownNat m, KnownNat n, KnownNat p)
      => DualNumber d (OS.Array '[m, n] r)
      -> DualNumber d (OS.Array '[n, p] r)
      -> DualNumber d (OS.Array '[m, p] r)
(<>$) d e = from2S $ fromS2 d <>! fromS2 e

conv2S :: forall d r kheight_minus_1 kwidth_minus_1 in_height in_width.
          ( KnownNat kheight_minus_1, KnownNat kwidth_minus_1
          , KnownNat in_height, KnownNat in_width
          , IsScalar d r )
       => DualNumber d (OS.Array '[kheight_minus_1 + 1, kwidth_minus_1 + 1] r)
       -> DualNumber d (OS.Array '[in_height, in_width] r)
       -> DualNumber d (OS.Array '[ in_height + kheight_minus_1
                                 , in_width + kwidth_minus_1 ] r)
conv2S ker x = from2S $ conv2' (fromS2 ker) (fromS2 x)

-- Convolution of many matrices at once. Some of the names of dimensions
-- are from https://www.tensorflow.org/api_docs/python/tf/nn/conv2d
conv24 :: forall kheight_minus_1 kwidth_minus_1
                 out_channels in_height in_width batch_size in_channels d r.
          ( KnownNat kheight_minus_1, KnownNat kwidth_minus_1
          , KnownNat out_channels, KnownNat in_height, KnownNat in_width
          , KnownNat batch_size, KnownNat in_channels
          , IsScalar d r )
       => DualNumber d (OS.Array '[ out_channels, in_channels
                                 , kheight_minus_1 + 1, kwidth_minus_1 + 1 ] r)
       -> DualNumber d (OS.Array '[batch_size, in_channels, in_height, in_width] r)
       -> DualNumber d (OS.Array '[ batch_size, out_channels
                                 , in_height + kheight_minus_1
                                 , in_width + kwidth_minus_1 ] r)
conv24 ker = mapS conv23 where
  conv23 :: DualNumber d (OS.Array '[in_channels, in_height, in_width] r)
         -> DualNumber d (OS.Array '[ out_channels
                                   , in_height + kheight_minus_1
                                   , in_width + kwidth_minus_1 ] r)
  conv23 x = mapS (convFilters x) ker
  convFilters
    :: DualNumber d (OS.Array '[in_channels, in_height, in_width] r)
    -> DualNumber d (OS.Array '[ in_channels
                              , kheight_minus_1 + 1, kwidth_minus_1 + 1 ] r)
    -> DualNumber d (OS.Array '[ in_height + kheight_minus_1
                              , in_width + kwidth_minus_1 ] r)
  convFilters x ker1 = sumOutermost $ zipWithS conv2S ker1 x
  sumOutermost :: DualNumber d (OS.Array '[ in_channels
                                         , in_height + kheight_minus_1
                                         , in_width + kwidth_minus_1 ] r)
               -> DualNumber d (OS.Array '[ in_height + kheight_minus_1
                                         , in_width + kwidth_minus_1 ] r)
  sumOutermost = sum . unravelToListS
    -- slow; should go through Tensor2, or the Num instance should when possible

maxPool24
  :: forall ksize_minus_1 stride in_height in_width batch_size channels d r m.
     ( KnownNat ksize_minus_1, KnownNat stride
     , KnownNat in_height, KnownNat in_width
     , KnownNat batch_size, KnownNat channels
     , 1 <= stride
     , ksize_minus_1 <= in_height
     , ksize_minus_1 <= in_width
     , 1 <= in_height - ksize_minus_1 + stride
     , 1 <= in_width - ksize_minus_1 + stride
     , DualMonad d r m )
     => DualNumber d (OS.Array '[batch_size, channels, in_height, in_width] r)
     -> m (DualNumber d
             (OS.Array '[ batch_size, channels
                         , (in_height - ksize_minus_1) `DivRoundUp` stride
                         , (in_width - ksize_minus_1) `DivRoundUp` stride ] r))
maxPool24 d = do
  res <- mapMS (mapMS (fmap from2S
                       . maxPool2 (valueOf @ksize_minus_1 + 1)
                                  (valueOf @stride)
                       . fromS2)) d
  returnLet res

-- * Operations creating delayed/outlined derivatives

-- | The version of the @D@ constructor lazy in the second argument.
-- To be used as in
--
-- > sinDelayed :: (Floating a, IsPrimal d a) => DualNumber d a -> DualNumber d a
-- > sinDelayed (D u u') = delayD (sin u) (dScale (cos u) u')
-- >
-- > plusDelayed :: (Floating a, IsPrimal d a)
-- >             => DualNumber d a -> DualNumber d a -> DualNumber d a
-- > plusDelayed (D u u') (D v v') = delayD (u + v) (dAdd u' v')
-- >
-- > x ** (sinDelayed x `plusDelayed` (id2 $ id2 $ id2 $ konst1 (sumElements0 x) 2))
--
-- The outlining is lost when serializing or logging, unlike with @Out@,
-- @Outline0@, etc.
--
-- Yet another incomparable variant that can't be serialized would be
-- (illustrating with an example of a constructor at rank 0)
--
-- > FromParams0 (Domains -> Delta0 r)
--
-- that expects the initial parameters. But it's more troublesome
-- than @Delay0@ both in implementation and usage.
delayD :: IsPrimal d a => a -> Dual d a -> DualNumber d a
delayD u ~u' = D u (dDelay u')

-- | A wrapper type to delay/outline computation of the derivatives of the given
-- primitive numeric function inside the dual component of the created dual
-- number. The rule is that if all arguments of a function are wrapped
-- in @Out@ then the function gets delayed. Inconsistent wrapping,
-- e.g., only one of the arguments, leads to early type errors.
--
-- To be used as in
--
-- > x ** unOut (sin (Out x) + Out (id $ id $ id $ konst1 (sumElements0 x) 2))
--
-- which delays computing the dual component of sine and of addition
-- (both in rank 1), but not of power, konst and sumElements. The last two
-- can't be currently delayed, because only primitive numeric functions
-- are supported (an attempt would not type-check).
newtype Out a = Out {unOut :: a}
  deriving (Eq, Ord)

returnOut :: (DualMonad d r m, IsPrimalWithScalar d a r)
          => Out (DualNumber d a) -> m (Out (DualNumber d a))
returnOut dOut = do
  dvar <- returnLet $ unOut dOut
  return $ Out dvar

instance (Num a, IsPrimal 'DModeGradient a, HasVariables a)
         => Num (Out (DualNumber 'DModeGradient a)) where
  Out (D u u') + Out (D v v') =
    Out $ D (u + v) (dOutline PlusOut [u, v] [u', v'])
  Out (D u u') - Out (D v v') =
    Out $ D (u - v) (dOutline MinusOut [u, v] [u', v'])
  Out (D u u') * Out (D v v') =
    Out $ D (u * v) (dOutline TimesOut [u, v] [u', v'])
  negate (Out (D v v')) = Out $ D (negate v) (dOutline NegateOut [v] [v'])
  abs (Out (D v v')) = Out $ D (abs v) (dOutline AbsOut [v] [v'])
  signum (Out (D v v')) = Out $ D (signum v) (dOutline SignumOut [v] [v'])
  fromInteger = Out . constant . fromInteger

instance (Real a, IsPrimal 'DModeGradient a, HasVariables a)
         => Real (Out (DualNumber 'DModeGradient a)) where
  toRational = undefined  -- TODO?

instance (Fractional a, IsPrimal 'DModeGradient a, HasVariables a)
         => Fractional (Out (DualNumber 'DModeGradient a)) where
  Out (D u u') / Out (D v v') =
    Out $ D (u / v) (dOutline DivideOut [u, v] [u', v'])
  recip (Out (D v v')) = Out $ D (recip v) (dOutline RecipOut [v] [v'])
  fromRational = Out . constant . fromRational

instance (Floating a, IsPrimal 'DModeGradient a, HasVariables a)
         => Floating (Out (DualNumber 'DModeGradient a)) where
  pi = Out $ constant pi
  exp (Out (D u u')) = Out $ D (exp u) (dOutline ExpOut [u] [u'])
  log (Out (D u u')) = Out $ D (log u) (dOutline LogOut [u] [u'])
  sqrt (Out (D u u')) = Out $ D (sqrt u) (dOutline SqrtOut [u] [u'])
  Out (D u u') ** Out (D v v') =
    Out $ D (u ** v) (dOutline PowerOut [u, v] [u', v'])
  logBase (Out (D u u')) (Out (D v v')) = Out $ D (logBase u v) (dOutline LogBaseOut [u, v] [u', v'])
  sin (Out (D u u')) = Out $ D (sin u) (dOutline SinOut [u] [u'])
  cos (Out (D u u')) = Out $ D (cos u) (dOutline CosOut [u] [u'])
  tan (Out (D u u')) = Out $ D (tan u) (dOutline TanOut [u] [u'])
  asin (Out (D u u')) = Out $ D (asin u) (dOutline AsinOut [u] [u'])
  acos (Out (D u u')) = Out $ D (acos u) (dOutline AcosOut [u] [u'])
  atan (Out (D u u')) = Out $ D (atan u) (dOutline AtanOut [u] [u'])
  sinh (Out (D u u')) = Out $ D (sinh u) (dOutline SinhOut [u] [u'])
  cosh (Out (D u u')) = Out $ D (cosh u) (dOutline CoshOut [u] [u'])
  tanh (Out (D u u')) = Out $ D (tanh u) (dOutline TanhOut [u] [u'])
  asinh (Out (D u u')) = Out $ D (asinh u) (dOutline AsinhOut [u] [u'])
  acosh (Out (D u u')) = Out $ D (acosh u) (dOutline AcoshOut [u] [u'])
  atanh (Out (D u u')) = Out $ D (atanh u) (dOutline AtanhOut [u] [u'])

instance (RealFrac a, IsPrimal 'DModeGradient a, HasVariables a)
         => RealFrac (Out (DualNumber 'DModeGradient a)) where
  properFraction = undefined
    -- very low priority, since these are all extremely not continuous

instance (RealFloat a, IsPrimal 'DModeGradient a, HasVariables a)
         => RealFloat (Out (DualNumber 'DModeGradient a)) where
  atan2 (Out (D u u')) (Out (D v v')) =
    Out $ D (atan2 u v) (dOutline Atan2Out [u, v] [u', v'])
      -- we can be selective here and omit the other methods,
      -- most of which don't even have a differentiable codomain

-- * Busywork to let the derivatives mode ignore all outlining

-- | Note that this should apply only when @d@ is @'DModeDerivative@.
-- However, GHC can't tell that @d@ has only two cases. Therefore, we need
-- to overgeneralize these definitions and mark them with @OVERLAPPABLE@
-- or else GHC complains that not enough instances are given
-- whenever type-checking code polymorphic on @d@.
instance {-# OVERLAPPABLE #-} (Num a, IsPrimal d a)
                              => Num (Out (DualNumber d a)) where
  Out d + Out e = Out (d + e)
  Out d - Out e = Out (d - e)
  Out d * Out e = Out (d * e)
  negate (Out e) = Out (negate e)
  abs (Out e) = Out (abs e)
  signum (Out e) = Out (signum e)
  fromInteger = Out . constant . fromInteger

instance {-# OVERLAPPABLE #-} (Real a, IsPrimal d a)
                              => Real (Out (DualNumber d a)) where
  toRational = undefined  -- TODO?

instance {-# OVERLAPPABLE #-} (Fractional a, IsPrimal d a)
                              => Fractional (Out (DualNumber d a)) where
  Out d / Out e = Out (d / e)
  recip (Out e) = Out (recip e)
  fromRational = Out . constant . fromRational

instance {-# OVERLAPPABLE #-} (Floating a, IsPrimal d a)
                              => Floating (Out (DualNumber d a)) where
  pi = Out $ constant pi
  exp (Out d) = Out (exp d)
  log (Out d) = Out (log d)
  sqrt (Out d) = Out (sqrt d)
  Out d ** Out e = Out (d ** e)
  logBase (Out x) (Out y) = Out (logBase x y)
  sin (Out d) = Out (sin d)
  cos (Out d) = Out (cos d)
  tan (Out d) = Out (tan d)
  asin (Out d) = Out (asin d)
  acos (Out d) = Out (acos d)
  atan (Out d) = Out (atan d)
  sinh (Out d) = Out (sinh d)
  cosh (Out d) = Out (cosh d)
  tanh (Out d) = Out (tanh d)
  asinh (Out d) = Out (asinh d)
  acosh (Out d) = Out (acosh d)
  atanh (Out d) = Out (atanh d)

instance {-# OVERLAPPABLE #-} (RealFrac a, IsPrimal d a)
                              => RealFrac (Out (DualNumber d a)) where
  properFraction = undefined
    -- very low priority, since these are all extremely not continuous

instance {-# OVERLAPPABLE #-} (RealFloat a, IsPrimal d a)
                              => RealFloat (Out (DualNumber d a)) where
  atan2 (Out d) (Out e) = Out (atan2 d e)

max_stars_repo_name: Mikolaj/horde-ad

max_issues_repo_name: Mikolaj/horde-ad

max_forks_repo_name: Mikolaj/horde-ad


text:
{-# LANGUAGE DataKinds             #-}
{-# LANGUAGE GADTs                 #-}
{-# LANGUAGE TypeFamilies          #-}
{-# LANGUAGE TypeOperators         #-}
{-# LANGUAGE FlexibleContexts      #-}
{-# LANGUAGE ScopedTypeVariables   #-}
{-# LANGUAGE RankNTypes            #-}

module Grenade.Utils.OneHot (
    oneHot
  , hotMap
  , makeHot
  , unHot
  , sample
  ) where

import qualified Control.Monad.Random as MR

import           Data.List ( group, sort )

import           Data.Map ( Map )
import qualified Data.Map as M

import           Data.Proxy
import           Data.Singletons.TypeLits

import           Data.Vector ( Vector )
import qualified Data.Vector as V
import qualified Data.Vector.Storable as VS

import           Numeric.LinearAlgebra ( maxIndex )
import           Numeric.LinearAlgebra.Devel
import           Numeric.LinearAlgebra.Static

import           Grenade.Core.Shape

-- | From an int which is hot, create a 1D Shape
--   with one index hot (1) with the rest 0.
--   Rerurns Nothing if the hot number is larger
--   than the length of the vector.
oneHot :: forall n. (KnownNat n)
       => Int -> Maybe (S ('D1 n))
oneHot hot =
  let len = fromIntegral $ natVal (Proxy :: Proxy n)
  in if hot < len
      then
        fmap S1D . create $ runSTVector $ do
        vec    <- newVector 0 len
        writeVector vec hot 1
        return vec
      else Nothing

-- | Create a one hot map from any enumerable.
--   Returns a map, and the ordered list for the reverse transformation
hotMap :: (Ord a, KnownNat n) => Proxy n -> [a] -> Either String (Map a Int, Vector a)
hotMap n as =
  let len  = fromIntegral $ natVal n
      uniq = [ c | (c:_) <- group $ sort as]
      hotl = length uniq
  in if hotl == len
      then
        Right (M.fromList $ zip uniq [0..], V.fromList uniq)
      else
        Left ("Couldn't create hotMap of size " ++ show len ++ " from vector with " ++ show hotl ++ " unique characters")

-- | From a map and value, create a 1D Shape
--   with one index hot (1) with the rest 0.
--   Rerurns Nothing if the hot number is larger
--   than the length of the vector or the map
--   doesn't contain the value.
makeHot :: forall a n. (Ord a, KnownNat n)
        => Map a Int -> a -> Maybe (S ('D1 n))
makeHot m x = do
  hot    <- M.lookup x m
  let len = fromIntegral $ natVal (Proxy :: Proxy n)
  if hot < len
      then
        fmap S1D . create $ runSTVector $ do
        vec    <- newVector 0 len
        writeVector vec hot 1
        return vec
      else Nothing

unHot :: forall a n. KnownNat n
      => Vector a -> S ('D1 n) -> Maybe a
unHot v (S1D xs)
  = (V.!?) v
  $ maxIndex (extract xs)

sample :: forall a n m. (KnownNat n, MR.MonadRandom m)
       => Double -> Vector a -> S ('D1 n) -> m a
sample temperature v (S1D xs) = do
  ix <- MR.fromList . zip [0..] . fmap (toRational . exp . (/ temperature) . log) . VS.toList . extract $ xs
  return $ v V.! ix

ext: hs

max_issues_repo_name: Alien-Inc/grenade

max_forks_repo_name: Alien-Inc/grenade


text:
{-# LANGUAGE ScopedTypeVariables #-}

module Qubism.AlgebraTests
  ( module Test.QuickCheck.Property.Generic
  , isVectorSpace, isHilbertSpace, isAlgebra)
  where

import Test.Hspec
import Test.QuickCheck
import Test.QuickCheck.Property.Generic
import Test.QuickCheck.Property.Common.Internal
import Data.Complex

import Qubism.Algebra

tupled :: T a -> T (a,a)
tupled _ = T

prop_Distributive
  :: (a -> b -> b) -> (b -> b -> b) -> T a -> T b -> a -> b -> b -> Equal b
prop_Distributive f g _ _ a b c = Equal left right
  where left  = a `f` (b `g` c)
        right = (a `f` b) `g` (a `f` c)

isVectorSpace :: (Show v, Arbitrary v, VectorSpace v) => T v -> Spec
isVectorSpace t = do
  it "+: is associative"          $ property . eq $ prop_Associative (+:) t
  it "+: is commutative"          $ property . eq $ prop_Commutative (+:) (tupled t)
  it "zero is the identity of +:" $ property . eq $ prop_Identity zero (+:) t
  it "neg produces the inverse"   $ property $
    let test (_ :: T v) (v :: v) = neg v +: v == zero
    in  test t
  it ".: distrbutes over +:" $ property . eq $
    prop_Distributive (.:) (+:) (T :: T C) t

prop_InnerLinear :: VectorSpace v => (v -> C) -> T v -> C -> v -> v -> Bool
prop_InnerLinear f _ a u v = magnitude (left - right) < 0.00001
  where left  = f $ (a .: u) +: v
        right = a * f u + f v

isHilbertSpace :: (Show v, Arbitrary v, HilbertSpace v) => T v -> Spec
isHilbertSpace t = do
  it "<.> is linear in the second argument" $ property $
    \w a u v -> prop_InnerLinear (w <.>) t a u v
  it "<a.b> = conj <b.a>" $ property $
    let test (_ :: T v)(a :: v)(b :: v) = a <.> b `shouldBe` conjugate (b <.> a)
    in  test t

prop_AlgLinear :: (Show v, Algebra v) => (v -> v) -> T v -> C -> v -> v -> Expectation
prop_AlgLinear f _ a u v = left `shouldBe` right
  where left  = f $ (a .: u) +: v
        right = (a .: f u) +: f v

isAlgebra :: (Show v, Arbitrary v, Algebra v) => T v -> Spec
isAlgebra t = do
  it "*: is linear in the first argument" $ property $
    \w a u v -> prop_AlgLinear (*: w) t a u v
  it "*: is linear in the second argument" $ property $
    \w a u v -> prop_AlgLinear (w *:) t a u v

max_stars_repo_name: qubitrot/qubism

max_issues_repo_name: qubitrot/qubism

max_forks_repo_name: qubitrot/qubism


text:
module LispVal where

import Control.Monad.Cont (liftIO, liftM)
import Control.Monad.Error.Class (MonadError, catchError)
import Control.Monad.Except (ExceptT, runExceptT, throwError)
import Data.Complex (Complex)
import Data.IORef
import GHC.IO.Handle.Types (Handle)
import Text.ParserCombinators.Parsec (ParseError)

instance Show LispVal where show = showVal

data LispVal
  = Atom String
  | List [LispVal]
  | DottedList [LispVal] LispVal
  | Number Integer
  | Float Float
  | Rational Rational
  | Complex (Complex Float)
  | String String
  | Character String
  | Bool Bool
  | PrimitiveFunc ([LispVal] -> ThrowsError LispVal)
  | Func
      { params :: [String],
        vararg :: Maybe String,
        body :: [LispVal],
        closure :: Environment
      }
  | IOFunc ([LispVal] -> IOThrowsError LispVal)
  | Port Handle

showVal :: LispVal -> String
showVal (List contents) = "(" ++ wordsList contents ++ ")"
showVal (DottedList x xs) = "(" ++ wordsList x ++ " . " ++ showVal xs ++ ")"
showVal (String contents) = "\"" ++ contents ++ "\""
showVal (Atom name) = name
showVal (Number contents) = show contents
showVal (Float contents) = show contents
showVal (Rational contents) = show contents
showVal (Complex contents) = show contents
showVal (Character contents) = show contents
showVal (Bool True) = "True"
showVal (Bool False) = "False"
showVal (PrimitiveFunc _) = "<primitive>"
showVal (Func {params = args, vararg = varargs, body = _, closure = _}) =
  "(lambda (" ++ unwords (map show args)
    ++ ( case varargs of
           Nothing -> ""
           Just arg -> " . " ++ arg
       )
    ++ ") ...)"
showVal (Port _) = "<IO port>"
showVal (IOFunc _) = "<IO primitive>"

type IOThrowsError = ExceptT LispError IO

liftThrows :: ThrowsError a -> IOThrowsError a
liftThrows (Left err) = throwError err
liftThrows (Right val) = return val

runIOThrows :: IOThrowsError String -> IO String
runIOThrows action = runExceptT (trapError action) >>= return . extractValue

type ThrowsError = Either LispError

instance Show LispError where show = showError

data LispError
  = NumArgs Integer [LispVal]
  | TypeMismatch String LispVal
  | Parser ParseError
  | BadSpecialForm String LispVal
  | NotFunction String String
  | UnboundVar String String
  | Default String

extractValue :: ThrowsError a -> a
extractValue (Right val) = val

trapError :: (MonadError a m, Show a) => m String -> m String
trapError action = catchError action (return . show)

showError :: LispError -> String
showError (UnboundVar message varName) = message ++ ": " ++ varName
showError (BadSpecialForm message form) = message ++ ": " ++ show form
showError (NotFunction message func) = message ++ ": " ++ show func
showError (NumArgs expected found) =
  "Expected " ++ show expected
    ++ " args; found values "
    ++ wordsList found
showError (TypeMismatch expected found) =
  "Invalid type: expected " ++ expected
    ++ ", found "
    ++ show found
showError (Parser parseErr) = "Parse error at " ++ show parseErr
showError _ = "Error"

wordsList :: [LispVal] -> String
wordsList = unwords . map showVal

type Environment = IORef [(String, IORef LispVal)]

getVar :: Environment -> String -> IOThrowsError LispVal
getVar envRef var = do
  env <- liftIO $ readIORef envRef
  maybe
    (throwError $ UnboundVar "Getting an unbound variable" var)
    (liftIO . readIORef)
    (lookup var env)

setVar :: Environment -> String -> LispVal -> IOThrowsError LispVal
setVar envRef var value = do
  env <- liftIO $ readIORef envRef
  maybe
    (throwError $ UnboundVar "Setting an unbound variable" var)
    (liftIO . (flip writeIORef value))
    (lookup var env)
  return value

defineVar :: Environment -> String -> LispVal -> IOThrowsError LispVal
defineVar envRef var value = do
  alreadyDefined <- liftIO $ isBound envRef var
  if alreadyDefined
    then setVar envRef var value >> return value
    else liftIO $ do
      valueRef <- newIORef value
      env <- readIORef envRef
      writeIORef envRef ((var, valueRef) : env)
      return value

bindVars :: Environment -> [(String, LispVal)] -> IO Environment
bindVars envRef binds = readIORef envRef >>= extendEnv binds >>= newIORef
  where
    extendEnv binds' env = liftM (++ env) (mapM addBind binds')
    addBind (var, value) = do
      ref <- newIORef value
      return (var, ref)

isBound :: Environment -> String -> IO Bool
isBound envRef var = readIORef envRef >>= return . maybe False (const True) . lookup var

nullEnv :: IO Environment
nullEnv = newIORef []


text:
{-# LANGUAGE BangPatterns #-}

module Rover
  (
    someFunc
  ) where

import Models.POI
import Models.RoverDomain
import Models.Rover
import Models.State
import Sim.Simulation
import Util

import CCEA
import NN.NeuralNetwork hiding (train)
import Numeric.LinearAlgebra.HMatrix hiding (corr)

import System.IO.Unsafe
import System.Random
import Control.Monad.Writer
import Data.List (sortOn, mapAccumL, sort)

type Bounds = (Int,Int)
type POICount = Int
type LoyalCount = Int
type TraitorCount = Int
type DomainInfo = (Bounds, POICount, LoyalCount, TraitorCount)
type Domain = RoverDomain Rover POI

-- Given a domain (which may or may not have other agents with neural nets), trains
-- a population of neural networks to either be loyal or disloyal
trainX ::  Domain -> Pop -> Int -> SimSteps -> Bool -> IO Pop
trainX domain nets n steps trainLoyal = do
  let fit = if trainLoyal then dFit steps else traitorFit steps
  gen <- newStdGen
  let ccea :: CCEA (Network Double) (StdGen, Domain) StdGen
      cceaBreeding = elitist (mutateWeights nnVars)
      ccea = CCEA nets fit cceaBreeding (tournament 2) (gen, domain)
      ccea' = evolveNCCEA n gen ccea
      finalNets = _pop (snd ccea')
  return $! finalNets

testPopWithT :: DomainInfo -> NNVars -> NNVars -> [Network Double] -> [Network Double] -> SimSteps -> IO Double
testPopWithT (b,p,l,t) varsL varsT loyal traitor n = do
  roverDom <- randomRoverDomain varsL varsT b p l t
  gen      <- newStdGen
  let roverDom' = assignNets (assignNetsT roverDom traitor) loyal
      afterSim  = simulation n roverDom'
  return $! getGlobalScore afterSim

type Pop = Population (Network Double)
type Team = [Network Double]
type Scores = [Double]

mean :: [Double] -> Double
mean !xs = sum xs / (fromIntegral (length xs))

stddev :: (Floating a) => [a] -> a
stddev xs = sqrt $ var xs

var xs = (var' 0 0 0 xs) / (fromIntegral $ length xs - 1)
    where
      var' _ _ s [] = s
      var' m n s (x:xs) = var' nm (n + 1) (s + delta * (x - nm)) xs
         where
           delta = x - m
           nm = m + delta/(fromIntegral $ n + 1)

type NGens = Int
type NTrials = Int
type NExps = Int
type SimSteps = Int

-- Trains a population of neural networks in an adversarial setting against another team
-- of neural networks in some random (fixed throughout training) domain.
-- Then tests the best team across random domains.
-- Returns the population, the best team, and the scores of the best team
evalTrain :: NTrials -> DomainInfo -> Pop -> [Network Double] -> NNVars -> NNVars -> NGens -> SimSteps -> Bool -> IO (Pop, [Network Double], [Double])
evalTrain nTrials info@(b,p,l,t) pop bestOthers varsL varsT nGens steps isLoyal = do
  dom <- randomRoverDomain varsL varsT b p l t
  let (bestF, assignF) = if isLoyal then (bestLoyal, assignNetsT)
                                 else (bestTraitor, assignNets)
  let domWithOthers = assignF dom bestOthers
  trained <- trainX domWithOthers pop nGens steps isLoyal
  bestTrained <- bestF domWithOthers steps trained

  let (bestLoyal, bestTraitor) = if isLoyal then (bestTrained, bestOthers)
                                 else (bestOthers, bestTrained)

  scores <- repIO nTrials (\_ -> testPopWithT info varsL varsT bestLoyal bestTraitor steps) :: IO [Double]
  return $! (trained, bestTrained, scores)

-- Next step. Make pointfree.
repEvalTrain :: NExps -> IO (Pop, Team, Scores) -> IO (Pop, Team, Scores)
repEvalTrain n f = repIO n (\_ -> f) >>= return . last . sortOn (\(a,b,c) -> mean c)

repEvalTrainT n f = repIO n (\_ -> f) >>= return . head . sortOn (\(a,b,c) -> mean c)

getTrainingData :: NExps -> IO [(Pop, Team, Scores)] -> IO [(Double, Double)]
getTrainingData n fs = do
  allGens <- fs
  repEachGen <- sequence $ map (repEvalTrain n . return) allGens
  return $! map (\(a,b,c) -> (mean c, stddev c)) repEachGen

getTrainingDataT n fs = do
  allGens <- fs
  repEachGen <- sequence $ map (repEvalTrainT n . return) allGens
  return $! map (\(a,b,c) -> (mean c, stddev c)) repEachGen

someFunc :: IO ()
someFunc = do
  let numL = 2
      numT = 4
      numP = 10
      b    = (10,10)
      domInfo  = (b, numP, numL, 0)
      domInfo2 = (b, numP, numL, numT)
      nTrials  = 30
      nGens    = 1
      nExps    = 15
      simStep  = 15
      varsN  = nnVarsNaive
      vars   = nnVars

  gen   <- newStdGen

  -- Strong foundation (Pre trained loyal)
  loyalNaive <- createPopulation 10 numL nnVarsNaive :: IO Pop
  (ln1,_, _) <- repEvalTrain 20 (evalTrain nTrials domInfo loyalNaive [] varsN varsN nGens simStep True)
  (ln2,_, _) <- repEvalTrain 20 (evalTrain nTrials domInfo ln1 [] varsN varsN nGens simStep True)
  (ln3,_, _) <- repEvalTrain 20 (evalTrain nTrials domInfo ln2 [] varsN varsN nGens simStep True)

  -- Loyal Agents Trained With 8 Input Neural Networks
  -- Step 1: Train loyal with 8
  let tNL x = evalTrain nTrials domInfo ln1 [] varsN varsN x simStep True
  (lnp, lnt, lns) <- repEvalTrain 20 (tNL nGens)
  putStrLn "Loyal Agents"
  getTrainingData 10 (sequence $ map (\n -> tNL n) [1,10..nGens]) >>= splitP

  newTraitors <- createPopulation 5 numT vars :: IO Pop
  let tTNL x = evalTrain nTrials domInfo2 newTraitors lnt varsN vars x simStep False
  (tp1,tt1,ts1) <- repEvalTrainT 20 (tTNL nGens)
  (tp2,tt2,ts2) <- repEvalTrainT 20 (evalTrain nTrials domInfo2 tp1 lnt varsN vars nGens simStep False)
  (tp, tt, ts) <- repEvalTrainT 20 (evalTrain nTrials domInfo2 tp2 lnt varsN vars nGens simStep False)

  putStrLn "Treasonous agents training with static loyal agents"
  getTrainingDataT 10 (sequence $ map (\n -> tTNL n) [1,10..nGens]) >>= print

  -- Loyal Agents Trained With 8 Input Neural Networks with Traitors
  -- Step 1: Train agents with above step traitors
  let tTNL' x = evalTrain nTrials domInfo2 lnp tt varsN vars x simStep True
  (lnp', lnt', lns') <- repEvalTrain 20 (tTNL' nGens)

  putStrLn "Loyal Agents trained with traitors"
  getTrainingData 10 (sequence $ map (\n -> tTNL' n) [1,10..nGens]) >>= print

  putStrLn "Now Agents can recognize traitors"

  -- Strong foundation (Pre trained loyal with traitor awareness)
  l <- createPopulation 10 numL nnVars :: IO Pop
  (l1, _, _) <- repEvalTrain 20 (evalTrain nTrials domInfo l [] vars vars nGens simStep True)
  (l2, _, _) <- repEvalTrain 20 (evalTrain nTrials domInfo l1 [] vars vars nGens simStep True)
  (l3, _, _) <- repEvalTrain 20 (evalTrain nTrials domInfo l2 [] vars vars nGens simStep True)

  -- -- Loyal Agents Trained With 12 Input Neural Networks

  let tL x = evalTrain nTrials domInfo l3 [] nnVars nnVars x simStep True
  (lp, lt, ls) <-  repEvalTrain 20 (tL nGens)

  putStrLn "Loyal agents that can recognize traitors"
  getTrainingData 10 (sequence $ map (\n -> tL n) [1,10..nGens]) >>= print

  -- Loyal Agents Trained With 12 Input Neural Networks with Traitors
  let tTL x = evalTrain nTrials domInfo2 lp tt nnVars nnVars x simStep True
  (lp', lt', ls') <-  repEvalTrain 20 (tTL nGens)

  putStrLn "Loyal Agents that can recognize traitors with traitors"
  getTrainingData 20 (sequence $ map (\n -> tTL n) [1,10..nGens]) >>= print

splitP as = matlabPrint (map fst as) >> matlabPrint (map snd as)
matlabPrint :: (Show a) => [a] -> IO ()
matlabPrint !as = putStr "[" >> matlabPrintHelper as
  where
    matlabPrintHelper [] = putStrLn "];"
    matlabPrintHelper (a:as) = putStr ((show a) ++ " ") >> matlabPrintHelper as

  --lnData          <- getTrainingData 20 (sequence $ map (\n -> evalTrainX nTrials domInfo loyalNaive [] nnVarsNaive nnVarsNaive n simStep True) [0,10..nGens])

  --  putStrLn $ show lnData
  --putStr "x = "
  --matlabPrint [0,10..nGens]
  --putStr "ln = "
  --matlabPrint (map fst lnData)
  --putStr "lne = "
  --matlabPrint (map snd lnData)

  -- Traitors trained with 12 Input Neural Networks vs Loyal Agents with 8 Input
  -- Step 1: Train with agents in above step

max_stars_repo_name: eklinkhammer/rover-domain

max_issues_repo_name: eklinkhammer/rover-domain

max_forks_repo_name: eklinkhammer/rover-domain


text:
{-# LANGUAGE RankNTypes #-}
{-# LANGUAGE TypeApplications #-}

-- | Flexible numeric parsers for real-world programming languages.
-- These parsers aim to be a superset of the numeric syntaxes across
-- the most popular programming languages.
--
-- All parsers assume any trailing whitespace has already been
-- consumed, and places no requirement for an @endOfInput@ at the end
-- of a literal. Be sure to handle these in a calling context. These
-- parsers do not use 'Text.Parser.Token.TokenParsing', and therefore
-- may fail while consuming input, depending on if you use a parser
-- that automatically backtracks or not. Apply 'try' if needed.
module Numeric.Parse.Flexible
  ( integer,
    natural,
    decimal,
    hexadecimal,
    octal,
    binary,
    floating,
    signed,
    imaginary,
  )
where

import Control.Applicative
import Control.Monad hiding (fail)
import Data.Scientific hiding (scientific)
import Numeric
import Text.Parser.Char
import Text.Parser.Combinators
import Text.Read (readMaybe)
import Prelude hiding (exponent, fail, takeWhile)
import Data.Complex
import Numeric.Natural (Natural)

-- | Parse an integer in 'decimal', 'hexadecimal', 'octal', or 'binary', with optional leading sign.
--
-- Note that because the 'octal' parser takes primacy over 'decimal', numbers with a leading
-- @0@ will be parsed as octal. This is unfortunate, but matches the behavior
-- of C, Python, and Ruby.
integer :: (CharParsing m, Monad m) => m Integer
integer = signed (choice [try hexadecimal, try octal, try binary, decimal])

-- | Parse a natural number in 'decimal', 'hexadecimal', 'octal', or 'binary'. As with 'integer',
-- a leading @0@ is interpreted as octal. Leading signs are not accepted.
natural :: (CharParsing m, Monad m) => m Natural
natural = fromIntegral <$> choice [try hexadecimal, try octal, try binary, decimal]

-- | Parse an integer in base 10.
--
-- Accepts @0..9@ and underscore separators. No leading signs are accepted.
decimal :: (CharParsing m, Monad m) => m Integer
decimal = do
  contents <- withUnder digit
  attempt contents

-- | Parse a number in hexadecimal.
--
-- Requires a @0x@ or @0X@ prefix. No leading signs are accepted.
-- Accepts @A..F@, @a..f@, @0..9@ and underscore separators.
hexadecimal :: forall a m. (Eq a, Num a, CharParsing m, Monad m) => m a
hexadecimal = do
  void (string "0x" <|> string "0X")
  contents <- withUnder hexDigit
  let res = readHex contents
  case res of
    [] -> unexpected ("unparsable hex literal " <> contents)
    [(x, "")] -> pure x
    _ -> unexpected ("ambiguous hex literal " <> contents)

-- | Parse a number in octal.
--
-- Requires a @0@, @0o@ or @0O@ prefix. No leading signs are accepted.
-- Accepts @0..7@ and underscore separators.
octal :: forall a m. (Num a, CharParsing m, Monad m) => m a
octal = do
  void (char '0' *> optional (oneOf "oO"))
  digs <- withUnder octDigit
  fromIntegral <$> attempt @Integer ("0o" <> digs)

-- | Parse a number in binary.
--
-- Requires a @0b@ or @0B@ prefix. No leading signs are accepted.
-- Accepts @0@, @1@, and underscore separators.
binary :: forall a m. (Show a, Num a, CharParsing m, Monad m) => m a
binary = do
  void (char '0')
  void (optional (oneOf "bB"))
  digs <- withUnder (oneOf "01")
  let c2b c = case c of
        '0' -> 0
        '1' -> 1
        x -> error ("Invariant violated: both Attoparsec and readInt let a bad digit through: " <> [x])
  let res = readInt 2 (`elem` "01") c2b digs
  case res of
    [] -> unexpected ("No parse of binary literal: " <> digs)
    [(x, "")] -> pure x
    others -> unexpected ("Too many parses of binary literal: " <> show others)

-- | Parse an arbitrary-precision number with an optional decimal part.
--
-- Unlike 'scientificP' or Scientific's 'Read' instance, this handles:
--
--   * omitted whole parts, e.g. @.5@
--   * omitted decimal parts, e.g. @5.@
--   * exponential notation, e.g. @3.14e+1@
--   * numeric parts, in whole or decimal or exponent parts, with @_@ characters
--   * hexadecimal, octal, and binary integer literals, without a decimal part.
--
-- You may either omit the whole or the leading part, not both; this parser also rejects the empty string.
-- It does /not/ handle hexadecimal floating-point numbers.
floating :: (CharParsing m, Monad m) => m Scientific
floating = signed (choice [hexadecimal, octal, binary, dec])
  where
    -- Compared to the binary parser, this is positively breezy.
    dec = do
      -- Try getting the whole part of a floating literal.
      leadings <- stripUnder <$> many (digit <|> char '_')
      -- Try reading a dot.
      void (optional (char '.'))
      -- The trailing part...
      trailings <- stripUnder <$> many (digit <|> char '_')
      -- ...and the exponent.
      exponent <- stripUnder <$> many (oneOf "eE_0123456789+-")
      -- Ensure we don't read an empty string, or one consisting only of a dot and/or an exponent.
      when (null trailings && null leadings) (unexpected "Does not accept a single dot")
      -- Replace empty parts with a zero.
      let leads = if null leadings then "0" else leadings
      let trail = if null trailings then "0" else trailings
      attempt (leads <> "." <> trail <> exponent)

-- | Converts a numeric parser to one that accepts an optional leading sign.
signed :: forall a m . (CharParsing m, Num a) => m a -> m a
signed p =
  (negate <$> (char '-' *> p))
    <|> (char '+' *> p)
    <|> p

-- | Converts a numeric parser to one that accepts a trailing imaginary specifier
-- @i@ or @j@. This does not add facilities for two-valued literals, i.e. @1+4j@,
-- as those are generally best left to high-level expression facilities.
imaginary :: forall a m . (CharParsing m, Monad m, Num a) => m a -> m (Complex a)
imaginary num = do
  real <- num
  void (oneOf "ij")
  pure (0 :+ real)

stripUnder :: String -> String
stripUnder = Prelude.filter (/= '_')

attempt :: (Read a, CharParsing m) => String -> m a
attempt str = maybe (unexpected ("No parse: " <> str)) pure (readMaybe str)

withUnder :: CharParsing m => m Char -> m String
withUnder p = stripUnder <$> ((:) <$> p <*> many (p <|> char '_'))

max_stars_repo_name: patrickt/flexible-numeric-parsers

max_issues_repo_name: patrickt/flexible-numeric-parsers

max_forks_repo_name: patrickt/flexible-numeric-parsers


text:
-- Lib.hs

module Lib
    ( symbol
    , readExp
    , spaces
    ) where

import Text.ParserCombinators.Parsec hiding (spaces)
import Numeric (readHex, readOct, readFloat)
import Data.Char (digitToInt)
import qualified Data.Complex as C
import Data.Ratio ((%))

data LispVal = Atom String
             | List [LispVal]
             | DottedList [LispVal] LispVal
             | Complex (C.Complex Double)
             | Integer Integer
             | Real Double
             | Rational Rational
             | String String
             | Char Char
             | Bool Bool

readExp :: String -> String
readExp input =
  case parse parseExpr "lisp" input of
    Left err -> "No match: " ++ show err
    Right _  -> "Found value"

parseExpr :: Parser LispVal
parseExpr = parseAtom
        <|> parseString
        <|> parseNumber
        <|> try parseChar
        <|> parseQuoted
        <|> do _ <- char '('
               x <- try parseList <|> try parseDottedList
               _ <- char ')'
               return x
        <|> parseQuasiQuoted
        <|> parseUnquote

parseAtom :: Parser LispVal
parseAtom =
  do first <- letter <|> symbol
     rest <- many $ letter <|> digit <|> symbol
     let atom = first:rest
     return $ case atom of
                "#t" -> Bool True
                "#f" -> Bool False
                _    -> Atom atom

parseString :: Parser LispVal
parseString =
  do _ <- char '"'
     x <- many $ escapedChars <|> noneOf "\""
     _ <- char '"'
     return $ String x

parseNumber :: Parser LispVal
parseNumber = try parseComplex
          <|> try parseReal
          <|> try parseRational
          <|> try parseInteger

parseInteger :: Parser LispVal
parseInteger = parseDec
           <|> parseDec2
           <|> parseOct
           <|> parseHex
           <|> parseBin

parseOct :: Parser LispVal
parseOct = try (string "#o") >> many1 octDigit >>= return . Integer . octToDig
  where octToDig = fst . head . readOct

parseDec :: Parser LispVal
parseDec = Integer . read <$> many1 digit

parseDec2 :: Parser LispVal
parseDec2 = try (string "#d") >> Integer . read <$> many1 digit

parseHex :: Parser LispVal
parseHex = try (string "#x") >> many1 hexDigit >>= return . Integer . hexToDig
  where hexToDig = fst . head . readHex

parseBin :: Parser LispVal
parseBin =
  try (string "#b") >>
  many1 (oneOf "10") >>=
  return . Integer . binToDig
-- x0 + 2 * (x1 + 2 * (x2 + 2 * (x3 + 2 * x4)))
  where binToDig "" = 0
        binToDig s = foldr f 0 ds
          where ds      = (toInteger . digitToInt) <$> reverse s
                f x acc = x + 2 * acc

parseChar :: Parser LispVal
parseChar = parseChar1 <|> parseChar2

parseChar1 :: Parser LispVal
parseChar1 =
  do _ <- try (string "#\\")
     x <- anyChar >>= \c -> notFollowedBy alphaNum >> return c
     return . Char $ x

parseChar2 :: Parser LispVal
parseChar2 =
  try (string "#\\") >>
  try (string "newline" <|> string "space") >>= \x ->
  return . Char $ case x of
                    "newline" -> '\n'
                    "space"   -> ' '

parseReal :: Parser LispVal
parseReal =
  do x <- many1 digit
     _ <- char '.'
     y <- many1 digit
     return . Real . fst . head . readFloat $ x ++ ['.'] ++ y

parseRational :: Parser LispVal
parseRational =
  do n <- many1 digit
     _ <- char '/'
     d <- many1 digit
     return . Rational $ read d % read n

parseComplex :: Parser LispVal
parseComplex =
  do r <- try $ parseReal <|> parseDec
     _ <- char '+'
     i <- try $ parseReal <|> parseDec
     _ <- char 'i'
     return . Complex $ (toDouble r) C.:+ (toDouble i)
  where toDouble (Real f)    = realToFrac f
        toDouble (Integer f) = fromIntegral f

parseList :: Parser LispVal
parseList = List <$> sepBy parseExpr spaces

parseDottedList :: Parser LispVal
parseDottedList =
  do h <- endBy parseExpr spaces
     t <- char '.' >> spaces >> parseExpr
     return $ DottedList h t

parseQuoted :: Parser LispVal
parseQuoted =
  do _ <- char '\''
     x <- parseExpr
     return $ List [Atom "quote", x]

parseQuasiQuoted :: Parser LispVal
parseQuasiQuoted =
  do _ <- char '`'
     e <- parseExpr
     return $ List [Atom "quasiquote", e]

parseUnquote :: Parser LispVal
parseUnquote =
  do _ <- char ','
     e <- parseExpr
     return $ List [Atom "unquote", e]

symbol :: Parser Char
symbol = oneOf "!#$%&|*+-/:<=>?@^_~"

spaces :: Parser ()
spaces = skipMany1 space

escapedChars :: Parser Char
escapedChars = char '\\' >>
               oneOf "\\\"\n\r\t" >>= \x ->
               return $ case x of
                          '\\' -> x
                          '"'  -> x
                          'n'  -> '\n'
                          'r'  -> '\r'

max_stars_repo_name: edgarlepe/scheme

max_issues_repo_name: edgarlepe/scheme

max_forks_repo_name: edgarlepe/scheme


text:
-- H24.12.02 22:30 Close it.
module Main where
import System.Environment
import Text.ParserCombinators.Parsec hiding (spaces)
import Control.Monad
import Numeric
import Data.Ratio
import Data.Complex
import Data.Array
import Control.Monad.Error
import System.IO
import Data.IORef

main :: IO ()
main = do args <- getArgs
          case length args of
            0 -> runRepl
            1 -> runOnce $ args !! 0
            otherwise -> putStrLn "Program takes only 0 or 1 argument"

flushStr :: String -> IO ()
flushStr str = putStr str >> hFlush stdout

readPrompt :: String -> IO String
readPrompt prompt = flushStr prompt >> getLine

evalAndPrint :: Env -> String -> IO ()
evalAndPrint env expr = evalString env expr >>= putStrLn

evalString :: Env -> String -> IO String
evalString env expr = runIOThrows $ liftM show $
    (liftThrows $ readExpr expr) >>= eval env

until_ :: Monad m => (a -> Bool) -> m a -> (a -> m ()) -> m ()
until_ pred prompt action = do
  result <- prompt
  if pred result
     then return ()
     else action result >> until_ pred prompt action

runOnce :: String -> IO ()
runOnce expr = primitiveBindings >>= flip evalAndPrint expr

runRepl :: IO ()
runRepl = primitiveBindings >>=
    until_ (== "quit") (readPrompt "Lisp>>> ") . evalAndPrint

data LispVal = Atom String
             | List [LispVal]
             | DottedList [LispVal] LispVal
             | Number Integer
             | String String
             | Bool Bool
             | Character Char
             | Float Double
             | Ratio Rational
             | Complex (Complex Double)
             | Vector (Array Int LispVal)
             | PrimitiveFunc ([LispVal] -> ThrowsError LispVal)
             | Func {params :: [String], vararg :: (Maybe String),
                      body :: [LispVal], closure :: Env}

spaces :: Parser ()
spaces = skipMany1 space

symbol :: Parser Char
symbol = oneOf "!#$%&|*+-/:<=>?@^_~"
symbol1 :: Parser Char
symbol1 = oneOf "!$%&|*+-/:<=>?@^_~"

parseAtom :: Parser LispVal
parseAtom = do first <- letter <|> symbol1
               rest <- many (letter <|> digit <|> symbol)
               return . Atom $ first : rest

parseBool :: Parser LispVal
parseBool = liftM Bool ((char 't' >> return True)
                         <|> (char 'f' >> return False))

parseSigned :: (Real r) => ReadS r -> Parser [Char] -> Parser r
parseSigned reader lexer = liftM (fst . head . readSigned reader) (do
  sign <- char '-' <|> return '0'
  num <- lexer
  return (sign : num))

lexFloat :: Parser [Char]
lexFloat = do x <- many1 digit
              char '.'
              y <- many digit
              return $ x ++ "." ++ y

parseFloat :: Parser LispVal
parseFloat = (liftM Float $ parseSigned readFloat lexFloat)

parseDigits :: Parser LispVal
parseDigits = (liftM Number $ parseSigned readDec (many1 digit))

parseRatio :: Parser LispVal
parseRatio = do
  x <- parseSigned readDec (many1 digit)
  char '/'
  y <- parseSigned readDec (many1 digit)
  return $ Ratio (x % y)

toDouble :: LispVal -> Double
toDouble (Number n) = fromIntegral n
toDouble (Float f) = f

parseComplex :: Parser LispVal
parseComplex = do x <- (try parseFloat <|> parseDigits)
                  char '+'
                  y <- (try parseFloat <|> parseDigits)
                  char 'i'
                  return $ Complex (toDouble x :+ toDouble y)

parseNum :: Parser LispVal
parseNum = try parseComplex
           <|> try parseRatio
           <|> try parseFloat
           <|> parseDigits

parseNumber :: Parser LispVal
parseNumber = liftM Number
                (parseBin <|> parseOct <|> parseDig <|> parseHex)
  where
        parseBin = char 'b' >> parseSigned readBin (many1 $ oneOf "01")
        parseOct = char 'o' >> parseSigned readOct (many1 octDigit)
        parseDig = char 'd' >> parseSigned readDec (many1 digit)
        parseHex = char 'x' >> parseSigned readHex (many1 hexDigit)
        readBin  = readInt 2 (\x->x=='0'||x=='1') (\x->if x=='0' then 0 else 1)

parseSharpSyntax :: Parser LispVal
parseSharpSyntax = char '#' >> (parseBool
                     <|> parseNumber
                     <|> parseCharacter
                     <|> parseVector)

escapedChars :: Parser Char
escapedChars = char '\\' >> oneOf "\\\"nrt" >>= \x ->
               return $ case x of
                 'n' -> '\n'
                 'r' -> '\r'
                 't' -> '\t'
                 _ -> x

parseString :: Parser LispVal
parseString = do char '"'
                 x <- many (escapedChars <|> noneOf "\"\\")
                 char '"'
                 return $ String x

parseCharacter :: Parser LispVal
parseCharacter = liftM Character (char '\\' >>
                   option ' ' (try (string "space" >> return ' ')
                                <|> try (string "newline" >> return '\n')
                                <|> do { x <- anyChar;
                                         notFollowedBy alphaNum;
                                         return x}))

parseList :: Parser LispVal
parseList = do
  body <- sepEndBy parseExpr spaces
  if length body == 0 then return $ List []
  else parseTail >>= \tail-> return $
    case tail of
      Just (List l)         -> List (body ++ l)
      Just (DottedList h t) -> DottedList (body ++ h) t
      Just val              -> DottedList body val
      Nothing               -> List body

parseTail :: Parser (Maybe LispVal)
parseTail = liftM Just (char '.' >> skipMany space >> parseExpr)
            <|> return Nothing

makeSyntaxSugarParser :: Parser a -> String -> Parser LispVal
makeSyntaxSugarParser ch symbol =
  ch >> parseExpr >>= (\x -> return $ List [Atom symbol, x])

parseQuoted :: Parser LispVal
parseQuoted = makeSyntaxSugarParser (char '\'') "quote"

parseQusaiquote :: Parser LispVal
parseQusaiquote = makeSyntaxSugarParser (char '`') "qusaiquote"

parseUnquote :: Parser LispVal
parseUnquote = makeSyntaxSugarParser (char ',') "unquote"

parseUnquoteS :: Parser LispVal
parseUnquoteS = makeSyntaxSugarParser (string ",@") "unquote-splicing"

parens :: Parser a -> Parser a
parens = between (char '(' >> skipMany space) (char ')')

parseVector :: Parser LispVal
parseVector = do
  arrayValues <- parens (sepBy parseExpr spaces)
  return $ Vector (listArray (0,(length arrayValues - 1)) arrayValues)

parseExpr :: Parser LispVal
parseExpr = parseSharpSyntax
        <|> try parseNum
        <|> parseAtom
        <|> parseString
        <|> parseQuoted <|> parseQusaiquote
        <|> try parseUnquoteS <|> parseUnquote
        <|> parens parseList

readExpr :: String -> ThrowsError LispVal
readExpr input = case parse parseExpr "lisp" input of
                   Left err -> throwError $ Parser err
                   Right val -> return val

showVal :: LispVal -> String
showVal (String contents) = "\"" ++ contents ++ "\""
showVal (Atom name) = name
showVal (Number contents) = show contents
showVal (Float contents) = showFFloat Nothing contents ""
showVal (Ratio contents) = show (numerator contents) ++ "/"
                         ++ show (denominator contents)
showVal (Complex x) = showFFloat Nothing (realPart x) $ "+"
                    ++ showFFloat Nothing (imagPart x) "i"
showVal (Bool True) = "#t"
showVal (Bool False) = "#f"
showVal (Character '\n') = "#\\newline"
showVal (Character ' ') = "#\\space"
showVal (Character char) = "#\\" ++ [char]
showVal (List contents) = "(" ++ unwordsList contents ++ ")"
showVal (DottedList head tail) = "(" ++ unwordsList head
                               ++ " . " ++ showVal tail ++ ")"
showVal (Vector contents) = "#(" ++ unwordsList (elems contents) ++ ")"
showVal (PrimitiveFunc _) = "<primitive>"
showVal (Func {params = args, vararg = varargs, body = body, closure = env}) =
  "(lambda (" ++ unwords (map show args) ++
    (case varargs of
       Nothing -> ""
       Just arg -> " . " ++ arg) ++ ") ...)"

unwordsList :: [LispVal] -> String
unwordsList = unwords . map showVal

instance Show LispVal where show = showVal

eval :: Env -> LispVal -> IOThrowsError LispVal
eval env val@(String _) = return val
eval env val@(Number _) = return val
eval env val@(Bool _) = return val
eval env (Atom id) = getVar env id
eval env (List [Atom "quote", val]) = return val
eval env (List [Atom "if", pred, conseq, alt]) = do
  result <- eval env pred
  case result of
    Bool False -> eval env alt
    otherwise -> eval env conseq
eval env (List [Atom "set!", Atom var, form]) =
    eval env form >>= setVar env var
eval env (List [Atom "define", Atom var, form]) =
    eval env form >>= defineVar env var
eval env (List (Atom "define" : List (Atom var : params) : body)) =
    makeNormalFunc env params body >>= defineVar env var
eval env (List (Atom "define" : DottedList(Atom var : params) varargs : body)) =
    makeVarargs varargs env params body >>= defineVar env var
eval env (List (Atom "lambda" : List params : body)) =
    makeNormalFunc env params body
eval env (List (Atom "lambda" : DottedList params varargs : body)) =
    makeVarargs varargs env params body
eval env (List (Atom "lambda" : varargs@(Atom _) : body)) =
    makeVarargs varargs env [] body
eval env (List (function : args)) = do
    func <- eval env function
    argVars <- mapM (eval env) args
    apply func argVars
eval env badForm =
    throwError $ BadSpecialForm "Unrecognized special form" badForm

apply :: LispVal -> [LispVal] -> IOThrowsError LispVal
apply (PrimitiveFunc func) args = liftThrows $ func args
apply (Func params varargs body closure) args =
    if num params /= num args && varargs == Nothing
       then throwError $ NumArgs (num params) args
       else (liftIO $ bindVars closure $ zip params args)
                >>= bindVarArgs varargs >>= evalBody
    where remainingArgs = drop (length params) args
          num = toInteger . length
          evalBody env = liftM last $ mapM (eval env) body
          bindVarArgs arg env = case arg of
              Just argName -> liftIO $
                  bindVars env [(argName, List $ remainingArgs)]
              Nothing -> return env

makeFunc varargs env params body =
    return $ Func (map showVal params) varargs body env

makeNormalFunc = makeFunc Nothing

makeVarargs = makeFunc . Just . showVal

primitives :: [(String, [LispVal] -> ThrowsError LispVal)]
primitives = [ ("+", numericBinop (+)),
               ("-", numericBinop (-)),
               ("*", numericBinop (*)),
               ("/", numericBinop div),
               ("mod", numericBinop mod),
               ("quotient", numericBinop quot),
               ("remainder", numericBinop rem),
               ("boolean?", unaryOp booleanp),
               ("symbol?", unaryOp symbolp),
{-               ("char?", unaryOp charp),
               ("vector?", unaryOp vectorp),
--               ("procedure?", unaryOp procedurep),-}
               ("pair?", unaryOp pairp),
               ("number?", unaryOp numberp),{-
               ("complex?", unaryOp complexp),
               ("real?", unaryOp realp),
               ("rational?", unaryOp rationalp),
               ("integer?", unaryOp integerp),-}
               ("string?", unaryOp stringp),
--               ("port?", unaryOp portp),-}
               ("symbol->string", unaryOp symbol2string),
               ("string->symbol", unaryOp string2symbol),
               ("=", numBoolBinop (==)),
               ("/=", numBoolBinop (/=)),
               ("<", numBoolBinop (<)),
               (">", numBoolBinop (>)),
               ("<=", numBoolBinop (<=)),
               (">=", numBoolBinop (>=)),
               ("&&", boolBoolBinop (&&)),
               ("||", boolBoolBinop (||)),
               ("string=?", strBoolBinop (==)),
               ("string<?", strBoolBinop (<)),
               ("string>?", strBoolBinop (>)),
               ("string<=?", strBoolBinop (<=)),
               ("string>=?", strBoolBinop (>=)),
               ("car", car),
               ("cdr", cdr),
               ("cons", cons),
               ("eq?", eqv),
               ("eqv?", eqv),
               ("equal?", equal)
             ]

numericBinop :: (Integer -> Integer -> Integer) -> [LispVal]
                -> ThrowsError LispVal
numericBinop op singleVal@[_] = throwError $ NumArgs 2 singleVal
numericBinop op params = mapM unpackNum params >>= return . Number . foldl1 op

unpackNum :: LispVal -> ThrowsError Integer
unpackNum (Number n) = return n
unpackNum nan = throwError $ TypeMismatch "number" nan

unaryOp :: (LispVal -> LispVal) -> [LispVal] -> ThrowsError LispVal
unaryOp f [v] = return $ f v

symbolp, numberp, stringp, booleanp, pairp :: LispVal -> LispVal
symbolp (Atom _)   = Bool True
symbolp _          = Bool False
numberp (Number _) = Bool True
numberp _          = Bool False
stringp (String _) = Bool True
stringp _          = Bool False
booleanp (Bool _)  = Bool True
booleanp   _       = Bool False
pairp (List _)     = Bool True
pairp (DottedList _ _) = Bool True
pairp _            = Bool False

symbol2string, string2symbol :: LispVal -> LispVal
symbol2string (Atom s)   = String s
symbol2string _          = String ""
string2symbol (String s) = Atom s
string2symbol _          = Atom ""
--booleanp, symbolp, charp, vectorp, {-procedurep,-}
--  pairp, numberp, complexp, realp, rationalp, integerp, stringp{-, portp-} :: LispVal -> Bool
--booleanp

boolBinop :: (LispVal -> ThrowsError a) -> (a -> a -> Bool)
             -> [LispVal] -> ThrowsError LispVal
boolBinop unpacker op args = if length args /= 2
                             then throwError $ NumArgs 2 args
                             else do left <- unpacker $ args !! 0
                                     right <- unpacker $ args !! 1
                                     return $ Bool $ left `op` right

numBoolBinop = boolBinop unpackNum
strBoolBinop = boolBinop unpackStr
boolBoolBinop = boolBinop unpackBool

unpackStr :: LispVal -> ThrowsError String
unpackStr (String s) = return s
unpackStr (Number s) = return $ show s
unpackStr (Bool s) = return $ show s
unpackStr notStr = throwError $ TypeMismatch "string" notStr

unpackBool :: LispVal -> ThrowsError Bool
unpackBool (Bool b) = return b
unpackBool (Number 0) = return False
unpackBool (Number 1) = return True
unpackBool notBool = throwError $ TypeMismatch "boolean" notBool

car :: [LispVal] -> ThrowsError LispVal
car [List (x : xs)] = return x
car [DottedList (x : xs) _] = return x
car [badArg] = throwError $ TypeMismatch "pair" badArg
car badArgList = throwError $ NumArgs 1 badArgList

cdr :: [LispVal] -> ThrowsError LispVal
cdr [List (x : xs)] = return $ List xs
cdr [DottedList [xs] x] = return x
cdr [DottedList (_ : xs) x] = return $ DottedList xs x
cdr [badArg] = throwError $ TypeMismatch "pair" badArg
cdr badArgList = throwError $ NumArgs 1 badArgList

cons :: [LispVal] -> ThrowsError LispVal
cons [x, List []] = return $ List [x]
cons [x, List xs] = return $ List (x : xs)
cons [x, DottedList xs xlast] = return $ DottedList (x : xs) xlast
cons [x1, x2] = return $ DottedList [x1] x2
cons badArgList = throwError $ NumArgs 2 badArgList

eqv :: [LispVal] -> ThrowsError LispVal
eqv [(Bool x), (Bool y)] = return $ Bool $ x == y
eqv [(Number x), (Number y)] = return $ Bool $ x == y
eqv [(String x), (String y)] = return $ Bool $ x == y
eqv [(Atom x), (Atom y)] = return $ Bool $ x == y
eqv [(DottedList xs x), (DottedList ys y)] =
    eqv [List $ xs ++ [x], List $ ys ++ [y]]
eqv [(List x), (List y)] =
    return $ Bool $ (length x == length y) && (all eqvPair $ zip x y)
  where eqvPair (x1, x2) = case eqv [x1, x2] of
                             Left err -> False
                             Right (Bool val) -> val
eqv [_, _] = return $ Bool False
eqv badArgList = throwError $ NumArgs 2 badArgList

data Unpacker = forall a. Eq a => AnyUnpacker (LispVal -> ThrowsError a)
unpackEquals :: LispVal -> LispVal -> Unpacker -> ThrowsError Bool
unpackEquals x y (AnyUnpacker unpacker) =
  do
    unpacked1 <- unpacker x
    unpacked2 <- unpacker y
    return $ unpacked1 == unpacked2
  `catchError` (const $ return False)

equal :: [LispVal] -> ThrowsError LispVal
equal [x, y] = do
  primitiveEquals <- liftM or $ mapM (unpackEquals x y)
    [AnyUnpacker unpackNum, AnyUnpacker unpackStr, AnyUnpacker unpackBool]
  eqvEquals <- eqv [x, y]
  return $ Bool $ (primitiveEquals || let (Bool x) = eqvEquals in x)
equal badArgList = throwError $ NumArgs 2 badArgList

data LispError = NumArgs Integer [LispVal]
               | TypeMismatch String LispVal
               | Parser ParseError
               | BadSpecialForm String LispVal
               | NotFunction String String
               | UnboundVar String String
               | Default String

showError :: LispError -> String
showError (UnboundVar message varname) = message ++ ": " ++ varname
showError (BadSpecialForm message form) = message ++ ": " ++ show form
showError (NotFunction message func) = message ++ ": " ++ show func
showError (NumArgs expected found) = "Expected: " ++ show expected
                                   ++ ", found " ++ show found
showError (TypeMismatch expected found) = "Invalid type: expected " ++ expected
                                        ++ ", found " ++ show found
showError (Parser parseErr) = "Parse error at " ++ show parseErr

instance Show LispError where show = showError
instance Error LispError where
  noMsg = Default "An error has occurred"
  strMsg = Default
type ThrowsError = Either LispError

--trapError ::
trapError action = catchError action (return . show)

extractValue :: ThrowsError a -> a
extractValue (Right val) = val

type Env = IORef [(String, IORef LispVal)]

nullEnv :: IO Env
nullEnv = newIORef []

primitiveBindings :: IO Env
primitiveBindings =
    nullEnv >>= (flip bindVars $ map makePrimitiveFunc primitives)
  where makePrimitiveFunc (var, func) = (var, PrimitiveFunc func)

type IOThrowsError = ErrorT LispError IO

liftThrows :: ThrowsError a -> IOThrowsError a
liftThrows (Left err) = throwError err
liftThrows (Right val) = return val

runIOThrows :: IOThrowsError String -> IO String
runIOThrows action = runErrorT (trapError action) >>= return . extractValue

isBound :: Env -> String -> IO Bool
isBound envRef var =
    readIORef envRef >>= return . maybe False (const True) . lookup var

getVar :: Env -> String -> IOThrowsError LispVal
getVar envRef var = do
  env <- liftIO $ readIORef envRef
  maybe (throwError $ UnboundVar "Getting an unbound variable" var)
        (liftIO . readIORef)
        (lookup var env)

setVar :: Env -> String -> LispVal -> IOThrowsError LispVal
setVar envRef var value = do
  env <- liftIO $ readIORef envRef
  maybe (throwError $ UnboundVar "Setting an unbound variable" var)
        (liftIO . (flip writeIORef value))
        (lookup var env)
  return value

defineVar :: Env -> String -> LispVal -> IOThrowsError LispVal
defineVar envRef var value = do
  alreadyDefined <- liftIO $ isBound envRef var
  if alreadyDefined
     then setVar envRef var value >> return value
     else liftIO $ do
       valueRef <- newIORef value
       env <- readIORef envRef
       writeIORef envRef ((var, valueRef) : env)
       return value

bindVars :: Env -> [(String, LispVal)] -> IO Env
bindVars envRef bindings = readIORef envRef >>= extendEnv bindings >>= newIORef
  where extendEnv bindings env = liftM (++ env) (mapM addBinding bindings)
        addBinding (var, value) = do ref <- newIORef value
                                     return (var, ref)

-- ghc -package parsec -fglasgow-exts -o 8function.exe --make 8function.hs
-- 8function.exe
-- Lisp>>> (define (f x y) (+ x y))
-- Lisp>>> (f 1 2)
-- Lisp>>> (f 1 2 3)
-- Lisp>>> (f 1)
-- Lisp>>> (define (factorial x) (if (= x 1) 1 (* x (factorial (- x 1)))))
-- Lisp>>> (factorial 10)
-- Lisp>>> (define (counter inc) (lambda (x) (set! inc (+ x inc)) inc))
-- Lisp>>> (define my-count (counter 5))
-- Lisp>>> (my-count 3)
-- Lisp>>> (my-count 6)
-- Lisp>>> (my-count 5)
-- Lisp>>> quit