content: -- 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 GeneralizedNewtypeDeriving #-}
  {-# LANGUAGE OverloadedStrings #-}
  {-# LANGUAGE Rank2Types #-}
  {-# LANGUAGE ScopedTypeVariables #-}
  
  module TensorFlow.Output
      ( ControlNode(..)
      , Device(..)
      -- * Ops
      , NodeName(..)
      , OpDef(..)
      , opName
      , opType
      , opAttr
      , opInputs
      , opControlInputs
      , OpType(..)
      , OutputIx(..)
      , Output(..)
      , output
      , PendingNodeName(..)
      )  where
  
  import Data.ProtoLens.Message(defMessage)
  import qualified Data.Map.Strict as Map
  import Data.String (IsString(..))
  import Data.Text (Text)
  import qualified Data.Text as Text
  import Lens.Family2 (Lens')
  import Lens.Family2.Unchecked (lens)
  import Proto.Tensorflow.Core.Framework.AttrValue (AttrValue)
  import TensorFlow.Types (Attribute, attrLens)
  
  -- | A type of graph node which has no outputs. These nodes are
  -- valuable for causing side effects when they are run.
  newtype ControlNode = ControlNode { unControlNode :: NodeName }
  
  -- | The type of op of a node in the graph.  This corresponds to the proto field
  -- NodeDef.op.
  newtype OpType = OpType { unOpType :: Text }
      deriving (Eq, Ord, Show)
  
  instance IsString OpType where
      fromString = OpType . Text.pack
  
  -- | An output of a TensorFlow node.
  data Output = Output {outputIndex :: !OutputIx, outputNodeName :: !NodeName}
      deriving (Eq, Ord, Show)
  
  output :: OutputIx -> NodeName -> Output
  output = Output
  
  newtype OutputIx = OutputIx { unOutputIx :: Int }
      deriving (Eq, Ord, Num, Enum, Show)
  
  -- | A device that a node can be assigned to.
  -- There's a naming convention where the device names
  -- are constructed from job and replica names.
  newtype Device = Device {deviceName :: Text}
      deriving (Eq, Ord, IsString)
  
  instance Show Device where
      show (Device d) = show d
  
  -- | Op definition. This corresponds somewhat to the 'NodeDef' proto.
  data OpDef = OpDef
      { _opName :: !PendingNodeName
      , _opType :: !OpType
      , _opAttrs :: !(Map.Map Text AttrValue)
      , _opInputs :: [Output]
      , _opControlInputs :: [NodeName]
      }  deriving (Eq, Ord)
  
  -- | The name specified for an unrendered Op.  If an Op has an
  -- ImplicitName, it will be assigned based on the opType plus a
  -- unique identifier.  Does not contain the "scope" prefix.
  data PendingNodeName = ExplicitName !Text | ImplicitName
      deriving (Eq, Ord, Show)
  
  instance IsString PendingNodeName where
      fromString = ExplicitName . fromString
  
  -- | The name of a node in the graph.  This corresponds to the proto field
  -- NodeDef.name.  Includes the scope prefix (if any) and a unique identifier
  -- (if the node was implicitly named).
  newtype NodeName = NodeName { unNodeName :: Text }
      deriving (Eq, Ord, Show)
  
  opName :: Lens' OpDef PendingNodeName
  opName = lens _opName (\o x -> o {_opName = x})
  
  opType :: Lens' OpDef OpType
  opType = lens _opType (\o x -> o { _opType = x})
  
  opAttr :: Attribute a => Text -> Lens' OpDef a
  opAttr n = lens _opAttrs (\o x -> o {_opAttrs = x})
                . lens (Map.findWithDefault defMessage n) (flip (Map.insert n))
                . attrLens
  
  opInputs :: Lens' OpDef [Output]
  opInputs = lens _opInputs (\o x -> o {_opInputs = x})
  
  opControlInputs :: Lens' OpDef [NodeName]
  opControlInputs = lens _opControlInputs (\o x -> o {_opControlInputs = x})
  
  -- TODO(gnezdo): IsString instance is weird and we should move that
  -- code into a Build function
  instance IsString Output where
      fromString s = case break (==':') s of
          (n, ':':ixStr) | ix <- read ixStr
            -> Output (fromInteger ix) $ assigned n
          _ -> Output 0 $ assigned s
       where assigned = NodeName . Text.pack
  