module Feldspar.Frontend where

import Prelude (Integral, Ord, RealFloat, RealFrac)
import qualified Prelude as P
import Prelude.EDSL

import Control.Monad.Identity
import Data.Bits (Bits, FiniteBits)
import qualified Data.Bits as Bits
import Data.Complex (Complex)
import Data.Int
import Data.List (genericLength)

import Language.Syntactic (Internal)
import Language.Syntactic.Functional
import qualified Language.Syntactic as Syntactic

import qualified Control.Monad.Operational.Higher as Oper
import Language.Embedded.Imperative (IxRange)
import qualified Language.Embedded.Imperative as Imp

import qualified Data.Inhabited as Inhabited
import Data.TypedStruct
import Feldspar.Primitive.Representation
import Feldspar.Representation
import Feldspar.Sugar ()

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

----------------------------------------
-- ** General constructs
----------------------------------------

-- | Force evaluation of a value and share the result. Note that due to common
-- sub-expression elimination, this function is rarely needed in practice.
share :: (Syntax a, Syntax b)
    => a         -- ^ Value to share
    -> (a -> b)  -- ^ Body in which to share the value
    -> b
share = shareTag ""
  -- Explicit sharing can be useful e.g. when the value to share contains a
  -- function or when the code motion algorithm for some reason doesn't work
  -- find opportunities for sharing.

-- | Explicit tagged sharing
shareTag :: (Syntax a, Syntax b)
    => String    -- ^ A tag (that may be empty). May be used by a back end to
                 --   generate a sensible variable name.
    -> a         -- ^ Value to share
    -> (a -> b)  -- ^ Body in which to share the value
    -> b
shareTag tag = sugarSymFeld (Let tag)

-- | For loop
forLoop :: Syntax st => Data Length -> st -> (Data Index -> st -> st) -> st
forLoop = sugarSymFeld ForLoop

-- | Conditional expression
cond :: Syntax a
    => Data Bool  -- ^ Condition
    -> a          -- ^ True branch
    -> a          -- ^ False branch
    -> a
cond = sugarSymFeld Cond

-- | Condition operator; use as follows:
--
-- @
-- cond1 `?` a $
-- cond2 `?` b $
-- cond3 `?` c $
--         default
-- @
(?) :: Syntax a
    => Data Bool  -- ^ Condition
    -> a          -- ^ True branch
    -> a          -- ^ False branch
    -> a
(?) = cond

infixl 1 ?

-- | Multi-way conditional expression
--
-- The first association @(a,b)@ in the list of cases for which @a@ is equal to
-- the scrutinee is selected, and the associated @b@ is returned as the result.
-- If no case matches, the default value is returned.
switch :: (Syntax a, Syntax b, PrimType (Internal a))
    => b        -- ^ Default result
    -> [(a,b)]  -- ^ Cases (match, result)
    -> a        -- ^ Scrutinee
    -> b        -- ^ Result
switch def [] _ = def
switch def cs s = P.foldr
    (\(c,a) b -> desugar c == desugar s ? a $ b)
    def
    cs

----------------------------------------
-- ** Literals
----------------------------------------

-- | Literal
value :: Syntax a => Internal a -> a
value = sugarSymFeld . Lit

false :: Data Bool
false = value False

true :: Data Bool
true = value True

instance Syntactic.Syntactic ()
  where
    type Domain ()   = FeldDomain
    type Internal () = Int32
    desugar () = unData 0
    sugar   _  = ()

-- | Example value
--
-- 'example' can be used similarly to 'undefined' in normal Haskell, i.e. to
-- create an expression whose value is irrelevant.
--
-- Note that it is generally not possible to use 'undefined' in Feldspar
-- expressions, as this will crash the compiler.
example :: Syntax a => a
example = value Inhabited.example

----------------------------------------
-- ** Primitive functions
----------------------------------------

instance (Bounded a, Type a) => Bounded (Data a)
  where
    minBound = value minBound
    maxBound = value maxBound

instance (Num a, PrimType a) => Num (Data a)
  where
    fromInteger = value . fromInteger
    (+)         = sugarSymFeld Add
    (-)         = sugarSymFeld Sub
    (*)         = sugarSymFeld Mul
    negate      = sugarSymFeld Neg
    abs         = sugarSymFeld Abs
    signum      = sugarSymFeld Sign

instance (Fractional a, PrimType a) => Fractional (Data a)
  where
    fromRational = value . fromRational
    (/) = sugarSymFeld FDiv

instance (Floating a, PrimType a) => Floating (Data a)
  where
    pi    = sugarSymFeld Pi
    exp   = sugarSymFeld Exp
    log   = sugarSymFeld Log
    sqrt  = sugarSymFeld Sqrt
    (**)  = sugarSymFeld Pow
    sin   = sugarSymFeld Sin
    cos   = sugarSymFeld Cos
    tan   = sugarSymFeld Tan
    asin  = sugarSymFeld Asin
    acos  = sugarSymFeld Acos
    atan  = sugarSymFeld Atan
    sinh  = sugarSymFeld Sinh
    cosh  = sugarSymFeld Cosh
    tanh  = sugarSymFeld Tanh
    asinh = sugarSymFeld Asinh
    acosh = sugarSymFeld Acosh
    atanh = sugarSymFeld Atanh

-- | Alias for 'pi'
π :: (Floating a, PrimType a) => Data a
π = pi

-- | Integer division truncated toward zero
quot :: (Integral a, PrimType a) => Data a -> Data a -> Data a
quot = sugarSymFeld Quot

-- | Integer remainder satisfying
--
-- > (x `quot` y)*y + (x `rem` y) == x
rem :: (Integral a, PrimType a) => Data a -> Data a -> Data a
rem = sugarSymFeld Rem

-- | Simultaneous @quot@ and @rem@
quotRem :: (Integral a, PrimType a) => Data a -> Data a -> (Data a, Data a)
quotRem a b = (q,r)
  where
    q = quot a b
    r = a - b * q

-- | Integer division truncated toward negative infinity
div :: (Integral a, PrimType a) => Data a -> Data a -> Data a
div = sugarSymFeld Div

-- | Integer modulus, satisfying
--
-- > (x `div` y)*y + (x `mod` y) == x
mod :: (Integral a, PrimType a) => Data a -> Data a -> Data a
mod = sugarSymFeld Mod

-- | Integer division assuming `unsafeBalancedDiv x y * y == x` (i.e. no
-- remainder)
--
-- The advantage of using 'unsafeBalancedDiv' over 'quot' or 'div' is that the
-- above assumption can be used for simplifying the expression.
unsafeBalancedDiv :: (Integral a, PrimType a) => Data a -> Data a -> Data a
unsafeBalancedDiv a b = guardValLabel
    InternalAssertion
    (rem a b == 0)
    "unsafeBalancedDiv: division not balanced"
    (sugarSymFeld DivBalanced a b)
  -- Note: We can't check that `result * b == a`, because `result * b` gets
  -- simplified to `a`.

-- | Construct a complex number
complex :: (Num a, PrimType a, PrimType (Complex a))
    => Data a  -- ^ Real part
    -> Data a  -- ^ Imaginary part
    -> Data (Complex a)
complex = sugarSymFeld Complex

-- | Construct a complex number
polar :: (Floating a, PrimType a, PrimType (Complex a))
    => Data a  -- ^ Magnitude
    -> Data a  -- ^ Phase
    -> Data (Complex a)
polar = sugarSymFeld Polar

-- | Extract the real part of a complex number
realPart :: (PrimType a, PrimType (Complex a)) => Data (Complex a) -> Data a
realPart = sugarSymFeld Real

-- | Extract the imaginary part of a complex number
imagPart :: (PrimType a, PrimType (Complex a)) => Data (Complex a) -> Data a
imagPart = sugarSymFeld Imag

-- | Extract the magnitude of a complex number's polar form
magnitude :: (RealFloat a, PrimType a, PrimType (Complex a)) =>
    Data (Complex a) -> Data a
magnitude = sugarSymFeld Magnitude

-- | Extract the phase of a complex number's polar form
phase :: (RealFloat a, PrimType a, PrimType (Complex a)) =>
    Data (Complex a) -> Data a
phase = sugarSymFeld Phase

-- | Complex conjugate
conjugate :: (RealFloat a, PrimType (Complex a)) =>
    Data (Complex a) -> Data (Complex a)
conjugate = sugarSymFeld Conjugate
  -- `RealFloat` could be replaced by `Num` here, but it seems more consistent
  -- to use `RealFloat` for all functions.

-- | Integral type casting
i2n :: (Integral i, Num n, PrimType i, PrimType n) => Data i -> Data n
i2n = sugarSymFeld I2N

-- | Cast integer to 'Bool'
i2b :: (Integral a, PrimType a) => Data a -> Data Bool
i2b = sugarSymFeld I2B

-- | Cast 'Bool' to integer
b2i :: (Integral a, PrimType a) => Data Bool -> Data a
b2i = sugarSymFeld B2I

-- | Round a floating-point number to an integer
round :: (RealFrac a, Num b, PrimType a, PrimType b) => Data a -> Data b
round = sugarSymFeld Round

-- | Boolean negation
not :: Data Bool -> Data Bool
not = sugarSymFeld Not

-- | Boolean conjunction
(&&) :: Data Bool -> Data Bool -> Data Bool
(&&) = sugarSymFeld And

infixr 3 &&

-- | Boolean disjunction
(||) :: Data Bool -> Data Bool -> Data Bool
(||) = sugarSymFeld Or

infixr 2 ||

-- | Equality
(==) :: PrimType a => Data a -> Data a -> Data Bool
(==) = sugarSymFeld Eq

-- | Inequality
(/=) :: PrimType a => Data a -> Data a -> Data Bool
a /= b = not (a==b)

-- | Less than
(<) :: (Ord a, PrimType a) => Data a -> Data a -> Data Bool
(<) = sugarSymFeld Lt

-- | Greater than
(>) :: (Ord a, PrimType a) => Data a -> Data a -> Data Bool
(>) = sugarSymFeld Gt

-- | Less than or equal
(<=) :: (Ord a, PrimType a) => Data a -> Data a -> Data Bool
(<=) = sugarSymFeld Le

-- | Greater than or equal
(>=) :: (Ord a, PrimType a) => Data a -> Data a -> Data Bool
(>=) = sugarSymFeld Ge

infix 4 ==, /=, <, >, <=, >=

