-- 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 FlexibleContexts #-}
{-# LANGUAGE FlexibleInstances #-}
{-# LANGUAGE MultiParamTypeClasses #-}
{-# LANGUAGE RankNTypes #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE TypeFamilies #-}
module TensorFlow.Nodes where

import Control.Applicative (liftA2, liftA3)
import Data.Map.Strict (Map)
import Data.Monoid ((<>))
import Data.Set (Set)
import Data.Text (Text)
import Lens.Family2 ((^.))
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 ControlNode where
    getNodes (ControlNode o) = Set.singleton <$> getOrAddOp 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 (Tensor v a) where
    getNodes t = Set.singleton <$> getOrAddOp (t ^. tensorOutput . outputOp)

fetchTensorVector :: forall a v . TensorType a
                  => Tensor v a -> Build (Fetch (TensorData a))
fetchTensorVector (Tensor _ o) = do
    outputName <- renderOutput o
    return $ 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