{-# LANGUAGE DataKinds #-}
module Example where

import AST
import AST.Pretty
import CHAD
import Simplify


bin :: SOp (TPair a b) c -> Ex env a -> Ex env b -> Ex env c
bin op a b = EOp ext op (EPair ext a b)

senv1 :: SList STy [TScal TF32, TScal TF32]
senv1 = STScal STF32 `SCons` STScal STF32 `SCons` SNil

-- x y |- x * y + x
--
-- let x3 = (x1, x2)
--     x4 = ((*) x3, x1)
--  in ( (+) x4
--     , let x5 = 1.0
--           x6 = Inr (x5, x5)
--        in case x6 of
--             Inl x7 -> return ()
--             Inr x8 ->
--               let x9 = fst x8
--                   x10 = Inr (snd x3 * x9, fst x3 * x9)
--                in case x10 of
--                     Inl x11 -> return ()
--                     Inr x12 ->
--                       let x13 = fst x12
--                        in one "v1" x13 >>= \x14 ->
--                             let x15 = snd x12
--                              in one "v2" x15 >>= \x16 ->
--                     let x17 = snd x8
--                      in one "v1" x17)
--
-- ( (x1 * x2) + x1
-- , let x5 = 1.0
--   in do one "v1" (x2 * x5)
--         one "v2" (x1 * x5)
--         one "v1" x5)
ex1 :: Ex [TScal TF32, TScal TF32] (TScal TF32)
ex1 =
  bin (OAdd STF32)
    (bin (OMul STF32)
      (EVar ext (STScal STF32) (IS IZ))
      (EVar ext (STScal STF32) IZ))
    (EVar ext (STScal STF32) (IS IZ))

-- x y |- let z = x + y in z * (z + x)
ex2 :: Ex [TScal TF32, TScal TF32] (TScal TF32)
ex2 =
  ELet ext (bin (OAdd STF32) (EVar ext (STScal STF32) (IS IZ))
                             (EVar ext (STScal STF32) IZ)) $
    bin (OMul STF32)
      (EVar ext (STScal STF32) IZ)
      (bin (OAdd STF32)
        (EVar ext (STScal STF32) IZ)
        (EVar ext (STScal STF32) (IS (IS IZ))))