-- | Return the smallest of two values
min :: (Ord a, PrimType a) => Data a -> Data a -> Data a
min a b = a<=b ? a $ b
  -- There's no standard definition of min/max in C:
  -- <http://stackoverflow.com/questions/3437404/min-and-max-in-c>
  --
  -- There is `fmin`/`fminf` for floating-point numbers, but these are
  -- implemented essentially as above (except that they handle `NaN`
  -- specifically:
  -- <https://sourceware.org/git/?p=glibc.git;a=blob;f=math/s_fmin.c;hb=HEAD>

-- | Return the greatest of two values
max :: (Ord a, PrimType a) => Data a -> Data a -> Data a
max a b = a>=b ? a $ b

----------------------------------------
-- ** Bit manipulation
----------------------------------------

-- | Bit-wise \"and\"
(.&.) :: (Bits a, PrimType a) => Data a -> Data a -> Data a
(.&.) = sugarSymFeld BitAnd

-- | Bit-wise \"or\"
(.|.) :: (Bits a, PrimType a) => Data a -> Data a -> Data a
(.|.) = sugarSymFeld BitOr

-- | Bit-wise \"xor\"
xor :: (Bits a, PrimType a) => Data a -> Data a -> Data a
xor = sugarSymFeld BitXor

-- | Bit-wise \"xor\"
(⊕) :: (Bits a, PrimType a) => Data a -> Data a -> Data a
(⊕) = xor

-- | Bit-wise complement
complement :: (Bits a, PrimType a) => Data a -> Data a
complement = sugarSymFeld BitCompl

-- | Left shift
shiftL :: (Bits a, PrimType a)
    => Data a      -- ^ Value to shift
    -> Data Int32  -- ^ Shift amount (negative value gives right shift)
    -> Data a
shiftL = sugarSymFeld ShiftL

-- | Right shift
shiftR :: (Bits a, PrimType a)
    => Data a      -- ^ Value to shift
    -> Data Int32  -- ^ Shift amount (negative value gives left shift)
    -> Data a
shiftR = sugarSymFeld ShiftR

-- | Left shift
(.<<.) :: (Bits a, PrimType a)
    => Data a      -- ^ Value to shift
    -> Data Int32  -- ^ Shift amount (negative value gives right shift)
    -> Data a
(.<<.) = shiftL

-- | Right shift
(.>>.) :: (Bits a, PrimType a)
    => Data a      -- ^ Value to shift
    -> Data Int32  -- ^ Shift amount (negative value gives left shift)
    -> Data a
(.>>.) = shiftR

infixl 8 `shiftL`, `shiftR`, .<<., .>>.
infixl 7 .&.
infixl 6 `xor`
infixl 5 .|.

bitSize :: forall a . FiniteBits a => Data a -> Length
bitSize _ = P.fromIntegral $ Bits.finiteBitSize (a :: a)
  where
    a = P.error "finiteBitSize evaluates its argument"

-- | Set all bits to one
allOnes :: (Bits a, Num a, PrimType a) => Data a
allOnes = complement 0

-- | Set the @n@ lowest bits to one
oneBits :: (Bits a, Num a, PrimType a) => Data Int32 -> Data a
oneBits n = complement (allOnes .<<. n)

-- | Extract the @k@ lowest bits
lsbs :: (Bits a, Num a, PrimType a) => Data Int32 -> Data a -> Data a
lsbs k i = i .&. oneBits k

-- | Integer logarithm in base 2. Returns \(\lfloor log_2(x) \rfloor\).
-- Assumes \(x>0\).
ilog2 :: (FiniteBits a, Integral a, PrimType a) => Data a -> Data a
ilog2 a = guardValLabel InternalAssertion (a >= 1) "ilog2: argument < 1" $
    snd $ P.foldr (\ffi vr -> share vr (step ffi)) (a,0) ffis
  where
    step (ff,i) (v,r) =
        share (b2i (v > fromInteger ff) .<<. value i) $ \shift ->
          (v .>>. i2n shift, r .|. shift)

    -- [(0x1, 0), (0x3, 1), (0xF, 2), (0xFF, 3), (0xFFFF, 4), ...]
    ffis
        = (`P.zip` [0..])
        $ P.takeWhile (P.<= (2 P.^ (bitSize a `P.div` 2) - 1 :: Integer))
        $ P.map ((subtract 1) . (2 P.^) . (2 P.^))
        $ [(0::Integer)..]
  -- Based on this algorithm:
  -- <http://graphics.stanford.edu/~seander/bithacks.html#IntegerLog>

----------------------------------------
-- ** Arrays
----------------------------------------

-- | Index into an array
arrIx :: Syntax a => IArr a -> Data Index -> a
arrIx arr i = resugar $ mapStruct ix $ unIArr arr
  where
    ix :: forall b . PrimType' b => Imp.IArr Index b -> Data b
    ix arr' = sugarSymFeldPrim
      (GuardVal InternalAssertion "arrIx: index out of bounds")
      (i < length arr)
      (sugarSymFeldPrim (ArrIx arr') (i + iarrOffset arr) :: Data b)

class Indexed a
  where
    type IndexedElem a

    -- | Indexing operator. If @a@ is 'Finite', it is assumed that
    -- @i < `length` a@ in any expression @a `!` i@.
    (!) :: a -> Data Index -> IndexedElem a

infixl 9 !

-- | Linear structures with a length. If the type is also 'Indexed', the length
-- is the successor of the maximal allowed index.
class Finite a
  where
    -- | The length of a finite structure
    length :: a -> Data Length

instance Finite (Arr a)  where length = arrLength
instance Finite (IArr a) where length = iarrLength

-- | Linear structures that can be sliced
class Slicable a
  where
    -- | Take a slice of a structure
    slice
        :: Data Index   -- ^ Start index
        -> Data Length  -- ^ Slice length
        -> a            -- ^ Structure to slice
        -> a

instance Syntax a => Indexed (IArr a)
  where
    type IndexedElem (IArr a) = a
    (!) = arrIx

instance Slicable (Arr a)
  where
    slice from len (Arr o l arr) = Arr o' l' arr
      where
        o' = guardValLabel InternalAssertion (from<=l) "invalid Arr slice" (o+from)
        l' = guardValLabel InternalAssertion (from+len<=l) "invalid Arr slice" len

instance Slicable (IArr a)
  where
    slice from len (IArr o l arr) = IArr o' l' arr
      where
        o' = guardValLabel InternalAssertion (from<=l) "invalid IArr slice" (o+from)
        l' = guardValLabel InternalAssertion (from+len<=l) "invalid IArr slice" len

----------------------------------------
-- ** Syntactic conversion
----------------------------------------

desugar :: Syntax a => a -> Data (Internal a)
desugar = Data . Syntactic.desugar

sugar :: Syntax a => Data (Internal a) -> a
sugar = Syntactic.sugar . unData

-- | Cast between two values that have the same syntactic representation
resugar :: (Syntax a, Syntax b, Internal a ~ Internal b) => a -> b
resugar = Syntactic.resugar

----------------------------------------
-- ** Assertions
----------------------------------------

-- | Guard a value by an assertion (with implicit label @`UserAssertion` ""@)
guardVal :: Syntax a
    => Data Bool  -- ^ Condition that is expected to be true
    -> String     -- ^ Error message
    -> a          -- ^ Value to attach the assertion to
    -> a
guardVal = guardValLabel $ UserAssertion ""

-- | Like 'guardVal' but with an explicit assertion label
guardValLabel :: Syntax a
    => AssertionLabel  -- ^ Assertion label
    -> Data Bool       -- ^ Condition that is expected to be true
    -> String          -- ^ Error message
    -> a               -- ^ Value to attach the assertion to
    -> a
guardValLabel c cond msg = sugarSymFeld (GuardVal c msg) cond

----------------------------------------
-- ** Unsafe operations
----------------------------------------

-- | Turn a 'Comp' computation into a pure value. For this to be safe, the
-- computation should be free of side effects and independent of its
-- environment.
unsafePerform :: Syntax a => Comp a -> a
unsafePerform = sugarSymFeld . UnsafePerform . fmap desugar

--------------------------------------------------------------------------------
-- * Programs with computational effects
--------------------------------------------------------------------------------

-- | Monads that support computational effects: mutable data structures and
-- control flow
class Monad m => MonadComp m
  where
    -- | Lift a 'Comp' computation
    liftComp :: Comp a -> m a
    -- | Conditional statement
    iff :: Data Bool -> m () -> m () -> m ()
    -- | For loop
    for :: (Integral n, PrimType n) =>
        IxRange (Data n) -> (Data n -> m ()) -> m ()
    -- | While loop
    while :: m (Data Bool) -> m () -> m ()

instance MonadComp Comp
  where
    liftComp        = id
    iff c t f       = Comp $ Imp.iff c (unComp t) (unComp f)
    for range body  = Comp $ Imp.for range (unComp . body)
    while cont body = Comp $ Imp.while (unComp cont) (unComp body)

----------------------------------------
-- ** References
----------------------------------------

-- | Create an uninitialized reference
newRef :: (Syntax a, MonadComp m) => m (Ref a)
newRef = newNamedRef "r"

-- | Create an uninitialized named reference
--
-- The provided base name may be appended with a unique identifier to avoid name
-- collisions.
newNamedRef :: (Syntax a, MonadComp m)
    => String  -- ^ Base name
    -> m (Ref a)
newNamedRef base = liftComp $ fmap Ref $
    mapStructA (const $ Comp $ Imp.newNamedRef base) typeRep

-- | Create an initialized named reference
initRef :: (Syntax a, MonadComp m) => a -> m (Ref a)
initRef = initNamedRef "r"

-- | Create an initialized reference
--
-- The provided base name may be appended with a unique identifier to avoid name
-- collisions.
initNamedRef :: (Syntax a, MonadComp m)
    => String  -- ^ Base name
    -> a       -- ^ Initial value
    -> m (Ref a)
initNamedRef base =
    liftComp . fmap Ref . mapStructA (Comp . Imp.initNamedRef base) . resugar

-- | Get the contents of a reference.
getRef :: (Syntax a, MonadComp m) => Ref a -> m a
getRef = liftComp . fmap resugar . mapStructA (Comp . Imp.getRef) . unRef

-- | Set the contents of a reference.
setRef :: (Syntax a, MonadComp m) => Ref a -> a -> m ()
setRef r
    = liftComp
    . sequence_
    . zipListStruct (\r' a' -> Comp $ Imp.setRef r' a') (unRef r)
    . resugar

-- | Modify the contents of reference.
modifyRef :: (Syntax a, MonadComp m) => Ref a -> (a -> a) -> m ()
modifyRef r f = setRef r . f =<< unsafeFreezeRef r

-- | Freeze the contents of reference (only safe if the reference is not updated
--   as long as the resulting value is alive).
unsafeFreezeRef :: (Syntax a, MonadComp m) => Ref a -> m a
unsafeFreezeRef
    = liftComp
    . fmap resugar
    . mapStructA (Comp . Imp.unsafeFreezeRef)
    . unRef

----------------------------------------
-- ** Arrays
----------------------------------------

-- | Create an uninitialized array
newArr :: (Type (Internal a), MonadComp m) => Data Length -> m (Arr a)
newArr = newNamedArr "a"

-- | Create an uninitialized named array
--
-- The provided base name may be appended with a unique identifier to avoid name
-- collisions.
newNamedArr :: (Type (Internal a), MonadComp m)
    => String  -- ^ Base name
    -> Data Length
    -> m (Arr a)
newNamedArr base l = liftComp $ fmap (Arr 0 l) $
    mapStructA (const (Comp $ Imp.newNamedArr base l)) typeRep

-- | Create an array and initialize it with a constant list
constArr :: (PrimType (Internal a), MonadComp m)
    => [Internal a]  -- ^ Initial contents
    -> m (Arr a)
constArr = constNamedArr "a"

-- | Create a named array and initialize it with a constant list
--
-- The provided base name may be appended with a unique identifier to avoid name
-- collisions.
constNamedArr :: (PrimType (Internal a), MonadComp m)
    => String        -- ^ Base name
    -> [Internal a]  -- ^ Initial contents
    -> m (Arr a)
constNamedArr base as =
    liftComp $ fmap (Arr 0 len . Single) $ Comp $ Imp.constNamedArr base as
  where
    len = value $ genericLength as

-- | Get an element of an array
getArr :: (Syntax a, MonadComp m) => Arr a -> Data Index -> m a
getArr arr i = do
    assertLabel
      InternalAssertion
      (i < length arr)
      "getArr: index out of bounds"
    liftComp
      $ fmap resugar
      $ mapStructA (Comp . flip Imp.getArr (i + arrOffset arr))
      $ unArr arr

-- | Set an element of an array
setArr :: forall m a . (Syntax a, MonadComp m) =>
    Arr a -> Data Index -> a -> m ()
setArr arr i a = do
    assertLabel
      InternalAssertion
      (i < length arr)
      "setArr: index out of bounds"
    liftComp
      $ sequence_
      $ zipListStruct
          (\a' arr' -> Comp $ Imp.setArr arr' (i + arrOffset arr) a') rep
      $ unArr arr
  where
    rep = resugar a :: Struct PrimType' Data (Internal a)

-- | Copy the contents of an array to another array. The length of the
-- destination array must not be less than that of the source array.
--
-- In order to copy only a part of an array, use 'slice' before calling
-- 'copyArr'.
copyArr :: MonadComp m
    => Arr a  -- ^ Destination
    -> Arr a  -- ^ Source
    -> m ()
copyArr arr1 arr2 = do
    assertLabel
      InternalAssertion
      (length arr1 >= length arr2)
      "copyArr: destination too small"
    liftComp $ sequence_ $
      zipListStruct
        (\a1 a2 ->
            Comp $ Imp.copyArr
              (a1, arrOffset arr1)
              (a2, arrOffset arr2)
              (length arr2)
        )
        (unArr arr1)
        (unArr arr2)

-- | Freeze a mutable array to an immutable one. This involves copying the array
-- to a newly allocated one.
freezeArr :: (Type (Internal a), MonadComp m) => Arr a -> m (IArr a)
freezeArr arr = liftComp $ do
    arr2 <- newArr (length arr)
    copyArr arr2 arr
    unsafeFreezeArr arr2
  -- This is better than calling `freezeArr` from imperative-edsl, since that
  -- one copies without offset.

-- | A version of 'freezeArr' that slices the array from 0 to the given length
freezeSlice :: (Type (Internal a), MonadComp m) =>
    Data Length -> Arr a -> m (IArr a)
freezeSlice len = fmap (slice 0 len) . freezeArr

-- | Freeze a mutable array to an immutable one without making a copy. This is
-- generally only safe if the the mutable array is not updated as long as the
-- immutable array is alive.
unsafeFreezeArr :: MonadComp m => Arr a -> m (IArr a)
unsafeFreezeArr arr
    = liftComp
    $ fmap (IArr (arrOffset arr) (length arr))
    $ mapStructA (Comp . Imp.unsafeFreezeArr)
    $ unArr arr

-- | A version of 'unsafeFreezeArr' that slices the array from 0 to the given
-- length
unsafeFreezeSlice :: MonadComp m => Data Length -> Arr a -> m (IArr a)
unsafeFreezeSlice len = fmap (slice 0 len) . unsafeFreezeArr

-- | Thaw an immutable array to a mutable one. This involves copying the array
-- to a newly allocated one.
thawArr :: (Type (Internal a), MonadComp m) => IArr a -> m (Arr a)
thawArr arr = liftComp $ do
    arr2 <- unsafeThawArr arr
    arr3 <- newArr (length arr)
    copyArr arr3 arr2
    return arr3

-- | Thaw an immutable array to a mutable one without making a copy. This is
-- generally only safe if the the mutable array is not updated as long as the
-- immutable array is alive.
unsafeThawArr :: MonadComp m => IArr a -> m (Arr a)
unsafeThawArr arr
    = liftComp
    $ fmap (Arr (iarrOffset arr) (length arr))
    $ mapStructA (Comp . Imp.unsafeThawArr)
    $ unIArr arr

-- | Create an immutable array and initialize it with a constant list
constIArr :: (PrimType (Internal a), MonadComp m) =>
    [Internal a] -> m (IArr a)
constIArr = constArr >=> unsafeFreezeArr

----------------------------------------
-- ** Control-flow
----------------------------------------

-- | Conditional statement that returns an expression
ifE :: (Syntax a, MonadComp m)
    => Data Bool  -- ^ Condition
    -> m a        -- ^ True branch
    -> m a        -- ^ False branch
    -> m a
ifE c t f = do
    res <- newRef
    iff c (t >>= setRef res) (f >>= setRef res)
    unsafeFreezeRef res

-- | Break out from a loop
break :: MonadComp m => m ()
break = liftComp $ Comp Imp.break

-- | Assertion (with implicit label @`UserAssertion` ""@)
assert :: MonadComp m
    => Data Bool  -- ^ Expression that should be true
    -> String     -- ^ Message in case of failure
    -> m ()
assert = assertLabel $ UserAssertion ""

-- | Like 'assert' but tagged with an explicit assertion label
assertLabel :: MonadComp m
    => AssertionLabel  -- ^ Assertion label
    -> Data Bool       -- ^ Expression that should be true
    -> String          -- ^ Message in case of failure
    -> m ()
assertLabel c cond msg =
    liftComp $ Comp $ Oper.singleInj $ Assert c cond msg

----------------------------------------
-- ** Misc.
----------------------------------------

-- | Force evaluation of a value and share the result (monadic version of
-- 'share')
shareM :: (Syntax a, MonadComp m) => a -> m a
shareM = initRef >=> unsafeFreezeRef
  -- This function is more commonly needed than `share`, since code motion
  -- doesn't work across monadic binds.


module Main where

import System.Environment
import System.Random

import Data.Maybe

import Options.Applicative

import Control.Monad.Trans.Maybe

import DMCTS.Client.State
import DMCTS.Client.Sampling
import DMCTS.Client.Searches
import DMCTS.Types
import DMCTS.Random

import DMCTSKnapsack.KnapTree
import DMCTSKnapsack.Parsers

import Criterion.Main
import Criterion.Main.Options
import qualified Criterion.Types as CT

import Statistics.Types

main :: IO ()
main = validator =<< execParser opts
  where
    opts = info (cmdParser <**> helper)
      ( fullDesc
     <> progDesc "Estimate solutions for knapsack problems"
     <> header "dmcts-knapsack-client - a client for estimating Knapsack Problems with DMCTS" )

-- Validates all input arguments before calling execMtx
validator :: KnapArgs -> IO ()
validator args@SingleRun {..} = do
  knapNode' <- newKnapCfg knapCfgPath
  clientState' <- runMaybeT $ newClientState lambdaCfgPath
  -- Same dirty error checking as above.
  -- TODO: Tweak this to print a friendly error message.
  let clientState = fromJust clientState'
  let knapNode = fromJust knapNode'
  execMtx args clientState knapNode

validator args@Benchmark {..} = do
  knapNode' <- newKnapCfg knapCfgPath
  clientState' <- runMaybeT $ newClientState lambdaCfgPath
  -- Same dirty error checking as above.
  -- TODO: Tweak this to print a friendly error message.
  let clientState = fromJust clientState'
  let knapNode = fromJust knapNode'
  execMtx args clientState knapNode

-- Given a set of arguments, calls the appropriate functions to evaluate them
-- and print the results.
--
-- NOTE: This _might_ be an anti-pattern. I'll have to see how it pans out, I haven't
--       figured out the proper design idioms for this kind of functional code yet.
--       I do like being able to implement switching logic through pattern matching
--       though.
execMtx :: KnapArgs -> DMCTSClientState -> KnapNode -> IO ()

-- Single run cases.
execMtx SingleRun {local=True, ..} cs node = do
  gen <- getStdGen
  putStrLn "Running a search locally..."
  (minNode, _) <- return $ minLeafL KnapLogicWR node mcSamples compare gen
  putStrLn $ "Solution: " ++ show minNode

-- Cases where computation should be offloaded to AWS lambda instances
execMtx SingleRun {local=False, ..} cs node = do
  -- Run the calculation.
  gen <- getStdGen
  putStrLn "Running a search remotely..."
  (minNode, _) <- minLeafR cs KnapLogicWR node mcSamples compare gen
  putStrLn $ "Solution: " ++ show minNode

-- Benchmark case.
execMtx args@Benchmark {..} cs node = do
  putStrLn "Running a benchmark..."
  gen <- getStdGen
  let runBench = buildBench args cs node gen
  runBench

buildBench :: KnapArgs -> DMCTSClientState -> KnapNode -> StdGen -> IO ()
buildBench Benchmark {..} cs node gen = runBench
  where
    -- Lambdas to assist in building the benchmarks
    localSamp  = \sampCnt gen -> fst $ minLeafL KnapLogicWR node sampCnt compare gen
    remoteSamp = \sampCnt gen -> fst <$> minLeafR cs KnapLogicWR node sampCnt compare gen
    -- Build the list of sample #'s for the benchmarks.
    exponents = [0..(nBenchmarks-1)] :: [Int]
    benchSamples = filter (/= 0) $ map ((minSamp+).(sampBase^)) exponents
    benchNames = map (\a -> "benchmark_samp_" ++ show a) benchSamples

    localFuncs  = map localSamp benchSamples
    remoteFuncs = map remoteSamp benchSamples

    localBenches = zipWith whnf localFuncs lSeeds :: [Benchmarkable]
    remoteBenches = zipWith whnfAppIO remoteFuncs lSeeds :: [Benchmarkable]

    -- Generate seeds
    (lSeeds, newGen) = getNGens gen (length benchSamples)
    (rSeeds, _)      = getNGens newGen (length benchSamples)

    localBenchmarks = bgroup "Local" $ zipWith bench benchNames localBenches
    remoteBenchmarks = bgroup "Remote" $ zipWith bench benchNames remoteBenches

    -- Replace benchmark configuration
    criterionArgs =
      defaultConfig
        { --CT.confInterval = cl90,
          --CT.resamples = nRuns,
          --CT.timeLimit = 600,
          CT.reportFile = Just "./bench-report.html",
          CT.csvFile    = Just "./bench-data.csv",
          CT.verbosity  = CT.Verbose}
    critMode = Run criterionArgs IPattern []
    --args = criterionArgs { CT.resamples = nRuns}
    runBench = runMode critMode [localBenchmarks, remoteBenchmarks]


{-# LANGUAGE GADTs #-}
{-# LANGUAGE MultiParamTypeClasses #-}
{-# LANGUAGE FlexibleInstances #-}
{-# LANGUAGE DataKinds #-}

module SimpleR.Interpreter.Commons.Vector
  ( Type (..)
  -- , SInt (..)
  -- , SDouble (..)
  -- , SComplex (..)
  -- , SString (..)
  -- , SBool (..)
  , Atom (..)
  , Vector (..)
  , idFromAtomString
  , vecFromConst
  , vecInts
  , vecDoubles
  , vecComplexes
  , vecStrings
  , vecBools
  , vecLength
  , vecResize
  , vecType
  , vecToType
  , joinType
  , vecPairJoinType
  , vecPairLengthen
  , vecPairLengthenJoinType
  ) where

import Data.Maybe
import qualified Data.Complex as C

import SimpleR.Language
import SimpleR.Smt

data Type = BoolTy | IntTy | DoubleTy | ComplexTy | StringTy | NullTy
  deriving (Eq, Show, Read)

data Atom a where
  Atom :: a -> Atom a
  NAtom :: Atom a

instance (Eq a) => Eq (Atom a) where
  (Atom a) == (Atom b) = a == b
  (NAtom) == (NAtom) = True
  _ == _ = False

instance (Show a) => Show (Atom a) where
  show (Atom a) = "Atom (" ++ show a ++ ")"
  show (NAtom) = "NAtom"

instance (Read a) => Read (Atom a) where
  readsPrec _ = error "readsPrec: Atom read unimplemented"

instance Functor Atom where
  fmap _ NAtom = NAtom
  fmap f (Atom a) = Atom (f a)

instance Applicative Atom where
  pure a = Atom a

  (NAtom) <*> _ = NAtom
  (Atom f) <*> atom = fmap f atom

instance Monad Atom where
  return a = Atom a

  (NAtom) >>= _ = NAtom
  (Atom a) >>= f = f a

data Vector =
    IntVec [Atom Int]
  | DoubleVec [Atom Double]
  | ComplexVec [Atom Complex]
  | BoolVec [Atom Bool]
  | StringVec [Atom String]
  | SymVec SmtIdent Type Int
  | NilVec
  deriving (Eq, Show, Read)

idFromAtomString :: Atom String -> Ident
idFromAtomString (Atom str) = idFromString str

-- Vector conversion
vecFromConst :: Const -> Vector
vecFromConst (IntConst int) = IntVec [Atom int]
vecFromConst (DoubleConst double) = DoubleVec [Atom double]
vecFromConst (ComplexConst complex) = ComplexVec [Atom complex]
vecFromConst (BoolConst bool) = BoolVec [Atom bool]
vecFromConst (StringConst str) = StringVec [Atom str]
vecFroMConst (NaConst) = NilVec

vecLength :: Vector -> Int
vecLength (IntVec xs) = length xs
vecLength (DoubleVec xs) = length xs
vecLength (ComplexVec xs) = length xs
vecLength (BoolVec xs) = length xs
vecLength (StringVec xs) = length xs
vecLength (SymVec _ _ len) = len -- error $ "vecLength: called with SymVec"

vecType :: Vector -> Type
vecType (IntVec _) = IntTy
vecType (DoubleVec _) = DoubleTy
vecType (ComplexVec _) = ComplexTy
vecType (BoolVec _) = BoolTy
vecType (StringVec _) = StringTy
vecType (SymVec _ ty _) = ty
vecType (NilVec) = NullTy

vecInts :: Vector -> Maybe [Atom Int]
vecInts (IntVec xs) = Just xs
vecInts _ = Nothing

vecDoubles :: Vector -> Maybe [Atom Double]
vecDoubles (DoubleVec xs) = Just xs
vecDoubles _ = Nothing

vecComplexes :: Vector -> Maybe [Atom Complex]
vecComplexes (ComplexVec xs) = Just xs
vecComplexes _ = Nothing

vecStrings :: Vector -> Maybe [Atom String]
vecStrings (StringVec xs) = Just xs
vecStrings _ = Nothing

vecBools :: Vector -> Maybe [Atom Bool]
vecBools (BoolVec xs) = Just xs
vecBools _ = Nothing

------
-- Resizing

snip :: Int -> [a] -> [a]
snip n xs = take n $ concat $ repeat xs

vecResize :: Int -> Vector -> Vector
vecResize n (IntVec (x : xs)) = IntVec $ snip n (x : xs)
vecResize n (DoubleVec (x : xs)) = DoubleVec $ snip n (x : xs)
vecResize n (ComplexVec (x : xs)) = ComplexVec $ snip n (x : xs)
vecResize n (StringVec (x : xs)) = StringVec $ snip n (x : xs)
vecResize n (BoolVec (x : xs)) = BoolVec $ snip n (x : xs)
vecResize _ (IntVec []) = error $ "vecLength: IntVec has length 0"
vecResize _ (DoubleVec []) = error $ "vecLength: DoubleVec has length 0"
vecResize _ (ComplexVec []) = error $ "vecLength: ComplexVec has length 0"
vecResize _ (StringVec []) = error $ "vecLength: StringVec has length 0"
vecResize _ (BoolVec []) = error $ "vecLength: BoolVec has length 0"
vecResize _ (NilVec) = error $ "vecLength: called with NilVec"
vecResize n (SymVec id ty len) =
  SymVec id ty n
  -- error $ "vecLength: called with SymVec"

-----
-- Conversion

maybeRead :: (Read a) => String -> Maybe a
maybeRead str = fmap fst $ listToMaybe $ reads str

-- Int
intFromInt :: Int -> Int
intFromInt int = int

intFromDouble :: Double -> Int
intFromDouble double = round double

intFromComplex :: Complex -> Int
intFromComplex complex = round $ C.realPart complex

intFromString :: String -> int
intFromString = error $ "intFromString: TODO"

intFromBool :: Bool -> Int
intFromBool bool = if bool then 1 else 0

-- Double
doubleFromInt :: Int -> Double
doubleFromInt int = fromIntegral int

doubleFromDouble :: Double -> Double
doubleFromDouble double = double

doubleFromComplex :: Complex -> Double
doubleFromComplex complex = C.realPart complex

doubleFromString :: String -> Double
doubleFromString = error "doubleFromString: TODO"

doubleFromBool :: Bool -> Double
doubleFromBool bool = if bool then 1.0 else 0.0

-- Complex
complexFromInt :: Int -> Complex
complexFromInt int = (fromIntegral int) C.:+ 0.0

complexFromDouble :: Double -> Complex
complexFromDouble double = double C.:+ 0.0

complexFromComplex :: Complex -> Complex
complexFromComplex complex = complex

complexFromString :: String -> Complex
complexFromString = error "complexFromString: TODO"

complexFromBool :: Bool -> Complex
complexFromBool bool = if bool then 1.0 C.:+ 0.0 else 0.0 C.:+ 0.0

-- String
stringFromInt :: Int -> String
stringFromInt int = show int

stringFromDouble :: Double -> String
stringFromDouble double = show double

stringFromComplex :: Complex -> String
stringFromComplex complex = show complex

stringFromString :: String -> String
stringFromString string = string

stringFromBool :: Bool -> String
stringFromBool bool = show bool

-- Bool
boolFromInt :: Int -> Bool
boolFromInt int = int /= 0

boolFromDouble :: Double -> Bool
boolFromDouble double = double /= 0.0

boolFromComplex :: Complex -> Bool
boolFromComplex complex = complex /= (0.0 C.:+ 0.0)

boolFromString :: String -> Bool
boolFromString = error $ "boolFromString: TODO"

boolFromBool :: Bool -> Bool
boolFromBool bool = bool

vecToType :: Vector -> Type -> Vector
vecToType (IntVec xs) IntTy = IntVec $ (map . fmap) intFromInt xs
vecToType (IntVec xs) DoubleTy = DoubleVec $ (map . fmap) doubleFromInt xs
vecToType (IntVec xs) ComplexTy = ComplexVec $ (map . fmap) complexFromInt xs
vecToType (IntVec xs) StringTy = StringVec $ (map . fmap) stringFromInt xs
vecToType (IntVec xs) BoolTy = BoolVec $ (map . fmap) boolFromInt xs

vecToType (DoubleVec xs) IntTy = IntVec $ (map . fmap) intFromDouble xs
vecToType (DoubleVec xs) DoubleTy =
  DoubleVec $ (map . fmap) doubleFromDouble xs
vecToType (DoubleVec xs) ComplexTy =
  ComplexVec $ (map . fmap) complexFromDouble xs
vecToType (DoubleVec xs) StringTy =
  StringVec $ (map . fmap) stringFromDouble xs
vecToType (DoubleVec xs) BoolTy = BoolVec $ (map . fmap) boolFromDouble xs

vecToType (ComplexVec xs) IntTy = IntVec $ (map . fmap) intFromComplex xs
vecToType (ComplexVec xs) DoubleTy =
  DoubleVec $ (map . fmap) doubleFromComplex xs
vecToType (ComplexVec xs) ComplexTy =
  ComplexVec $ (map . fmap) complexFromComplex xs
vecToType (ComplexVec xs) StringTy =
  StringVec $ (map . fmap) stringFromComplex xs
vecToType (ComplexVec xs) BoolTy = BoolVec $ (map . fmap) boolFromComplex xs

vecToType (StringVec xs) IntTy = IntVec $ (map . fmap) intFromString xs
vecToType (StringVec xs) DoubleTy =
  DoubleVec $ (map . fmap) doubleFromString xs
vecToType (StringVec xs) ComplexTy =
  ComplexVec $ (map . fmap) complexFromString xs
vecToType (StringVec xs) StringTy =
  StringVec $ (map . fmap) stringFromString xs
vecToType (StringVec xs) BoolTy = BoolVec $ (map . fmap) boolFromString xs

vecToType (BoolVec xs) IntTy = IntVec $ (map . fmap) intFromBool xs
vecToType (BoolVec xs) DoubleTy = DoubleVec $ (map . fmap) doubleFromBool xs
vecToType (BoolVec xs) ComplexTy =
  ComplexVec $ (map . fmap) complexFromBool xs
vecToType (BoolVec xs) StringTy = StringVec $ (map . fmap) stringFromBool xs
vecToType (BoolVec xs) BoolTy = BoolVec $ (map . fmap) boolFromBool xs

vecToType (SymVec sym _ len) ty = SymVec sym ty len
vecToType (NilVec) _ = NilVec

-- Type lattice
--  BoolTy <= IntTy <= Double <= Complex <= String <= Null
instance Ord Type where
  BoolTy `compare` BoolTy = EQ
  BoolTy `compare` _ = LT

  IntTy `compare` IntTy = EQ
  IntTy `compare` _ = LT

  DoubleTy `compare` DoubleTy = EQ
  DoubleTy `compare` _ = LT

  ComplexTy `compare` ComplexTy = EQ
  ComplexTy `compare` _ = LT

  StringTy `compare` StringTy = EQ
  StringTy `compare` _ = LT

  NullTy `compare` NullTy = EQ
  NullTy `compare` _ = GT

joinType :: Type -> Type -> Type
joinType ty1 ty2 = if ty1 < ty2 then ty2 else ty1

vecPairJoinType :: (Vector, Vector) -> (Vector, Vector)
vecPairJoinType (vec1, vec2) =
  let joinTy = joinType (vecType vec1) (vecType vec2) in
    (vecToType vec1 joinTy, vecToType vec2 joinTy)

vecPairLengthen :: (Vector, Vector) -> (Vector, Vector)
vecPairLengthen (vec1, vec2) =
  let len1 = vecLength vec1 in
  let len2 = vecLength vec2 in
    if len1 > len2 then
      (vec1, vecResize len1 vec2)
    else
      (vecResize len2 vec1, vec2)

vecPairLengthenJoinType :: (Vector, Vector) -> (Vector, Vector)
vecPairLengthenJoinType (vec1, vec2) =
  vecPairLengthen $ vecPairJoinType (vec1, vec2)


{-# LANGUAGE CPP                       #-}
{-# LANGUAGE DataKinds                 #-}
{-# LANGUAGE DeriveAnyClass            #-}
{-# LANGUAGE DeriveGeneric             #-}
{-# LANGUAGE ExistentialQuantification #-}
{-# LANGUAGE FlexibleContexts          #-}
{-# LANGUAGE FlexibleInstances         #-}
{-# LANGUAGE GADTs                     #-}
{-# LANGUAGE KindSignatures            #-}
{-# LANGUAGE RankNTypes                #-}
{-# LANGUAGE ScopedTypeVariables       #-}
{-# LANGUAGE TypeOperators             #-}
{-# LANGUAGE UndecidableInstances      #-}

{-|
Module      : Grenade.Dynamic.Specification
Description : Specifying network architectures for dynamically creating nets.
Copyright   : (c) Manuel Schneckenreither, 2020
License     : BSD2
Stability   : experimental
-}
module Grenade.Dynamic.Specification
  ( FromDynamicLayer (..)
  , ToDynamicLayer (..)
  , SpecNetwork (..)
  , SpecNet (..)
  , GNum (..)
  -- Layer specific data types
  , Dimensions
  , SpecFullyConnected(..)
  , SpecConvolution(..)
  , SpecDeconvolution(..)
  , SpecDropout(..)
  , SpecElu(..)
  , SpecGelu(..)
  , SpecLeakyRelu(..)
  , SpecLogit(..)
  , SpecRelu(..)
  , SpecReshape(..)
  , SpecSinusoid(..)
  , SpecSoftmax(..)
  , SpecTanh(..)
  , SpecTrivial(..)
  ) where

import           Control.DeepSeq
import           Data.Constraint                       (Dict (..))
import           Data.Proxy
import           Data.Reflection                       (reifyNat)
import           Data.Serialize
import           Data.Singletons
import           Data.Singletons.Prelude.List
import           Data.Typeable                         as T
import           GHC.Generics                          hiding (R)
import           GHC.TypeLits
import           Unsafe.Coerce                         (unsafeCoerce)
#if MIN_VERSION_base(4,9,0)
import           Data.Kind                             (Type)
#endif
#ifndef FLYCHECK
import {-# SOURCE #-} Grenade.Dynamic.Layers.Convolution    ()
import {-# SOURCE #-} Grenade.Dynamic.Layers.Deconvolution  ()
import {-# SOURCE #-} Grenade.Dynamic.Layers.Dropout        ()
import {-# SOURCE #-} Grenade.Dynamic.Layers.Elu            ()
import {-# SOURCE #-} Grenade.Dynamic.Layers.FullyConnected ()
import {-# SOURCE #-} Grenade.Dynamic.Layers.Gelu           ()
import {-# SOURCE #-} Grenade.Dynamic.Layers.LeakyRelu      ()
import {-# SOURCE #-} Grenade.Dynamic.Layers.Logit          ()
import {-# SOURCE #-} Grenade.Dynamic.Layers.Relu           ()
import {-# SOURCE #-} Grenade.Dynamic.Layers.Reshape        ()
import {-# SOURCE #-} Grenade.Dynamic.Layers.Sinusoid       ()
import {-# SOURCE #-} Grenade.Dynamic.Layers.Softmax        ()
import {-# SOURCE #-} Grenade.Dynamic.Layers.Tanh           ()
import {-# SOURCE #-} Grenade.Dynamic.Layers.Trivial        ()
#endif

import           Control.Monad.Primitive               (PrimBase, PrimState)
import           Control.Parallel.Strategies
import           Numeric.LinearAlgebra.Static          (L, R)
import qualified Numeric.LinearAlgebra.Static          as H
import           System.Random.MWC                     (Gen)

import           Grenade.Core
import           Grenade.Types
import           Grenade.Utils.ListStore

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

-- | Grenade Num class.
--
-- This allows for instance scalar multiplication of the weights, which is useful for slowly adapting networks, e.g. NN'
-- <- \tau * NN' + (1-\tau) * NN. Or one could sum up some gradients in parallel and apply them at once: @applyUpdate lp
-- net $ foldl1 (|+) ...@aq.
class GNum a where
  (|*) :: Rational -> a -> a
  (|+) :: a -> a -> a
  gFromRational :: Rational -> a

infixl 7 |*
infixr 5 |+

instance (SingI i) => GNum (Network '[] '[ i]) where
  _ |* NNil = NNil
  _ |+ NNil = NNil
  gFromRational _ = NNil

instance (SingI i, SingI o, Layer x i o, NFData (Network xs (o ': rs)), GNum x, GNum (Network xs (o ': rs))) => GNum (Network (x ': xs) (i ': o ': rs)) where
  s |* (x :~> xs) =
    let x' = (s |* x)
        xs' = (s |* xs) `using` rdeepseq
     in x' :~> xs'
  (x :~> xs) |+ (y :~> ys) =
    let x' = (x |+ y)
        xs' = (xs |+ ys) `using` rdeepseq
     in x' :~> xs'
  gFromRational r = gFromRational r :~> gFromRational r

instance GNum (Gradients '[]) where
  _ |* GNil = GNil
  _ |+ GNil = GNil
  gFromRational _ = GNil

instance (GNum a) => GNum [a] where
  r |* xs = fmap (r |*) xs
  xs |+ ys = zipWith (|+) xs ys
  gFromRational _ = error "Cannot create a list of elements using gFromRational"

instance (UpdateLayer x, GNum (Gradient x), GNum (Gradients xs), NFData (Gradients xs)) => GNum (Gradients (x ': xs)) where
  s |* (x :/> xs) =
    let x' = (s |* x)
        xs' = (s |* xs) `using` rdeepseq
     in x' :/> xs'
  (x :/> xs) |+ (y :/> ys) =
    let x' = (x |+ y)
        xs' = (xs |+ ys) `using` rdeepseq
     in x' :/> xs'
  gFromRational r = gFromRational r :/> gFromRational r

instance GNum () where
  _ |* () = ()
  _ |+ () = ()
  gFromRational _ = ()

instance (GNum a, GNum b) => GNum (a, b) where
  s |* (a, b) = (s |* a, s |* b)
  (a1, b1) |+ (a2, b2) = (a1 |+ a2, b1 |+ b2)
  gFromRational v = (gFromRational v, gFromRational v)

instance (KnownNat m) => GNum (R m) where
  s |* vec = H.dvmap (fromRational s *) vec
  (|+) = (+)
  gFromRational = fromRational

instance (KnownNat m, KnownNat n) => GNum (L m n) where
  s |* mat = H.dmmap (fromRational s *) mat
  (|+) = (+)
  gFromRational = fromRational

instance (GNum a) => GNum (ListStore a) where
  s |* (ListStore n xs) = ListStore n $ map (fmap (s |*)) xs
  (ListStore n1 xs) |+ (ListStore n2 ys) = ListStore (max n1 n2) $ zipWith (\x y -> (|+) <$> x <*> y) xs ys
  gFromRational _ = mkListStore

-- | Create a runtime dynamic specification of a network. Dynamic layers (and networks), for storing and restoring specific network structures (e.g. in saving the network structures to a DB and
-- restoring it from there) or simply generating them at runtime. This does not store the weights and biases! They have to be handled separately (see Serialize)!
class FromDynamicLayer x where
  fromDynamicLayer :: SomeSing Shape -> SomeSing Shape -> x -> SpecNet

-- | Class for generating layers from a specification. See the function @networkFromSpecification@.
class (Show spec) => ToDynamicLayer spec where
  toDynamicLayer :: (PrimBase m) => WeightInitMethod -> Gen (PrimState m) -> spec -> m SpecNetwork

----------------------------------------
-- Return value of toDynamicLayer

-- | Specification of a network or layer.
data SpecNetwork :: Type where
  SpecNetwork
    :: ( FromDynamicLayer (Network layers shapes)
       , FoldableGradient (Gradients layers)
       , GNum (Gradients layers)
       , GNum (Network layers shapes)
       , Layer (Network layers shapes) (Head shapes) (Last shapes)
       , NFData (Gradients layers)
       , NFData (Network layers shapes)
       , NFData (Tapes layers shapes)
       , RandomLayer (Network layers shapes)
       , Serialize (Network layers shapes)
       , Serialize (Gradients layers)
       , Show (Network layers shapes)
       , SingI (Head shapes)
       , SingI (Last shapes)
       , SingI shapes
       , Typeable (Head shapes)
       , Typeable layers
       , Typeable shapes
       , CreatableNetwork layers shapes
       , RunnableNetwork layers shapes
       )
    => !(Network layers shapes)
    -> SpecNetwork
  SpecLayer
    :: ( FromDynamicLayer x
       , FoldableGradient (Gradient x)
       , GNum (Gradient x)
       , GNum x
       , Layer x i o
       , NFData (Gradient x)
       , NFData (Tape x i o)
       , NFData x
       , RandomLayer x
       , Serialize x
       , Serialize (Gradient x)
       , Show x
       , Typeable i
       , Typeable x
       , CreatableNetwork '[x] '[i,o]
       )
    => !x
    -> !(Sing i)
    -> !(Sing o)
    -> SpecNetwork

instance Show SpecNetwork where
  show (SpecNetwork net) = show net
  show (SpecLayer x _ _) = show x

instance FromDynamicLayer SpecNetwork where
  fromDynamicLayer inp out (SpecNetwork net) = fromDynamicLayer inp out net
  fromDynamicLayer inp out (SpecLayer x _ _) = fromDynamicLayer inp out x

-- | Data structure for holding specifications for networks. Networks can be built dynamically with @toDynamicLayer@. Further see the functions @|=>@, @specNil1D@, @specNil2D@, @specNil3D@, and
-- possibly any other layer implementation of @ToDynamicLayer@ for building specifications. See functions @networkFromSpecification@ to create a network from a specification and
-- @networkToSpecification@ to create a specification from a network.
data SpecNet
  = SpecNNil1D !Integer                   -- ^ 1D network output
  | SpecNNil2D !Integer !Integer          -- ^ 2D network output
  | SpecNNil3D !Integer !Integer !Integer -- ^ 3D network output
  | SpecNCons !SpecNet !SpecNet           -- ^ x :~> xs, where x can also be a network
  | forall spec . (ToDynamicLayer spec, Typeable spec, Ord spec, Eq spec, Show spec, Serialize spec, NFData spec) => SpecNetLayer !spec -- ^ Specification of a layer

instance Eq SpecNet where
  SpecNNil1D r1 == SpecNNil1D r2 = r1 == r2
  SpecNNil2D r1 c1 == SpecNNil2D r2 c2 = r1 == r2 && c1 == c2
  SpecNNil3D r1 c1 d1 == SpecNNil3D r2 c2 d2 = r1 == r2 && c1 == c2 && d1 == d2
  SpecNCons l1 r1 == SpecNCons l2 r2 = l1 == l2 && r1 == r2
  SpecNetLayer spec1 == SpecNetLayer spec2 = T.typeOf spec1 == T.typeOf spec2 && Just spec1 == cast spec2
  _ == _ = False

instance Ord SpecNet where
  SpecNNil1D r1 `compare` SpecNNil1D r2 = r1 `compare` r2
  SpecNNil2D r1 c1 `compare` SpecNNil2D r2 c2 = compare (r1, c1) (r2, c2)
  SpecNNil3D r1 c1 d1 `compare` SpecNNil3D r2 c2 d2 = compare (r1, c1, d1) (r2, c2, d2)
  SpecNCons l1 r1 `compare` SpecNCons l2 r2 = compare (l1, r1) (l2, r2)
  SpecNetLayer spec1 `compare` SpecNetLayer spec2
    | T.typeOf spec1 == T.typeOf spec2 = compare (Just spec1) (cast spec2)
    | otherwise = LT
  x `compare` y = compare x y

instance NFData SpecNet where
  rnf (SpecNNil1D r1)       = rnf r1
  rnf (SpecNNil2D r1 c1)    = rnf r1 `seq` rnf c1
  rnf (SpecNNil3D r1 c1 d1) = rnf r1 `seq` rnf c1 `seq` rnf d1
  rnf (SpecNCons l1 r1)     = rnf l1 `seq` rnf r1
  rnf (SpecNetLayer spec1)  = rnf spec1

instance Serialize SpecNet where
  put (SpecNNil1D r1)       = put (0 :: Int) >> put r1
  put (SpecNNil2D r1 c1)    = put (1 :: Int) >> put r1 >> put c1
  put (SpecNNil3D r1 c1 d1) = put (2 :: Int) >> put r1 >> put c1 >> put d1
  put (SpecNCons l1 r1)     = put (3 :: Int) >> put l1 >> put r1
  put (SpecNetLayer spec1)  = put (4 :: Int) >> put (show $ typeOf spec1) >> put spec1
  get = do
    (nr :: Int) <- get
    case nr of
      0 -> SpecNNil1D <$> get
      1 -> SpecNNil2D <$> get <*> get
      2 -> SpecNNil3D <$> get <*> get <*> get
      3 -> SpecNCons <$> get <*> get
      4 -> get >>= parseSpecDataConstructor
      _ -> error "unexpected input in Serialize instance of DynamicNetwork"

instance Show SpecNet where
  show (SpecNNil1D o) = "SpecNNil1D " ++ show o
  show (SpecNNil2D r c) = "SpecNNil2D " ++ show r ++ "x" ++ show c
  show (SpecNNil3D x y z) = "SpecNNil3D " ++ show x ++ "x" ++ show y ++ "x" ++ show z
  show (SpecNCons layer rest) = case layer of
    SpecNCons{} -> "[ " ++ show layer ++ " ] :=> " ++ show rest
    _           -> show layer ++ " :=> " ++ show rest
  show (SpecNetLayer layer) = show layer

-- | Function to parse the known implemented types. Can only deserialise these!
parseSpecDataConstructor :: String -> Get SpecNet
#ifdef FLYCHECK
parseSpecDataConstructor = error "disable -DFLYCHECK while compiling!"
#else
parseSpecDataConstructor repStr
  | repStr == show (typeRep (Proxy :: Proxy SpecFullyConnected)) = SpecNetLayer <$> (get :: Get SpecFullyConnected)
  | repStr == show (typeRep (Proxy :: Proxy SpecConvolution))    = SpecNetLayer <$> (get :: Get SpecConvolution)
  | repStr == show (typeRep (Proxy :: Proxy SpecDeconvolution))  = SpecNetLayer <$> (get :: Get SpecDeconvolution)
  | repStr == show (typeRep (Proxy :: Proxy SpecDropout))        = SpecNetLayer <$> (get :: Get SpecDropout)
  | repStr == show (typeRep (Proxy :: Proxy SpecElu))            = SpecNetLayer <$> (get :: Get SpecElu)
  | repStr == show (typeRep (Proxy :: Proxy SpecGelu))           = SpecNetLayer <$> (get :: Get SpecGelu)
  | repStr == show (typeRep (Proxy :: Proxy SpecLeakyRelu))      = SpecNetLayer <$> (get :: Get SpecLeakyRelu)
  | repStr == show (typeRep (Proxy :: Proxy SpecLogit))          = SpecNetLayer <$> (get :: Get SpecLogit)
  | repStr == show (typeRep (Proxy :: Proxy SpecRelu))           = SpecNetLayer <$> (get :: Get SpecRelu)
  | repStr == show (typeRep (Proxy :: Proxy SpecReshape))        = SpecNetLayer <$> (get :: Get SpecReshape)
  | repStr == show (typeRep (Proxy :: Proxy SpecSinusoid))       = SpecNetLayer <$> (get :: Get SpecSinusoid)
  | repStr == show (typeRep (Proxy :: Proxy SpecSoftmax))        = SpecNetLayer <$> (get :: Get SpecSoftmax)
  | repStr == show (typeRep (Proxy :: Proxy SpecTanh))           = SpecNetLayer <$> (get :: Get SpecTanh)
  | repStr == show (typeRep (Proxy :: Proxy SpecTrivial))        = SpecNetLayer <$> (get :: Get SpecTrivial)
  | repStr == show (typeRep (Proxy :: Proxy SpecNet))            = SpecNetLayer <$> (get :: Get SpecNet)
  | otherwise = error $ "unexpected input parseSpecDataConstructor: " ++ repStr
#endif

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

-- Network instances

instance (KnownNat rows) => FromDynamicLayer (Network '[] '[ 'D1 rows ]) where
  fromDynamicLayer _ _ _ = SpecNNil1D nr
    where nr = natVal (Proxy :: Proxy rows)
instance (KnownNat rows, KnownNat cols) => FromDynamicLayer (Network '[] '[ 'D2 rows cols ]) where
  fromDynamicLayer _ _ _ = SpecNNil2D (natVal (Proxy :: Proxy rows)) (natVal (Proxy :: Proxy cols))

instance (KnownNat rows, KnownNat cols, KnownNat depth) => FromDynamicLayer (Network '[] '[ 'D3 rows cols depth ]) where
  fromDynamicLayer _ _ _ = SpecNNil3D (natVal (Proxy :: Proxy rows)) (natVal (Proxy :: Proxy cols)) (natVal (Proxy :: Proxy depth))

instance (FromDynamicLayer x, FromDynamicLayer (Network xs (h : rs)), SingI h) => FromDynamicLayer (Network (x ': xs) (i ': h ': rs)) where
  fromDynamicLayer inp out ((x :: x) :~> (xs :: Network xs (h ': rs))) = SpecNCons (fromDynamicLayer inp hShape x) (fromDynamicLayer hShape out xs)
    where
      hShape = SomeSing (sing :: Sing h)

instance ToDynamicLayer SpecNet where
  toDynamicLayer _ _ (SpecNNil1D nrOut) =
    reifyNat nrOut $ \(_ :: (KnownNat o') => Proxy o') ->
      case (unsafeCoerce (Dict :: Dict ()) :: Dict (layers ~ '[]), unsafeCoerce (Dict :: Dict ()) :: Dict (shapes ~ '[ 'D1 o'])) of
        (Dict, Dict) -> return $ SpecNetwork (NNil :: Network '[] '[ 'D1 o'])

  toDynamicLayer _ _ (SpecNNil2D rows cols) =
    reifyNat rows $ \(_ :: (KnownNat r) => Proxy r) -> reifyNat cols $ \(_ :: (KnownNat c) => Proxy c) ->
      case (unsafeCoerce (Dict :: Dict ()) :: Dict (layers ~ '[]), unsafeCoerce (Dict :: Dict ()) :: Dict (shapes ~ '[ 'D2 r c ])) of
        (Dict, Dict) -> return $ SpecNetwork (NNil :: Network '[] '[ 'D2 r c])

  toDynamicLayer _ _ (SpecNNil3D rows cols depth) =
    reifyNat rows $ \(_ :: (KnownNat r) => Proxy r) -> reifyNat cols $ \(_ :: (KnownNat c) => Proxy c) -> reifyNat depth $ \(_ :: (KnownNat d) => Proxy d) ->
      case ( unsafeCoerce (Dict :: Dict ()) :: Dict (layers ~ '[])
           , unsafeCoerce (Dict :: Dict ()) :: Dict (shapes ~ '[ 'D3 r c d ])
           , unsafeCoerce (Dict :: Dict ()) :: Dict (KnownNat (r GHC.TypeLits.* d))) of
        (Dict, Dict, Dict) -> return $ SpecNetwork (NNil :: Network '[] '[ 'D3 r c d])

  toDynamicLayer wInit gen (SpecNetLayer layer) = toDynamicLayer wInit gen layer

  toDynamicLayer wInit gen (SpecNCons left right) = do
    restNet <- toDynamicLayer wInit gen right
    case restNet of
      SpecLayer{} -> error "Expected a network, but received a layer in Specification"
      SpecNetwork (xs :: Network restLayers restShapes) ->
        case (sing :: Sing restShapes) of
          SNil -> error "unexpected empty network (shapes where SNil) in Specification"
          (SCons (h :: Sing h) (hs :: Sing hs)) -> withSingI h $ withSingI hs $ do
            layer <- toDynamicLayer wInit gen left
            case layer of
              SpecLayer (x :: xType) singIn singOut ->
                case (singIn, singOut) of
                  (pxIn :: Sing inShape, pxOut :: Sing outShape) ->
                    withSingI pxIn $ withSingI h $ withSingI pxOut $
                     case (unsafeCoerce (Dict :: Dict ()) :: Dict (outShape ~ h)) of
                       Dict -> return $ SpecNetwork (x :~> xs :: Network (xType ': restLayers) (inShape ': restShapes))
              SpecNetwork (x :: Network xLayers xShapes) ->
                case (sing :: Sing xShapes) of
                  SNil -> error "unexpected empty network (SpecNNil) as layer in specification. Cannot proceed."
                  SCons (i :: Sing i) (xHs :: Sing xHs) -> withSingI i $ withSingI xHs $
                    case ( unsafeCoerce (Dict :: Dict ()) :: Dict (Head restShapes ~ Last xShapes)) of
                      Dict -> return $ SpecNetwork (x :~> xs :: Network (Network xLayers xShapes ': restLayers) (i ': restShapes))

-- Data structures stances for Layers (needs to be defined here)

type Dimensions = (Integer, Integer, Integer)

-- | Data Structure to save a fully connected layer. Saves number of input and output nodes.
data SpecFullyConnected = SpecFullyConnected !Integer !Integer
  deriving (Show, Eq, Ord, Serialize, Generic, NFData)

-- data SpecConcat = SpecConcat SpecNet SpecNet
--   deriving (Show, Eq, Ord, Serialize, Generic, NFData)

-- | Specifiation of a convolutional layer. Saves input and channels, filters, kernelRows, kernelColumns, strideRows, strideColumns, and kernelFlattened.
data SpecConvolution =
  SpecConvolution !Dimensions !Integer !Integer !Integer !Integer !Integer !Integer
  deriving (Show, Read, Eq, Ord, Serialize, Generic, NFData)

-- | Specifiation of a deconvolutional layer. Saves input and channels, filters, kernelRows, kernelColumns, strideRows, strideColumns, and kernelFlattened.
data SpecDeconvolution =
  SpecDeconvolution !Dimensions !Integer !Integer !Integer !Integer !Integer !Integer
  deriving (Show, Read, Eq, Ord, Serialize, Generic, NFData)

-- | Specification of a dropout filter. Saves number of input nodes, ratio and maybe a seed.
data SpecDropout = SpecDropout !Integer !RealNum !(Maybe Int)
  deriving (Show, Read, Eq, Ord, Serialize, Generic, NFData)

-- | Specification of Elu, saves input dimensions as triple, where every element >= 1.
newtype SpecElu = SpecElu Dimensions
  deriving (Show, Read, Eq, Ord, Serialize, Generic, NFData)

-- | Specification of Leaky Relu (with 0.01), saves input dimensions as triple, where every element >= 1.
newtype SpecGelu = SpecGelu Dimensions
  deriving (Show, Read, Eq, Ord, Serialize, Generic, NFData)

-- | Specification of Leaky Relu, saves alpha input dimensions as triple, where every element >= 1.
data SpecLeakyRelu = SpecLeakyRelu Dimensions RealNum
  deriving (Show, Read, Eq, Ord, Serialize, Generic, NFData)

-- | Specification of Logit, saves input dimensions as triple, where every element >= 1.
newtype SpecLogit = SpecLogit Dimensions
  deriving (Show, Read, Eq, Ord, Serialize, Generic, NFData)

-- | Specification of Relu, saves input dimensions as triple, where every element >= 1.
newtype SpecRelu = SpecRelu Dimensions
  deriving (Show, Read, Eq, Ord, Serialize, Generic, NFData)

-- | Specification of a reshape layer, saves input and output dimensions, where every element >= 1.
data SpecReshape = SpecReshape !Dimensions !Dimensions
  deriving (Show, Read, Eq, Ord, Serialize, Generic, NFData)

-- | Specification of Sinusoid, saves input dimensions as triple, where every element >= 1.
newtype SpecSinusoid = SpecSinusoid Dimensions
  deriving (Show, Read, Eq, Ord, Serialize, Generic, NFData)

-- | Specification of Softmax, saves input dimensions as triple, where every element >= 1.
newtype SpecSoftmax = SpecSoftmax Integer
  deriving (Show, Read, Eq, Ord, Serialize, Generic, NFData)

-- | Specification of Tanh, saves input dimensions as triple, where every element >= 1.
newtype SpecTanh = SpecTanh Dimensions
  deriving (Show, Read, Eq, Ord, Serialize, Generic, NFData)

-- | Specification of Trivial, saves input dimensions as triple, where every element >= 1.
newtype SpecTrivial = SpecTrivial Dimensions
  deriving (Show, Read, Eq, Ord, Serialize, Generic, NFData)


module IOUtils where

import System.IO (FilePath, readFile, writeFile)
import Data.Complex (Complex)

parseLines :: [String] -> [Double]
parseLines ls = map read ls :: [Double]

readData :: FilePath -> IO [Double]
readData filePath = do
    content <- readFile filePath
    return $ parseLines $ lines content

writeData :: FilePath -> [Double] -> IO ()
writeData filePath d = writeFile filePath $ unlines $ map show d

printData :: Show a => [a] -> IO ()
printData d = putStrLn $ unlines $ map (++ ", ") $ map show d


{-# LANGUAGE BangPatterns          #-}
{-# LANGUAGE CPP                   #-}
{-# LANGUAGE ConstraintKinds       #-}
{-# LANGUAGE DeriveDataTypeable    #-}
{-# LANGUAGE FlexibleContexts      #-}
{-# LANGUAGE FlexibleInstances     #-}
{-# LANGUAGE MultiParamTypeClasses #-}
{-# LANGUAGE ScopedTypeVariables   #-}
{-# LANGUAGE TypeFamilies          #-}
{-# LANGUAGE UndecidableInstances  #-}
{-# OPTIONS_GHC -fno-warn-orphans #-}
#if __GLASGOW_HASKELL__ >= 800
    {-# OPTIONS_GHC -fno-warn-redundant-constraints #-}
#endif
-- |
-- Module      : Graphics.Image.Interface.Repa.Generic
-- Copyright   : (c) Alexey Kuleshevich 2017
-- License     : BSD3
-- Maintainer  : Alexey Kuleshevich <lehins@yandex.ru>
-- Stability   : experimental
-- Portability : non-portable
--
module Graphics.Image.Interface.Repa.Generic
  ( RImage
  , Strategy(..)
  , computeR
  , makeImageR
  , dimsR
  , scalarR
  , index00R
  , makeImageWindowedR
  , mapR
  , imapR
  , eqR
  , zipWithR
  , izipWithR
  , unsafeTraverseR
  , traverseR
  , unsafeTraverse2R
  , traverse2R
  , transposeR
  , backpermuteR
  , fromListsR
  , foldR
  , foldIxR
  , toVectorUnboxedR
  , fromVectorUnboxedR
  , toVectorStorableR
  , fromVectorStorableR
  , fromRepaArrayR
  , multR
  , ix2sh
  , sh2ix
  )where

import           Data.Array.Repa                     as R
import           Data.Array.Repa.Eval                as R (Elt (..), Target,
                                                           fromList,
                                                           suspendedComputeP)
import           Data.Array.Repa.Operators.Traversal as R
import           Data.Array.Repa.Repr.ForeignPtr
import           Data.Array.Repa.Repr.Partitioned    (P, Range (..))
import           Data.Array.Repa.Repr.Unboxed        (Unbox)
import           Data.Array.Repa.Repr.Undefined      (X)
import           Data.Complex                        as C
import           Data.Maybe                          (listToMaybe)
import qualified Data.Vector.Storable                as VS
import qualified Data.Vector.Unboxed                 as VU
import           Graphics.Image.ColorSpace.Binary    (Bit (..))
import           Graphics.Image.Interface            (ColorSpace (..), Pixel,
                                                      checkDims, toIx)
import           Prelude                             as P

data Strategy
  = Parallel
  | Sequential

data RImage r p = RTImage !(R.Array r R.DIM2 p)
                | RDImage !(R.Array R.D R.DIM2 p)

sh2ix :: DIM2 -> (Int, Int)
sh2ix (Z :. i :. j) = (i, j)
{-# INLINE sh2ix #-}

ix2sh :: (Int, Int) -> DIM2
ix2sh !(i, j) = Z :. i :. j
{-# INLINE ix2sh #-}

errorR :: String -> String -> a
errorR fName errMsg =
  error $ "Graphics.Image.Interface.Repa.Generic." P.++ fName P.++ ": " P.++ errMsg

checkDimsR :: String -> (Int, Int) -> DIM2
checkDimsR fName = ix2sh . checkDims ("Graphics.Image.Interface.Repa.Generic." P.++ fName)
{-# INLINE checkDimsR #-}

dimsR :: Source r p => RImage r p -> (Int, Int)
dimsR (RTImage arr) = sh2ix $ R.extent arr
dimsR (RDImage arr) = sh2ix $ R.extent arr
{-# INLINE dimsR #-}

makeImageR :: (Int, Int) -> ((Int, Int) -> px) -> RImage t px
makeImageR sz f =
  RDImage $ R.fromFunction (checkDimsR "makeImageR" sz) (f . sh2ix)
{-# INLINE makeImageR #-}

makeImageWindowedR
  :: (Int, Int)
  -> (Int, Int)
  -> (Int, Int)
  -> ((Int, Int) -> p)
  -> ((Int, Int) -> p)
  -> RImage r p
makeImageWindowedR !sz !wIx !wSz getWindowPx getBorderPx =
  RDImage $ R.delay $ makeWindowed sh wIx wSz wArr bArr
  where
    wArr = R.fromFunction sh (getWindowPx . sh2ix)
    bArr = R.fromFunction sh (getBorderPx . sh2ix)
    !sh = checkDimsR "makeImageWindowedR" sz
{-# INLINE makeImageWindowedR #-}

makeWindowed
  :: (Source r1 p, Source r2 p)
  => DIM2 -- ^ Extent of array.
  -> (Int, Int) -- ^ Window starting index
  -> (Int, Int) -- ^ Window size.
  -> Array r1 DIM2 p -- ^ Array for internal elements.
  -> Array r2 DIM2 p -- ^ Array for border elements.
  -> Array (P r1 (P r2 (P r2 (P r2 (P r2 X))))) DIM2 p
makeWindowed sh@(Z :. m :. n) !(it, jt) !(wm, wn) arrWindow arrBorder =
  let inInternal (Z :. i :. j) = i >= it && i < ib && j >= jt && j < jb
      {-# INLINE inInternal #-}
      inBorder = not . inInternal
      {-# INLINE inBorder #-}
      !(ib, jb) = (wm + it, wn + jt)
  in APart sh (Range (Z :. it :. jt) (Z :. wm       :. wn      ) inInternal) arrWindow $
     APart sh (Range (Z :. 0  :. 0 ) (Z :. it       :. n       ) inBorder  ) arrBorder $
     APart sh (Range (Z :. it :. 0 ) (Z :. wm       :. jt      ) inBorder  ) arrBorder $
     APart sh (Range (Z :. it :. jb) (Z :. wm       :. (n - jb)) inBorder  ) arrBorder $
     APart sh (Range (Z :. ib :. 0 ) (Z :. (m - ib) :. n       ) inBorder  ) arrBorder $
     AUndefined sh
{-# INLINE makeWindowed #-}

fromRepaArrayR :: (Source r1 p) => R.Array r1 DIM2 p -> RImage r p
fromRepaArrayR = RDImage . R.delay
{-# INLINE fromRepaArrayR #-}

scalarR :: p -> RImage t p
scalarR = RDImage . fromFunction (Z :. 1 :. 1) . const
{-# INLINE scalarR #-}

index00R :: Source r e => RImage r e -> e
index00R (RTImage arr) = R.index arr (Z :. 0 :. 0)
index00R (RDImage arr) = R.index arr (Z :. 0 :. 0)
{-# INLINE index00R #-}

mapR :: Source r1 a => (a -> p) -> RImage r1 a -> RImage r p
mapR f (RTImage arr) = RDImage (R.map f arr)
mapR f (RDImage arr) = RDImage (R.map f arr)
{-# INLINE mapR #-}

imapR :: Source r2 b =>
         ((Int, Int) -> b -> p) -> RImage r2 b -> RImage r p
imapR f (RTImage arr) = RDImage (imapArr (f . sh2ix) arr)
imapR f (RDImage arr) = RDImage (imapArr (f . sh2ix) arr)
{-# INLINE imapR #-}

imapArr :: R.Source r2 b =>
           (DIM2 -> b -> c) -> R.Array r2 DIM2 b -> R.Array R.D DIM2 c
imapArr f !arr = R.traverse arr id (\ getPx !sh -> f sh (getPx sh))
{-# INLINE imapArr #-}

getDelayed :: Source r e => RImage r e -> Array D DIM2 e
getDelayed (RTImage arr) = delay arr
getDelayed (RDImage arr) = arr
{-# INLINE getDelayed #-}

zipWithR
  :: (Source r2 t, Source r1 e) =>
     (t -> e -> c) -> RImage r2 t -> RImage r1 e -> RImage r c
zipWithR f !img1 !img2 = zipWithArr f (getDelayed img1) (getDelayed img2)
{-# INLINE zipWithR #-}

zipWithArr
  :: (Source r1 a, Source r2 b)
  => (a -> b -> c) -> Array r1 DIM2 a -> Array r2 DIM2 b -> RImage r c
zipWithArr f arr1 arr2 =
  RDImage $
  case (extent arr1, extent arr2) of
    (Z :. 1 :. 1, _) -> R.map (f (unsafeIndex arr1 (Z :. 0 :. 0))) arr2
    (_, Z :. 1 :. 1) -> R.map (`f` (unsafeIndex arr2 (Z :. 0 :. 0))) arr1
    _                -> R.zipWith f arr1 arr2
{-# INLINE zipWithArr #-}

izipWithR
  :: (Source r2 a, Source r1 b) =>
     ((Int, Int) -> a -> b -> c)
     -> RImage r2 a -> RImage r1 b -> RImage r c
izipWithR f !img1 !img2 = izipWithArr f (getDelayed img1) (getDelayed img2)
{-# INLINE izipWithR #-}

izipWithArr
  :: (Source r1 a, Source r2 b)
  => ((Int, Int) -> a -> b -> c)
  -> Array r1 DIM2 a
  -> Array r2 DIM2 b
  -> RImage r c
izipWithArr f arr1 arr2 =
  RDImage $
  case (extent arr1, extent arr2) of
    (Z :. 1 :. 1, _) ->
      imapArr (\ !sh -> f (sh2ix sh) (unsafeIndex arr1 (Z :. 0 :. 0))) arr2
    (_, Z :. 1 :. 1) ->
      imapArr
        (\ !sh !px1 -> f (sh2ix sh) px1 (unsafeIndex arr2 (Z :. 0 :. 0))) arr1
    (Z :. m1 :. n1, Z :. m2 :. n2) -> traverse2 arr1 arr2 getNewDims getNewPx
      where getNewDims _ _ = Z :. min m1 m2 :. min n1 n2
            {-# INLINE getNewDims #-}
            getNewPx getPx1 getPx2 !sh = f (sh2ix sh) (getPx1 sh) (getPx2 sh)
            {-# INLINE getNewPx #-}
{-# INLINE izipWithArr #-}

traverseR
  :: Source r1 c
  => RImage r1 c
  -> ((Int, Int) -> (Int, Int))
  -> (((Int, Int) -> c) -> (Int, Int) -> p)
  -> RImage r p
traverseR img getNewDims getNewPx =
  RDImage $
  R.traverse
    (getDelayed img)
    (checkDimsR "traverseR" . getNewDims . sh2ix)
    (\ getPx -> getNewPx (getPx . ix2sh) . sh2ix)
{-# INLINE traverseR #-}

unsafeTraverseR
  :: Source r1 c
  => RImage r1 c
  -> ((Int, Int) -> (Int, Int))
  -> (((Int, Int) -> c) -> (Int, Int) -> p)
  -> RImage r p
unsafeTraverseR img getNewDims getNewPx =
  RDImage $
  R.unsafeTraverse
    (getDelayed img)
    (checkDimsR "traverseR" . getNewDims . sh2ix)
    (\ getPx -> getNewPx (getPx . ix2sh) . sh2ix)
{-# INLINE unsafeTraverseR #-}

traverse2R
  :: (Source r1 a, Source r2 b)
  => RImage r1 a
  -> RImage r2 b
  -> ((Int, Int) -> (Int, Int) -> (Int, Int))
  -> (((Int, Int) -> a) -> ((Int, Int) -> b) -> (Int, Int) -> c)
  -> RImage r c
traverse2R img1 img2 getNewDims getNewPx =
  RDImage $
  R.traverse2
    (getDelayed img1)
    (getDelayed img2)
    (\ !sh1 !sh2 -> checkDimsR "traverse2R" (getNewDims (sh2ix sh1) (sh2ix sh2)))
    (\getPx1 getPx2 -> getNewPx (getPx1 . ix2sh) (getPx2 . ix2sh) . sh2ix)
{-# INLINE traverse2R #-}

unsafeTraverse2R
  :: (Source r1 a, Source r2 b)
  => RImage r1 a
  -> RImage r2 b
  -> ((Int, Int) -> (Int, Int) -> (Int, Int))
  -> (((Int, Int) -> a) -> ((Int, Int) -> b) -> (Int, Int) -> c)
  -> RImage r c
unsafeTraverse2R img1 img2 getNewDims getNewPx =
  RDImage $
  R.unsafeTraverse2
    (getDelayed img1)
    (getDelayed img2)
    (\ !sh1 !sh2 -> checkDimsR "traverse2R" (getNewDims (sh2ix sh1) (sh2ix sh2)))
    (\getPx1 getPx2 -> getNewPx (getPx1 . ix2sh) (getPx2 . ix2sh) . sh2ix)
{-# INLINE unsafeTraverse2R #-}

transposeR :: Source r1 p => RImage r1 p -> RImage r p
transposeR (RDImage arr) = RDImage (R.transpose arr)
transposeR (RTImage arr) = RDImage (R.transpose arr)
{-# INLINE transposeR #-}

backpermuteR
  :: R.Source r e
  => (Int, Int) -> ((Int, Int) -> (Int, Int)) -> RImage r e -> RImage r e
backpermuteR !newDims g !img =
  RDImage $ R.backpermute
    (ix2sh (checkDims "backpermuteR" newDims))
    (ix2sh . g . sh2ix)
    (getDelayed img)
{-# INLINE backpermuteR #-}

fromListsR :: Target r p => [[p]] -> RImage r p
fromListsR ls =
  if all (== n) (P.map length ls)
    then RTImage $ R.fromList (Z :. m :. n) . concat $ ls
    else error "fromListsRepa: Inner lists do not all have an equal length."
  where
    !(m, n) = checkDims "fromListsRepa" (length ls, maybe 0 length $ listToMaybe ls)
{-# INLINE fromListsR #-}

multR :: (Num p, Elt p, Unbox p, Source r2 p, Source r1 p, Source r p, Target r p,
           Target r2 p, Target r1 p) =>
         Strategy -> RImage r2 p -> RImage r1 p -> RImage r p
multR strategy (RTImage arr1) (RTImage arr2) =
  if n1 /= m2
    then errorR "multR" $
         "Inner dimensions of multiplied images must be the same, but received: " P.++
         show (m1, n1) P.++
         " X " P.++
         show (m2, n2)
    else computeR strategy $ RDImage $ fromFunction (Z :. m1 :. n2) getPx
  where
    (Z :. m1 :. n1) = R.extent arr1
    (Z :. m2 :. n2) = R.extent arr2
    getPx (Z :. i :. j) =
      R.sumAllS
        (R.slice arr1 (R.Any :. (i :: Int) :. R.All) R.*^
         R.slice arr2 (R.Any :. (j :: Int)))
    {-# INLINE getPx #-}
multR strategy img1 img2 =
  multR strategy (computeR strategy img1) (computeR strategy img2)
{-# INLINE multR #-}

computeR :: (Target r p, Source r p) => Strategy -> RImage r p -> RImage r p
computeR Parallel (RDImage arr) = arrManifest `deepSeqArray` RTImage arrManifest
  where arrManifest = suspendedComputeP arr
computeR Sequential (RDImage arr) = RTImage $ computeS arr
computeR _ img = img
{-# INLINE computeR #-}

foldR :: (Source r p, Target r p, Elt p, Unbox p)
       => Strategy -> (p -> p -> p) -> p -> RImage r p -> p
foldR Parallel f !px0 !img =
  let RTImage arr = computeR Parallel img
  in case foldAllP f px0 arr of
       Just e  -> e
       Nothing -> errorR "foldPR" "impossible happened."
foldR Sequential f !px0 img =
  let RTImage arr = computeR Sequential img
  in foldAllS f px0 arr
{-# INLINE foldR #-}

addIxArr
  :: Source r2 b =>
     Array r2 DIM2 b -> Array R.D DIM2 (Int, b)
addIxArr !arr = R.zipWith (,) arrIx arr
  where
    !sh = extent arr
    !arrIx = R.fromFunction (R.extent arr) (toIndex sh)
{-# INLINE addIxArr #-}

foldIxR :: (Source r b, Elt b, Unbox b) =>
           Strategy -> (b -> (Int, Int) -> b -> b) -> b -> RImage r b -> b
foldIxR strategy f !acc' !img =
  case strategy of
    Parallel ->
      case foldAllP accumWithIx (-1, acc') arr of
        Just (_, acc) -> acc
        Nothing       -> error $ "foldIxPR: impossible happened."
    Sequential -> snd $ R.foldAllS accumWithIx (-1, acc') arr
  where
    (Z :. _ :. n) = extent arr
    !arr = addIxArr (getDelayed img)
    accumWithIx !(-1, acc) !(k, px) = (-1, f acc (toIx n k) px)
    accumWithIx !(k, px) (-1, acc) = (-1, f acc (toIx n k) px)
    accumWithIx (acc1Ix, _) (acc2Ix, _) =
      errorR "accumWithx" $
      "Impossible happened. Received: " P.++ show acc1Ix P.++ " " P.++
      show acc2Ix
    {-# INLINE accumWithIx #-}
{-# INLINE foldIxR #-}

eqR :: (Eq a, Source r1 a, Source r a) =>
        Strategy -> RImage r1 a -> RImage r a -> Bool
eqR strategy !img1 !img2 =
  case strategy of
    Parallel ->
      let eqArr = R.zipWith (==) (getDelayed img1) (getDelayed img2)
      in dimsR img1 == dimsR img2 && mCompute "eqR" (foldAllP (&&) True eqArr)
    Sequential -> R.equalsS (getDelayed img1) (getDelayed img2)
{-# INLINE eqR #-}

mCompute :: String -> Maybe t -> t
mCompute _ (Just res)  = res
mCompute fName Nothing = errorR fName "impossible happened"
{-# INLINE mCompute #-}

fromVectorUnboxedR :: Unbox e => (Int, Int) -> VU.Vector e -> RImage U e
fromVectorUnboxedR !sz = RTImage . fromUnboxed (ix2sh sz)
{-# INLINE fromVectorUnboxedR #-}

toVectorR :: (Target r2 e, Source r2 e) =>
             (R.Array r2 DIM2 e -> t) -> Strategy -> RImage r2 e -> t
toVectorR toVector Parallel (RDImage arr) =
 toVector $ mCompute "toVectorR" (computeP arr)
toVectorR toVector Sequential (RDImage arr) = toVector $ computeS arr
toVectorR toVector _ (RTImage arr) = toVector arr
{-# INLINE toVectorR #-}

toVectorUnboxedR :: Unbox e => Strategy -> RImage U e -> VU.Vector e
toVectorUnboxedR = toVectorR toUnboxed
{-# INLINE toVectorUnboxedR #-}

toVectorStorableR :: VS.Storable e => Strategy -> RImage F e -> VS.Vector e
toVectorStorableR = toVectorR toStorableR
{-# INLINE toVectorStorableR #-}

toStorableR :: VS.Storable a => Array F DIM2 a -> VS.Vector a
toStorableR arr = VS.unsafeFromForeignPtr0 (toForeignPtr arr) (m * n)
  where
    (Z :. m :. n) = R.extent arr
{-# INLINE toStorableR #-}

fromVectorStorableR
  :: VS.Storable px
  => (Int, Int) -> VS.Vector px -> RImage F px
fromVectorStorableR !(m, n) !v
  | sz == sz' = RTImage $ fromForeignPtr (ix2sh (m, n)) fp
  | otherwise =
    errorR "fromVectorStorableR" $
    "(impossible) Vector size mismatch: " P.++ show sz P.++ " vs " P.++ show sz'
  where
    !(fp, sz) = VS.unsafeToForeignPtr0 v
    !sz' = m * n
{-# INLINE fromVectorStorableR #-}

instance (Num e, R.Elt e) => Elt (C.Complex e) where
  touch (r :+ i) = R.touch r >> R.touch i
  {-# INLINE touch #-}

  zero     = 0 :+ 0
  {-# INLINE zero #-}

  one      = 1 :+ 0
  {-# INLINE one #-}

instance Elt Bit where
  touch (Bit w) = R.touch w
  {-# INLINE touch #-}

  zero     = 0
  {-# INLINE zero #-}

  one      = 1
  {-# INLINE one #-}

instance (ColorSpace cs e, Elt e, Num (Pixel cs e)) => Elt (Pixel cs e) where
  touch = foldlPx (\ a !e -> a >> touch e) (return ())
  {-# INLINE touch #-}

  zero     = 0
  {-# INLINE zero #-}

  one      = 1
  {-# INLINE one #-}


module Render where

import           Control.Monad

import           Data.Complex
import           Data.Matrix                  (Matrix)
import qualified Data.Matrix                  as M
import           Data.Vector                  (Vector, (!), (//))
import qualified Data.Vector                  as V

import qualified Graphics.UI.Threepenny       as UI
import           Graphics.UI.Threepenny.Core
import           Graphics.UI.Threepenny.Timer

type GridElem = UI ()
type Grid = Matrix GridElem

render :: (Grid -> Grid) -> IO ()
render xformer = startGUI defaultConfig { jsStatic = Nothing } (setup xformer)

canvasSize, clockInterval :: Num n => n
canvasSize = 2000
clockInterval = 50

constMatrix :: Int -> Int -> a -> Matrix a
constMatrix x y val = M.matrix x y (const val)

(///) :: Matrix a -> [(Int, Int, a)] -> Matrix a
m /// []             = m
m /// ((a, b, c):xs) = M.setElem c (a, b) m /// xs

setup :: (Grid -> Grid) -> Window -> UI ()
setup stateTransformer window = do
    _ <- return window # set title "Canvas - Examples"

    canvas <- UI.canvas
        # set UI.height canvasSize
        # set UI.width  canvasSize
        # set style [("border", "solid black 1px"), ("background", "#eee")]

    clock <- timer # set interval clockInterval
    start clock

    _ <- getBody window #+ [ column [element canvas] ]

    canvas # UI.clearCanvas
    set UI.fillStyle (UI.htmlColor "#23364A") (return canvas)
    UI.fillRect (0, 0) canvasSize canvasSize canvas
    let fillProp color = do
          set UI.fillStyle (UI.htmlColor color) (return canvas)
          UI.fill canvas
          return ()
    let defaultProp = fillProp "#A2C5E8"
    let darkProp    = fillProp "#D2691E"
    let gridSize    = 8
    let defaultData = constMatrix (2 * gridSize) gridSize defaultProp
    -- example of modifying defaultData:
    -- let customData  = (defaultData /// [ (5, 0, darkProp)
    --                                    , (5, 3, darkProp)
    --                                    , (5, 5, darkProp)
    --                                    , (5, 7, darkProp)
    --                                    ])

    clockEv <- accumE defaultData $ stateTransformer <$ UI.tick clock
    onEvent clockEv $ genGrid canvas

genGrid :: UI.Canvas
        -> Matrix GridElem
        -> UI ()
genGrid canvas aGrid =
  forM_ [1, 3 .. (M.nrows aGrid)] $ \n -> do
    genRow canvas (s1 n) (M.getRow n       aGrid) sc
    genRow canvas (s2 n) (M.getRow (n + 1) aGrid) sc
  where
    sc          = 50
    psc         = 10 + sc
    s1 n        = (1.00 * psc, (fromIntegral n)     * psc)
    s2 n        = (2.25 * psc, (fromIntegral n + 1) * psc)

genRow :: UI.Canvas
       -> (Double, Double)
       -> Vector GridElem
       -> Double
       -> UI ()
genRow canvas (startX, startY) aRow scale =
  forM_ [0 .. V.length aRow - 1] $ \n ->
    genHex canvas (startX + offset n, startY) (aRow ! n) scale
  where
    offset n = 3 * scale * fromIntegral n

genHex :: UI.Canvas
       -> (Double, Double)
       -> GridElem
       -> Double
       -> UI ()
genHex canvas (startX, startY) props scale = do
  canvas # UI.beginPath
  movePos $ vertex 0
  forM_ [1 .. 6] $ linePos . vertex
  canvas # UI.closePath
  props
  canvas # UI.stroke
  where
    fromComp (x :+ y) = (x, y)
    startPos          = startX :+ startY
    sixthROU          = cis $ pi / 3

    vertex :: Integer -> Complex Double
    vertex n = ((sixthROU ^ n) * (scale :+ 0)) + startPos

    movePos pos = flip UI.moveTo canvas $ fromComp pos
    linePos pos = flip UI.lineTo canvas $ fromComp pos


module Syntax
  ( LispVal (..)
  , unwordsList
  , ThrowsError
  , trapError
  , extractValue,
  LispError (..)
  ) where
import           Data.Complex (Complex (..))
import Text.ParserCombinators.Parsec
import Control.Monad.Except

instance Show LispVal where
  show (Atom s)          = s
  show (List xs)         = "(" ++ unwordsList xs ++ ")"
  show (DottedList xs x) = "(" ++ unwordsList xs ++ " . " ++ show x ++ ")"
  show (Number x)        = show x
  show (String s)        = "\"" ++ s ++ "\""
  show (Bool True)       = "#t"
  show (Bool False)      = "#f"
  show (Character c)     = [c]
  show (Float x)         = show x
  show (Complex x)       = show x
  show (Rational x)      = show x

data LispVal = Atom String
             | List [LispVal]
             | DottedList [LispVal] LispVal
             | Number Int
             | Bool Bool
             | Character Char
             | String String
             | Float Double
             | Complex (Complex Double)
             | Rational Rational
             | PrimitiveFunc [LispVal -> ThrowsError LispVal]

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

-- Error
type ThrowsError = Either LispError

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

instance Show LispError where
  show (NumArgs expected found) = "Expected " ++ show expected ++ " args; found values " ++ unwordsList found
  show (TypeMismatch expected found) = "Invalid type: expected " ++ expected ++ ", found " ++ show found
  show (Parser parseErr) = "Parse error at " ++ show parseErr
  show (BadSpecialForm message form) = message ++ ": " ++ show form
  show (NotFunction message func) = message ++ ": " ++ show func
  show (UnboundVar message varname) = message ++ ": " ++ varname
  show (Default message) = message

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

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


{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE FlexibleInstances #-}
{-# LANGUAGE GADTs #-}
{-# LANGUAGE GeneralizedNewtypeDeriving #-}
{-# LANGUAGE InstanceSigs #-}
{-# LANGUAGE MultiParamTypeClasses #-}
{-# LANGUAGE OverloadedStrings #-}
{-# LANGUAGE RankNTypes #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE StandaloneDeriving #-}
{-# LANGUAGE TypeFamilies #-}
{-# LANGUAGE TypeOperators #-}

-- |
-- Module      :  Spiral.Program.Monad
-- Copyright   :  (c) 2017 Drexel University
-- License     :  BSD-style
-- Maintainer  :  mainland@drexel.edu

module Spiral.Program.Monad (
    P(..),
    evalP,

    alwaysUnroll,
    shouldUnroll,

    comment,

    (.:=.),
    assignP,
    forP,
    tempP,

    cacheExp,

    newArray,
    cacheArray
  ) where

import Control.Monad.Exception (MonadException(..))
import Control.Monad.Fail (MonadFail)
import Control.Monad.IO.Class (MonadIO)
import Control.Monad.Primitive (PrimMonad(..))
import Control.Monad.Ref (MonadRef(..))
import Control.Monad.Reader (MonadReader(..),
                             ReaderT,
                             asks,
                             runReaderT)
import Control.Monad.State (MonadState(..),
                            StateT,
                            execStateT,
                            gets,
                            modify)
import Control.Monad.Trans.Class (MonadTrans(..))
import Data.Complex
import Data.IORef (IORef)
import Data.Map (Map)
import qualified Data.Map as Map
import Data.Ratio
import Data.Sequence ((|>))
import Text.PrettyPrint.Mainland

import Data.Heterogeneous
import Spiral.Array
import Spiral.Array.Repr.Hidden
import Spiral.Array.Repr.Complex
import Spiral.Array.Repr.Concrete
import Spiral.Config
import Spiral.Exp
import Spiral.Monad (MonadSpiral)
import Spiral.Program.Syntax
import Spiral.Util.Trace
import Spiral.Util.Uniq

data PEnv = PEnv { unroll :: Bool }

defaultPEnv :: PEnv
defaultPEnv = PEnv { unroll = False }

data PState = PState
    { decls  :: Decls
    , stms   :: Stms
    , ecache :: Map (Some Exp) Var
    }

defaultPState :: PState
defaultPState = PState
    { decls  = mempty
    , stms   = mempty
    , ecache = mempty
    }

newtype P m a = P { unP :: StateT PState (ReaderT PEnv m) a}
    deriving (Functor, Applicative, Monad, MonadFail, MonadIO,
              MonadException,
              MonadReader PEnv,
              MonadState PState,
              MonadUnique,
              MonadTrace,
              MonadConfig)

instance MonadTrans P where
    lift = P . lift . lift

deriving instance MonadRef IORef m => MonadRef IORef (P m)

instance PrimMonad m => PrimMonad (P m) where
    type PrimState (P m) = PrimState m
    primitive = P . primitive

instance MonadSpiral m => MonadSpiral (P m) where

evalP :: Monad m => P m () -> m Block
evalP m = do
    s <- runReaderT (execStateT (unP m) defaultPState) defaultPEnv
    return $ Block (decls s) (stms s)

appendDecl :: Monad m => Decl -> P m ()
appendDecl d = modify $ \s -> s { decls = decls s |> d }

appendStm :: Monad m => Stm -> P m ()
appendStm stm = modify $ \s -> s { stms = stms s |> stm }

inNewBlock :: Monad m => P m a -> P m (Block, a)
inNewBlock k = do
    old_decls <- gets decls
    old_stms  <- gets stms
    modify $ \s -> s { decls = mempty, stms = mempty }
    x <- k
    block <- Block <$> gets decls <*> gets stms
    modify $ \s -> s { decls = old_decls, stms = old_stms }
    return (block, x)

inNewBlock_ :: Monad m => P m a -> P m Block
inNewBlock_ m = fst <$> inNewBlock m

-- | Look up an expression in the expression cache.
lookupCachedExp :: Monad m => Exp a -> P m (Maybe Var)
lookupCachedExp e = Map.lookup (Some e) <$> gets ecache

-- | Insert an expression into the expression cache.
insertCachedExp :: Monad m => Var -> Exp a -> P m ()
insertCachedExp v e =
    modify $ \s -> s { ecache = Map.insert (Some e) v (ecache s) }

-- | Always unroll loops in the given computation.
alwaysUnroll :: Monad m => P m a -> P m a
alwaysUnroll = local $ \env -> env { unroll = True }

-- | Return 'True' if loops of the given size should be unrolled.
shouldUnroll :: MonadSpiral m => Int -> P m Bool
shouldUnroll n = do
    always <- asks unroll
    maxun  <- asksConfig maxUnroll
    return $ always || n <= maxun

-- | Comment on the code being generated.
comment :: MonadSpiral m => Doc -> P m ()
comment doc = whenDynFlag GenComments $ appendStm (CommentS doc)

-- | @'forP' begin end k@ generates code for a fro loop with bounds @begin@
-- and @end@, where the body of the loop is generated by calling theVarDVarD
-- continuation @k@ with an expression representing the loop variable.
forP :: MonadSpiral m => Int -> Int -> (Exp Int -> P m ()) -> P m ()
forP lo hi k = do
    should <- shouldUnroll (hi - lo)
    if should
      then mapM_ k [intE i | i <- [lo..hi-1::Int]]
      else do
        i     <- gensym "i"
        block <- inNewBlock_ $ k (VarE i)
        appendStm $ ForS i lo hi block

-- | Generate a temporary of the given type.
tempP :: forall a m . (Typed a, MonadSpiral m) => P m (Exp a)
tempP = do
    useComplexType <- asksConfig $ testDynFlag UseComplex
    go useComplexType tau
  where
    tau :: Type a
    tau = typeOf (undefined :: a)

    go :: Bool -> Type b -> P m (Exp b)
    go False (ComplexT tau') = do
        er <- go False tau'
        ei <- go False tau'
        return $ ComplexE er ei

    go _ tau = do
        v <- gensym "t"
        appendDecl $ VarD v tau
        return $ VarE v

-- | An alias for 'assignP'.
infix 4 .:=.
(.:=.) :: (Typed a, Num (Exp a), MonadSpiral m) => Exp a -> Exp a -> P m ()
(.:=.) = assignP

-- | Assign one expression to another.
assignP :: forall a m . (Typed a, Num (Exp a), MonadSpiral m) => Exp a -> Exp a -> P m ()
assignP e1 e2 = do
    e2' <- cacheExp e2
    assignS e1 e2'

assignS :: (Typed a, MonadSpiral m) => Exp a -> Exp a -> P m ()
assignS e1@(VarE v) e2 = do
    insertCachedExp v e2
    appendStm $ AssignS e1 e2

assignS (ComplexE e1r e1i) (ComplexE e2r e2i) = do
    assignS e1r e2r
    assignS e1i e2i

assignS e1 e2 =
    appendStm $ AssignS e1 e2

-- | Cache the given expression. This serves as a hint that it will be used
-- more than once.
cacheExp :: forall a m . (Typed a, Num (Exp a), MonadSpiral m) => Exp a -> P m (Exp a)
cacheExp e = do
    config <- askConfig
    cacheWithConfig config e

cacheWithConfig :: forall a m . (Typed a, Num (Exp a), MonadSpiral m)
                => Config
                -> Exp a
                -> P m (Exp a)
cacheWithConfig config = cache
  where
    useComplex, splitComplex, doCSE, doRewrite, difRewrite :: Bool
    useComplex   = UseComplex `testDynFlag` config
    splitComplex = SplitComplex `testDynFlag` config
    doCSE        = CSE `testDynFlag` config
    doRewrite    = Rewrite `testDynFlag` config
    difRewrite   = DifRewrite `testDynFlag` config

    cache :: forall a . (Typed a, Num (Exp a)) => Exp a -> P m (Exp a)
    cache e
      | doRewrite = rewrite e
      | otherwise = mustCache e

    rewrite :: forall a . (Typed a, Num (Exp a)) => Exp a -> P m (Exp a)
    rewrite (ConstE (FloatC x)) | x < 0 =
        return $ UnopE Neg (ConstE (FloatC (-x)))

    rewrite (ConstE (DoubleC x)) | x < 0 =
        return $ UnopE Neg (ConstE (DoubleC (-x)))

    rewrite e@ConstE{} =
        return e

    rewrite e@VarE{} =
        return e

    rewrite (ComplexE er ei) =
        ComplexE <$> rewrite er <*> rewrite ei

    -- Don't cache negated constants. We don't want to fall through to the next
    -- case because the call to 'negate' will re-negate the constant!
    rewrite e@(UnopE Neg ConstE{}) =
        return e

    -- Cache the term we are negating, not the negated term.
    rewrite (UnopE Neg e) =
        negate <$> rewrite e

    rewrite (UnopE op e) = do
        e' <- rewrite e
        mustCache (UnopE op e')

    rewrite (BinopE op e1 e2) = do
        e1' <- rewrite e1
        e2' <- rewrite e2
        go op e1' e2'
      where
        go :: Binop -> Exp a -> Exp a -> P m (Exp a)
        -- Choose canonical variable ordering
        go Add e1@VarE{} e2@VarE{} | e2 < e1 =
            mustCache (BinopE Add e2 e1)

        go Add e1@IdxE{} e2@IdxE{} | e2 < e1 =
            mustCache (BinopE Add e2 e1)

        go Sub e1@VarE{} e2@VarE{} | e2 < e1 =
            UnopE Neg <$> mustCache (BinopE Sub e2 e1)

        go Sub e1@IdxE{} e2@IdxE{} | e2 < e1 =
            UnopE Neg <$> mustCache (BinopE Sub e2 e1)

        go Mul e1@VarE{} e2@VarE{} | e2 < e1 =
            mustCache (BinopE Mul e2 e1)

        go Mul e1@IdxE{} e2@IdxE{} | e2 < e1 =
            mustCache (BinopE Mul e2 e1)

        -- Constants always come first
        go Mul e1 e2@ConstE{} | not (isConstE e1) =
            go Mul e2 e1

        -- Push negation out
        go Add e1 (UnopE Neg e2) =
            rewrite $ e1 - e2

        go Add (UnopE Neg e1) e2 =
            rewrite $ e2 - e1

        go Sub e1 (UnopE Neg e2) =
            rewrite $ e1 + e2

        go Sub (UnopE Neg e1) e2 =
            rewrite $ -(e1 + e2)

        go Mul e1 (UnopE Neg e2) =
            rewrite $ -(e1 * e2)

        go Mul (UnopE Neg e1) e2 =
            rewrite $ -(e1 * e2)

        -- Rewrite @w1*w2*t@ to @(w1*w2)*t@
        go Mul w1@(ConstE W{}) (BinopE Mul w2@(ConstE W{}) e) | difRewrite =
            cache $ (w1*w2) * e

        go Mul e1@(ConstE W{}) e2 =
            return $ BinopE Mul e1 e2

        -- Rewrite @w1*t1 + w2*t2@ to @w1*(t1 + (w2/w1)*t2)@ where either @w1@
        -- or @w2/w1@ is a 4th root of unity.
        go op e1@(BinopE Mul w1@(ConstE W{}) t1) e2@(BinopE Mul w2@(ConstE W{}) t2) | difRewrite, Just (+/-) <- fromAddSub op =
            case tau of
              ComplexT{} | isTrivialRoot w1 || isTrivialRoot (w2/w1) ->
                cache $ w1 * (t1 +/- ((w2/w1) * t2))
              _ -> mustCache (BinopE op e1 e2)
          where
            tau :: Type a
            tau = typeOf (undefined :: a)

        go op e1 e2 =
            mustCache (BinopE op e1 e2)

        fromAddSub :: Binop -> Maybe (Exp a -> Exp a -> Exp a)
        fromAddSub Add = return (+)
        fromAddSub Sub = return (-)
        fromAddSub _   = fail "isAddSub: neither (+) nor (-)"

    rewrite (ReE e) =
        ReE <$> rewrite e >>= mustCache

    rewrite (ImE e) =
        ImE <$> rewrite e >>= mustCache

    rewrite e =
        mustCache e

    -- When we are not using complex types, always separately cache the real and
    -- imaginary parts of a complex expression.
    mustCache :: forall a . (Typed a, Num (Exp a)) => Exp a -> P m (Exp a)
    mustCache (ComplexE er ei) =
        ComplexE <$> mustCache er <*> mustCache ei

    mustCache (ReE e) | not useComplex, splitComplex = do
        let (er, _ei) = forceUnComplexE e
        cache er

    mustCache (ImE e) | not useComplex, splitComplex = do
        let (_er, ei) = forceUnComplexE e
        cache ei

    mustCache e | not useComplex, splitComplex, ComplexT{} <- tau = do
        let (er, ei) = forceUnComplexE e
        ComplexE <$> cache er <*> cache ei >>= cache
      where
        tau :: Type a
        tau = typeOf (undefined :: a)

    mustCache e | doCSE = do
        maybe_v <- lookupCachedExp e
        case maybe_v of
          Just v  -> return $ VarE v
          Nothing -> do temp <- tempP
                        assignS temp e
                        return temp

    mustCache e =
        return e

-- | @isNthRoot n r@ returns 'True' if @r@ is an nth root of unity.
isNthRoot :: Int -> Exp (Complex a) -> Bool
isNthRoot n (ConstE (W n' k _)) = denominator ((k % n') * fromIntegral n) == 1
isNthRoot _ _                   = False

-- | 'isTrivialRoot' returns 'True' if its argument is a trivial root, i.e., a
-- fourth root of unity.
isTrivialRoot :: Exp (Complex a) -> Bool
isTrivialRoot = isNthRoot 4

-- | Create a new concrete array of the given shape.
newArray :: forall sh a m . (Shape sh, Typed a, Num (Exp a), MonadSpiral m)
         => (sh :. Int)
         -> P m (Array H (sh :. Int) (Exp a))
newArray (sh :. i) = do
    useComplexType <- asksConfig $ testDynFlag UseComplex
    v <- gensym "V"
    go useComplexType v (typeOf (undefined :: a))
  where
    go :: Bool -> Var -> Type a -> P m (Array H (sh :. Int) (Exp a))
    go False v (ComplexT tau) = do
        appendDecl $ ArrD v (sh :. 2*i) tau
        return $ H $ CMPLX $ C (sh :. 2*i) v

    go _ v tau = do
        appendDecl $ ArrD v (sh :. i) tau
        return $ H $ C (sh :. i) v

-- | Cache the contents of a matrix, returning it as a concrete matrix.
cacheArray :: forall a r sh m . (Typed a, Num (Exp a), Shape sh, IArray r (sh :. Int) (Exp a), MonadSpiral m)
           => Array r (sh :. Int) (Exp a)
           -> P m (Array H (sh :. Int) (Exp a))
cacheArray arr = do
    useComplexType <- asksConfig $ testDynFlag UseComplex
    v <- gensym "K"
    go useComplexType v (typeOf (undefined :: a))
  where
    go :: Bool -> Var -> Type a -> P m (Array H (sh :. Int) (Exp a))
    go False v ComplexT{} = do
        appendDecl $ ConstArrD v (RE arr)
        return $ H $ CMPLX $ C (extent (RE arr)) v

    go _ v _ = do
        appendDecl $ ConstArrD v arr
        return $ H $ C (extent arr) v