  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
--
--     -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 CPP #
# LANGUAGE FlexibleContexts #
# LANGUAGE LambdaCase #
{-# LANGUAGE OverloadedStrings #-}
# LANGUAGE TypeFamilies #
  | Rendering of TensorFlow operations as  functions . 

 The basic type signature generated for each op is : 

 > { constraints } = > { mandatory attrs } - > { input tensors } - > { output tensors } 

 where : 

 * @{mandatory attrs}@ is of the form @A_1 - > ... - >  , where each @A@ is an 
  op attribute that does n't have a default and ca n't be inferred from other 
  inputs . 

 * @{constraints}@ restrict the type parameters of the input and output tensors 
  ( for example : ' TensorType ' or '  ' ) . 

 * @{input tensors}@ is of the form @T_1 - > ... - > T_N@ , where each @T@ is of 
 the form @Tensor Ref a@ or  v a@ ( or a list of one of those types ) , 
 and @a@ is either a concrete type or a ( constrained ) type variable . 

 * @{output tensors}@ is of the form @(T_1, ... ,T_N)@ for " pure " ops , and 
  ( T_1, ... ,T_N)@ for " stateful " ops .   An op is considered " stateful " if 
 it takes a @Tensor Ref@ or @Tensor v ResourceHandle@ as input , or if it 's 
 explicitly marked \"Stateful\ " in its @REGISTER_OP@ definition .   ( If there 
 are no outputs , it is either @ControlNode@ or @Build ControlNode@. ) 


The basic type signature generated for each op is:

> {constraints} => {mandatory attrs} -> {input tensors} -> {output tensors}

where:

* @{mandatory attrs}@ is of the form @A_1 -> ... -> A_N@, where each @A@ is an
 op attribute that doesn't have a default and can't be inferred from other
 inputs.

* @{constraints}@ restrict the type parameters of the input and output tensors
 (for example: 'TensorType' or 'OneOf').

* @{input tensors}@ is of the form @T_1 -> ... -> T_N@, where each @T@ is of
the form @Tensor Ref a@ or @Tensor v a@ (or a list of one of those types),
and @a@ is either a concrete type or a (constrained) type variable.

* @{output tensors}@ is of the form @(T_1,...,T_N)@ for "pure" ops, and
@Build (T_1,...,T_N)@ for "stateful" ops.  An op is considered "stateful" if
it takes a @Tensor Ref@ or @Tensor v ResourceHandle@ as input, or if it's
explicitly marked \"Stateful\" in its @REGISTER_OP@ definition.  (If there
are no outputs, it is either @ControlNode@ or @Build ControlNode@.)
-}

module TensorFlow.OpGen
  ( OpGenFlags(..)
  , docOpList
  , flagParser)
  where

import Data.Foldable (toList)
import Data.Maybe (fromMaybe)
import Data.ProtoLens.Default(def)
import Data.ProtoLens (showMessage)
import Data.List (sortOn)
import Data.List.NonEmpty (NonEmpty)
import qualified Data.List.NonEmpty as NE
import Lens.Family2 ((^.), (.~), (&), view)
import Options.Applicative (Parser, help, long, strOption, value)
import Proto.Tensorflow.Core.Framework.OpDef
  ( OpList
  , OpDef
  )
import Proto.Tensorflow.Core.Framework.OpDef_Fields
  ( attr
  , inputArg
  , name
  , op
  , outputArg
  )
import Proto.Tensorflow.Core.Framework.Types (DataType(..))
import System.FilePath (takeBaseName)
import TensorFlow.OpGen.ParsedOp
import Text.PrettyPrint.Mainland
  ( Doc
  , (<+>)
  , (</>)
  , (<+/>)
  , brackets
  , comma
  , commasep
  , dquotes
  , empty
  , enclose
  , flatten
  , folddoc
  , hang
  , indent
  , parens
  , sep
  , stack
  , strictText
  , tuple
  )
import qualified Data.Set as Set
import qualified Data.Text as Text

data OpGenFlags = OpGenFlags
     { outputFile :: String
     , prefix :: String
     , excludeList :: String
     }

flagParser :: Parser OpGenFlags
flagParser = OpGenFlags
     <$> strOption (mconcat [ long "output"
                            , help "File to write."
                            ])
     <*> strOption (mconcat [ long "prefix"
                            , help "Haskell package prefix to use"
                            ])
     <*> strOption (mconcat [ long "exclude_list"
                            , value ""
                            , help "Comma separated Ops names to ignore"
                            ])


