-- | Code generation for sequential Python.
module Futhark.CodeGen.Backends.SequentialPython
  ( compileProg,
  )
where

import Data.Text qualified as T
import Futhark.CodeGen.Backends.GenericPython qualified as Py
import Futhark.CodeGen.Backends.GenericPython.AST
import Futhark.CodeGen.ImpCode.Sequential qualified as Imp
import Futhark.CodeGen.ImpGen.Sequential qualified as ImpGen
import Futhark.IR.SeqMem (Prog, SeqMem)
import Futhark.MonadFreshNames

-- | Compile the program to Python.
compileProg ::
  (MonadFreshNames m) =>
  Py.CompilerMode ->
  String ->
  Prog SeqMem ->
  m (ImpGen.Warnings, T.Text)
compileProg :: forall (m :: * -> *).
MonadFreshNames m =>
CompilerMode -> String -> Prog SeqMem -> m (Warnings, Text)
compileProg CompilerMode
mode String
class_name Prog SeqMem
prog = do
  (ws, defs) <- Prog SeqMem -> m (Warnings, Program)
forall (m :: * -> *).
MonadFreshNames m =>
Prog SeqMem -> m (Warnings, Program)
ImpGen.compileProg Prog SeqMem
prog
  (ws,)
    <$> Py.compileProg
      mode
      class_name
      constructor
      imports
      defines
      operations
      ()
      []
      []
      defs
  where
    imports :: [PyStmt]
imports =
      [ String -> Maybe String -> PyStmt
Import String
"sys" Maybe String
forall a. Maybe a
Nothing,
        String -> Maybe String -> PyStmt
Import String
"numpy" (Maybe String -> PyStmt) -> Maybe String -> PyStmt
forall a b. (a -> b) -> a -> b
$ String -> Maybe String
forall a. a -> Maybe a
Just String
"np",
        String -> Maybe String -> PyStmt
Import String
"ctypes" (Maybe String -> PyStmt) -> Maybe String -> PyStmt
forall a b. (a -> b) -> a -> b
$ String -> Maybe String
forall a. a -> Maybe a
Just String
"ct",
        String -> Maybe String -> PyStmt
Import String
"time" Maybe String
forall a. Maybe a
Nothing
      ]
    constructor :: Constructor
constructor =
      [String] -> [PyStmt] -> Constructor
Py.Constructor
        [String
"self", String
"user_sizes=user_sizes"]
        [PyExp -> PyStmt
Exp (PyExp -> PyStmt) -> PyExp -> PyStmt
forall a b. (a -> b) -> a -> b
$ String -> [PyExp] -> PyExp
Py.simpleCall String
"set_user_params" [String -> PyExp
Var String
"self.sizes", String -> PyExp
Var String
"user_sizes"]]
    defines :: [a]
defines = []
    operations :: Py.Operations Imp.Sequential ()
    operations :: Operations Sequential ()
operations =
      Operations Sequential ()
forall op s. Operations op s
Py.defaultOperations
        { Py.opsCompiler = const $ pure ()
        }