content: -- Copyright 2023 Google LLC
  --
  -- Use of this source code is governed by a BSD-style
  -- license that can be found in the LICENSE file or at
  -- https://developers.google.com/open-source/licenses/bsd
  
  {-# LANGUAGE DuplicateRecordFields #-}
  
  module JAX.Concrete where
  
  import GHC.Generics (Generic (..))
  
  import Control.Category hiding ((.))
  import Prelude hiding (id)
  import Data.Aeson
  import Data.Aeson.Types qualified as A
  
  import IRVariants
  import Name
  
  data Primitive =
      Sin | Add
    | Scan ScanParams
    | ConvertElementType ConvertElementTypeParams
    -- others!
    deriving (Generic)
  
  data ScanParams = ScanParams
    { reverse :: Bool
    , length :: Int
    , jaxpr :: (ClosedJaxpr VoidS) -- The scan body
    , num_consts :: Int
    , num_carry :: Int
    , linear :: [Bool] -- which arguments are we linear wrt?
    , unroll :: Int -- how many copies of the body to replicate
    }
    deriving (Generic)
  
  data ConvertElementTypeParams = ConvertElementTypeParams
    { new_dtype :: DType
    , weak_type :: Bool
    }
    deriving (Generic)
  
  data DType = F64 | F32 | I64 | I32 -- others
    deriving (Generic)
  
  data DimSizeDeBrujin = LitDimSize Int
    | RefDimSizeInput Int  -- "de Brujin" index but counted from the left of the list
    | RefDimSizeOutput Int -- same
    deriving (Generic)
  
  data JArgType = JArrayDeBrujin DType [DimSizeDeBrujin]
  -- On the output, can refer to input binders and preceding things in the output
    deriving (Generic)
  
  data JFuncType = JFunc [JArgType] JEffects [JArgType]
  -- The ints just count from the beginning of the list of inputs of the jaxpr
  -- being typed by this JEffects datum
    deriving (Generic)
  
  -- TODO Variable references, polynomials, etc
  data DimSizeName = DimSize Int
    deriving (Generic)
  
  data JVarType = JArrayName
    { shape :: [DimSizeName]
    , dtype :: DType
    }
    deriving (Generic)
  
  data JEffects = IO | Read Int | Write Int | Append Int
    deriving (Generic)
  
  data JAtom (n::S) =
      JVariable (JVar n)
    | JLiteral JLit
             -- DropVar?
    deriving (Generic)
  
  data JSourceName = JSourceName
    { name :: Int
    , count :: Int
    , suffix :: String
    }
    deriving (Eq, Ord, Show)
  
  instance HasNameHint JSourceName where
    getNameHint JSourceName{suffix=suffix} = getNameHint suffix
  
  data JSourceNameOr (a::E) (n::S) where
    SourceName :: JSourceName -> JSourceNameOr a 'VoidS
    InternalName :: JSourceName -> a n -> JSourceNameOr a n
  
  data JVar (n::S) = JVar
    { sourceName :: JSourceNameOr (Name (AtomNameC SimpIR)) n
    , ty :: JVarType
    }
    deriving (Generic)
  
  data JLit = JLit
    { val :: A.Value -- TODO What's the real type?
    , ty :: JVarType
    }
    deriving (Generic)
  
  data JBinder (n::S) (l::S) where
    JBindSource :: JSourceName -> JVarType -> JBinder n n
    JBind :: JSourceName -> JVarType -> NameBinder (AtomNameC SimpIR) n l -> JBinder n l
  
  instance ProvesExt JBinder where
  
  instance BindsNames JBinder where
    toScopeFrag = \case
      JBindSource _ _ -> id
      JBind _ _ b -> toScopeFrag b
  
  instance BindsAtMostOneName JBinder (AtomNameC SimpIR) where
    b @> x = case b of
      JBindSource _ _ -> error "Unexpected source binder after parsing"
      JBind _ _ b' -> singletonSubst b' x
    {-# INLINE (@>) #-}
  
  jBinderVar :: JBinder n l -> JVar 'VoidS
  jBinderVar = \case
    JBindSource src ty -> JVar (SourceName src) ty
    JBind src ty _ -> JVar (SourceName src) ty
  
  data JEqn (n::S) (l::S) = JEqn
    { outvars :: Nest JBinder n l
    , primitive :: Primitive
    , invars :: [JAtom n]
    }
    deriving (Generic)
  
  -- examples of primitive parameters
  --  * broadcast - data describing axes
  --  * conv -- strides
  --  * contraction dimensions etc in matmul-ish things
  --  * scan -- direction, "unroll", num_carry, num_consts
  --  * transpose -- permutation
  
  -- other things
  --  * effects
  --  * dynamic shapes
  
  data Jaxpr (n::S) where
    Jaxpr ::
      Nest JBinder n n' -- invars
      -> Nest JBinder n' l -- constvars
      -> Nest JEqn l o -- eqns
      -> [JAtom o] -- outvars
      -> Jaxpr n
  
  data ClosedJaxpr (n::S) = ClosedJaxpr
    { jaxpr :: Jaxpr n
    , consts :: [A.Value]
    }
    deriving (Generic)
  
  instance ToJSON (JVar n) where
    toJSON JVar {sourceName, ty} = object
      [ "name" .= name
      , "count" .= count
      , "suffix" .= suffix
      , "ty" .= ty
      ] where JSourceName {name, count, suffix} = case sourceName of
                SourceName nm -> nm
                InternalName nm _ -> nm
  
  instance FromJSON (JVar VoidS) where
    parseJSON = \case
      (A.Object obj) -> do
        name <- obj .: "name"
        count <- obj .: "count"
        suffix <- obj .: "suffix"
        ty <- obj .: "ty"
        return JVar { sourceName = SourceName $ JSourceName {..}
                    , ty = ty
                    }
      invalid -> A.prependFailure "parsing var failed, " $
        A.typeMismatch "Object" invalid
  
  instance ToJSON (JEqn n l) where
    toJSON JEqn {..} = object
      [ "primitive" .= name
      , "params" .= params
      , "invars" .= A.toJSON invars
      , "outvars" .= A.toJSON (nestToList' jBinderVar outvars)
      ] where (name, params) = dumpPrimitive primitive
  
  dumpPrimitive :: Primitive -> (String, A.Value)
  dumpPrimitive = \case
    Add -> ("add", object [])
    Sin -> ("sin", object [])
    Scan params -> ("scan", A.toJSON params)
    ConvertElementType params -> ("convert_element_type", A.toJSON params)
  
  instance FromJSON (JEqn 'VoidS 'VoidS) where
    parseJSON = \case
      (A.Object obj) -> do
        invars <- obj .: "invars"
        outvars <- obj .: "outvars"
        primName <- obj .: "primitive"
        prim <- parsePrimitive primName =<< obj .: "params"
        return $ JEqn (varsAsBinders outvars) prim invars
      invalid -> A.prependFailure "parsing eqn failed, " $
        A.typeMismatch "Object" invalid
  
  voidNest :: [b 'VoidS 'VoidS] -> Nest b 'VoidS 'VoidS
  voidNest = foldr Nest Empty
  
  varsAsBinders :: [JVar 'VoidS] -> Nest JBinder 'VoidS 'VoidS
  varsAsBinders = voidNest . map varAsBinder
  
  varAsBinder :: JVar 'VoidS -> JBinder 'VoidS 'VoidS
  varAsBinder JVar{sourceName, ty} = case sourceName of
    SourceName srcName -> JBindSource srcName ty
    InternalName _ _ -> error "Unexpected internal name during parsing"
  
  parsePrimitive :: String -> A.Value -> A.Parser Primitive
  parsePrimitive name params = case name of
    "add" -> return Add
    "sin" -> return Sin
    "scan" -> Scan <$> parseJSON params
    "convert_element_type" -> ConvertElementType <$> parseJSON params
    _ -> fail $ "Unknown primitive " ++ name
  
  instance ToJSON (JAtom n) where
    toJSON (JVariable var) = object ["var" .= var]
    toJSON (JLiteral  lit) = object ["lit" .= lit]
  
  instance FromJSON (JAtom VoidS) where
    parseJSON = \case
      o@(A.Object obj) -> do
        obj .:? "var" >>= \case
          Just var -> JVariable <$> A.parseJSON var
          Nothing -> do
            obj .:? "lit" >>= \case
              Just lit -> JLiteral <$> A.parseJSON lit
              Nothing -> wrong o
      invalid -> wrong invalid
      where
        wrong invalid = A.prependFailure "parsing atom failed, " $
          A.typeMismatch "Object with a var or lit key" invalid
  
  instance ToJSON DType where
    toJSON F64 = A.String "f64"
    toJSON F32 = A.String "f32"
    toJSON I64 = A.String "i64"
    toJSON I32 = A.String "i32"
  
  instance FromJSON DType where
    parseJSON (A.String "f64") = pure F64
    parseJSON (A.String "f32") = pure F32
    parseJSON (A.String "i64") = pure I64
    parseJSON (A.String "i32") = pure I32
    parseJSON invalid = A.prependFailure "parsing dtype failed, " $
      A.typeMismatch "dtype string" invalid
  
  instance ToJSON (Jaxpr n) where
    toJSON (Jaxpr invars constvars eqns outvars) = object
      [ "invars" .= nestToList' jBinderVar invars
      , "constvars" .= nestToList' jBinderVar constvars
      , "eqns" .= nestToList' voidify eqns
      , "outvars" .= outvars
      ] where
      voidify :: forall n' l. JEqn n' l -> JEqn VoidS VoidS
      voidify = unsafeCoerceB
  
  instance FromJSON (Jaxpr VoidS) where
    parseJSON = \case
      (A.Object obj) -> do
        invars <- obj .: "invars"
        constvars <- obj .: "constvars"
        eqns <- obj .: "eqns"
        outvars <- obj .: "outvars"
        return $ Jaxpr (varsAsBinders invars) (varsAsBinders constvars)
          (voidNest eqns) outvars
      invalid -> A.prependFailure "parsing jaxpr failed, " $
        A.typeMismatch "Object" invalid
  
  instance ToJSON JLit
  instance ToJSON JVarType
  instance ToJSON JEffects
  instance ToJSON DimSizeName
  instance ToJSON JFuncType
  instance ToJSON Primitive
  instance ToJSON DimSizeDeBrujin
  instance ToJSON JArgType
  instance ToJSON (ClosedJaxpr n)
  instance ToJSON ScanParams
  instance ToJSON ConvertElementTypeParams
  
  instance FromJSON JLit
  instance FromJSON JVarType
  instance FromJSON JEffects
  instance FromJSON DimSizeName
  instance FromJSON JFuncType
  instance FromJSON Primitive
  instance FromJSON DimSizeDeBrujin
  instance FromJSON JArgType
  instance FromJSON (ClosedJaxpr VoidS)
  instance FromJSON ScanParams
  instance FromJSON ConvertElementTypeParams
  