docOpList :: OpGenFlags -> OpList -> Doc
docOpList flags opList =
  stack [ "{-# LANGUAGE ConstraintKinds #-}"
        , "{-# LANGUAGE DataKinds #-}"
        , "{-# LANGUAGE FlexibleContexts #-}"
        , "{-# LANGUAGE FlexibleInstances #-}"
        , "{-# LANGUAGE OverloadedStrings #-}"
        , "{-# LANGUAGE ScopedTypeVariables #-}"
          -- Avoids reports about shadowing standard library names.
        , "{-# OPTIONS_GHC -fno-warn-name-shadowing #-}"
          -- eqLengthGuard never returns false and dies instead.
        , "{-# OPTIONS_GHC -fno-warn-incomplete-patterns #-}"
        , "module" <+> strictText moduleName <+> "where"
        , empty
        , imports
        , empty
        , folddoc (\x y -> x </> empty </> y)
                  (map renderOpAndExtras $
                   sortOn (view name) $
                   filter (not . flip elem exclusions . view name) $
                   toList $ opList ^. op)
        ]
  where moduleName =
            Text.pack (prefix flags) <> "." <> camelCase
             -- Discards the optional trailing _ops_op_lib
            (fromMaybe shortName (Text.stripSuffix "_ops_op_lib" shortName))
        shortName = Text.pack (takeBaseName $ outputFile flags)
        exclusions = Text.splitOn "," $ Text.pack $ excludeList flags
        renderOpAndExtras o = renderOp (parseOp o) </> extras o

imports :: Doc
imports = stack [
      "import Data.ByteString (ByteString)"
    , "import Data.Complex (Complex)"
    , "import Data.Int (Int8, Int16, Int32, Int64)"
    , "import Data.Proxy (Proxy(Proxy))"
    , "import Data.Word (Word8, Word16, Word32, Word64)"
    , "import Lens.Family2 ((.~), (&))"
    , "import TensorFlow.Build"
    , "import TensorFlow.BuildOp"
    , "import TensorFlow.Tensor"
    , "import TensorFlow.Types"
    ]

renderHaskellName, renderTFName, renderQuotedTFName :: Name -> Doc
renderHaskellName = strictText . unHaskellName . haskellName
renderTFName = strictText . unTFName . tfName
renderQuotedTFName = dquotes . renderTFName


-- | Generate the source code for a single op.
-- For example:
--
-- -- | {haddock comment}
-- foo :: {type sig}
  foo attr1 attr2 input1 input2 | eqLengthGuard [ ... ] = { function body }
renderOp :: ParsedOp -> Doc
renderOp pOp = stack $
    [ haddocks
    -- Prevent unreasonably long compilation times on ghc-7.10, due
    -- to stack calling "-dump-hi" which (unnecessarily) includes the
    -- inlining information, and is large for ops with many arguments.
#if __GLASGOW_HASKELL__ < 800
    , "{-# NOINLINE" <+> n <+> "#-}"
#endif
    , n <+> "::" <+> hang 0 (typeSig empty pOp)
    , n <+> "=" <+> n <> "' id"
    , n' <+> "::" <+> hang 0 (typeSig "OpParams ->" pOp)
    , n' <+> hang 0 args <+> "|" <+> funcGuard listSizeAttrs
                <+> "=" </>  -- args are indented
                    -- the body needs to be indented wrt the name
                    indent indentation (functionBody pOp)
    ] ++ whereClause listSizeAttrs
  where
    n = renderHaskellName $ parsedOpName pOp
    n' = n <> "'"
    listSizeAttrs = inferredListSizeAttrs pOp
    args = sep $ "op'options"
               : (map renderHaskellName
                    $ map attrName (explicitInputAttrs pOp)
                    ++ map parsedArgName (parsedInputs pOp))
    haddocks = "-- |" <+> multilineComment (parsedOpSummary pOp) (parsedOpDescription pOp)

-- | A check that all lists of the given size have the given length.
-- For example:
    eqLengthGuard [ ( " N " , [ ( " input1 " , length input1 ) , ( " input2 " , length input2 ) ] ) ]
funcGuard :: [Attr (NonEmpty Name)] -> Doc
funcGuard attrs = "eqLengthGuard" <+> brackets (commasep entries)
      where
        entries =
            [ parens $ nAttr <> comma <+>
              brackets (commasep $ toList $
                            map renderTensorName (toList $ attrInfo a))
            | a <- attrs
            , let nAttr = renderQuotedTFName (attrName a)
            ]
        renderTensorName x = parens $ renderQuotedTFName x <> comma <+>
                        "length" <+> renderHaskellName x

-- | Define the implicit list length attributes.
-- For example:
--   where
      n1 = fromIntegral ( length input1 ) : : Int64
      n2 = fromIntegral ( length input2 ) : : Int64
whereClause :: [Attr (NonEmpty Name)] -> [Doc]
whereClause [] = []
whereClause as = [indent 2 $ "where" </> indent 2 (stack $ map defineLengthAttr as)]
  where
    defineLengthAttr a = renderHaskellAttrName a <+> "="
                            <+> "fromIntegral (length"
                            <+> renderHaskellName (NE.head $ attrInfo a)
                            <> ") :: Int64"

renderHaskellAttrName :: Attr a -> Doc
renderHaskellAttrName = renderHaskellName . attrName

