{-# LANGUAGE PolyKinds #-}
{-# LANGUAGE TemplateHaskell #-}
{-# LANGUAGE UndecidableInstances #-}
{-# LANGUAGE DataKinds #-}

-- | Internal representation of Feldspar programs

module Feldspar.Representation where

import Control.Monad.Reader
import Data.Complex
import Data.Hash (hashInt)
import Data.Int
import Data.List (genericTake)
import Data.Proxy
import Data.Typeable (Typeable)
import Data.Word

import Data.Constraint (Dict (..))

import Language.Syntactic
import Language.Syntactic.Functional
import Language.Syntactic.Functional.Tuple
import Language.Syntactic.TH

import qualified Control.Monad.Operational.Higher as Operational

import qualified Language.Embedded.Expression as Imp
import qualified Language.Embedded.Imperative.CMD as Imp
import qualified Language.Embedded.Backend.C.Expression as Imp

import Data.Inhabited
import Data.Selection
import Data.TypedStruct
import Feldspar.Primitive.Representation
import Feldspar.Primitive.Backend.C ()

--------------------------------------------------------------------------------
-- * Object-language types
--------------------------------------------------------------------------------

-- | Representation of all supported types
type TypeRep = Struct PrimType' PrimTypeRep

-- | Supported types
class (Eq a, Show a, Typeable a, Inhabited a) => Type a
  where
    -- | Reify a type
    typeRep :: TypeRep a

instance Type Bool             where typeRep = Single BoolT
instance Type Int8             where typeRep = Single Int8T
instance Type Int16            where typeRep = Single Int16T
instance Type Int32            where typeRep = Single Int32T
instance Type Int64            where typeRep = Single Int64T
instance Type Word8            where typeRep = Single Word8T
instance Type Word16           where typeRep = Single Word16T
instance Type Word32           where typeRep = Single Word32T
instance Type Word64           where typeRep = Single Word64T
instance Type Float            where typeRep = Single FloatT
instance Type Double           where typeRep = Single DoubleT
instance Type (Complex Float)  where typeRep = Single ComplexFloatT
instance Type (Complex Double) where typeRep = Single ComplexDoubleT
instance (Type a, Type b) => Type (a,b) where typeRep = Two typeRep typeRep

-- | Alias for the conjunction of 'PrimType'' and 'Type'
class    (PrimType' a, Type a) => PrimType a
instance (PrimType' a, Type a) => PrimType a

instance Imp.CompTypeClass PrimType
  where
    compType _ = Imp.compType (Proxy :: Proxy PrimType')
    compLit _  = Imp.compLit (Proxy :: Proxy PrimType')
  -- This instance is used by <https://github.com/kmate/raw-feldspar-mcs>

-- | Convert any 'Struct' with a 'PrimType' constraint to a 'TypeRep'
toTypeRep :: Struct PrimType' c a -> TypeRep a
toTypeRep = mapStruct (const primTypeRep)

-- | Check whether two type representations are equal
typeEq :: TypeRep a -> TypeRep b -> Maybe (Dict (a ~ b))
typeEq (Single t) (Single u) = primTypeEq t u
typeEq (Two t1 t2) (Two u1 u2) = do
    Dict <- typeEq t1 u1
    Dict <- typeEq t2 u2
    return Dict
typeEq _ _ = Nothing

-- | Reflect a 'TypeRep' to a 'Typeable' constraint
witTypeable :: TypeRep a -> Dict (Typeable a)
witTypeable (Single t)
    | Dict <- witPrimType t
    = Dict
witTypeable (Two ta tb)
    | Dict <- witTypeable ta
    , Dict <- witTypeable tb
    = Dict

-- | Representation of supported value types + N-ary functions over such types
data TypeRepFun a
  where
    ValT :: TypeRep a -> TypeRepFun a
    FunT :: TypeRep a -> TypeRepFun b -> TypeRepFun (a -> b)
  -- Another option would have been to make `FunT` a constructor in `TypeRep`.
  -- That would have got rid of the extra layer at the expense of less accurate
  -- types (functions would be allowed in pairs, etc.). The current design was
  -- chosen in order to be able to reuse `Struct` instead of making `TypeRep` a
  -- new data type.

-- | Check whether two type representations are equal
typeEqFun :: TypeRepFun a -> TypeRepFun b -> Maybe (Dict (a ~ b))
typeEqFun (ValT t)     (ValT u)     = typeEq t u
typeEqFun (FunT ta tb) (FunT ua ub) = do
    Dict <- typeEq ta ua
    Dict <- typeEqFun tb ub
    return Dict
typeEqFun _ _ = Nothing

-- | Mutable variable
newtype Ref a = Ref { unRef :: Struct PrimType' Imp.Ref (Internal a) }
  -- A reference to a tuple is a struct of smaller references. This means that
  -- creating a reference to a tuple will generate several calls to generate new
  -- references. This must be done already in the front end, which means that
  -- the work in the back end becomes simpler.
  --
  -- Another option would be to allow a single reference to refer to a tuple,
  -- and then turn that into smaller references in the back end. However, this
  -- would complicate the back end for no obvious reason. (One way to do it
  -- would be to run the back end in a store that maps each front end reference
  -- to a struct of small references. Among other things, that would require
  -- dynamic typing.)

-- | Reference specialized to 'Data' elements
type DRef a = Ref (Data a)

instance Eq (Ref a)
  where
    Ref a == Ref b = and $ zipListStruct (==) a b

-- | Mutable array
data Arr a = Arr
    { arrOffset :: Data Index
    , arrLength :: Data Length
    , unArr     :: Struct PrimType' (Imp.Arr Index) (Internal a)
    }
  -- `arrOffset` gives the offset into the internal arrays. The user should not
  -- be able to access the offset or the internal arrays.
  --
  -- `arrLength` gives the maximum index+1 according to the view provided by
  -- `arrIx`. Thus, the maximum index in the internal arrays is
  -- `arrLength + arrOffset - 1`.
  -- `arrLength` should not be exported directly to prevent the user from using
  -- record syntax to change the length of an array. But an equivalent function
  -- can be exported.

  -- An array of tuples is represented as a struct of smaller arrays. See
  -- comment to `Ref`.

-- | Array specialized to 'Data' elements
type DArr a = Arr (Data a)

-- | '==' checks if two 'Arr' use the same physical array. The length and offset
-- are ignored.
instance Eq (Arr a)
  where
    Arr _ _ arr1 == Arr _ _ arr2 = and (zipListStruct (==) arr1 arr2)

-- | Immutable array
data IArr a = IArr
    { iarrOffset :: Data Index
    , iarrLength :: Data Length
    , unIArr     :: Struct PrimType' (Imp.IArr Index) (Internal a)
    }

-- | Immutable array specialized to 'Data' elements
type DIArr a = IArr (Data a)

-- | Check if an 'Arr' and and 'IArr' use the same physical array. The length
-- and offset are ignored. This operation may give false negatives, but never
-- false positives. Whether or not false negatives occur may also depend on the
-- interpretation of the program.
--
-- Due to this unreliability, the function should only be used to affect the
-- non-functional properties of a program (e.g. to avoid unnecessary array
-- copying).
unsafeEqArrIArr :: Arr a -> IArr a -> Bool
unsafeEqArrIArr (Arr _ _ arr1) (IArr _ _ arr2) =
    and (zipListStruct sameId arr1 arr2)
  where
    sameId :: Imp.Arr Index a -> Imp.IArr Index a -> Bool
    sameId (Imp.ArrComp a) (Imp.IArrComp i) = a==i
    sameId _ _ = False

--------------------------------------------------------------------------------
-- * Pure expressions
--------------------------------------------------------------------------------

-- | Assertion labels
data AssertionLabel
    = InternalAssertion
        -- ^ Internal assertion to guarantee meaningful results
    | LibraryAssertion String
        -- ^ Assertion related to a specific library
    | UserAssertion String
        -- ^ Assertion in user code. The default label for user assertions is
        --   @`UserAssertion` ""@
  deriving (Eq, Show)

-- | A selection that includes all labels defined as 'UserAssertion'
onlyUserAssertions :: Selection AssertionLabel
onlyUserAssertions = selectBy $ \l -> case l of
    UserAssertion _ -> True
    _ -> False

data ExtraPrimitive sig
  where
    -- Integer division assuming `divBalanced x y * y == x` (i.e. no remainder).
    -- The purpose of the assumption is to use it in simplifications.
    DivBalanced :: (Integral a, PrimType' a) =>
        ExtraPrimitive (a :-> a :-> Full a)

    -- Guard a value by an assertion
    GuardVal ::
        AssertionLabel -> String -> ExtraPrimitive (Bool :-> a :-> Full a)

instance Eval ExtraPrimitive
  where
    evalSym DivBalanced = \a b -> if a `mod` b /= 0
      then error $ unwords ["divBalanced", show a, show b, "is not balanced"]
      else div a b
    evalSym (GuardVal _ msg) = \cond a ->
        if cond then a else error $ "Feldspar assertion failure: " ++ msg

instance Equality ExtraPrimitive
  where
    equal DivBalanced    DivBalanced    = True
    equal (GuardVal _ _) (GuardVal _ _) = True
    equal _ _ = False

    hash DivBalanced    = hashInt 1
    hash (GuardVal _ _) = hashInt 2

-- | For loop
data ForLoop sig
  where
    ForLoop :: Type st =>
        ForLoop (Length :-> st :-> (Index -> st -> st) :-> Full st)

instance Eval ForLoop
  where
    evalSym ForLoop = \len init body ->
        foldl (flip body) init $ genericTake len [0..]

-- | Interaction with the IO layer
data Unsafe sig
  where
    -- Turn a program into a pure value
    UnsafePerform :: Comp (Data a) -> Unsafe (Full a)

instance Render Unsafe
  where
    renderSym (UnsafePerform _) = "UnsafePerform ..."

instance Eval Unsafe
  where
    evalSym s = error $ "eval: cannot evaluate unsafe operation " ++ renderSym s

-- | 'equal' always returns 'False'
instance Equality Unsafe
  where
    equal _ _ = False

type FeldConstructs
    =   BindingT
    :+: Let
    :+: Tuple
    :+: Primitive
    :+: ExtraPrimitive
    :+: ForLoop
    :+: Unsafe

type FeldDomain = FeldConstructs :&: TypeRepFun

newtype Data a = Data { unData :: ASTF FeldDomain a }

-- | Declaring 'Data' as syntactic sugar
instance Syntactic (Data a)
  where
    type Domain (Data a)   = FeldDomain
    type Internal (Data a) = a
    desugar = unData
    sugar   = Data

instance Syntactic (Struct PrimType' Data a)
    -- Note that this instance places no constraints on `a`. This is crucial in
    -- the way it is used in the rest of the code. It would be possible to
    -- define `desugar` and `sugar` in terms of the instance for pairs; however,
    -- that would require constraining `a`.
  where
    type Domain   (Struct PrimType' Data a) = FeldDomain
    type Internal (Struct PrimType' Data a) = a

    desugar (Single a) = unData a
    desugar (Two a b)  = sugarSymDecor (ValT $ Two ta tb) Pair a' b'
      where
        a' = desugar a
        b' = desugar b
        ValT ta = getDecor a'
        ValT tb = getDecor b'

    sugar a = case getDecor a of
        ValT (Single _)  -> Single $ Data a
        ValT (Two ta tb) ->
            Two (sugarSymDecor (ValT ta) Fst a) (sugarSymDecor (ValT tb) Snd a)

-- | Specialization of the 'Syntactic' class for the Feldspar domain
class    (Syntactic a, Domain a ~ FeldDomain, Type (Internal a)) => Syntax a
instance (Syntactic a, Domain a ~ FeldDomain, Type (Internal a)) => Syntax a

-- | Make a smart constructor for a symbol
sugarSymFeld
    :: ( Signature sig
       , fi         ~ SmartFun FeldDomain sig
       , sig        ~ SmartSig fi
       , FeldDomain ~ SmartSym fi
       , SyntacticN f fi
       , sub :<: FeldConstructs
       , Type (DenResult sig)
       )
    => sub sig -> f
sugarSymFeld = sugarSymDecor $ ValT typeRep

-- | Make a smart constructor for a symbol
sugarSymFeldPrim
    :: ( Signature sig
       , fi         ~ SmartFun FeldDomain sig
       , sig        ~ SmartSig fi
       , FeldDomain ~ SmartSym fi
       , SyntacticN f fi
       , sub :<: FeldConstructs
       , PrimType' (DenResult sig)
       )
    => sub sig -> f
sugarSymFeldPrim = sugarSymDecor $ ValT $ Single primTypeRep

instance Imp.FreeExp Data
  where
    type FreePred Data = PrimType'
    constExp = sugarSymFeldPrim . Lit
    varExp   = sugarSymFeldPrim . FreeVar

--------------------------------------------------------------------------------
-- * Monadic computations
--------------------------------------------------------------------------------

data AssertCMD fs a
  where
    Assert
        :: AssertionLabel
        -> exp Bool
        -> String
        -> AssertCMD (Operational.Param3 prog exp pred) ()

instance Operational.HFunctor AssertCMD
  where
    hfmap _ (Assert c cond msg) = Assert c cond msg

instance Operational.HBifunctor AssertCMD
  where
    hbimap _ g (Assert c cond msg) = Assert c (g cond) msg

instance Operational.InterpBi AssertCMD IO (Operational.Param1 PrimType')
  where
    interpBi (Assert _ cond msg) = do
        cond' <- cond
        unless cond' $ error $ "Assertion failed: " ++ msg

type CompCMD
  =               Imp.RefCMD
  Operational.:+: Imp.ArrCMD
  Operational.:+: Imp.ControlCMD
  Operational.:+: AssertCMD

-- | Monad for computational effects: mutable data structures and control flow
newtype Comp a = Comp
    { unComp ::
        Operational.Program CompCMD (Operational.Param2 Data PrimType') a
    }
  deriving (Functor, Applicative, Monad)

--------------------------------------------------------------------------------
-- Template Haskell instances
--------------------------------------------------------------------------------

deriveSymbol    ''ExtraPrimitive
deriveRender id ''ExtraPrimitive

deriveSymbol    ''ForLoop
deriveRender id ''ForLoop
deriveEquality  ''ForLoop

deriveSymbol ''Unsafe

--------------------------------------------------------------------------------
-- * Interpretation
--------------------------------------------------------------------------------

-- The stuff below depends on the TH-generated instances. From GHC 9, it seems
-- that generated instances are only in scope after the splice. The alternative
-- to move the splices up into the code is not possible (at least not easily)
-- because of the mutual dependencies between the different types. (Definitions
-- that occur after a splice are not in scope before the splice.)

instance EvalEnv ExtraPrimitive env
instance StringTree ExtraPrimitive

instance EvalEnv ForLoop env
instance StringTree ForLoop

instance EvalEnv Unsafe env
instance StringTree Unsafe

-- | Evaluate a closed expression
eval :: (Syntactic a, Domain a ~ FeldDomain) => a -> Internal a
eval = evalClosed . desugar
  -- Note that a `Syntax` constraint would rule out evaluating functions

instance Imp.EvalExp Data
  where
    evalExp = eval


{-# LANGUAGE BangPatterns #-}
{-# LANGUAGE DataKinds #-}
{-# LANGUAGE DefaultSignatures #-}
{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE FlexibleInstances #-}
{-# LANGUAGE GeneralizedNewtypeDeriving #-}
{-# LANGUAGE MultiParamTypeClasses #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE TypeFamilies #-}
{-# LANGUAGE TypeOperators #-}
{-# LANGUAGE UndecidableInstances #-}
{-# OPTIONS_HADDOCK hide #-}
module Data.Vector.Unboxing.Internal
  (Unboxable(Rep)
  ,Vector(UnboxingVector)
  ,MVector(UnboxingMVector)
  ,Generics(..)
  ,Enum(..)
  ,EnumRep(..)
  ,coerceVector
  ,liftCoercion
  ,vectorCoercion
  ,toUnboxedVector
  ,fromUnboxedVector
  ,coercionWithUnboxedVector
  ,coerceMVector
  ,liftCoercionM
  ,mVectorCoercion
  ,toUnboxedMVector
  ,fromUnboxedMVector
  ,coercionWithUnboxedMVector
  ) where
import Prelude hiding (Enum)
import qualified Prelude
import qualified Data.Vector.Generic as G
import qualified Data.Vector.Generic.Mutable as GM
import qualified Data.Vector.Unboxed as U
import qualified Data.Vector.Unboxed.Mutable as UM
import qualified Data.Vector.Fusion.Bundle as Bundle
import Data.Coerce
import Data.Type.Coercion
import Data.Int
import Data.Word
import qualified GHC.Generics
import Data.Type.Bool
import qualified Data.Complex
import qualified Data.Functor.Identity
import qualified Data.Functor.Const
import qualified Data.Functor.Compose
import qualified Data.Ord
import qualified Data.Semigroup
import qualified Data.Monoid
import GHC.Exts (IsList(..))
import Control.DeepSeq (NFData(..))
import Text.Read (Read(..),readListPrecDefault)
import GHC.TypeLits (TypeError,ErrorMessage(Text))

newtype Vector a = UnboxingVector (U.Vector (Rep a))
newtype MVector s a = UnboxingMVector (UM.MVector s (Rep a))

type instance G.Mutable Vector = MVector

-- | Types that can be stored in unboxed vectors ('Vector' and 'MVector').
--
-- You can define instances of this class like:
--
-- > newtype Foo = Foo Int
-- > instance Unboxable Foo where
-- >   type Rep Foo = Int
--
-- The type specified by 'Rep' needs to be an instance of 'U.Unbox',
-- and coercion must be possible between the two types.
--
-- Instances can also be derived with @GeneralizedNewtypeDeriving@.
-- GND always works if the base type is an instance of 'Unboxable'.
--
-- If you want to have non-trivial correspondence between the type and the representation,
-- use 'Generics' wrapper with @DerivingVia@.
--
-- Note that @UndecidableInstances@ is needed if you use GND or @DerivingVia@ to derive instances.
class U.Unbox (Rep a) => Unboxable a where
  -- | The underlying type of @a@.  Must be an instance of 'U.Unbox'.
  type Rep a

  -- Hidden members:

  -- Used by 'coerceVector'
  type CoercibleRep a
  type CoercibleRep a = Rep a

  -- True if both 'unboxingFrom and 'unboxingTo' are just 'coerce'
  type IsTrivial a :: Bool
  type IsTrivial a = 'True
  -- TODO: Use ConstraintKinds?

  unboxingFrom :: a -> Rep a
  default unboxingFrom :: Coercible a (Rep a) => a -> Rep a
  unboxingFrom = coerce
  {-# INLINE unboxingFrom #-}

  unboxingTo :: Rep a -> a
  default unboxingTo :: Coercible a (Rep a) => Rep a -> a
  unboxingTo = coerce
  {-# INLINE unboxingTo #-}

coerceVector :: (Coercible a b, Unboxable a, Unboxable b, CoercibleRep a ~ CoercibleRep b, Rep a ~ Rep b) => Vector a -> Vector b
coerceVector = coerce
{-# INLINE coerceVector #-}

liftCoercion :: (Unboxable a, Unboxable b, CoercibleRep a ~ CoercibleRep b, Rep a ~ Rep b) => Coercion a b -> Coercion (Vector a) (Vector b)
liftCoercion Coercion = Coercion
{-# INLINE liftCoercion #-}

vectorCoercion :: (Coercible a b, Unboxable a, Unboxable b, CoercibleRep a ~ CoercibleRep b, Rep a ~ Rep b) => Coercion (Vector a) (Vector b)
vectorCoercion = Coercion
{-# INLINE vectorCoercion #-}

toUnboxedVector :: (Unboxable a, Rep a ~ a, IsTrivial a ~ 'True) => Vector a -> U.Vector a
toUnboxedVector = coerce
{-# INLINE toUnboxedVector #-}

fromUnboxedVector :: (Unboxable a, Rep a ~ a, IsTrivial a ~ 'True) => U.Vector a -> Vector a
fromUnboxedVector = coerce
{-# INLINE fromUnboxedVector #-}

coerceMVector :: (Coercible a b, Unboxable a, Unboxable b, CoercibleRep a ~ CoercibleRep b, Rep a ~ Rep b) => MVector s a -> MVector s b
coerceMVector = coerce
{-# INLINE coerceMVector #-}

liftCoercionM :: (Unboxable a, Unboxable b, CoercibleRep a ~ CoercibleRep b, Rep a ~ Rep b) => Coercion a b -> Coercion (MVector s a) (MVector s b)
liftCoercionM Coercion = Coercion
{-# INLINE liftCoercionM #-}

mVectorCoercion :: (Coercible a b, Unboxable a, Unboxable b, CoercibleRep a ~ CoercibleRep b, Rep a ~ Rep b) => Coercion (MVector s a) (MVector s b)
mVectorCoercion = Coercion
{-# INLINE mVectorCoercion #-}

toUnboxedMVector :: (Unboxable a, Rep a ~ a, IsTrivial a ~ 'True) => MVector s a -> UM.MVector s a
toUnboxedMVector = coerce
{-# INLINE toUnboxedMVector #-}

fromUnboxedMVector :: (Unboxable a, Rep a ~ a, IsTrivial a ~ 'True) => UM.MVector s a -> MVector s a
fromUnboxedMVector = coerce
{-# INLINE fromUnboxedMVector #-}

coercionWithUnboxedVector :: (Unboxable a, Rep a ~ a, IsTrivial a ~ 'True) => Coercion (Vector a) (U.Vector a)
coercionWithUnboxedVector = Coercion
{-# INLINE coercionWithUnboxedVector #-}

coercionWithUnboxedMVector :: (Unboxable a, Rep a ~ a, IsTrivial a ~ 'True) => Coercion (MVector s a) (U.MVector s a)
coercionWithUnboxedMVector = Coercion
{-# INLINE coercionWithUnboxedMVector #-}

--
-- Generics
--

-- | A newtype wrapper to be used with @DerivingVia@.
--
-- Usage:
--
-- > data Bar = Bar !Int !Int
-- >   deriving Generic
-- >   deriving Unboxable via Generics Bar
newtype Generics a = Generics a

instance (GHC.Generics.Generic a, U.Unbox (Rep' (GHC.Generics.Rep a)), Unboxable' (GHC.Generics.Rep a)) => Unboxable (Generics a) where
  type Rep (Generics a) = Rep' (GHC.Generics.Rep a)
  type CoercibleRep (Generics a) = a
  type IsTrivial (Generics a) = 'False
  unboxingFrom (Generics x) = from' (GHC.Generics.from x)
  {-# INLINE unboxingFrom #-}
  unboxingTo y = Generics (GHC.Generics.to (to' y))
  {-# INLINE unboxingTo #-}

class Unboxable' f where
  type Rep' f
  from' :: f x -> Rep' f
  to' :: Rep' f -> f x
instance Unboxable' GHC.Generics.U1 where
  type Rep' GHC.Generics.U1 = ()
  from' _ = ()
  to' _ = GHC.Generics.U1
  {-# INLINE from' #-}
  {-# INLINE to' #-}
instance Unboxable c => Unboxable' (GHC.Generics.K1 i c) where
  type Rep' (GHC.Generics.K1 i c) = Rep c
  from' = {- from . GHC.Generics.unK1 -} coerce (unboxingFrom :: c -> Rep c)
  to' = {- GHC.Generics.K1 . to -} coerce (unboxingTo :: Rep c -> c)
  {-# INLINE from' #-}
  {-# INLINE to' #-}
instance Unboxable' f => Unboxable' (GHC.Generics.M1 i c f) where
  type Rep' (GHC.Generics.M1 i c f) = Rep' f
  from' = from' . GHC.Generics.unM1
  to' = GHC.Generics.M1 . to'
  {-# INLINE from' #-}
  {-# INLINE to' #-}
instance (Unboxable' f, Unboxable' g) => Unboxable' (f GHC.Generics.:*: g) where
  type Rep' (f GHC.Generics.:*: g) = (Rep' f, Rep' g)
  from' (x GHC.Generics.:*: y) = (from' x, from' y)
  to' (x, y) = (to' x GHC.Generics.:*: to' y)
  {-# INLINE from' #-}
  {-# INLINE to' #-}
instance Unboxable' (f GHC.Generics.:+: g) where
  type Rep' (f GHC.Generics.:+: g) = TypeError ('Text "Cannot derive Unboxable instance for a sum type.")
  from' = undefined
  to' = undefined

--
-- Enum
--

-- | A newtype wrapper to be used with @DerivingVia@.
-- The value will be stored as 'Int', via `fromEnum`/`toEnum`.
--
-- Usage:
--
-- > data Direction = North | South | East | West
-- >   deriving Enum
-- >   deriving Data.Vector.Unboxing.Unboxable via Data.Vector.Unboxing.Enum Bar
newtype Enum a = Enum a
instance (Prelude.Enum a) => Unboxable (Enum a) where
  type Rep (Enum a) = Int
  type CoercibleRep (Enum a) = a
  type IsTrivial (Enum a) = 'False
  unboxingFrom (Enum x) = fromEnum x
  {-# INLINE unboxingFrom #-}
  unboxingTo y = Enum (toEnum y)
  {-# INLINE unboxingTo #-}

-- | A newtype wrapper to be used with @DerivingVia@.
--
-- Usage:
--
-- > data Direction = North | South | East | West
-- >   deriving Enum
-- >   deriving Data.Vector.Unboxing.Unboxable via Data.Vector.Unboxing.EnumRep Int8 Bar
newtype EnumRep rep a = EnumRep a
instance (Prelude.Enum a, Integral rep, U.Unbox rep) => Unboxable (EnumRep rep a) where
  type Rep (EnumRep rep a) = rep
  type CoercibleRep (EnumRep rep a) = a
  type IsTrivial (EnumRep rep a) = 'False
  unboxingFrom (EnumRep x) = fromIntegral (fromEnum x)
  {-# INLINE unboxingFrom #-}
  unboxingTo y = EnumRep (toEnum (fromIntegral y))
  {-# INLINE unboxingTo #-}

--
-- Instances
--

instance (Unboxable a) => IsList (Vector a) where
  type Item (Vector a) = a
  fromList = G.fromList
  fromListN = G.fromListN
  toList = G.toList
  {-# INLINE fromList #-}
  {-# INLINE fromListN #-}
  {-# INLINE toList #-}

instance (Eq a, Unboxable a) => Eq (Vector a) where
  xs == ys = Bundle.eq (G.stream xs) (G.stream ys)
  xs /= ys = not (Bundle.eq (G.stream xs) (G.stream ys))
  {-# INLINE (==) #-}
  {-# INLINE (/=) #-}

instance (Ord a, Unboxable a) => Ord (Vector a) where
  compare xs ys = Bundle.cmp (G.stream xs) (G.stream ys)
  {-# INLINE compare #-}

instance (Show a, Unboxable a) => Show (Vector a) where
  showsPrec = G.showsPrec
  {-# INLINE showsPrec #-}

instance (Read a, Unboxable a) => Read (Vector a) where
  readPrec = G.readPrec
  readListPrec = readListPrecDefault
  {-# INLINE readPrec #-}
  {-# INLINE readListPrec #-}

instance (Unboxable a) => Data.Semigroup.Semigroup (Vector a) where
  (<>) = (G.++)
  sconcat = G.concatNE
  {-# INLINE (<>) #-}
  {-# INLINE sconcat #-}

instance (Unboxable a) => Data.Monoid.Monoid (Vector a) where
  mempty = G.empty
  mappend = (Data.Semigroup.<>)
  mconcat = G.concat
  {-# INLINE mempty #-}
  {-# INLINE mappend #-}
  {-# INLINE mconcat #-}

instance NFData (Vector a) where
  rnf !_ = () -- the content is unboxed

instance (Unboxable a) => GM.MVector MVector a where
  basicLength (UnboxingMVector mv)                           = GM.basicLength mv
  basicUnsafeSlice i l (UnboxingMVector mv)                  = UnboxingMVector (GM.basicUnsafeSlice i l mv)
  basicOverlaps (UnboxingMVector mv) (UnboxingMVector mv')   = GM.basicOverlaps mv mv'
  basicUnsafeNew l                                           = UnboxingMVector <$> GM.basicUnsafeNew l
  basicInitialize (UnboxingMVector mv)                       = GM.basicInitialize mv
  basicUnsafeReplicate i x                                   = UnboxingMVector <$> GM.basicUnsafeReplicate i (unboxingFrom x)
  basicUnsafeRead (UnboxingMVector mv) i                     = unboxingTo <$> GM.basicUnsafeRead mv i
  basicUnsafeWrite (UnboxingMVector mv) i x                  = GM.basicUnsafeWrite mv i (unboxingFrom x)
  basicClear (UnboxingMVector mv)                            = GM.basicClear mv
  basicSet (UnboxingMVector mv) x                            = GM.basicSet mv (unboxingFrom x)
  basicUnsafeCopy (UnboxingMVector mv) (UnboxingMVector mv') = GM.basicUnsafeCopy mv mv'
  basicUnsafeMove (UnboxingMVector mv) (UnboxingMVector mv') = GM.basicUnsafeMove mv mv'
  basicUnsafeGrow (UnboxingMVector mv) n                     = UnboxingMVector <$> GM.basicUnsafeGrow mv n
  {-# INLINE basicLength #-}
  {-# INLINE basicUnsafeSlice #-}
  {-# INLINE basicOverlaps #-}
  {-# INLINE basicUnsafeNew #-}
  {-# INLINE basicInitialize #-}
  {-# INLINE basicUnsafeRead #-}
  {-# INLINE basicUnsafeWrite #-}
  {-# INLINE basicClear #-}
  {-# INLINE basicSet #-}
  {-# INLINE basicUnsafeCopy #-}
  {-# INLINE basicUnsafeMove #-}
  {-# INLINE basicUnsafeGrow #-}

instance (Unboxable a) => G.Vector Vector a where
  basicUnsafeFreeze (UnboxingMVector mv)                  = UnboxingVector <$> G.basicUnsafeFreeze mv
  basicUnsafeThaw (UnboxingVector v)                      = UnboxingMVector <$> G.basicUnsafeThaw v
  basicLength (UnboxingVector v)                          = G.basicLength v
  basicUnsafeSlice i l (UnboxingVector v)                 = UnboxingVector (G.basicUnsafeSlice i l v)
  basicUnsafeIndexM (UnboxingVector v) i                  = unboxingTo <$> G.basicUnsafeIndexM v i
  basicUnsafeCopy (UnboxingMVector mv) (UnboxingVector v) = G.basicUnsafeCopy mv v
  elemseq (UnboxingVector v) x y                          = G.elemseq v (unboxingFrom x) y -- ?
  {-# INLINE basicUnsafeFreeze #-}
  {-# INLINE basicUnsafeThaw #-}
  {-# INLINE basicLength #-}
  {-# INLINE basicUnsafeSlice #-}
  {-# INLINE basicUnsafeIndexM #-}
  {-# INLINE basicUnsafeCopy #-}
  {-# INLINE elemseq #-}

--
-- Unboxable instances
--

instance Unboxable Bool where   type Rep Bool = Bool
instance Unboxable Char where   type Rep Char = Char
instance Unboxable Double where type Rep Double = Double
instance Unboxable Float where  type Rep Float = Float
instance Unboxable Int where    type Rep Int = Int
instance Unboxable Int8 where   type Rep Int8 = Int8
instance Unboxable Int16 where  type Rep Int16 = Int16
instance Unboxable Int32 where  type Rep Int32 = Int32
instance Unboxable Int64 where  type Rep Int64 = Int64
instance Unboxable Word where   type Rep Word = Word
instance Unboxable Word8 where  type Rep Word8 = Word8
instance Unboxable Word16 where type Rep Word16 = Word16
instance Unboxable Word32 where type Rep Word32 = Word32
instance Unboxable Word64 where type Rep Word64 = Word64
instance Unboxable () where     type Rep () = ()

instance (Unboxable a) => Unboxable (Data.Complex.Complex a) where
  type Rep (Data.Complex.Complex a) = Data.Complex.Complex (Rep a)
  type CoercibleRep (Data.Complex.Complex a) = Data.Complex.Complex (CoercibleRep a)
  type IsTrivial (Data.Complex.Complex a) = IsTrivial a
  unboxingFrom = fmap unboxingFrom
  unboxingTo = fmap unboxingTo
  {-# INLINE unboxingFrom #-}
  {-# INLINE unboxingTo #-}

instance (Unboxable a, Unboxable b) => Unboxable (a, b) where
  type Rep (a, b) = (Rep a, Rep b)
  type CoercibleRep (a, b) = (CoercibleRep a, CoercibleRep b)
  type IsTrivial (a, b) = IsTrivial a && IsTrivial b
  unboxingFrom (a, b) = (unboxingFrom a, unboxingFrom b)
  unboxingTo (a, b) = (unboxingTo a, unboxingTo b)
  {-# INLINE unboxingFrom #-}
  {-# INLINE unboxingTo #-}

instance (Unboxable a, Unboxable b, Unboxable c) => Unboxable (a, b, c) where
  type Rep (a, b, c) = (Rep a, Rep b, Rep c)
  type CoercibleRep (a, b, c) = (CoercibleRep a, CoercibleRep b, CoercibleRep c)
  type IsTrivial (a, b, c) = IsTrivial a && IsTrivial b && IsTrivial c
  unboxingFrom (a, b, c) = (unboxingFrom a, unboxingFrom b, unboxingFrom c)
  unboxingTo (a, b, c) = (unboxingTo a, unboxingTo b, unboxingTo c)
  {-# INLINE unboxingFrom #-}
  {-# INLINE unboxingTo #-}

instance (Unboxable a, Unboxable b, Unboxable c, Unboxable d) => Unboxable (a, b, c, d) where
  type Rep (a, b, c, d) = (Rep a, Rep b, Rep c, Rep d)
  type CoercibleRep (a, b, c, d) = (CoercibleRep a, CoercibleRep b, CoercibleRep c, CoercibleRep d)
  type IsTrivial (a, b, c, d) = IsTrivial a && IsTrivial b && IsTrivial c && IsTrivial d
  unboxingFrom (a, b, c, d) = (unboxingFrom a, unboxingFrom b, unboxingFrom c, unboxingFrom d)
  unboxingTo (a, b, c, d) = (unboxingTo a, unboxingTo b, unboxingTo c, unboxingTo d)
  {-# INLINE unboxingFrom #-}
  {-# INLINE unboxingTo #-}

instance (Unboxable a, Unboxable b, Unboxable c, Unboxable d, Unboxable e) => Unboxable (a, b, c, d, e) where
  type Rep (a, b, c, d, e) = (Rep a, Rep b, Rep c, Rep d, Rep e)
  type CoercibleRep (a, b, c, d, e) = (CoercibleRep a, CoercibleRep b, CoercibleRep c, CoercibleRep d, CoercibleRep e)
  type IsTrivial (a, b, c, d, e) = IsTrivial a && IsTrivial b && IsTrivial c && IsTrivial d && IsTrivial e
  unboxingFrom (a, b, c, d, e) = (unboxingFrom a, unboxingFrom b, unboxingFrom c, unboxingFrom d, unboxingFrom e)
  unboxingTo (a, b, c, d, e) = (unboxingTo a, unboxingTo b, unboxingTo c, unboxingTo d, unboxingTo e)
  {-# INLINE unboxingFrom #-}
  {-# INLINE unboxingTo #-}

instance (Unboxable a, Unboxable b, Unboxable c, Unboxable d, Unboxable e, Unboxable f) => Unboxable (a, b, c, d, e, f) where
  type Rep (a, b, c, d, e, f) = (Rep a, Rep b, Rep c, Rep d, Rep e, Rep f)
  type CoercibleRep (a, b, c, d, e, f) = (CoercibleRep a, CoercibleRep b, CoercibleRep c, CoercibleRep d, CoercibleRep e, CoercibleRep f)
  type IsTrivial (a, b, c, d, e, f) = IsTrivial a && IsTrivial b && IsTrivial c && IsTrivial d && IsTrivial e && IsTrivial f
  unboxingFrom (a, b, c, d, e, f) = (unboxingFrom a, unboxingFrom b, unboxingFrom c, unboxingFrom d, unboxingFrom e, unboxingFrom f)
  unboxingTo (a, b, c, d, e, f) = (unboxingTo a, unboxingTo b, unboxingTo c, unboxingTo d, unboxingTo e, unboxingTo f)
  {-# INLINE unboxingFrom #-}
  {-# INLINE unboxingTo #-}

{-
GND for type classes with associated type families is not available until GHC 8.2.
With GHC 8.2 or later, one can derive these instances like:

> deriving instance Unboxable a => Unboxable (Data.Functor.Identity.Identity a)
> deriving instance Unboxable a => Unboxable (Data.Functor.Const.Const a b)
> deriving instance Unboxable a => Unboxable (Data.Semigroup.Min a)
> deriving instance Unboxable a => Unboxable (Data.Semigroup.Max a)
> deriving instance Unboxable a => Unboxable (Data.Semigroup.First a)
> deriving instance Unboxable a => Unboxable (Data.Semigroup.Last a)
> deriving instance Unboxable a => Unboxable (Data.Semigroup.WrappedMonoid a)
> deriving instance Unboxable a => Unboxable (Data.Monoid.Dual a)
> deriving instance Unboxable Data.Monoid.All
> deriving instance Unboxable Data.Monoid.Any
> deriving instance Unboxable a => Unboxable (Data.Monoid.Sum a)
> deriving instance Unboxable a => Unboxable (Data.Monoid.Product a)
> deriving instance Unboxable a => Unboxable (Data.Ord.Down a)
-}

instance Unboxable a => Unboxable (Data.Functor.Identity.Identity a) where
  type Rep (Data.Functor.Identity.Identity a) = Rep a
  unboxingFrom = unboxingFrom .# Data.Functor.Identity.runIdentity
  unboxingTo = Data.Functor.Identity.Identity #. unboxingTo
  {-# INLINE unboxingFrom #-}
  {-# INLINE unboxingTo #-}

instance Unboxable a => Unboxable (Data.Functor.Const.Const a b) where
  type Rep (Data.Functor.Const.Const a b) = Rep a
  unboxingFrom = unboxingFrom .# Data.Functor.Const.getConst
  unboxingTo = Data.Functor.Const.Const #. unboxingTo
  {-# INLINE unboxingFrom #-}
  {-# INLINE unboxingTo #-}

instance Unboxable a => Unboxable (Data.Semigroup.Min a) where
  type Rep (Data.Semigroup.Min a) = Rep a
  unboxingFrom = unboxingFrom .# Data.Semigroup.getMin
  unboxingTo = Data.Semigroup.Min #. unboxingTo
  {-# INLINE unboxingFrom #-}
  {-# INLINE unboxingTo #-}

instance Unboxable a => Unboxable (Data.Semigroup.Max a) where
  type Rep (Data.Semigroup.Max a) = Rep a
  unboxingFrom = unboxingFrom .# Data.Semigroup.getMax
  unboxingTo = Data.Semigroup.Max #. unboxingTo
  {-# INLINE unboxingFrom #-}
  {-# INLINE unboxingTo #-}

instance Unboxable a => Unboxable (Data.Semigroup.First a) where
  type Rep (Data.Semigroup.First a) = Rep a
  unboxingFrom = unboxingFrom .# Data.Semigroup.getFirst
  unboxingTo = Data.Semigroup.First #. unboxingTo
  {-# INLINE unboxingFrom #-}
  {-# INLINE unboxingTo #-}

instance Unboxable a => Unboxable (Data.Semigroup.Last a) where
  type Rep (Data.Semigroup.Last a) = Rep a
  unboxingFrom = unboxingFrom .# Data.Semigroup.getLast
  unboxingTo = Data.Semigroup.Last #. unboxingTo
  {-# INLINE unboxingFrom #-}
  {-# INLINE unboxingTo #-}

instance Unboxable a => Unboxable (Data.Semigroup.WrappedMonoid a) where
  type Rep (Data.Semigroup.WrappedMonoid a) = Rep a
  unboxingFrom = unboxingFrom .# Data.Semigroup.unwrapMonoid
  unboxingTo = Data.Semigroup.WrapMonoid #. unboxingTo
  {-# INLINE unboxingFrom #-}
  {-# INLINE unboxingTo #-}

instance Unboxable a => Unboxable (Data.Monoid.Dual a) where
  type Rep (Data.Monoid.Dual a) = Rep a
  unboxingFrom = unboxingFrom .# Data.Monoid.getDual
  unboxingTo = Data.Monoid.Dual #. unboxingTo
  {-# INLINE unboxingFrom #-}
  {-# INLINE unboxingTo #-}

instance Unboxable Data.Monoid.All where
  type Rep Data.Monoid.All = Bool

instance Unboxable Data.Monoid.Any where
  type Rep Data.Monoid.Any = Bool

instance Unboxable a => Unboxable (Data.Monoid.Sum a) where
  type Rep (Data.Monoid.Sum a) = Rep a
  unboxingFrom = unboxingFrom .# Data.Monoid.getSum
  unboxingTo = Data.Monoid.Sum #. unboxingTo
  {-# INLINE unboxingFrom #-}
  {-# INLINE unboxingTo #-}

instance Unboxable a => Unboxable (Data.Monoid.Product a) where
  type Rep (Data.Monoid.Product a) = Rep a
  unboxingFrom = unboxingFrom .# Data.Monoid.getProduct
  unboxingTo = Data.Monoid.Product #. unboxingTo
  {-# INLINE unboxingFrom #-}
  {-# INLINE unboxingTo #-}

instance Unboxable a => Unboxable (Data.Ord.Down a) where
  type Rep (Data.Ord.Down a) = Rep a
  unboxingFrom = unboxingFrom .# (\(Data.Ord.Down x) -> x)
  unboxingTo = Data.Ord.Down #. unboxingTo
  {-# INLINE unboxingFrom #-}
  {-# INLINE unboxingTo #-}

instance Unboxable Ordering where
  type Rep Ordering = Int8
  unboxingFrom x = fromIntegral (fromEnum x)
  unboxingTo y = toEnum (fromIntegral y)
  {-# INLINE unboxingFrom #-}
  {-# INLINE unboxingTo #-}

instance (Unboxable a, Unboxable b) => Unboxable (Data.Semigroup.Arg a b) where
  type Rep (Data.Semigroup.Arg a b) = (Rep a, Rep b)
  unboxingFrom (Data.Semigroup.Arg x y) = (unboxingFrom x, unboxingFrom y)
  unboxingTo (x, y) = Data.Semigroup.Arg (unboxingTo x) (unboxingTo y)
  {-# INLINE unboxingFrom #-}
  {-# INLINE unboxingTo #-}

instance Unboxable (f a) => Unboxable (Data.Monoid.Alt f a) where
  type Rep (Data.Monoid.Alt f a) = Rep (f a)
  unboxingFrom = unboxingFrom .# Data.Monoid.getAlt
  unboxingTo = Data.Monoid.Alt #. unboxingTo
  {-# INLINE unboxingFrom #-}
  {-# INLINE unboxingTo #-}

{-
-- Since base-4.12.0.0
instance Unboxable (f a) => Unboxable (Data.Monoid.Ap f a) where
  type Rep (Data.Monoid.Ap f a) = Rep (f a)
  unboxingFrom = unboxingFrom .# Data.Monoid.getAp
  unboxingTo = Data.Monoid.Ap #. unboxingTo
  {-# INLINE unboxingFrom #-}
  {-# INLINE unboxingTo #-}
-}

instance Unboxable (f (g a)) => Unboxable (Data.Functor.Compose.Compose f g a) where
  type Rep (Data.Functor.Compose.Compose f g a) = Rep (f (g a))
  unboxingFrom = unboxingFrom .# Data.Functor.Compose.getCompose
  unboxingTo = Data.Functor.Compose.Compose #. unboxingTo
  {-# INLINE unboxingFrom #-}
  {-# INLINE unboxingTo #-}

infixr 9 #.
infixl 8 .#

(#.) :: Coercible b c => (b -> c) -> (a -> b) -> a -> c
_ #. f = coerce f
{-# INLINE (#.) #-}

(.#) :: Coercible a b => (b -> c) -> (a -> b) -> a -> c
f .# _ = coerce f
{-# INLINE (.#) #-}


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

  {-  LANGUAGE CPP #-}  -- specified in .cabal default-extensions

#define USE_NFDATA_SUPERCLASS 0

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

  -- Later: I'm not so sure about this, actually; is the arithmetic
  -- on n actually not piling up thunks, without the bang-patterns?
  --
  -- It would be easy to get rid of the bang-patterns.
  -- The Complex instance is done as an example.
  -- If you do go with the case's, probably want -fno-warn-name-shadowing.
  {-# LANGUAGE BangPatterns #-}
  {-# OPTIONS_GHC -fno-warn-name-shadowing #-}

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

-- |
-- Module      :  Control.DeepSeq.Bounded.NFDataN
-- Copyright   :  Andrew G. Seniuk 2014-2015
-- License     :  BSD-style (see the file LICENSE)
--
-- Maintainer  :  Andrew Seniuk <rasfar@gmail.com>
-- Stability   :  provisional
-- Portability :  portable
--
-- This module provides an overloaded function, 'deepseqn', for partially
-- (or fully) evaluating data structures to a given depth.
--

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

  module Control.DeepSeq.Bounded.NFDataN
  (

      -- * Depth-bounded analogues of 'deepseq' and 'force'

      deepseqn
    , forcen

      -- * Class of things that can be evaluated to an arbitrary finite depth

#if JUST_ALIAS_GNFDATAN
    , rnfn
#else
    , NFDataN(..)
#endif

  )
  where

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

--import qualified Data.Generics.Shape as S  -- not actually used

#if JUST_ALIAS_GNFDATAN
  import Generics.SOP ( Generic )
  import Control.DeepSeq.Bounded.Generic.GNFDataN ( grnfn )
#endif

#if USE_NFDATA_SUPERCLASS
  import Control.DeepSeq
#endif

--import Data.Data

  import Data.Int
  import Data.Word
  import Data.Ratio
  import Data.Complex
  import Data.Array
  import Data.Fixed
  import Data.Version

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

#if JUST_ALIAS_GNFDATAN
  type NFDataN = Generic
#endif

#if JUST_ALIAS_GNFDATAN
  rnfn = grnfn
#endif

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

-- infixr 0 $!!

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

  -- | @'deepseqn' n x y@ evaluates @x@ to depth n, before returning @y@.
  --
  -- This is used when expression @x@ also appears in @y@, but mere
  -- evaluation of @y@ does not force the embedded @x@ subexpression
  -- as deeply as we wish.
  --
  -- The name 'deepseqn' is used to illustrate the relationship to 'seq':
  -- where 'seq' is shallow in the sense that it only evaluates the top
  -- level of its argument, @'deepseqn' n@ traverses (evaluates) the entire
  -- top @n@ levels of the data structure.
  --
  -- A typical use is to ensure any exceptions hidden within lazy
  -- fields of a data structure do not leak outside the scope of the
  -- exception handler; another is to force evaluation of a data structure in
  -- one thread, before passing it to another thread (preventing work moving
  -- to the wrong threads). Unlike <http://hackage.haskell.org/package/deepseq/docs/Control-DeepSeq.html DeepSeq>, potentially infinite values of coinductive
  -- data types are supported by principled bounding of deep evaluation.
  --
  -- It is also useful for diagnostic purposes when trying to understand
  -- and manipulate space\/time trade-offs in lazy code,
  -- and as a less indiscriminate substitute for 'deepseq'.
  --
  -- Furthermore, 'deepseqn' can sometimes be a better solution than
  -- using stict fields in your data structures, because the
  -- latter will behave strictly everywhere that its constructors
  -- are used, instead of just where its laziness is problematic.
  --
  -- There may be possible applications to the prevention of resource leaks
  -- in lazy streaming, but I'm not certain.
  --
  -- One of the special qualities of 'NFDataN' is shape oblivity:
  -- it doesn't care about any details of the shape of the term
  -- it's forcing, it only cares about stratifying levels of
  -- recursion depth.
  -- (I would say \"as contrasted with <http://hackage.haskell.org/package/deepseq-bounded-0.8.0.0/docs/Control-DeepSeq-Bounded-NFDataP.html NFDataP>\" but cannot, because
  -- <http://hackage.haskell.org/package/deepseq-bounded-0.8.0.0/docs/Control-DeepSeq-Bounded-NFDataP.html NFDataP>
  -- was extended to include
  -- <http://hackage.haskell.org/package/deepseq-bounded-0.8.0.0/docs/Control-DeepSeq-Bounded-NFDataN.html NFDataN>
  -- syntax\/capabilities, precisely to ammend this deficiency.)

  deepseqn :: NFDataN a => Int -> a -> b -> b
  deepseqn n a b = rnfn n a `seq` b
#if 1
  -- XXX Need to double-check that this makes sense; didn't think
  -- it through -- when b != a, things are not so simple.
  -- XXX Partially-applied; is that okay in GHC RULES?
  {-# RULES
    "deepseqn/composition"    forall n1 n2 x.  (.) (deepseqn n2) (deepseqn n1) x = deepseqn ( max n1 n2 ) x
      #-}
#endif

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

#if 0
-- As per DeepSeq:
  -- | the deep analogue of '$!'.  In the expression @f $!! x@, @x@ is
  -- fully evaluated before the function @f@ is applied to it.
  ($!!) :: NFData a => (a -> b) -> a -> b
  f $!! x = x `deepseq` f x
#endif

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

  -- | 'forcen' is a variant of 'deepseqn' that is useful in some circumstances.
  --
  -- > forcen n x = deepseqn n x x
  --
  -- @forcen n x@ evaluates @x@ to depth @n@, and then returns it.
  -- Recall that, in common with all Haskell functions, @forcen@ only
  -- performs evaluation when upstream demand actually occurs,
  -- so essentially it turns shallow evaluation into depth-@n@ evaluation.

  forcen :: NFDataN a => Int -> a -> a
  forcen n x = deepseqn n x x
#if 1
  {-# RULES
    "forcen/composition"    forall n1 n2 x.  (.) (forcen n2) (forcen n1) x = forcen ( max n1 n2 ) x
      #-}
#endif

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

#if ! JUST_ALIAS_GNFDATAN

  -- | A class of types that can be evaluated to arbitrary depth.
#if USE_NFDATA_SUPERCLASS
  class NFData a => NFDataN a where
#else
  class NFDataN a where
#endif
#if 0
    -- | rnf should reduce its argument to normal form (that is, fully
    -- evaluate all sub-components), and then return '()'.
    --
    -- The default implementation of 'rnf' is
    --
    -- > rnf a = a `seq` ()
    --
    -- which may be convenient when defining instances for data types with
    -- no unevaluated fields (e.g. enumerations).
    rnf :: a -> ()
    rnf a = a `seq` ()
#endif
    rnfn :: Int -> a -> ()
#if USE_NFDATA_SUPERCLASS
    -- ?? not sure about this (whatever, we'll override it)
#if 1
    rnfn n a | n <= 0     = ()  -- case needed? (seems to be!)
             | otherwise  = rnf a
#else
    rnfn _ a = rnf a
#endif
#else
#if 0
#elif 1
    rnfn n x | n <= 0     = ()
             | otherwise  = x `seq` ()
#elif 0
    rnfn _ x = x `seq` ()
#elif 0
    rnfn _ _ = ()
#endif
#endif
#if 1
  {-# RULES
    "rnfn/composition"    forall n1 n2 x.  (.) (rnfn n2) (rnfn n1) x = rnfn ( max n1 n2 ) x
      #-}
--   "rnfp/composition"    forall p1 p2 x.  compose (rnfp p2) (rnfp p1) x = rnfp ( unionPat [p1, p2] ) x
--   "rnfp/composition"    forall p1 p2 x.  ( rnfp p2 . rnfp p1 ) x = rnfp ( unionPat [p1, p2] ) x
#endif

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

  instance NFDataN Int
  instance NFDataN Word
  instance NFDataN Integer
  instance NFDataN Float
  instance NFDataN Double

  instance NFDataN Char
  instance NFDataN Bool
  instance NFDataN ()

  instance NFDataN Int8
  instance NFDataN Int16
  instance NFDataN Int32
  instance NFDataN Int64

  instance NFDataN Word8
  instance NFDataN Word16
  instance NFDataN Word32
  instance NFDataN Word64

  instance NFDataN (Fixed a)

  -- [Quoted from deepseq:]
  -- This instance is for convenience and consistency with 'seq'.
  -- This assumes that WHNF is equivalent to NF for functions.
  instance NFDataN (a -> b)

#if 0
  instance NFDataN a => NFDataN (IO a) where
    rnfn !n x = ()  -- XXX ignore if in IO (cludge easing an experiment...)
#endif

  -- not taken to be a level of depth
  instance (Integral a, NFDataN a) => NFDataN (Ratio a) where
    rnfn !n x = rnfn n (numerator x, denominator x)

  instance (RealFloat a, NFDataN a) => NFDataN (Complex a) where
#if 1
    rnfn n (x:+y) = case n of
     n | n <= 0  -> ()
     _           -> let n' = -1+n in
                    rnfn n' x `seq`
                    rnfn n' y `seq`
                    ()
#else
    rnfn !n (x:+y)
     | n <= 0    = ()
     | otherwise = let n' = -1+n in
                   rnfn n' x `seq`
                   rnfn n' y `seq`
                   ()
#endif

  instance NFDataN a => NFDataN (Maybe a) where
    rnfn _ Nothing  = ()
    rnfn !n (Just x)
     | n <= 0     = ()
     | otherwise  = rnfn (-1+n) x

  instance (NFDataN a, NFDataN b) => NFDataN (Either a b) where
    rnfn !n (Left x)
     | n <= 0     = ()
     | otherwise  = rnfn (-1+n) x
    rnfn !n (Right y)
     | n <= 0     = ()
     | otherwise  = rnfn (-1+n) y

  instance NFDataN Data.Version.Version where
    rnfn !n (Data.Version.Version branch tags)
     | n <= 0     = ()
     | otherwise  = let n' = -1+n in
                    rnfn n' branch `seq` rnfn n' tags

  instance NFDataN a => NFDataN [a] where
    rnfn _ [] = ()
    rnfn !n (x:xs)
     | n <= 0     = ()
     | otherwise  = let n' = -1+n in
                    rnfn n' x `seq` rnfn n' xs

  -- not taken to be a level of depth
  instance (Ix a, NFDataN a, NFDataN b) => NFDataN (Array a b) where
    rnfn !n x = rnfn n (bounds x, Data.Array.elems x)

  instance (NFDataN a, NFDataN b) => NFDataN (a,b) where
    rnfn !n (x,y)
     | n <= 0     = ()
     | otherwise  = let n' = -1+n in
                    rnfn n' x `seq` rnfn n' y

  instance (NFDataN a, NFDataN b, NFDataN c) => NFDataN (a,b,c) where
    rnfn !n (x,y,z)
     | n <= 0     = ()
     | otherwise  = let n' = -1+n in
                    rnfn n' x `seq` rnfn n' y `seq` rnfn n' z

  instance (NFDataN a, NFDataN b, NFDataN c, NFDataN d) => NFDataN (a,b,c,d) where
    rnfn !n (x1,x2,x3,x4)
     | n <= 0     = ()
     | otherwise  = let n' = -1+n in
                    rnfn n' x1 `seq`
                    rnfn n' x2 `seq`
                    rnfn n' x3 `seq`
                    rnfn n' x4

  instance (NFDataN a1, NFDataN a2, NFDataN a3, NFDataN a4, NFDataN a5) =>
           NFDataN (a1, a2, a3, a4, a5) where
    rnfn !n (x1,x2,x3,x4,x5)
     | n <= 0     = ()
     | otherwise  = let n' = -1+n in
                    rnfn n' x1 `seq`
                    rnfn n' x2 `seq`
                    rnfn n' x3 `seq`
                    rnfn n' x4 `seq`
                    rnfn n' x5

  instance (NFDataN a1, NFDataN a2, NFDataN a3, NFDataN a4, NFDataN a5, NFDataN a6) =>
           NFDataN (a1, a2, a3, a4, a5, a6) where
    rnfn !n (x1,x2,x3,x4,x5,x6)
     | n <= 0     = ()
     | otherwise  = let n' = -1+n in
                    rnfn n' x1 `seq`
                    rnfn n' x2 `seq`
                    rnfn n' x3 `seq`
                    rnfn n' x4 `seq`
                    rnfn n' x5 `seq`
                    rnfn n' x6

  instance (NFDataN a1, NFDataN a2, NFDataN a3, NFDataN a4, NFDataN a5, NFDataN a6, NFDataN a7) =>
           NFDataN (a1, a2, a3, a4, a5, a6, a7) where
    rnfn !n (x1,x2,x3,x4,x5,x6,x7)
     | n <= 0     = ()
     | otherwise  = let n' = -1+n in
                    rnfn n' x1 `seq`
                    rnfn n' x2 `seq`
                    rnfn n' x3 `seq`
                    rnfn n' x4 `seq`
                    rnfn n' x5 `seq`
                    rnfn n' x6 `seq`
                    rnfn n' x7

  instance (NFDataN a1, NFDataN a2, NFDataN a3, NFDataN a4, NFDataN a5, NFDataN a6, NFDataN a7, NFDataN a8) =>
           NFDataN (a1, a2, a3, a4, a5, a6, a7, a8) where
    rnfn !n (x1,x2,x3,x4,x5,x6,x7,x8)
     | n <= 0     = ()
     | otherwise  = let n' = -1+n in
                    rnfn n' x1 `seq`
                    rnfn n' x2 `seq`
                    rnfn n' x3 `seq`
                    rnfn n' x4 `seq`
                    rnfn n' x5 `seq`
                    rnfn n' x6 `seq`
                    rnfn n' x7 `seq`
                    rnfn n' x8

  instance (NFDataN a1, NFDataN a2, NFDataN a3, NFDataN a4, NFDataN a5, NFDataN a6, NFDataN a7, NFDataN a8, NFDataN a9) =>
           NFDataN (a1, a2, a3, a4, a5, a6, a7, a8, a9) where
    rnfn !n (x1,x2,x3,x4,x5,x6,x7,x8,x9)
     | n <= 0     = ()
     | otherwise  = let n' = -1+n in
                    rnfn n' x1 `seq`
                    rnfn n' x2 `seq`
                    rnfn n' x3 `seq`
                    rnfn n' x4 `seq`
                    rnfn n' x5 `seq`
                    rnfn n' x6 `seq`
                    rnfn n' x7 `seq`
                    rnfn n' x8 `seq`
                    rnfn n' x9

#endif

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


{-# LANGUAGE BangPatterns, CPP, FlexibleContexts, GADTs, OverloadedStrings,
             RankNTypes #-}

{- |
   Module      : TestBench.Evaluate
   Description : Tree-based representation for Criterion and Weigh
   Copyright   : (c) Ivan Lazar Miljenovic
   License     : MIT
   Maintainer  : Ivan.Miljenovic@gmail.com

An extremely simple rose tree-based representation of criterion
benchmarks and weigh measurements.

 -}
module TestBench.Evaluate
  ( -- * Types
    EvalTree
  , EvalForest
  , Eval(..)
    -- ** Weights
  , GetWeight
  , getWeight
  , getWeightIO
    -- * Conversion
  , flattenBenchTree
  , flattenBenchForest
    -- * Running benchmarks
  , evalForest
    -- ** Weighing individual functions
  , weighIndex
  ) where

import TestBench.Commands  (resetUnusedConfig, weighFileArg, weighIndexArg)
import TestBench.LabelTree

import Criterion.Analysis          (OutlierVariance(ovFraction),
                                    SampleAnalysis(..))
import Criterion.Internal          (runAndAnalyseOne)
import Criterion.Measurement       (initializeTime, secs)
import Criterion.Measurement.Types (Benchmark, Benchmarkable, bench, bgroup)
import Criterion.Monad             (withConfig)
import Criterion.Types             (Config(..), DataRecord(..), Report(..))
import Statistics.Types            (ConfInt(..), Estimate(..))
import Weigh                       (weighAction, weighFunc)

import Data.Csv (DefaultOrdered(..), Field, Name, ToField, ToNamedRecord(..),
                 ToRecord(..), header, namedRecord, record, toField)

import qualified Data.DList        as DL
import           Streaming         (Of, Stream)
import           Streaming.Cassava (encodeByNameDefault)
import qualified Streaming.Prelude as S
import           Streaming.With    (writeBinaryFile)

import Control.Applicative              (liftA2)
import Control.DeepSeq                  (NFData)
import Control.Monad                    (join, when, zipWithM_)
import Control.Monad.IO.Class           (liftIO)
import Control.Monad.Trans.Class        (lift)
import Control.Monad.Trans.State.Strict (StateT, evalStateT, get, put)
import Data.Int                         (Int64)
import Data.List                        (intercalate)
import Data.Maybe                       (isJust, listToMaybe, mapMaybe)
import Data.String                      (IsString)
import System.Environment               (getExecutablePath)
import System.Exit                      (ExitCode(..))
import System.IO                        (hClose)
import System.IO.Temp                   (withSystemTempFile)
import System.Process                   (rawSystem)
import Text.Printf                      (printf)

#if MIN_VERSION_base (4,9,0)
import Data.Semigroup (Semigroup(..))
#endif

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

-- | A more explicit tree-like structure for benchmarks than using
--   Criterion's 'Benchmark' type.
type EvalTree = LabelTree Eval

type EvalForest = [EvalTree]

data Eval = Eval { eName  :: !String
                 , eBench :: !(Maybe Benchmarkable)
                 , eWeigh :: !(Maybe GetWeight)
                 }

-- | The results from measuring memory usage.
--
--   @since 0.2.0.0
data GetWeight where
  GetWeight   :: forall a b. (NFData b) => (a -> b)    -> a -> GetWeight
  GetWeightIO :: forall a b. (NFData b) => (a -> IO b) -> a -> GetWeight

runGetWeight :: GetWeight -> IO Weight
runGetWeight gw = mkWeight <$> case gw of
                                 GetWeight   f a -> weighFunc   f a
                                 GetWeightIO f a -> weighAction f a
  where
    mkWeight (b,gc,_,_) = Weight b gc

data Weight = Weight { bytesAlloc :: !Int64
                     , numGC      :: !Int64
                     }
            deriving (Eq, Ord, Show, Read)

-- | How to weigh a function.
--
--   @since 0.2.0.0
getWeight :: (NFData b) => (a -> b) -> a -> GetWeight
getWeight = GetWeight

-- | An IO-based variant of 'getWeight'.
--
--   @since 0.2.0.0
getWeightIO :: (NFData b) => (a -> IO b) -> a -> GetWeight
getWeightIO = GetWeightIO

flattenBenchTree :: EvalTree -> Maybe Benchmark
flattenBenchTree = fmap (foldLTree (const bgroup) (flip const))
                   . mapMaybeTree (liftA2 fmap (bench . eName) eBench)

-- | Remove the explicit tree-like structure into the implicit one
--   used by Criterion.
--
--   Useful for embedding the results into an existing benchmark
--   suite.
flattenBenchForest :: EvalForest -> [Benchmark]
flattenBenchForest = mapMaybe flattenBenchTree

-- | Run the specified benchmarks, printing the results (once they're
--   all complete) to stdout in a tabular format for easier
--   comparisons.
evalForest :: Config -> EvalForest -> IO ()
evalForest cfg ef = do when (hasBench ep) initializeTime
                       let ec = EC cfg ep
                       printHeaders ep
                       (`evalStateT` zeroIndex)
                         . maybeCSV
                         . S.mapM_ (liftIO . printRow ep)
                         . S.copy
                         $ toRows ec ef
  where
    ep = checkForest ef

    -- Ideally we would do the S.copy here only if we're writing to
    -- CSV; unfortunately, as we're not longer using ResourceT we
    -- don't have a MonadMask instance and thus can't use another
    -- Stream as the inner Monad; this /must/ be run after processing
    -- the stdout version.
    maybeCSV = maybe S.effects streamCSV (csvFile cfg)

    -- In reality, this type signature contains StateT, but that
    -- over-complicates understanding what it does, and to specify it
    -- generically requires bringing in MonadBaseControl and
    -- MonadThrow.

    -- streamCSV :: FilePath -> Stream (Of Row) IO () -> IO ()
    streamCSV fp = writeBinaryFile fp
                  . encodeByNameDefault
                  . S.filter isLeaf

data EvalParams = EP { hasBench  :: !Bool
                     , hasWeigh  :: !Bool
                     , nameWidth :: !Width
                     }
                deriving (Eq, Ord, Show, Read)

#if MIN_VERSION_base (4,9,0)
instance Semigroup EvalParams where
  (<>) = appendEP
#endif

instance Monoid EvalParams where
  mempty = EP { hasBench  = False
              , hasWeigh  = False
              , nameWidth = 0
              }
  mappend = appendEP

appendEP :: EvalParams -> EvalParams -> EvalParams
appendEP ec1 ec2 = EP { hasBench  = appendBy hasBench
                      , hasWeigh  = appendBy hasWeigh
                      , nameWidth = nameWidth ec1 `max` nameWidth ec2
                      }
  where
    appendBy f = f ec1 || f ec2

checkForest :: EvalForest -> EvalParams
checkForest = mconcat . map (foldLTree mergeNode calcConfig)
  where
    mergeNode d lbl ls = mempty { nameWidth = width d lbl }
                         `mappend` mconcat ls

    calcConfig d e = EP { hasBench  = isJust (eBench e)
                        , hasWeigh  = isJust (eWeigh e)
                        , nameWidth = width d (eName e)
                        }

    width d nm = indentPerLevel * d + length nm

data EvalConfig = EC { benchConfig :: {-# UNPACK #-}!Config
                     , evalParam   :: {-# UNPACK #-}!EvalParams
                     }
                deriving (Eq, Show, Read)

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

type PathList = DL.DList String

data Row = Row { rowLabel  :: !String
               , rowPath   :: !PathList -- ^ Invariant: length == rowDepth
               , rowDepth  :: {-# UNPACK #-} !Depth
               , isLeaf    :: !Bool
               , rowBench  :: !(Maybe BenchResults)
               , rowWeight :: !(Maybe Weight)
               }
  deriving (Eq, Show, Read)

pathLabel :: Row -> String
pathLabel row = intercalate "/" (DL.toList (DL.snoc (rowPath row) (rowLabel row)))

-- | Unlike terminal output, this instance creates columns for benchmarks, weighing, etc. even if they're not used.
instance ToRecord Row where
  toRecord row =
    record (toField fmtLabel : benchRecord ++ weightRecord)
    where
      fmtLabel = pathLabel row

      benchRecord = timed resMean ++ timed resStdDev ++ [bField (ovFraction . resOutVar)]
        where
          bField = mField . (. rowBench) . fmap
          timed f = map (bField . (. f)) [estPoint, estLowerBound, estUpperBound]

      weightRecord = [wField bytesAlloc, wField numGC]
        where
          wField = mField . (. rowWeight) . fmap

      mField :: (ToField a) => (Row -> Maybe a) -> Field
      mField f = maybe mempty toField (f row)

-- | Unlike terminal output, this instance creates columns for benchmarks, weighing, etc. even if they're not used.
instance ToNamedRecord Row where
  toNamedRecord row =
    namedRecord (fmtLabel : benchRecord ++ weightRecord)
    where
      fmtLabel = (labelName, toField (pathLabel row))

      benchRecord = zip benchNames
                        (timed resMean ++ timed resStdDev ++ [bField (ovFraction . resOutVar)])
        where
          bField = mField . (. rowBench) . fmap
          timed f = map (bField . (. f)) [estPoint, estLowerBound, estUpperBound]

      weightRecord = zip weighNames
                         [wField bytesAlloc, wField numGC]
        where
          wField = mField . (. rowWeight) . fmap

      mField :: (ToField a) => (Row -> Maybe a) -> Field
      mField f = maybe mempty toField (f row)

instance DefaultOrdered Row where
  headerOrder _ = header (labelName : benchNames ++ weighNames)

toRows :: EvalConfig -> EvalForest -> Stream (Of Row) (StateT Index IO) ()
toRows cfg = f2r DL.empty
  where
    f2r :: PathList -> EvalForest -> Stream (Of Row) (StateT Index IO) ()
    f2r pl = mapM_ (t2r pl)

    t2r :: PathList -> EvalTree -> Stream (Of Row) (StateT Index IO) ()
    t2r pl bt = case bt of
                  Leaf   d e      -> do i <- lift get
                                        lift (put (i+1))
                                        r <- liftIO (makeRow cfg pl i d e)
                                        S.yield r
                  Branch d lbl ts -> S.cons (Row lbl pl d False Nothing Nothing)
                                            (f2r (pl `DL.snoc` lbl) ts)

makeRow :: EvalConfig -> PathList -> Index -> Depth -> Eval -> IO Row
makeRow cfg pl idx d e = Row lbl pl d True
                         <$> tryRun hasBench eBench (getBenchResults (benchConfig cfg) lbl)
                         <*> tryRun hasWeigh eWeigh (const (tryGetWeight idx))
  where
    lbl = eName e
    ep = evalParam cfg

    tryRun :: (EvalParams -> Bool) -> (Eval -> Maybe a) -> (a -> IO (Maybe b)) -> IO (Maybe b)
    tryRun p f r =
      if p ep
         then maybe (return Nothing) r (f e)
         else return Nothing

data BenchResults = BenchResults { resMean   :: !(Estimate ConfInt Double)
                                 , resStdDev :: !(Estimate ConfInt Double)
                                 , resOutVar :: !OutlierVariance
                                 }
  deriving (Eq, Show, Read)

getBenchResults :: Config -> String -> Benchmarkable -> IO (Maybe BenchResults)
getBenchResults cfg lbl b = do dr <- withConfig cfg' (runAndAnalyseOne i lbl b)
                               return $ case dr of
                                          Measurement{} -> Nothing
                                          Analysed rpt  -> Just $
                                            let sa = reportAnalysis rpt
                                            in BenchResults { resMean   = anMean sa
                                                            , resStdDev = anStdDev sa
                                                            , resOutVar = anOutlierVar sa
                                                            }

  where
    -- Set this here just in case someone didn't use the top-level
    -- 'testBench' function.
    --
    -- Also, just in case a CSV file is being outputted, don't try and
    -- write to it.
    cfg' = resetUnusedConfig cfg { csvFile = Nothing }

    i = 0 -- We're ignoring this value anyway, so it should be OK to
          -- just set it.

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

type Index = Int

zeroIndex :: Index
zeroIndex = 0

tryGetWeight :: Index -> IO (Maybe Weight)
tryGetWeight idx = withSystemTempFile "testBench.weigh" $ \fp h -> do
  -- We use a temporary file in case the program prints something else
  -- out to stdout.
  hClose h -- We're not writing to it, just need the file
  exe <- getExecutablePath
  ec <- rawSystem exe ["--" ++ weighIndexArg, show idx, "--" ++ weighFileArg, fp, "+RTS", "-T", "-RTS"]
  case ec of
    ExitFailure{} -> return Nothing
    ExitSuccess   -> do
      out <- readFile fp
      case reads out of
        [(!mw,_)] -> return mw
        _         -> return Nothing

weighIndex :: EvalForest -> Index -> IO (Maybe Weight)
weighIndex ef = fmap join . mapM (mapM runGetWeight . eWeigh) . index es
  where
    es = concatMap leaves ef

index :: [a] -> Index -> Maybe a
index as n = listToMaybe . drop (n-1) $ as

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

printHeaders :: EvalParams -> IO ()
printHeaders ep = do putStr (replicate (nameWidth ep) ' ')
                     when (hasBench ep) (mapM_ toPrintf benchHeaders)
                     when (hasWeigh ep) (mapM_ toPrintf weighHeaders)
                     putStr "\n"
  where
    toPrintf (w,hdr) = printf "%s%*s" columnSpace w hdr

printRow :: EvalParams -> Row -> IO ()
printRow ep r = do printf "%-*s" (nameWidth ep) label
                   when (hasBench ep) (printBench (rowBench r))
                   when (hasWeigh ep) (printWeigh (rowWeight r))
                   putStr "\n"
  where
    label :: String
    label = printf "%s%s" (replicate (rowDepth r * indentPerLevel) ' ') (rowLabel r)

type Width = Int

indentPerLevel :: Width
indentPerLevel = 2

columnGap :: Width
columnGap = 2

columnSpace :: String
columnSpace = replicate columnGap ' '

labelName :: Name
labelName = "Label"

benchNames :: (IsString str) => [str]
benchNames = ["Mean", "MeanLB", "MeanUB", "Stddev", "StddevLB", "StddevUB", "OutlierVariance"]

benchHeaders :: [(Width, String)]
benchHeaders = map addWidth benchNames

weighNames :: (IsString str) => [str]
weighNames = ["AllocBytes", "NumGC"]

weighHeaders :: [(Width, String)]
weighHeaders = map addWidth weighNames

-- Maximum width a numeric field can take.  Might as well make them
-- all the same width.  All other formatters have been manually
-- adjusted to produce nothing longer than this.
secsWidth :: Width
secsWidth = length (secs ((-pi) / 000))

addWidth :: String -> (Width, String)
addWidth nm = (max (length nm) secsWidth, nm)

printBench :: Maybe BenchResults -> IO ()
printBench mr = zipWithM_ (printf "%s%*s" columnSpace) wdths cols
  where
    cols = maybe (repeat "")
                 (\r -> timed (resMean r) ++ timed (resStdDev r) ++ [ov r])
                 mr

    timed bs = [secIt estPoint, secIt estLowerBound, secIt estUpperBound]
      where
        secIt f = secs (f bs)

    ov r = percent (ovFraction (resOutVar r))

    wdths = map fst benchHeaders

estLowerBound :: (Num a) => Estimate ConfInt a -> a
estLowerBound e = estPoint e - confIntLDX (estError e)

estUpperBound :: (Num a) => Estimate ConfInt a -> a
estUpperBound e = estPoint e + confIntUDX (estError e)

printWeigh :: Maybe Weight -> IO ()
printWeigh mr = zipWithM_ (printf "%s%*s" columnSpace) wdths cols
  where
    cols = maybe (repeat "")
                 (\r -> [bytes (bytesAlloc r), count (numGC r)])
                 mr

    wdths = map fst weighHeaders

percent :: Double -> String
percent p = printf "%.3f%%" (p * 100)

-- | Human-readable description of the number of bytes used.  Assumed
--   non-negative.
bytes :: Int64 -> String
bytes b  = printf "%.3f %sB" val p
  where
    prefixes = [" ", "K", "M", "G", "T", "P", "E"] :: [String]

    base = 1024 :: Num a => a

    mult
      | b == 0    = 0
      | otherwise = floor (logBase (base :: Double) (fromIntegral b))

    val = fromIntegral b / (base ^ mult)

    p = prefixes !! mult

count :: Int64 -> String
count = printf "%.3e" . (`asTypeOf` (0::Double)) . fromIntegral


{-|
  Copyright   :  (C) 2019, Google Inc.
  License     :  BSD2 (see the file LICENSE)
  Maintainer  :  Christiaan Baaij <christiaan.baaij@gmail.com>
-}
{-# LANGUAGE CPP             #-}
{-# LANGUAGE TemplateHaskell #-}

{-# OPTIONS_GHC -Wno-orphans #-}

module Clash.Class.AutoReg.Instances where

import           Clash.Class.AutoReg.Internal
import           Clash.CPP                           (maxTupleSize)

#if MIN_VERSION_base(4,12,0)
import           Data.Complex (Complex)
import           Data.Ord (Down)
#endif

import           Data.Ratio (Ratio)

#if MIN_VERSION_base(4,12,0)
deriveAutoReg ''Complex
deriveAutoReg ''Down
#endif

deriveAutoReg ''Ratio

-- TODO: this needs NFDataX instances to maxTupleSize that in turn need
-- TODO: generic instances for large tuples, but this slows GHC down by a
-- TODO: lot. See: https://phabricator.haskell.org/D2899. Maybe we could add
-- TODO: -flarge-tuples to clash-prelude?
-- deriveAutoRegTuples [2..maxTupleSize]
deriveAutoRegTuples [2..min maxTupleSize 15]


{-# LANGUAGE BangPatterns          #-}
{-# LANGUAGE DeriveAnyClass        #-}
{-# LANGUAGE DeriveGeneric         #-}
{-# LANGUAGE DuplicateRecordFields #-}
{-# LANGUAGE RecordWildCards       #-}
{-# LANGUAGE ViewPatterns          #-}

module Main where

import Control.Parallel.Strategies
import Data.Char
import Data.Foldable
import Data.List.Split (splitOn)
import Data.Map.Strict (Map)
import qualified Data.Map.Strict as M
import Data.Maybe
import Data.Monoid
import Data.Set (Set)
import qualified Data.Set as S
import qualified Data.Vector as B
import GHC.Conc (numCapabilities)
import GHC.Generics (Generic)
import Numeric.GSL.SimulatedAnnealing
import qualified Numeric.LinearAlgebra as H

main :: IO ()
main = do
  !input <- parse <$> readFile "inputs/day-8.txt"

  -- The fun part is that doing this in parallel on a 12 core 24 thread CPU is
  -- only twice as fast as when not using any parallelization, and with 24
  -- threads it's 2.5 times slower. FFI calls require global locking, which is
  -- probably what's slowing this down.
  let chunkSize      = ceiling $ fromIntegral (length input) / (fromIntegral numCapabilities :: Float)
      mappings       = map (fromJust . solveMapping) input
      !mappedOutputs = map mapDisplayOutputs mappings `using` parListChunk chunkSize rdeepseq

  putStrLn "Part 1:"
  print . length . filter (`elem` [1, 4, 7, 8]) $ concat mappedOutputs

  putStrLn "\nPart 2:"
  print . sum $ map digitsToInt mappedOutputs

-- * Segments
--
-- These represent each value of the seven segment display in different ways.

-- | The expected lit segments for each of the ten digits that can be
-- represented by the four segment display. The segment indices for each number
-- are in sorted.
expectedSegments :: B.Vector (Set Char)
expectedSegments = B.fromList . map S.fromList $
  [ "abcefg"  -- 0
  , "cf"      -- 1
  , "acdeg"   -- 2
  , "acdfg"   -- 3
  , "bcdf"    -- 4
  , "abdfg"   -- 5
  , "abdefg"  -- 6
  , "acf"     -- 7
  , "abcdefg" -- 8
  , "abcdfg"  -- 9
  ]

-- * Solution

data Display = Display
  { -- | The ten digits from 'expectedSegments', but with the character ->
    -- segment assignments scrambled.
    patterns :: B.Vector (Set Char)
    -- | A four digit output value, with the same scrambled assignments as those
    -- in 'patterns'.
  , outputs  :: B.Vector (Set Char)
  }
  deriving (Show, Generic, NFData)

-- | The solver state for solving the mappings within a display.
data SolverState = SolverState
  { display        :: Display
    -- | The mapping from a pattern from 'patterns' to the character in
    -- 'expectedSegments'.
  , patternMapping :: !(Map Int Int)
    -- | The mapping from a character on the seven segment display as used in a
    -- 'Display' to the actual letter of the segment.
  , charMapping    :: !(Map Char Char)
  }
  deriving (Show, Generic, NFData)

-- | The actions we can take in our solver. Swapping a mapping means swapping
-- two values in either of the maps in 'SolverState'.
data SolverAction = SwapPatternMapping | SwapCharMapping
  deriving (Show, Enum, Bounded)

applyMapping :: Ord a => Map a a -> Set a -> Set a
applyMapping mapping = S.map (mapping M.!)

-- | The penalty/loss for a solver state. We want to minimize this.
solverPenalty :: SolverState -> Int
solverPenalty SolverState{..} = lengthPenalty + mappingPenalty
  where
    -- We need the pattern state from the solver permuted to match the expected
    -- segments, alongside the expected segments
    permutedPatterns = B.backpermute (patterns display) (B.fromList $ map (patternMapping M.!) [0..9])
    zippedPatterns   = B.zip expectedSegments permutedPatterns

    -- Because the set difference only tells you which elements from set A are not in set B
    symmetricDifference :: Ord a => Set a -> Set a -> Set a
    symmetricDifference a b = S.difference a b `S.union` S.difference b a

    -- | The sum of length errors when using the pattern mapping.
    lengthPenalty = getSum $
      foldMap (\(expected, actual) -> Sum . abs $ S.size expected - S.size actual)
              zippedPatterns

    -- | The number of difference characters in each segment set after applying
    -- the pattern mapping.
    mappingPenalty = getSum $
      foldMap (\(expected, actual) -> Sum . abs . S.size $ symmetricDifference expected (applyMapping charMapping actual))
              zippedPatterns

-- | Find the matching mapping from the scrambled pattern space to the expected
-- segments. There's a clever way to do this. This is not it.
solveMapping :: Display -> Maybe SolverState
solveMapping display'
  -- There should only be an exact solution
  | solverPenalty optimalSolution == 0 = Just optimalSolution
  | otherwise                          = Nothing
  where
    optimalSolution = simanSolve 420 3 solverParams initialState
      (fromIntegral . solverPenalty)
      -- This distance metric is not entirely correct
      (\x y -> fromIntegral . abs $ solverPenalty x - solverPenalty y)
      step
      Nothing -- (Just $ show . solverPenalty)

    -- This solver is WAY overkill, so the cooling rate and final temperature
    -- are relatively high to not waste too much time on the simpler to solve
    -- cases
    solverParams = SimulatedAnnealingParams
      100   -- Tries per step
      1000  -- Tries per temperature
      1     -- Maximum step size in random walk
      1.0   -- Boltzman constant for random walks
      1.0   -- Initial temperature
      2.0   -- Cooling rate
      0.069 -- Final temperature

    initialState = SolverState
      { display        = display'
      , patternMapping = M.fromList $ map (\idx -> (idx, idx)) [0..9]
      , charMapping    = M.fromList $ map (\c -> (c, c)) "abcdefg"
      }

    solverActionRange = fromIntegral (fromEnum (maxBound :: SolverAction))

    swapRange SwapPatternMapping = 9.0
    swapRange SwapCharMapping    = fromIntegral $ ord 'g' - ord 'a'

    swap :: Ord k => k -> k -> Map k v -> Map k v
    swap from to m = M.insert to (m M.! from) $ M.insert from (m M.! to) m

    -- Where n starts at 0, as generated by the RNG.
    char :: Int -> Char
    char n = chr (n + ord 'a')

    step :: H.Vector Double -> Double -> SolverState -> SolverState
    step rands _ state@SolverState{..} =
      let action   = toEnum (round $ (rands H.! 0) * solverActionRange) :: SolverAction
          -- These two can be the same
          swapFrom = round $ (rands H.! 1) * swapRange action
          swapTo   = round $ (rands H.! 2) * swapRange action
       in case action of
            SwapPatternMapping -> state { patternMapping = swap swapFrom swapTo patternMapping }
            SwapCharMapping    -> state { charMapping    = swap (char swapFrom) (char swapTo) charMapping }

-- | Grab the translated display numbers from a solver result.
mapDisplayOutputs :: SolverState -> [Int]
mapDisplayOutputs SolverState{display = Display{..}, ..}
  = map ((segmentsMapping M.!) . applyMapping charMapping)
  $ B.toList outputs
  where
    segmentsMapping :: Map (Set Char) Int
    segmentsMapping = M.fromList $ B.toList $ B.imap (\idx segment -> (segment, idx)) expectedSegments

-- | convert a list of digits to an integer.
digitsToInt :: [Int] -> Int
digitsToInt = foldl' (\acc n -> (acc * 10) + n) 0

parse :: String -> [Display]
parse = map parseLine . lines
  where
    parseLine :: String -> Display
    parseLine l = case splitOn " | " l of
      [words -> ps, words -> os] -> Display (parseSegments ps) (parseSegments os)
      _                          -> error "What is this even!?"

    parseSegments :: [[Char]] -> B.Vector (Set Char)
    parseSegments = B.fromList . map S.fromList


module Feldspar
  ( module Prelude.EDSL
  , module Control.Monad
    -- * Types
    -- ** Syntax
  , Data
  , Syntax
  , Comp
    -- ** Object-language types
  , module Data.Int
  , module Data.Word
  , Complex (..)
  , PrimType'
  , PrimType
  , Type
  , Length
  , Index
  , Ref, DRef
  , Arr, DArr
  , IArr, DIArr
  , Inhabited
  , Syntactic
  , Domain
  , Internal
    -- * Front end
  , eval
  , showAST
  , drawAST
  , module Feldspar.Frontend
  , Bits
  , FiniteBits
  , Integral
  , Ord
  , RealFloat
  , RealFrac
  , Border (..)
  , IxRange
  , AssertionLabel (..)
  ) where

import Prelude.EDSL

import Control.Monad hiding (foldM)

import Data.Bits (Bits, FiniteBits)
import Data.Complex (Complex (..))
import Data.Int
import Data.Word

import Language.Syntactic (Syntactic, Domain, Internal)
import qualified Language.Syntactic as Syntactic

import Language.Embedded.Imperative (Border (..), IxRange)

import Data.Inhabited
import Feldspar.Primitive.Representation
import Feldspar.Representation
import Feldspar.Frontend
import Feldspar.Optimize

-- | Show the syntax tree using Unicode art
--
-- Only user assertions will be included.
showAST :: (Syntactic a, Domain a ~ FeldDomain) => a -> String
showAST = Syntactic.showAST . optimize onlyUserAssertions . Syntactic.desugar

-- | Draw the syntax tree on the terminal using Unicode art
--
-- Only user assertions will be included.
drawAST :: (Syntactic a, Domain a ~ FeldDomain) => a -> IO ()
drawAST = Syntactic.drawAST . optimize onlyUserAssertions . Syntactic.desugar


{-# LANGUAGE DeriveDataTypeable #-}

module Text.Namelist.Types where

import Data.Complex
import Data.CaseInsensitive (CI)
import Data.Data(Data, Typeable)

data Index
    = Index Int
    | Range (Maybe Int) (Maybe Int) (Maybe Int)
    deriving (Show, Read, Data, Typeable, Eq)

data Key
    = Key     (CI String)
    | Indexed (CI String) [Index]
    | Sub     (CI String) Key
    deriving (Show, Read, Data, Typeable, Eq)

data Value
    = Integer Int
    | Real Double
    | Complex (Complex Double)
    | Logical Bool
    | String String
    | Array [Value]
    | Int :* Value
    | Null
    deriving (Show, Read, Data, Typeable, Eq)

data Pair = Key := Value
    deriving (Show, Read, Data, Typeable, Eq)

data Group = Group (CI String) [Pair]
    deriving (Show, Read, Data, Typeable, Eq)


module LatticeSymmetries.Group
  ( Permutation,
    unPermutation,
    mkPermutation,
    identityPermutation,
    fromGenerators,
    pgLength,
    Symmetry,
    mkSymmetry,
    SymmetriesHeader (..),
    mkSymmetries,
    symmetriesFromHeader,
    Symmetries (..),
    nullSymmetries,
    emptySymmetries,
    symmetriesGetNumberBits,
    symmetriesFromSpec,
  )
where

import Data.Aeson (FromJSON (..), ToJSON (..), object, withObject, (.:), (.=))
import Data.Complex
import qualified Data.List
import Data.Ratio
import qualified Data.Set as Set
import qualified Data.Vector as B
import qualified Data.Vector.Generic as G
import qualified Data.Vector.Storable as S
import Foreign.C.Types (CDouble)
import Foreign.ForeignPtr
import Foreign.Ptr
import Foreign.Storable
import GHC.Exts (IsList (..))
import LatticeSymmetries.Benes
import LatticeSymmetries.Dense
import LatticeSymmetries.FFI
import LatticeSymmetries.IO
import LatticeSymmetries.Utils
import System.IO.Unsafe (unsafePerformIO)
import Prelude hiding (identity, permutations, toList)

newtype PermutationGroup = PermutationGroup (B.Vector Permutation)
  deriving stock (Show, Eq)

getPeriodicity :: Permutation -> Int
getPeriodicity p₀ = go 1 p₀
  where
    identity = identityPermutation (G.length (unPermutation p₀))
    go !n !p
      | p == identity = n
      | otherwise = go (n + 1) (p₀ <> p)

nullPermutationGroup :: PermutationGroup -> Bool
nullPermutationGroup (PermutationGroup gs) = G.null gs

data Symmetry = Symmetry
  { symmetryPermutation :: !Permutation,
    symmetryPhase :: !(Ratio Int)
  }
  deriving stock (Show, Eq, Ord)

symmetrySector :: Symmetry -> Int
symmetrySector s
  | denominator sector == 1 = numerator sector
  | otherwise = error $ "this should not have happened: invalid symmetry: " <> show s
  where
    sector = symmetryPhase s * fromIntegral (getPeriodicity (symmetryPermutation s))

instance FromJSON Symmetry where
  parseJSON = withObject "Symmetry" $ \v -> do
    mkSymmetry <$> v .: "permutation" <*> v .: "sector"

instance ToJSON Symmetry where
  toJSON s = object ["permutation" .= symmetryPermutation s, "sector" .= symmetrySector s]

mkSymmetry :: Permutation -> Int -> Symmetry
mkSymmetry p sector
  | sector < 0 || sector >= periodicity =
    error $
      "invalid sector: " <> show sector <> "; permutation has periodicity " <> show periodicity
  | otherwise = Symmetry p (sector % periodicity)
  where
    periodicity = getPeriodicity p

symmetryNumberSites :: Symmetry -> Int
symmetryNumberSites (Symmetry p _) = G.length (unPermutation p)

modOne :: Integral a => Ratio a -> Ratio a
modOne x
  | x >= 1 = x - fromIntegral (numerator x `div` denominator x)
  | otherwise = x

instance Semigroup Symmetry where
  (<>) (Symmetry pa λa) (Symmetry pb λb) = Symmetry (pa <> pb) (modOne (λa + λb))

data SymmetriesHeader = SymmetriesHeader
  { symmHeaderGroup :: !PermutationGroup,
    symmHeaderNetwork :: !BatchedBenesNetwork,
    symmHeaderCharactersReal :: !(S.Vector Double),
    symmHeaderCharactersImag :: !(S.Vector Double)
  }
  deriving stock (Show, Eq)

getNumberBits :: SymmetriesHeader -> Int
getNumberBits (SymmetriesHeader (PermutationGroup gs) _ _ _)
  | G.null gs = 0
  | otherwise = G.length . unPermutation . G.head $ gs

emptySymmetriesHeader :: SymmetriesHeader
emptySymmetriesHeader = SymmetriesHeader (PermutationGroup G.empty) emptyBatchedBenesNetwork G.empty G.empty

data Symmetries = Symmetries {symmetriesHeader :: !SymmetriesHeader, symmetriesContents :: !(ForeignPtr Cpermutation_group)}
  deriving stock (Show)

instance Eq Symmetries where
  (==) a b = symmetriesHeader a == symmetriesHeader b

instance FromJSON Symmetries where
  parseJSON xs = mkSymmetries <$> parseJSON xs

instance ToJSON Symmetries where
  toJSON symmetries = toJSON (toSymmetryList symmetries)

toSymmetryList :: Symmetries -> [Symmetry]
toSymmetryList (Symmetries (SymmetriesHeader (PermutationGroup gs) _ λsRe λsIm) _) =
  Data.List.zipWith3 _toSymmetry (G.toList gs) (G.toList λsRe) (G.toList λsIm)
  where
    _toSymmetry :: Permutation -> Double -> Double -> Symmetry
    _toSymmetry g λRe λIm
      | r ≈ 1 && s' ≈ fromIntegral (round s' :: Int) = mkSymmetry g (round s')
      | otherwise = error $ "failed to reconstruct Symmetry from " <> show (toList g) <> " and λ = " <> show λRe <> " + " <> show λIm <> "j"
      where
        -- φ = sector / periodicity, but care needs to be taken because φ is approximate
        s' = φ * fromIntegral n
        n = getPeriodicity g
        -- polar returns the phase in (-π, π], we add 2π to get a positive value
        (r, _φ) = polar (λRe :+ λIm)
        φ = _φ / (2 * pi) + (if _φ < 0 then 1 else 0)

pgLength :: PermutationGroup -> Int
pgLength (PermutationGroup gs) = G.length gs

fromGenerators :: (Semigroup a, Ord a) => a -> [a] -> [a]
fromGenerators _ [] = []
fromGenerators identity gs = go Set.empty (Set.singleton identity)
  where
    go !interior !boundary
      | Set.null boundary = Set.toAscList $ interior
      | otherwise = go interior' boundary'
      where
        interior' = interior `Set.union` boundary
        boundary' = Set.fromList [h <> g | h <- Set.toList boundary, g <- gs] Set.\\ interior'

nullSymmetries :: Symmetries -> Bool
nullSymmetries = nullPermutationGroup . symmHeaderGroup . symmetriesHeader

emptySymmetries :: Symmetries
emptySymmetries = symmetriesFromHeader emptySymmetriesHeader

symmetriesGetNumberBits :: Symmetries -> Int
symmetriesGetNumberBits = getNumberBits . symmetriesHeader

mkSymmetries :: [Symmetry] -> Symmetries
mkSymmetries gs = case symmetriesFromHeader <$> mkSymmetriesHeader gs of
  Just s -> s
  Nothing -> error "incompatible symmetries"

mkSymmetriesHeader :: [Symmetry] -> Maybe SymmetriesHeader
mkSymmetriesHeader [] = Just emptySymmetriesHeader
mkSymmetriesHeader gs@(g : _)
  | all ((== n) . symmetryNumberSites) gs = case isConsistent of
    True ->
      let permutations = G.fromList $ symmetryPermutation <$> symmetries
          permGroup = PermutationGroup permutations
          benesNetwork = mkBatchedBenesNetwork $ G.map toBenesNetwork permutations
          charactersReal = G.fromList $ (\φ -> cos (-2 * pi * realToFrac φ)) <$> symmetryPhase <$> symmetries
          charactersImag = G.fromList $ (\φ -> sin (-2 * pi * realToFrac φ)) <$> symmetryPhase <$> symmetries
       in Just $ SymmetriesHeader permGroup benesNetwork charactersReal charactersImag
    False -> Nothing
  | otherwise = error "symmetries have different number of sites"
  where
    n = symmetryNumberSites g
    identity = mkSymmetry (identityPermutation n) 0
    symmetries = fromGenerators identity gs
    isConsistent = all id $ do
      s₁ <- symmetries
      s₂ <- symmetries
      let s₃@(Symmetry p₃ λ₃) = s₁ <> s₂
      pure $
        Set.member s₃ set
          && denominator (λ₃ * fromIntegral (getPeriodicity p₃)) == 1
      where
        set = Set.fromList symmetries

symmetriesFromHeader :: SymmetriesHeader -> Symmetries
symmetriesFromHeader x = unsafePerformIO $ do
  fp <- mallocForeignPtr
  let (DenseMatrix numberShifts numberMasks masks) = bbnMasks $ symmHeaderNetwork x
      shifts = bbnShifts $ symmHeaderNetwork x
      numberBits = getNumberBits x
  withForeignPtr fp $ \ptr ->
    S.unsafeWith masks $ \masksPtr ->
      S.unsafeWith shifts $ \shiftsPtr ->
        S.unsafeWith (symmHeaderCharactersReal x) $ \eigvalsRealPtr ->
          S.unsafeWith (symmHeaderCharactersImag x) $ \eigvalsImagPtr ->
            poke ptr $
              Cpermutation_group
                { cpermutation_group_refcount = 0,
                  cpermutation_group_number_bits = fromIntegral numberBits,
                  cpermutation_group_number_shifts = fromIntegral numberShifts,
                  cpermutation_group_number_masks = fromIntegral numberMasks,
                  cpermutation_group_masks = masksPtr,
                  cpermutation_group_shifts = shiftsPtr,
                  cpermutation_group_eigvals_re = (castPtr :: Ptr Double -> Ptr CDouble) eigvalsRealPtr,
                  cpermutation_group_eigvals_im = (castPtr :: Ptr Double -> Ptr CDouble) eigvalsImagPtr
                }
  pure $ Symmetries x fp
{-# NOINLINE symmetriesFromHeader #-}

symmetryFromSpec :: SymmetrySpec -> Symmetry
symmetryFromSpec (SymmetrySpec p k) =
  mkSymmetry (mkPermutation (G.fromList (toList p))) k

symmetriesFromSpec :: [SymmetrySpec] -> Symmetries
symmetriesFromSpec = mkSymmetries . fmap symmetryFromSpec