-- Copyright 2016 TensorFlow authors.
--
-- Licensed under the Apache License, Version 2.0 (the "License");
-- you may not use this file except in compliance with the License.
-- You may obtain a copy of the License at
--
--     http://www.apache.org/licenses/LICENSE-2.0
--
-- Unless required by applicable law or agreed to in writing, software
-- distributed under the License is distributed on an "AS IS" BASIS,
-- WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-- See the License for the specific language governing permissions and
-- limitations under the License.

{-# LANGUAGE DataKinds #-}
{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE FlexibleInstances #-}
{-# LANGUAGE MultiParamTypeClasses #-}
{-# LANGUAGE RankNTypes #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE TypeFamilies #-}
{-# LANGUAGE TypeOperators #-}
{-# LANGUAGE UndecidableInstances #-}  -- For Fetchable (TensorExpr a)
module TensorFlow.Nodes where

import Control.Applicative (liftA2, liftA3)
import Data.Functor.Identity (Identity)
import Data.Map.Strict (Map)
import Data.Set (Set)
import Data.Text (Text)
import qualified Data.Map.Strict as Map
import qualified Data.Set as Set

import TensorFlow.Build
import TensorFlow.Output
import TensorFlow.Tensor
import TensorFlow.Types
import qualified TensorFlow.Internal.FFI as FFI

-- | Types that contain ops which can be run.
class Nodes t where
    getNodes :: t -> Build (Set NodeName)

-- | Types that tensor representations (e.g. 'Tensor', 'ControlNode') can be
-- fetched into.
--
-- Includes collections of tensors (e.g. tuples).
class Nodes t => Fetchable t a where
    getFetch :: t -> Build (Fetch a)

-- | Fetch action. Keeps track of what needs to be fetched and how to decode
-- the fetched data.
data Fetch a = Fetch
          { -- | Nodes to fetch
            fetches :: Set Text
            -- | Function to create an 'a' from the fetched data.
          , fetchRestore :: Map Text FFI.TensorData -> a
          }

instance Functor Fetch where
    fmap f (Fetch fetch restore) = Fetch fetch (f . restore)

instance Applicative Fetch where
    pure x = Fetch Set.empty (const x)
    Fetch fetch restore <*> Fetch fetch' restore' =
        Fetch (fetch <> fetch') (restore <*> restore')

nodesUnion :: (Monoid b, Traversable t, Applicative f) => t (f b) -> f b
nodesUnion = fmap (foldMap id) . sequenceA

instance (Nodes t1, Nodes t2) => Nodes (t1, t2) where
    getNodes (x, y) = nodesUnion [getNodes x, getNodes y]

instance (Nodes t1, Nodes t2, Nodes t3) => Nodes (t1, t2, t3) where
    getNodes (x, y, z) = nodesUnion [getNodes x, getNodes y, getNodes z]

instance (Fetchable t1 a1, Fetchable t2 a2) => Fetchable (t1, t2) (a1, a2) where
    getFetch (x, y) = liftA2 (,) <$> getFetch x <*> getFetch y

instance (Fetchable t1 a1, Fetchable t2 a2, Fetchable t3 a3)
         => Fetchable (t1, t2, t3) (a1, a2, a3) where
    getFetch (x, y, z) =
        liftA3 (,,) <$> getFetch x <*> getFetch y <*> getFetch z

instance Nodes t => Nodes [t] where
    getNodes = nodesUnion . map getNodes

instance Fetchable t a => Fetchable [t] [a] where
    getFetch ts  = sequenceA <$> mapM getFetch ts

instance Nodes t => Nodes (Maybe t) where
    getNodes = nodesUnion . fmap getNodes

instance Fetchable t a => Fetchable (Maybe t) (Maybe a) where
    getFetch = fmap sequenceA . mapM getFetch

instance Nodes ControlNode where
    getNodes (ControlNode o) = pure $ Set.singleton o

-- We use the constraint @(a ~ ())@ to help with type inference.  For example,
-- if @t :: ControlNode@, then this constraint ensures that @run t :: Session
-- ()@.  If we used @instance Fetchable ControlNode ()@ instead, then that
-- expression would be ambiguous without explicitly specifying the return type.
instance a ~ () => Fetchable ControlNode a where
    getFetch _ = return $ pure ()

instance Nodes (ListOf f '[]) where
    getNodes _ = return Set.empty

instance (Nodes (f a), Nodes (ListOf f as)) => Nodes (ListOf f (a ': as)) where
    getNodes (x :/ xs) = liftA2 Set.union (getNodes x) (getNodes xs)

instance l ~ List '[] => Fetchable (ListOf f '[]) l where
    getFetch _ = return $ pure Nil

instance (Fetchable (f t) a, Fetchable (ListOf f ts) (List as), i ~ Identity)
    => Fetchable (ListOf f (t ': ts)) (ListOf i (a ': as)) where
    getFetch (x :/ xs) = liftA2 (\y ys -> y /:/ ys) <$> getFetch x <*> getFetch xs

instance Nodes (Tensor v a) where
    getNodes (Tensor o) = Set.singleton . outputNodeName <$> toBuild o

fetchTensorVector :: forall a v . (TensorType a)
                  => Tensor v a -> Build (Fetch (TensorData a))
fetchTensorVector (Tensor o) = do
    outputName <- encodeOutput <$> toBuild o
    pure $ Fetch (Set.singleton outputName) $ \tensors ->
        let tensorData = tensors Map.! outputName
            expectedType = tensorType (undefined :: a)
            actualType = FFI.tensorDataType tensorData
            badTypeError = error $ "Bad tensor type: expected "
                                   ++ show expectedType
                                   ++ ", got "
                                   ++ show actualType
        in if expectedType /= actualType
               then badTypeError
               else TensorData tensorData

-- The constraint "a ~ a'" means that the input/output of fetch can constrain
-- the TensorType of each other.
instance (TensorType a, a ~ a') => Fetchable (Tensor v a) (TensorData a') where
    getFetch = fetchTensorVector

instance (TensorType a, TensorDataType s a, a ~ a') => Fetchable (Tensor v a) (s a') where
    getFetch t = fmap decodeTensorData <$> fetchTensorVector t