functionBody :: ParsedOp -> Doc
functionBody pOp
    | parsedOpIsMonadic pOp
        = "build $ do"
            </> indent indentation (bindOpInputsVar
                        </> "buildOp" <+> outputListsSizes <+> opDef)
    | otherwise
        = "pureOp" <+> outputListsSizes <+> "$ do"
            </> indent indentation (bindOpInputsVar </> "return" <+> opDef)
  where
    outputListsSizes = brackets $ commasep
        [ renderHaskellName a
        | ParsedArg { parsedArgCase = ListArg { argLength = a } }
            <- parsedOutputs pOp
        ]
    opInputsVar = "op'inputs"
    bindOpInputsVar = opInputsVar <+> "<- fmap Prelude.concat $ Prelude.sequence"
                            <+> brackets (commasep $ map (\a -> "buildInputs" <+> a) tensorArgs)
    opDef = parens $ hang 0 $ stack $
        "opDef" <+> renderQuotedTFName (parsedOpName pOp) :
        -- Renders type parameter arguments.
        [ "& opAttr" <+> renderQuotedTFName n <+> ".~" <+> inferredTypeExpr a
        | a <- inferredTypeAttrs pOp, let n = attrName a
        ] ++
        -- Renders mandatory attributes as function parameters.
        [ "& opAttr" <+> renderQuotedTFName n <+> ".~" <+> renderHaskellName n
        | a <- explicitInputAttrs pOp, let n = attrName a
        ] ++
        -- Renders sizes of tensor list types having number_attr.
        [ "& opAttr" <+> renderQuotedTFName n <+> ".~" <+> renderHaskellName n
        | a <- inferredListSizeAttrs pOp, let n = attrName a
        ] ++
        ["& op'options & opInputs .~" <+> opInputsVar]
    tensorArgs = renderTensorArg <$> parsedInputs pOp
    renderTensorArg = renderHaskellName . parsedArgName
    inferredTypeExpr a
        | typeParamIsList $ attrInfo a
            = "fromTensorTypes (Proxy :: Proxy" <+> renderHaskellAttrName a
                    <> ")"
        | otherwise = "tensorType (undefined ::" <+> renderHaskellAttrName a
                            <> ")"

-- | Write a comment with the inputs/outputs/attributes in proto format, for
-- debugging.
extras :: OpDef -> Doc
extras d = enclose "{-\n" "\n-}" $
            strictText $ Text.pack $
            showMessage ((def :: OpDef)
                        & inputArg .~ (d ^. inputArg)
                        & outputArg .~ (d ^. outputArg)
                        & attr .~ (d ^. attr))

-- | The type signature for an op.
-- Of the form:
  forall t1 t2 v1 v2 . ( TensorType t1 , TensorType t2 )
--      => {pre} Float -> Tensor t1 v1 -> Tensor t2 v2
-- where "Float" is an explicit input attribute, "Tensor t1 v1" is an input, and
-- "Tensor t2 v2" is an output.
typeSig :: Doc -> ParsedOp -> Doc
typeSig pre pOp = constraints
            <+/> pre </> signatureFold (map attrInput (explicitInputAttrs pOp)
                                ++ map tensorArgAndComment (parsedInputs pOp)
                                ++ [outputs])
  where
    constraints
        | null classConstraints = empty
        | otherwise = "forall" <+> sep typeParams <+> "." <+> tuple classConstraints <+> "=>"
    typeParams = [strictText v | k <- parsedInputs pOp ++ parsedOutputs pOp,
                  ArgSomeTensor v <- [argKind $ parsedArgCase k]]
                ++ [renderHaskellAttrName n | n <- inferredTypeAttrs pOp]
                ++ if parsedOpIsMonadic pOp then ["m'"] else []
    -- Use m' as the type parameter to avoid clashing with an attribute name.
    monadConstraint
        | parsedOpIsMonadic pOp = ["MonadBuild m'"]
        | otherwise = []
    classConstraints = monadConstraint ++ map tensorArgConstraint
                                                    (inferredTypeAttrs pOp)
    signatureFold = folddoc (\x y -> x </> "->" <+> y)
    attrInput a = renderAttrType (attrInfo a) <+> hang 0 ("-- ^" <+> attrComment a)
    renderAttrType (AttrSingle a) = renderAttrBaseType a
    renderAttrType (AttrList a) = brackets $ renderAttrBaseType a
    renderAttrBaseType = \case
        AttrBytes -> "ByteString"
        AttrInt64 -> "Data.Int.Int64"
        AttrFloat -> "Float"
        AttrBool -> "Bool"
        AttrType -> "DataType"
        AttrShape -> "Shape"
        AttrTensor -> "TensorProto"
        AttrFunc -> error "AttrFunc not supported"

    tensorArgAndComment t = tensorArg t <+> hang 0 ("-- ^" <+> argComment t)
    outputs = case parsedOutputs pOp of
        [] -> wrapOutput "ControlNode"
        -- TODO(judahjacobson): To improve indentation: `tensorArgAndComment a`
        [a] -> wrapOutput (tensorArg a) <+> "-- ^" <+> argComment a
        as -> wrapOutput (tuple (map tensorArg as)) <+/> resultComment as
    wrapOutput o
        | parsedOpIsMonadic pOp = "m'" <+> parens o
        | otherwise = o

-- | Render an op input or output.
-- For example: "Tensor Ref Int64", "Tensor v t"
tensorArg :: ParsedArg -> Doc
tensorArg p = case parsedArgCase p of
    SimpleArg { argType = t, argKind = k } -> tensorType t k
    ListArg { argType = t, argKind = k } -> brackets $ tensorType t k
    MixedListArg {argTypeAttr = t, argKind = k}
        -> "TensorList" <+> parens (kind k) <+> renderHaskellName t
  where
    kind k = case k of
                ArgTensorRef -> "Ref"
                ArgTensorValue -> "Value"
                ArgTensorBuild -> "Build"
                ArgSomeTensor v -> strictText v
    tensorType t k = let
        a = case t of
                ArgTypeFixed dt -> strictText $ dtTypeToHaskell dt
                ArgTypeAttr n -> renderHaskellName n
        in "Tensor" <+> kind k <+> a

