{-# LANGUAGE DataKinds #-}
{-# LANGUAGE GADTs #-}
{-# LANGUAGE LambdaCase #-}
{-# LANGUAGE PolyKinds #-}
{-# LANGUAGE QuantifiedConstraints #-}
{-# LANGUAGE RankNTypes #-}
{-# LANGUAGE StandaloneDeriving #-}
{-# LANGUAGE TypeFamilies #-}
{-# LANGUAGE TypeOperators #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE TypeApplications #-}
module CHAD where

import AST


data Bindings f env env' where
  BTop :: Bindings f env env
  BPush :: Bindings f env env' -> (STy t, f env' t) -> Bindings f env (t : env')
deriving instance (forall e t. Show (f e t)) => Show (Bindings f env env')
infixl `BPush`

weakenBindings :: (forall e1 e2 t. e1 :> e2 -> f e1 t -> f e2 t)
               -> env1 :> env2 -> Bindings f env1 env'
               -> (forall env2'. Bindings f env2 env2' -> env' :> env2' -> r) -> r
weakenBindings _ w BTop k = k BTop w
weakenBindings wf w (BPush b (t, x)) k =
  weakenBindings wf w b $ \b' w' -> k (BPush b' (t, wf w' x)) (WCopy w')

sinkWithBindings :: Bindings f env env' -> env :> env'
sinkWithBindings BTop = WId
sinkWithBindings (BPush b _) = WSink .> sinkWithBindings b

bconcat :: Bindings f env1 env2 -> Bindings f env2 env3 -> Bindings f env1 env3
bconcat b1 BTop = b1
bconcat b1 (BPush b2 (t, x)) = BPush (bconcat b1 b2) (t, x)

bconcat' :: (forall e1 e2 t. e1 :> e2 -> f e1 t -> f e2 t)
         -> Bindings f env env1 -> Bindings f env env2
         -> (forall env12. Bindings f env env12 -> r) -> r
bconcat' wf b1 b2 k = weakenBindings wf (sinkWithBindings b1) b2 $ \b2' _ -> k (bconcat b1 b2')

bsnoc :: STy t -> f env t -> Bindings f (t : env) env' -> Bindings f env env'
bsnoc t x b = bconcat (BTop `BPush` (t, x)) b

data TupBindsReconstruct f env1 env2 env3 =
  forall env4.
    TupBindsReconstruct (Bindings f env3 env4)
                        (env2 :> env4)

data TupBinds f env1 env2 =
  forall tape.
    TupBinds (STy tape)
             (forall env2'. env2 :> env2' -> Ex env2' tape)
             (forall env3. env1 :> env3 -> Idx env3 tape -> TupBindsReconstruct f env1 env2 env3)

tupBinds :: Bindings Ex env1 env2 -> TupBinds Ex env1 env2
tupBinds BTop = TupBinds STNil (\_ -> ENil ext) (\w _ -> TupBindsReconstruct BTop w)
tupBinds (BPush binds (t, _))
  | TupBinds tape collect recon <- tupBinds binds
  = TupBinds (STPair tape t)
             (\w -> EPair ext (collect (w .> WSink))
                              (EVar ext t (w @> IZ)))
             (\w tapeidx ->
               case recon (WSink .> w) IZ of
                 TupBindsReconstruct rebinds wunder ->
                   let rebinds1 = bsnoc tape (EFst ext (EVar ext (STPair tape t) tapeidx)) rebinds
                   in TupBindsReconstruct
                        (rebinds1 `BPush`
                          (t, ESnd ext (EVar ext (STPair tape t)
                                             (sinkWithBindings rebinds1 @> tapeidx))))
                        (WCopy wunder))

letBinds :: Bindings Ex env env' -> Ex env' t -> Ex env t
letBinds BTop = id
letBinds (BPush b (_, rhs)) = letBinds b . ELet ext rhs

type family D1 t where
  D1 TNil = TNil
  D1 (TPair a b) = TPair (D1 a) (D1 b)
  D1 (TEither a b) = TEither (D1 a) (D1 b)
  D1 (TArr n t) = TArr n (D1 t)
  D1 (TScal t) = TScal t

type family D2 t where
  D2 TNil = TNil
  D2 (TPair a b) = TEither TNil (TPair (D2 a) (D2 b))
  D2 (TEither a b) = TEither TNil (TEither (D2 a) (D2 b))
  -- D2 (TArr n t) = _
  D2 (TScal t) = D2s t

type family D2s t where
  D2s TI32 = TNil
  D2s TI64 = TNil
  D2s TF32 = TScal TF32
  D2s TF64 = TScal TF64
  D2s TBool = TNil

type family D1E env where
  D1E '[] = '[]
  D1E (t : env) = D1 t : D1E env

type family D2E env where
  D2E '[] = '[]
  D2E (t : env) = D2 t : D2E env

d1 :: STy t -> STy (D1 t)
d1 STNil = STNil
d1 (STPair a b) = STPair (d1 a) (d1 b)
d1 (STEither a b) = STEither (d1 a) (d1 b)
d1 (STArr n t) = STArr n (d1 t)
d1 (STScal t) = STScal t
d1 STEVM{} = error "EVM not allowed in input program"

d2 :: STy t -> STy (D2 t)
d2 STNil = STNil
d2 (STPair a b) = STEither STNil (STPair (d2 a) (d2 b))
d2 (STEither a b) = STEither STNil (STEither (d2 a) (d2 b))
d2 STArr{} = error "TODO arrays"
d2 (STScal t) = case t of
  STI32 -> STNil
  STI64 -> STNil
  STF32 -> STScal STF32
  STF64 -> STScal STF64
  STBool -> STNil
d2 STEVM{} = error "EVM not allowed in input program"

d2e :: SList STy list -> SList STy (D2E list)
d2e SNil = SNil
d2e (SCons t list) = SCons (d2 t) (d2e list)

d2list :: SList STy env -> SList STy (D2E env)
d2list SNil = SNil
d2list (SCons x l) = SCons (d2 x) (d2list l)

conv1Idx :: Idx env t -> Idx (D1E env) (D1 t)
conv1Idx IZ = IZ
conv1Idx (IS i) = IS (conv1Idx i)

conv2Idx :: Idx env t -> Idx (D2E env) (D2 t)
conv2Idx IZ = IZ
conv2Idx (IS i) = IS (conv2Idx i)

zero :: STy t -> Ex env (D2 t)
zero STNil = ENil ext
zero (STPair t1 t2) = EInl ext (STPair (d2 t1) (d2 t2)) (ENil ext)
zero (STEither t1 t2) = EInl ext (STEither (d2 t1) (d2 t2)) (ENil ext)
zero STArr{} = error "TODO arrays"
zero (STScal t) = case t of
  STI32 -> ENil ext
  STI64 -> ENil ext
  STF32 -> EConst ext STF32 0.0
  STF64 -> EConst ext STF64 0.0
  STBool -> ENil ext
zero STEVM{} = error "EVM not allowed in input program"

data Ret env t =
  forall env'.
    Ret (Bindings Ex (D1E env) env')
        (Ex env' (D1 t))
        (Ex (D2 t : env') (TEVM (D2E env) TNil))
deriving instance Show (Ret env t)

data RetPair env0 env t =
    RetPair (Ex env (D1 t))
            (Ex (D2 t : env) (TEVM (D2E env0) TNil))
  deriving (Show)

data Rets env0 env list =
  forall env'.
    Rets (Bindings Ex env env')
         (SList (RetPair env0 env') list)
deriving instance Show (Rets env0 env list)

-- d1W :: env :> env' -> D1E env :> D1E env'
-- d1W WId = WId
-- d1W WSink = WSink
-- d1W (WCopy w) = WCopy (d1W w)
-- d1W (WPop w) = WPop (d1W w)
-- d1W (WThen u w) = WThen (d1W u) (d1W w)

weakenRetPair :: env :> env' -> RetPair env0 env t -> RetPair env0 env' t
weakenRetPair w (RetPair e1 e2) = RetPair (weakenExpr w e1) (weakenExpr (WCopy w) e2)

weakenRets :: env :> env' -> Rets env0 env list -> Rets env0 env' list
weakenRets w (Rets binds list) =
  weakenBindings weakenExpr w binds $ \binds' wbinds' ->
    Rets binds' (slistMap (weakenRetPair wbinds') list)

retConcat :: forall env list. SList (Ret env) list -> Rets env (D1E env) list
retConcat SNil = Rets BTop SNil
retConcat (SCons (Ret (b :: Bindings Ex (D1E env) env2) p d) list)
  | Rets binds pairs <- weakenRets (sinkWithBindings b) (retConcat list)
  = Rets (bconcat b binds)
         (SCons (RetPair (weakenExpr (sinkWithBindings binds) p)
                         (weakenExpr (WCopy (sinkWithBindings binds)) d))
                pairs)

d1op :: SOp a t -> Ex env (D1 a) -> Ex env (D1 t)
d1op (OAdd t) e = EOp ext (OAdd t) e
d1op (OMul t) e = EOp ext (OMul t) e
d1op (ONeg t) e = EOp ext (ONeg t) e
d1op (OLt t) e = EOp ext (OLt t) e
d1op (OLe t) e = EOp ext (OLe t) e
d1op (OEq t) e = EOp ext (OEq t) e
d1op ONot e = EOp ext ONot e

-- both primal and dual must be duplicable expressions
d2op :: SOp a t -> Ex env (D1 a) -> Ex env (D2 t) -> Ex env (D2 a)
d2op op e d = case op of
  OAdd _ -> EInr ext STNil (EPair ext d d)
  OMul t -> d2opBinArrangeInt t $
    EInr ext STNil (EPair ext (EOp ext (OMul t) (EPair ext (ESnd ext e) d))
                              (EOp ext (OMul t) (EPair ext (EFst ext e) d)))
  ONeg t -> d2opUnArrangeInt t $ EOp ext (ONeg t) d
  OLt t -> EInl ext (STPair (d2 (STScal t)) (d2 (STScal t))) (ENil ext)
  OLe t -> EInl ext (STPair (d2 (STScal t)) (d2 (STScal t))) (ENil ext)
  OEq t -> EInl ext (STPair (d2 (STScal t)) (d2 (STScal t))) (ENil ext)
  ONot -> ENil ext
  where
    d2opUnArrangeInt :: SScalTy a
                     -> (D2s a ~ TScal a => Ex env (D2 (TScal a)))
                     -> Ex env (D2 (TScal a))
    d2opUnArrangeInt ty float = case ty of
      STI32 -> ENil ext
      STI64 -> ENil ext
      STF32 -> float
      STF64 -> float
      STBool -> ENil ext

    d2opBinArrangeInt :: SScalTy a
                      -> (D2s a ~ TScal a => Ex env (D2 (TPair (TScal a) (TScal a))))
                      -> Ex env (D2 (TPair (TScal a) (TScal a)))
    d2opBinArrangeInt ty float = case ty of
      STI32 -> EInl ext (STPair STNil STNil) (ENil ext)
      STI64 -> EInl ext (STPair STNil STNil) (ENil ext)
      STF32 -> float
      STF64 -> float
      STBool -> EInl ext (STPair STNil STNil) (ENil ext)

drev :: SList STy env -> Ex env t -> Ret env t
drev senv = \case
  EVar _ t i ->
    Ret BTop
        (EVar ext (d1 t) (conv1Idx i))
        (EMOne (d2list senv) (conv2Idx i) (EVar ext (d2 t) IZ))

  ELet _ rhs body
    | Ret rhs0 rhs1 rhs2 <- drev senv rhs
    , Ret body0 body1 body2 <- drev (SCons (typeOf rhs) senv) body ->
    weakenBindings weakenExpr (WCopy (sinkWithBindings rhs0)) body0 $ \body0' wbody0' ->
    Ret (bconcat (rhs0 `BPush` (d1 (typeOf rhs), rhs1)) body0')
        (weakenExpr wbody0' body1)
        (EMBind (EMScope (weakenExpr (WCopy wbody0') body2))
                (ELet ext (ESnd ext (EVar ext (STPair STNil (d2 (typeOf rhs))) IZ)) $
                  weakenExpr (WCopy (wSinks @[_,_] .> WPop (sinkWithBindings body0'))) rhs2))

  EPair _ a b
    | Rets binds (RetPair a1 a2 `SCons` RetPair b1 b2 `SCons` SNil)
        <- retConcat $ drev senv a `SCons` drev senv b `SCons` SNil
    , let dt = STPair (d2 (typeOf a)) (d2 (typeOf b)) ->
    Ret binds
        (EPair ext a1 b1)
        (ECase ext (EVar ext (STEither STNil (STPair (d2 (typeOf a)) (d2 (typeOf b)))) IZ)
           (EMReturn (d2e senv) (ENil ext))
           (EMBind (ELet ext (EFst ext (EVar ext dt IZ))
                      (weakenExpr (WCopy (wSinks @[_,_])) a2))
                   (ELet ext (ESnd ext (EVar ext dt (IS IZ)))
                      (weakenExpr (WCopy (wSinks @[_,_,_])) b2))))

  EFst _ e
    | Ret e0 e1 e2 <- drev senv e
    , STPair t1 t2 <- typeOf e ->
    Ret e0
        (EFst ext e1)
        (ELet ext (EInr ext STNil (EPair ext (EVar ext (d2 t1) IZ) (zero t2))) $
           weakenExpr (WCopy WSink) e2)

  ESnd _ e
    | Ret e0 e1 e2 <- drev senv e
    , STPair t1 t2 <- typeOf e ->
    Ret e0
        (ESnd ext e1)
        (ELet ext (EInr ext STNil (EPair ext (zero t1) (EVar ext (d2 t2) IZ))) $
           weakenExpr (WCopy WSink) e2)

  ENil _ -> Ret BTop (ENil ext) (EMReturn (d2e senv) (ENil ext))

  EInl _ t2 e
    | Ret e0 e1 e2 <- drev senv e ->
    Ret e0
        (EInl ext (d1 t2) e1)
        (ECase ext (EVar ext (STEither STNil (STEither (d2 (typeOf e)) (d2 t2))) IZ)
           (EMReturn (d2e senv) (ENil ext))
           (ECase ext (EVar ext (STEither (d2 (typeOf e)) (d2 t2)) IZ)
              (weakenExpr (WCopy (wSinks @[_,_])) e2)
              (EError (STEVM (d2e senv) STNil) "inl<-dinr")))

  EInr _ t1 e
    | Ret e0 e1 e2 <- drev senv e ->
    Ret e0
        (EInr ext (d1 t1) e1)
        (ECase ext (EVar ext (STEither STNil (STEither (d2 t1) (d2 (typeOf e)))) IZ)
           (EMReturn (d2e senv) (ENil ext))
           (ECase ext (EVar ext (STEither (d2 t1) (d2 (typeOf e))) IZ)
              (EError (STEVM (d2e senv) STNil) "inr<-dinl")
              (weakenExpr (WCopy (wSinks @[_,_])) e2)))

  ECase _ e a b
    | STEither t1 t2 <- typeOf e
    , Ret e0 e1 e2 <- drev senv e
    , Ret a0 a1 a2 <- drev (SCons t1 senv) a
    , Ret b0 b1 b2 <- drev (SCons t2 senv) b
    , TupBinds tapeA collectA reconA <- tupBinds a0
    , TupBinds tapeB collectB reconB <- tupBinds b0
    , let tPrimal = STPair (d1 (typeOf a)) (STEither tapeA tapeB) ->
    weakenBindings weakenExpr (WCopy (WSink .> sinkWithBindings e0)) a0 $ \a0' wa0' ->
    weakenBindings weakenExpr (WCopy (WSink .> sinkWithBindings e0)) b0 $ \b0' wb0' ->
    Ret (e0 `BPush`
         (d1 (typeOf e), e1) `BPush`
         (tPrimal,
            ECase ext (EVar ext (d1 (typeOf e)) IZ)
              (letBinds a0' (EPair ext (weakenExpr wa0' a1) (EInl ext tapeB (collectA wa0'))))
              (letBinds b0' (EPair ext (weakenExpr wb0' b1) (EInr ext tapeA (collectB wb0'))))))
        (EFst ext (EVar ext tPrimal IZ))
        (EMBind
           (ECase ext (EVar ext (STEither (d1 t1) (d1 t2)) (IS (IS IZ)))
              (ECase ext (ESnd ext (EVar ext tPrimal (IS (IS IZ))))
                 (case reconA (WSink .> WCopy (wSinks @[_,_,_] .> sinkWithBindings e0)) IZ of
                    TupBindsReconstruct rebinds wrebinds ->
                      letBinds rebinds $
                        ELet ext (EVar ext (d2 (typeOf a)) (sinkWithBindings rebinds @> IS (IS IZ))) $
                          EMBind (weakenExpr (WCopy wrebinds) (EMScope a2))
                                 (EMReturn (d2e senv)
                                    (EInr ext STNil (EInl ext (d2 t2)
                                       (ESnd ext (EVar ext (STPair STNil (d2 t1)) IZ))))))
                 (EError (STEVM (d2e senv) (STEither STNil (STEither (d2 t1) (d2 t2)))) "dcase l/rtape"))
              (ECase ext (ESnd ext (EVar ext tPrimal (IS (IS IZ))))
                 (EError (STEVM (d2e senv) (STEither STNil (STEither (d2 t1) (d2 t2)))) "dcase r/ltape")
                 (case reconB (WSink .> WCopy (wSinks @[_,_,_] .> sinkWithBindings e0)) IZ of
                    TupBindsReconstruct rebinds wrebinds ->
                      letBinds rebinds $
                        ELet ext (EVar ext (d2 (typeOf a)) (sinkWithBindings rebinds @> IS (IS IZ))) $
                          EMBind (weakenExpr (WCopy wrebinds) (EMScope b2))
                                 (EMReturn (d2e senv)
                                    (EInr ext STNil (EInr ext (d2 t1)
                                       (ESnd ext (EVar ext (STPair STNil (d2 t2)) IZ))))))))
           (weakenExpr (WCopy (wSinks @[_,_,_])) e2))

  EConst _ t val ->
    Ret BTop
        (EConst ext t val)
        (EMReturn (d2e senv) (ENil ext))

  EOp _ op e
    | Ret e0 e1 e2 <- drev senv e ->
    Ret (e0 `BPush` (d1 (typeOf e), e1))
        (d1op op $ EVar ext (d1 (typeOf e)) IZ)
        (ELet ext (d2op op (EVar ext (d1 (typeOf e)) (IS IZ))
                           (EVar ext (d2 (opt2 op)) IZ))
           (weakenExpr (WCopy (wSinks @[_,_])) e2))

  e -> error $ "CHAD: unsupported " ++ takeWhile (/= ' ') (show e)