effect Python
    "Simantics/Python/Python"
    "org.simantics.pythonlink.PythonContext"

class PythonCompatible a where
    setPythonVariable :: String -> a -> <Python> ()
    getPythonVariable :: String -> <Python> a

importJava "org.simantics.pythonlink.NDArray" where
    data NDArray
    
    @JavaName "<init>"
    ndarray :: Vector Double -> NDArray
    @JavaName "<init>"
    ndarrayM :: Integer -> Integer -> Vector Double -> NDArray
    @JavaName "<init>"
    ndarrayN :: Vector Integer -> Vector Double -> NDArray
    
    @JavaName size
    ndarraySize :: NDArray -> Integer
    @JavaName dims
    ndarrayDims :: NDArray -> Vector Integer
    @JavaName getValues
    ndarrayValues :: NDArray -> Vector Double
    @JavaName getValue
    ndarrayElement :: NDArray -> Integer -> Double
    @JavaName getValue
    ndarrayElementM :: NDArray -> Integer -> Integer -> Double
    @JavaName getValue
    ndarrayElementN :: NDArray -> Vector Integer -> Double
    
    @JavaName toString
    ndarrayToString :: NDArray -> String
    @JavaName equals
    ndarrayEquals :: NDArray -> NDArray -> Boolean

class NDArrayInput a where
    ndarrayInput :: a -> ([Integer], [Double])
    
instance NDArrayInput [Double] where
    ndarrayInput v = ([length v], v)
    
instance (NDArrayInput [a]) => NDArrayInput [[a]] where
    ndarrayInput v = let
        inp = map ndarrayInput v
        dims = map fst inp
        vals = join $ map snd inp
      in
        if not (all (== (dims!0)) dims) then
            fail "List of lists not of uniform dimension"
          else
            ([length v] + (dims!0), vals)

ndarrayFromList :: NDArrayInput a => a -> NDArray
ndarrayFromList v = ndarrayN (vector dims) (vector vals) where
    (dims, vals) = ndarrayInput v

instance Show NDArray where
    show = ndarrayToString

instance Eq NDArray where
    (==) = ndarrayEquals

importJava "org.simantics.pythonlink.PythonContext" where
    data PythonContext
    
    @JavaName close
    closePythonContext :: PythonContext -> <Proc> ()
    
    executePythonStatement :: String -> <Python> ()
    
    setPythonBooleanVariable :: String -> Boolean -> <Python> ()
    setPythonIntegerVariable :: String -> Integer -> <Python> ()
    setPythonLongVariable :: String -> Long -> <Python> ()
    setPythonDoubleVariable :: String -> Double -> <Python> ()
    setPythonStringVariable :: String -> String -> <Python> ()
    
    setPythonBooleanArrayVariable :: String -> Vector Boolean -> <Python> ()
    setPythonIntegerArrayVariable :: String -> Vector Integer -> <Python> ()
    setPythonLongArrayVariable :: String -> Vector Long -> <Python> ()
    setPythonDoubleArrayVariable :: String -> Vector Double -> <Python> ()
    setPythonStringArrayVariable :: String -> Vector String -> <Python> ()
    setPythonNDArrayVariable :: String -> NDArray -> <Python> ()
    
    getPythonBooleanVariable :: String -> <Python> Boolean
    getPythonIntegerVariable :: String -> <Python> Integer
    getPythonLongVariable :: String -> <Python> Long
    getPythonDoubleVariable :: String -> <Python> Double
    getPythonStringVariable :: String -> <Python> String
    
    getPythonBooleanArrayVariable :: String -> <Python> Vector Boolean
    getPythonIntegerArrayVariable :: String -> <Python> Vector Integer
    getPythonLongArrayVariable :: String -> <Python> Vector Long
    getPythonDoubleArrayVariable :: String -> <Python> Vector Double
    getPythonStringArrayVariable  :: String -> <Python> Vector String
    getPythonNDArrayVariable :: String -> <Python> NDArray
    
    setPythonVariantVariable :: String -> Variant -> <Python> ()
    getPythonVariantVariable :: String -> <Python> Variant

importJava "org.simantics.pythonlink.Python" where
    openPythonContext :: <Proc> PythonContext
    
    runPythonF :: (() -> <Python, e> a) -> <Proc, e> a
    runWithPythonContextF :: PythonContext -> (() -> <Python, e> a) -> <Proc, e> a

runPython :: (<Python, e> a) -> <Proc, e> a
runPython v = runPythonF (\_ -> v)

runWithPythonContext :: PythonContext -> (<Python, e> a) -> <Proc, e> a
runWithPythonContext python v = runWithPythonContextF python (\_ -> v) 

instance (VecComp a, PythonCompatible (Vector a)) => PythonCompatible [a] where
    setPythonVariable n v = setPythonVariable n (vector v)
    getPythonVariable = vectorToList . getPythonVariable
    
instance PythonCompatible Double where
    setPythonVariable = setPythonDoubleVariable
    getPythonVariable = getPythonDoubleVariable
    
instance PythonCompatible (Vector Double) where
    setPythonVariable = setPythonDoubleArrayVariable
    getPythonVariable = getPythonDoubleArrayVariable

instance PythonCompatible Integer where
    setPythonVariable = setPythonIntegerVariable
    getPythonVariable = getPythonIntegerVariable
    
instance PythonCompatible (Vector Integer) where
    setPythonVariable = setPythonIntegerArrayVariable
    getPythonVariable = getPythonIntegerArrayVariable
    
instance PythonCompatible String where
    setPythonVariable = setPythonStringVariable
    getPythonVariable = getPythonStringVariable
    
instance PythonCompatible (Vector String) where
    setPythonVariable = setPythonStringArrayVariable
    getPythonVariable = getPythonStringArrayVariable
    
instance PythonCompatible NDArray where
    setPythonVariable = setPythonNDArrayVariable
    getPythonVariable = getPythonNDArrayVariable