attrComment :: Attr a -> Doc
attrComment a = argComment' (attrName a) (attrDescription a)

argComment :: ParsedArg -> Doc
argComment a = argComment' (parsedArgName a) (parsedArgDescription a)

argComment' :: Name -> Text.Text -> Doc
argComment' argName argDesc =
    bold (renderTFName argName) <> splitMultilineText (":" <+>) argDesc

bold :: Doc -> Doc
bold n = "__" <> n <> "__"

-- | Comment for the outputs of an op.
-- For example:
--   -- ^ (__output1__, __output2__)
--   --
    -- * _ _ output1 _ _ : 
--   --
    -- * _ _ output2 _ _ : 
resultComment :: [ParsedArg] -> Doc
resultComment os = stack $ flatten commentSummary : map commentDetails os
  where
    commentSummary = "-- ^" <+> tuple [bold (renderTFName $ parsedArgName o) | o <- os]
    commentDetails o =
        stack [ "--"
              , "-- *" <+> argComment o
              ]

-- | Constraints for a given type parameter.
  E.g. : " TensorType t " or " OneOf [ Int64 , Float ] t "
  or "  ts " or " OneOfs [ .. ] ts " .
tensorArgConstraint :: Attr TypeParam -> Doc
tensorArgConstraint a = case attrInfo a of
    TypeParam False Nothing -> "TensorType" <+> n
    TypeParam False (Just as) -> "OneOf" <+> typeList as <+> n
    TypeParam True Nothing -> "TensorTypes" <+> n
    TypeParam True (Just as) -> "OneOfs" <+> typeList as <+> n
  where
    n = renderHaskellAttrName a
    -- Produces a type-level list, e.g.: '[Int32,Int64,Float]
    typeList = ("'" <>) . brackets . commasep . map strictText .
                    Set.toList . Set.fromList .
                    map dtTypeToHaskell . toList

-- NOTE: The cases of this function should be kept in sync with
  TensorFlow . Types .  .
dtTypeToHaskell :: DataType -> Text.Text
dtTypeToHaskell DT_BOOL = "Bool"
dtTypeToHaskell DT_BFLOAT16 = "Data.Word.Word16"
dtTypeToHaskell DT_COMPLEX128 = "(Data.Complex.Complex Double)"
dtTypeToHaskell DT_COMPLEX64 = "(Data.Complex.Complex Float)"
dtTypeToHaskell DT_DOUBLE = "Double"
dtTypeToHaskell DT_FLOAT = "Float"
dtTypeToHaskell DT_INT16 = "Data.Int.Int16"
dtTypeToHaskell DT_INT32 = "Data.Int.Int32"
dtTypeToHaskell DT_INT64 = "Data.Int.Int64"
dtTypeToHaskell DT_INT8 = "Data.Int.Int8"
  TODO(gnezdo ): make unique
  TODO(gnezdo ): make unique
  TODO(gnezdo ): make unique
  TODO(gnezdo ): make unique
  TODO(gnezdo ): make unique
dtTypeToHaskell DT_STRING = "Data.ByteString.ByteString"
dtTypeToHaskell DT_UINT16 = "Data.Word.Word16"
dtTypeToHaskell DT_UINT32 = "Data.Word.Word32"
dtTypeToHaskell DT_UINT64 = "Data.Word.Word64"
  TODO(gnezdo ): make unique
dtTypeToHaskell DT_UINT8 = "Data.Word.Word8"
dtTypeToHaskell DT_RESOURCE = "ResourceHandle"
dtTypeToHaskell DT_VARIANT = "Variant"
dtTypeToHaskell x =
    Text.pack $ "Unsupported type in dtTypeToHaskell: " ++ show x

-- | haddockComment escapes TensorFlow doc strings into haddock.
  TODO(gnezdo ): deal with the markup .
haddockComment :: Text.Text -> Doc
haddockComment = strictText

-- | Generate a multiline comment.  For example:
--   summary'
--   --
--   -- detail_line1
--   -- detail_line2
--   -- ...
multilineComment :: Text.Text -> Text.Text -> Doc
multilineComment summary' detail =
    haddockComment summary' </>
    splitMultilineText insertParagraphAndComment detail
  where insertParagraphAndComment x = "--" </> "--" <+> x

-- | Converts the given multi-line detail string into
-- a multi-line haddock. Applies the given lead to the
  first line . Returns an empty document for empty detail .
splitMultilineText :: (Doc -> Doc) -> Text.Text -> Doc
splitMultilineText lead detail =
  case Text.lines detail of
    [] -> empty
    (l : ls) -> stack $ lead (haddockComment l)
                      : map (("--" <+>) . haddockComment) ls

indentation :: Int
indentation = 4


 you may not use this file except in compliance with the License.
 You may obtain a copy of the License at

     -2.0

 Unless required by applicable law or agreed to in writing, software
 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 OverloadedStrings #
 Avoids reports about shadowing standard library names.
 eqLengthGuard never returns false and dies instead.
 Discards the optional trailing _ops_op_lib
 | Generate the source code for a single op.
 For example:

 -- | {haddock comment}
 foo :: {type sig}
 Prevent unreasonably long compilation times on ghc-7.10, due
 to stack calling "-dump-hi" which (unnecessarily) includes the
 inlining information, and is large for ops with many arguments.
 args are indented
 the body needs to be indented wrt the name
 | A check that all lists of the given size have the given length.
 For example:
 | Define the implicit list length attributes.
 For example:
   where
 Renders type parameter arguments.
 Renders mandatory attributes as function parameters.
 Renders sizes of tensor list types having number_attr.
 | Write a comment with the inputs/outputs/attributes in proto format, for
 debugging.
 | The type signature for an op.
 Of the form:
      => {pre} Float -> Tensor t1 v1 -> Tensor t2 v2
 where "Float" is an explicit input attribute, "Tensor t1 v1" is an input, and
 "Tensor t2 v2" is an output.
 Use m' as the type parameter to avoid clashing with an attribute name.
 TODO(judahjacobson): To improve indentation: `tensorArgAndComment a`
 | Render an op input or output.
 For example: "Tensor Ref Int64", "Tensor v t"
 | Comment for the outputs of an op.
 For example:
   -- ^ (__output1__, __output2__)
   --
 * _ _ output1 _ _ : 
   --
 * _ _ output2 _ _ : 
 | Constraints for a given type parameter.
 Produces a type-level list, e.g.: '[Int32,Int64,Float]
 NOTE: The cases of this function should be kept in sync with
 | haddockComment escapes TensorFlow doc strings into haddock.
 | Generate a multiline comment.  For example:
   summary'
   --
   -- detail_line1
   -- detail_line2
   -- ...
 | Converts the given multi-line detail string into
 a multi-line haddock. Applies the given lead to the
  Copyright 2016 TensorFlow authors .
  Licensed under the Apache License , Version 2.0 ( the " License " ) ;
  distributed under the License is distributed on an " AS IS " BASIS ,

# LANGUAGE CPP #
# LANGUAGE FlexibleContexts #
# LANGUAGE LambdaCase #
# LANGUAGE TypeFamilies #
  | Rendering of TensorFlow operations as  functions . 

 The basic type signature generated for each op is : 

 > { constraints } = > { mandatory attrs } - > { input tensors } - > { output tensors } 

 where : 

 * @{mandatory attrs}@ is of the form @A_1 - > ... - >  , where each @A@ is an 
  op attribute that does n't have a default and ca n't be inferred from other 
  inputs . 

 * @{constraints}@ restrict the type parameters of the input and output tensors 
  ( for example : ' TensorType ' or '  ' ) . 

 * @{input tensors}@ is of the form @T_1 - > ... - > T_N@ , where each @T@ is of 
 the form @Tensor Ref a@ or  v a@ ( or a list of one of those types ) , 
 and @a@ is either a concrete type or a ( constrained ) type variable . 

 * @{output tensors}@ is of the form @(T_1, ... ,T_N)@ for " pure " ops , and 
  ( T_1, ... ,T_N)@ for " stateful " ops .   An op is considered " stateful " if 
 it takes a @Tensor Ref@ or @Tensor v ResourceHandle@ as input , or if it 's 
 explicitly marked \"Stateful\ " in its @REGISTER_OP@ definition .   ( If there 
 are no outputs , it is either @ControlNode@ or @Build ControlNode@. ) 


The basic type signature generated for each op is:

> {constraints} => {mandatory attrs} -> {input tensors} -> {output tensors}

where:

* @{mandatory attrs}@ is of the form @A_1 -> ... -> A_N@, where each @A@ is an
 op attribute that doesn't have a default and can't be inferred from other
 inputs.

* @{constraints}@ restrict the type parameters of the input and output tensors
 (for example: 'TensorType' or 'OneOf').

* @{input tensors}@ is of the form @T_1 -> ... -> T_N@, where each @T@ is of
the form @Tensor Ref a@ or @Tensor v a@ (or a list of one of those types),
and @a@ is either a concrete type or a (constrained) type variable.

* @{output tensors}@ is of the form @(T_1,...,T_N)@ for "pure" ops, and
@Build (T_1,...,T_N)@ for "stateful" ops.  An op is considered "stateful" if
it takes a @Tensor Ref@ or @Tensor v ResourceHandle@ as input, or if it's
explicitly marked \"Stateful\" in its @REGISTER_OP@ definition.  (If there
are no outputs, it is either @ControlNode@ or @Build ControlNode@.)
-}

module TensorFlow.OpGen
  ( OpGenFlags(..)
  , docOpList
  , flagParser)
  where

import Data.Foldable (toList)
import Data.Maybe (fromMaybe)
import Data.ProtoLens.Default(def)
import Data.ProtoLens (showMessage)
import Data.List (sortOn)
import Data.List.NonEmpty (NonEmpty)
import qualified Data.List.NonEmpty as NE
import Lens.Family2 ((^.), (.~), (&), view)
import Options.Applicative (Parser, help, long, strOption, value)
import Proto.Tensorflow.Core.Framework.OpDef
  ( OpList
  , OpDef
  )
import Proto.Tensorflow.Core.Framework.OpDef_Fields
  ( attr
  , inputArg
  , name
  , op
  , outputArg
  )
import Proto.Tensorflow.Core.Framework.Types (DataType(..))
import System.FilePath (takeBaseName)
import TensorFlow.OpGen.ParsedOp
import Text.PrettyPrint.Mainland
  ( Doc
  , (<+>)
  , (</>)
  , (<+/>)
  , brackets
  , comma
  , commasep
  , dquotes
  , empty
  , enclose
  , flatten
  , folddoc
  , hang
  , indent
  , parens
  , sep
  , stack
  , strictText
  , tuple
  )
import qualified Data.Set as Set
import qualified Data.Text as Text

data OpGenFlags = OpGenFlags
     { outputFile :: String
     , prefix :: String
     , excludeList :: String
     }

flagParser :: Parser OpGenFlags
flagParser = OpGenFlags
     <$> strOption (mconcat [ long "output"
                            , help "File to write."
                            ])
     <*> strOption (mconcat [ long "prefix"
                            , help "Haskell package prefix to use"
                            ])
     <*> strOption (mconcat [ long "exclude_list"
                            , value ""
                            , help "Comma separated Ops names to ignore"
                            ])


docOpList :: OpGenFlags -> OpList -> Doc
docOpList flags opList =
  stack [ "{-# LANGUAGE ConstraintKinds #-}"
        , "{-# LANGUAGE DataKinds #-}"
        , "{-# LANGUAGE FlexibleContexts #-}"
        , "{-# LANGUAGE FlexibleInstances #-}"
        , "{-# LANGUAGE OverloadedStrings #-}"
        , "{-# LANGUAGE ScopedTypeVariables #-}"
        , "{-# OPTIONS_GHC -fno-warn-name-shadowing #-}"
        , "{-# OPTIONS_GHC -fno-warn-incomplete-patterns #-}"
        , "module" <+> strictText moduleName <+> "where"
        , empty
        , imports
        , empty
        , folddoc (\x y -> x </> empty </> y)
                  (map renderOpAndExtras $
                   sortOn (view name) $
                   filter (not . flip elem exclusions . view name) $
                   toList $ opList ^. op)
        ]
  where moduleName =
            Text.pack (prefix flags) <> "." <> camelCase
            (fromMaybe shortName (Text.stripSuffix "_ops_op_lib" shortName))
        shortName = Text.pack (takeBaseName $ outputFile flags)
        exclusions = Text.splitOn "," $ Text.pack $ excludeList flags
        renderOpAndExtras o = renderOp (parseOp o) </> extras o

imports :: Doc
imports = stack [
      "import Data.ByteString (ByteString)"
    , "import Data.Complex (Complex)"
    , "import Data.Int (Int8, Int16, Int32, Int64)"
    , "import Data.Proxy (Proxy(Proxy))"
    , "import Data.Word (Word8, Word16, Word32, Word64)"
    , "import Lens.Family2 ((.~), (&))"
    , "import TensorFlow.Build"
    , "import TensorFlow.BuildOp"
    , "import TensorFlow.Tensor"
    , "import TensorFlow.Types"
    ]

renderHaskellName, renderTFName, renderQuotedTFName :: Name -> Doc
renderHaskellName = strictText . unHaskellName . haskellName
renderTFName = strictText . unTFName . tfName
renderQuotedTFName = dquotes . renderTFName


  foo attr1 attr2 input1 input2 | eqLengthGuard [ ... ] = { function body }
renderOp :: ParsedOp -> Doc
renderOp pOp = stack $
    [ haddocks
#if __GLASGOW_HASKELL__ < 800
    , "{-# NOINLINE" <+> n <+> "#-}"
#endif
    , n <+> "::" <+> hang 0 (typeSig empty pOp)
    , n <+> "=" <+> n <> "' id"
    , n' <+> "::" <+> hang 0 (typeSig "OpParams ->" pOp)
    , n' <+> hang 0 args <+> "|" <+> funcGuard listSizeAttrs
                    indent indentation (functionBody pOp)
    ] ++ whereClause listSizeAttrs
  where
    n = renderHaskellName $ parsedOpName pOp
    n' = n <> "'"
    listSizeAttrs = inferredListSizeAttrs pOp
    args = sep $ "op'options"
               : (map renderHaskellName
                    $ map attrName (explicitInputAttrs pOp)
                    ++ map parsedArgName (parsedInputs pOp))
    haddocks = "-- |" <+> multilineComment (parsedOpSummary pOp) (parsedOpDescription pOp)

    eqLengthGuard [ ( " N " , [ ( " input1 " , length input1 ) , ( " input2 " , length input2 ) ] ) ]
funcGuard :: [Attr (NonEmpty Name)] -> Doc
funcGuard attrs = "eqLengthGuard" <+> brackets (commasep entries)
      where
        entries =
            [ parens $ nAttr <> comma <+>
              brackets (commasep $ toList $
                            map renderTensorName (toList $ attrInfo a))
            | a <- attrs
            , let nAttr = renderQuotedTFName (attrName a)
            ]
        renderTensorName x = parens $ renderQuotedTFName x <> comma <+>
                        "length" <+> renderHaskellName x

      n1 = fromIntegral ( length input1 ) : : Int64
      n2 = fromIntegral ( length input2 ) : : Int64
whereClause :: [Attr (NonEmpty Name)] -> [Doc]
whereClause [] = []
whereClause as = [indent 2 $ "where" </> indent 2 (stack $ map defineLengthAttr as)]
  where
    defineLengthAttr a = renderHaskellAttrName a <+> "="
                            <+> "fromIntegral (length"
                            <+> renderHaskellName (NE.head $ attrInfo a)
                            <> ") :: Int64"

renderHaskellAttrName :: Attr a -> Doc
renderHaskellAttrName = renderHaskellName . attrName

functionBody :: ParsedOp -> Doc
functionBody pOp
    | parsedOpIsMonadic pOp
        = "build $ do"
            </> indent indentation (bindOpInputsVar
                        </> "buildOp" <+> outputListsSizes <+> opDef)
    | otherwise
        = "pureOp" <+> outputListsSizes <+> "$ do"
            </> indent indentation (bindOpInputsVar </> "return" <+> opDef)
  where
    outputListsSizes = brackets $ commasep
        [ renderHaskellName a
        | ParsedArg { parsedArgCase = ListArg { argLength = a } }
            <- parsedOutputs pOp
        ]
    opInputsVar = "op'inputs"
    bindOpInputsVar = opInputsVar <+> "<- fmap Prelude.concat $ Prelude.sequence"
                            <+> brackets (commasep $ map (\a -> "buildInputs" <+> a) tensorArgs)
    opDef = parens $ hang 0 $ stack $
        "opDef" <+> renderQuotedTFName (parsedOpName pOp) :
        [ "& opAttr" <+> renderQuotedTFName n <+> ".~" <+> inferredTypeExpr a
        | a <- inferredTypeAttrs pOp, let n = attrName a
        ] ++
        [ "& opAttr" <+> renderQuotedTFName n <+> ".~" <+> renderHaskellName n
        | a <- explicitInputAttrs pOp, let n = attrName a
        ] ++
        [ "& opAttr" <+> renderQuotedTFName n <+> ".~" <+> renderHaskellName n
        | a <- inferredListSizeAttrs pOp, let n = attrName a
        ] ++
        ["& op'options & opInputs .~" <+> opInputsVar]
    tensorArgs = renderTensorArg <$> parsedInputs pOp
    renderTensorArg = renderHaskellName . parsedArgName
    inferredTypeExpr a
        | typeParamIsList $ attrInfo a
            = "fromTensorTypes (Proxy :: Proxy" <+> renderHaskellAttrName a
                    <> ")"
        | otherwise = "tensorType (undefined ::" <+> renderHaskellAttrName a
                            <> ")"

extras :: OpDef -> Doc
extras d = enclose "{-\n" "\n-}" $
            strictText $ Text.pack $
            showMessage ((def :: OpDef)
                        & inputArg .~ (d ^. inputArg)
                        & outputArg .~ (d ^. outputArg)
                        & attr .~ (d ^. attr))

  forall t1 t2 v1 v2 . ( TensorType t1 , TensorType t2 )
typeSig :: Doc -> ParsedOp -> Doc
typeSig pre pOp = constraints
            <+/> pre </> signatureFold (map attrInput (explicitInputAttrs pOp)
                                ++ map tensorArgAndComment (parsedInputs pOp)
                                ++ [outputs])
  where
    constraints
        | null classConstraints = empty
        | otherwise = "forall" <+> sep typeParams <+> "." <+> tuple classConstraints <+> "=>"
    typeParams = [strictText v | k <- parsedInputs pOp ++ parsedOutputs pOp,
                  ArgSomeTensor v <- [argKind $ parsedArgCase k]]
                ++ [renderHaskellAttrName n | n <- inferredTypeAttrs pOp]
                ++ if parsedOpIsMonadic pOp then ["m'"] else []
    monadConstraint
        | parsedOpIsMonadic pOp = ["MonadBuild m'"]
        | otherwise = []
    classConstraints = monadConstraint ++ map tensorArgConstraint
                                                    (inferredTypeAttrs pOp)
    signatureFold = folddoc (\x y -> x </> "->" <+> y)
    attrInput a = renderAttrType (attrInfo a) <+> hang 0 ("-- ^" <+> attrComment a)
    renderAttrType (AttrSingle a) = renderAttrBaseType a
    renderAttrType (AttrList a) = brackets $ renderAttrBaseType a
    renderAttrBaseType = \case
        AttrBytes -> "ByteString"
        AttrInt64 -> "Data.Int.Int64"
        AttrFloat -> "Float"
        AttrBool -> "Bool"
        AttrType -> "DataType"
        AttrShape -> "Shape"
        AttrTensor -> "TensorProto"
        AttrFunc -> error "AttrFunc not supported"

    tensorArgAndComment t = tensorArg t <+> hang 0 ("-- ^" <+> argComment t)
    outputs = case parsedOutputs pOp of
        [] -> wrapOutput "ControlNode"
        [a] -> wrapOutput (tensorArg a) <+> "-- ^" <+> argComment a
        as -> wrapOutput (tuple (map tensorArg as)) <+/> resultComment as
    wrapOutput o
        | parsedOpIsMonadic pOp = "m'" <+> parens o
        | otherwise = o

tensorArg :: ParsedArg -> Doc
tensorArg p = case parsedArgCase p of
    SimpleArg { argType = t, argKind = k } -> tensorType t k
    ListArg { argType = t, argKind = k } -> brackets $ tensorType t k
    MixedListArg {argTypeAttr = t, argKind = k}
        -> "TensorList" <+> parens (kind k) <+> renderHaskellName t
  where
    kind k = case k of
                ArgTensorRef -> "Ref"
                ArgTensorValue -> "Value"
                ArgTensorBuild -> "Build"
                ArgSomeTensor v -> strictText v
    tensorType t k = let
        a = case t of
                ArgTypeFixed dt -> strictText $ dtTypeToHaskell dt
                ArgTypeAttr n -> renderHaskellName n
        in "Tensor" <+> kind k <+> a

attrComment :: Attr a -> Doc
attrComment a = argComment' (attrName a) (attrDescription a)

argComment :: ParsedArg -> Doc
argComment a = argComment' (parsedArgName a) (parsedArgDescription a)

argComment' :: Name -> Text.Text -> Doc
argComment' argName argDesc =
    bold (renderTFName argName) <> splitMultilineText (":" <+>) argDesc

bold :: Doc -> Doc
bold n = "__" <> n <> "__"

resultComment :: [ParsedArg] -> Doc
resultComment os = stack $ flatten commentSummary : map commentDetails os
  where
    commentSummary = "-- ^" <+> tuple [bold (renderTFName $ parsedArgName o) | o <- os]
    commentDetails o =
        stack [ "--"
              , "-- *" <+> argComment o
              ]

  E.g. : " TensorType t " or " OneOf [ Int64 , Float ] t "
  or "  ts " or " OneOfs [ .. ] ts " .
tensorArgConstraint :: Attr TypeParam -> Doc
tensorArgConstraint a = case attrInfo a of
    TypeParam False Nothing -> "TensorType" <+> n
    TypeParam False (Just as) -> "OneOf" <+> typeList as <+> n
    TypeParam True Nothing -> "TensorTypes" <+> n
    TypeParam True (Just as) -> "OneOfs" <+> typeList as <+> n
  where
    n = renderHaskellAttrName a
    typeList = ("'" <>) . brackets . commasep . map strictText .
                    Set.toList . Set.fromList .
                    map dtTypeToHaskell . toList

  TensorFlow . Types .  .
dtTypeToHaskell :: DataType -> Text.Text
dtTypeToHaskell DT_BOOL = "Bool"
dtTypeToHaskell DT_BFLOAT16 = "Data.Word.Word16"
dtTypeToHaskell DT_COMPLEX128 = "(Data.Complex.Complex Double)"
dtTypeToHaskell DT_COMPLEX64 = "(Data.Complex.Complex Float)"
dtTypeToHaskell DT_DOUBLE = "Double"
dtTypeToHaskell DT_FLOAT = "Float"
dtTypeToHaskell DT_INT16 = "Data.Int.Int16"
dtTypeToHaskell DT_INT32 = "Data.Int.Int32"
dtTypeToHaskell DT_INT64 = "Data.Int.Int64"
dtTypeToHaskell DT_INT8 = "Data.Int.Int8"
  TODO(gnezdo ): make unique
  TODO(gnezdo ): make unique
  TODO(gnezdo ): make unique
  TODO(gnezdo ): make unique
  TODO(gnezdo ): make unique
dtTypeToHaskell DT_STRING = "Data.ByteString.ByteString"
dtTypeToHaskell DT_UINT16 = "Data.Word.Word16"
dtTypeToHaskell DT_UINT32 = "Data.Word.Word32"
dtTypeToHaskell DT_UINT64 = "Data.Word.Word64"
  TODO(gnezdo ): make unique
dtTypeToHaskell DT_UINT8 = "Data.Word.Word8"
dtTypeToHaskell DT_RESOURCE = "ResourceHandle"
dtTypeToHaskell DT_VARIANT = "Variant"
dtTypeToHaskell x =
    Text.pack $ "Unsupported type in dtTypeToHaskell: " ++ show x

  TODO(gnezdo ): deal with the markup .
haddockComment :: Text.Text -> Doc
haddockComment = strictText

multilineComment :: Text.Text -> Text.Text -> Doc
multilineComment summary' detail =
    haddockComment summary' </>
    splitMultilineText insertParagraphAndComment detail
  where insertParagraphAndComment x = "--" </> "--" <+> x

  first line . Returns an empty document for empty detail .
splitMultilineText :: (Doc -> Doc) -> Text.Text -> Doc
splitMultilineText lead detail =
  case Text.lines detail of
    [] -> empty
    (l : ls) -> stack $ lead (haddockComment l)
                      : map (("--" <+>) . haddockComment) ls

indentation :: Int
indentation = 4
