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

import Data.Bifunctor (first, second)
import Data.Functor.Const
import Data.Kind (Type)
import Data.Proxy
import Data.Some
import GHC.TypeLits (Symbol)

import AST
import Data
import Lemmas


data Bindings f env binds where
  BTop :: Bindings f env '[]
  BPush :: Bindings f env binds -> (STy t, f (Append binds env) t) -> Bindings f env (t : binds)
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 binds
               -> (Bindings f env2 binds, Append binds env1 :> Append binds env2)
weakenBindings _ w BTop = (BTop, w)
weakenBindings wf w (BPush b (t, x)) =
  let (b', w') = weakenBindings wf w b
  in (BPush b' (t, wf w' x), WCopy w')

sinkOver :: SList STy ts -> env :> Append ts env
sinkOver SNil = WId
sinkOver (SCons _ ts) = WSink .> sinkOver ts

weakenOver :: SList STy ts -> env :> env' -> Append ts env :> Append ts env'
weakenOver SNil w = w
weakenOver (SCons _ ts) w = WCopy (weakenOver ts w)

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

bconcat :: forall f env binds1 binds2. Bindings f env binds1 -> Bindings f (Append binds1 env) binds2 -> Bindings f env (Append binds2 binds1)
bconcat b1 BTop = b1
bconcat b1 (BPush (b2 :: Bindings f (Append binds1 env) binds2C) (t, x))
  | Refl <- lemAppendAssoc @binds2C @binds1 @env
  = 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')

type family Snoc l x where
  Snoc '[] x = '[x]
  Snoc (y : ys) x = y : Snoc ys x

bsnoc :: (forall env1 env2 t'. env1 :> env2 -> f env1 t' -> f env2 t')
      -> STy t -> f env t -> Bindings f env binds -> (Bindings f env (Snoc binds t), Append binds env :> Append (Snoc binds t) env)
bsnoc _ t x BTop = (BPush BTop (t, x), WSink)
bsnoc wf t x (BPush b (t', y)) =
  let (b', w) = bsnoc wf t x b
  in (BPush b' (t', wf w y), WCopy w)

type family Tape binds where
  Tape '[] = TNil
  Tape (t : ts) = TPair t (Tape ts)

-- TODO: The problem is that in the 3rd field of TupBinds, we should reconstruct a stack of let bindings from the tape, but we can't directly without having quadratic code size (due to nested projections). Instead we should produce _two_ Bindings there: one to an existential intermediate 'tempbinds', and another that picks up from there and creates 'binds'.
data TupBinds f env binds =
  TupBinds (SList STy binds)
           (forall env2. Append binds env :> env2 -> Ex env2 (Tape binds))
           (forall env2. Idx env2 (Tape binds) -> Bindings f env2 binds)

bindingsBinds :: Bindings f env binds -> SList STy binds
bindingsBinds BTop = SNil
bindingsBinds (BPush binds (t, _)) = SCons t (bindingsBinds binds)

tapeTy :: SList STy binds -> STy (Tape binds)
tapeTy SNil = STNil
tapeTy (SCons t ts) = STPair t (tapeTy ts)

bindingsCollect :: Bindings f env binds -> Append binds env :> env2 -> Ex env2 (Tape binds)
bindingsCollect BTop _ = ENil ext
bindingsCollect (BPush binds (t, _)) w =
  EPair ext (EVar ext t (w @> IZ))
            (bindingsCollect binds (w .> WSink))

-- type family TapeUnfoldings binds where
--   TapeUnfoldings '[] = '[]
--   TapeUnfoldings (t : ts) = Snoc (TapeUnfoldings ts) (Tape (t : ts))

-- -- The input Ex must be duplicable.
-- -- This function is quadratic, and returns code whose internal representation is quadratic in size (due to types). It runtime should be linear, however.
-- tapeUnfoldings :: forall binds env. SList STy binds -> Ex env (Tape binds) -> Bindings Ex env (TapeUnfoldings binds)
-- tapeUnfoldings SNil _ = BTop
-- tapeUnfoldings (SCons t ts) e = fst $ bsnoc weakenExpr (tapeTy (SCons t ts)) e (tapeUnfoldings ts (ESnd ext e))

-- reconFromTape :: SList STy binds -> Bindings Ex env (TapeUnfoldings binds) -> Bindings Ex (Append (TapeUnfoldings binds) env) binds
-- reconFromTape SNil BTop = BTop
-- reconFromTape (SCons t ts) (BPush unfbinds (_, e)) = _
-- reconFromTape SCons{} BTop = error "unreachable"

-- TODO: This function produces quadratic code, but it must be linear. Need to fix this!
reconstructBindings :: SList STy binds -> Ex env (Tape binds) -> Bindings Ex env binds
reconstructBindings SNil _ = BTop
reconstructBindings (SCons t ts) e = BPush (reconstructBindings ts (ESnd ext e)) (t, weakenExpr (sinkOver ts) (EFst ext e))

letBinds :: Bindings Ex env binds -> Ex (Append binds 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

-- | Select only the types from the environment that have the specified storage
type family Select env sto s where
  Select '[] '[] _ = '[]
  Select (t : ts) (s : sto) s = t : Select ts sto s
  Select (_ : ts) (_ : sto) s = Select ts sto s

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"

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

conv2Idx :: Descr env sto -> Idx env t -> Either (Idx (D2E (Select env sto "accum")) (D2 t))
                                                 (Idx (Select env sto "merge") t)
conv2Idx (DPush _   (_, SAccum)) IZ = Left IZ
conv2Idx (DPush _   (_, SMerge)) IZ = Right IZ
conv2Idx (DPush des (_, SAccum)) (IS i) = first IS (conv2Idx des i)
conv2Idx (DPush des (_, SMerge)) (IS i) = second IS (conv2Idx des i)
conv2Idx DTop i = case i of {}

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"

plus :: STy t -> Ex env (D2 t) -> Ex env (D2 t) -> Ex env (D2 t)
plus STNil _ _ = ENil ext
plus (STPair t1 t2) a b =
  let t = STPair (d2 t1) (d2 t2)
  in plusSparse t a b $
       EPair ext (plus t1 (EFst ext (EVar ext t (IS IZ)))
                          (EFst ext (EVar ext t IZ)))
                 (plus t2 (ESnd ext (EVar ext t (IS IZ)))
                          (ESnd ext (EVar ext t IZ)))
plus (STEither t1 t2) a b =
  let t = STEither (d2 t1) (d2 t2)
  in plusSparse t a b $
       ECase ext (EVar ext t (IS IZ))
         (ECase ext (EVar ext t (IS IZ))
           (EInl ext (d2 t2) (plus t1 (EVar ext (d2 t1) (IS IZ)) (EVar ext (d2 t1) IZ)))
           (EError t "plus l+r"))
         (ECase ext (EVar ext t (IS IZ))
           (EError t "plus r+l")
           (EInr ext (d2 t1) (plus t2 (EVar ext (d2 t2) (IS IZ)) (EVar ext (d2 t2) IZ))))
plus STArr{} _ _ = error "TODO arrays"
plus (STScal t) a b = case t of
  STI32 -> ENil ext
  STI64 -> ENil ext
  STF32 -> EOp ext (OAdd STF32) (EPair ext a b)
  STF64 -> EOp ext (OAdd STF64) (EPair ext a b)
  STBool -> ENil ext
plus STEVM{} _ _ = error "EVM not allowed in input program"

plusSparse :: STy a
           -> Ex env (TEither TNil a) -> Ex env (TEither TNil a)
           -> Ex (a : a : env) a
           -> Ex env (TEither TNil a)
plusSparse t a b adder =
  ELet ext b $
    ECase ext (weakenExpr WSink a)
      (EVar ext (STEither STNil t) (IS IZ))
      (EInr ext STNil
        (ECase ext (EVar ext (STEither STNil t) (IS IZ))
          (EVar ext t (IS IZ))
          (weakenExpr (WCopy (WCopy WSink)) adder)))

type family Tup env where
  Tup '[] = TNil
  Tup (t : ts) = TPair (Tup ts) t

tTup :: SList STy env -> STy (Tup env)
tTup SNil = STNil
tTup (SCons t ts) = STPair (tTup ts) t

zeroTup :: SList STy env0 -> Ex env (Tup (D2E env0))
zeroTup SNil = ENil ext
zeroTup (SCons t env) = EPair ext (zeroTup env) (zero t)

onehotTup :: SList STy env0 -> Idx env0 t -> Ex env (D2 t) -> Ex env (Tup (D2E env0))
onehotTup (SCons _ env) IZ d = EPair ext (zeroTup env) d
onehotTup (SCons t env) (IS i) d = EPair ext (onehotTup env i d) (zero t)
onehotTup SNil i _ = case i of {}

plusTup :: SList STy env0 -> Ex env (Tup (D2E env0)) -> Ex env (Tup (D2E env0)) -> Ex env (Tup (D2E env0))
plusTup SNil _ _ = ENil ext
plusTup env0@(SCons t env) a b =
  ELet ext a $
  ELet ext (weakenExpr WSink b) $
    EPair ext (plusTup env (EFst ext (EVar ext (tTup (d2e env0)) (IS IZ)))
                           (EFst ext (EVar ext (tTup (d2e env0)) IZ)))
              (plus t (ESnd ext (EVar ext (tTup (d2e env0)) (IS IZ)))
                      (ESnd ext (EVar ext (tTup (d2e env0)) IZ)))

data Subenv env env' where
  SETop :: Subenv '[] '[]
  SEYes :: Subenv env env' -> Subenv (t : env) (t : env')
  SENo  :: Subenv env env' -> Subenv (t : env) env'
deriving instance Show (Subenv env env')

data Ret env0 sto t =
  forall shbinds env0Merge.
    Ret (Bindings Ex (D1E env0) shbinds)  -- shared binds
        (Ex (Append shbinds (D1E env0)) (D1 t))
        (Subenv (Select env0 sto "merge") env0Merge)
        (Ex (D2 t : shbinds) (TEVM (D2E (Select env0 sto "accum")) (Tup (D2E env0Merge))))
deriving instance Show (Ret env0 sto t)

data RetPair env0 sto env shbinds t =
  forall env0Merge.
    RetPair (Ex (Append shbinds env) (D1 t))
            (Subenv (Select env0 sto "merge") env0Merge)
            (Ex (D2 t : shbinds) (TEVM (D2E (Select env0 sto "accum")) (Tup (D2E env0Merge))))
deriving instance Show (RetPair env0 sto env shbinds t)

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

subList :: SList f env -> Subenv env env' -> SList f env'
subList SNil SETop = SNil
subList (SCons x xs) (SEYes sub) = SCons x (subList xs sub)
subList (SCons _ xs) (SENo sub) = subList xs sub

subenvNone :: SList STy env -> Subenv env '[]
subenvNone SNil = SETop
subenvNone (SCons _ env) = SENo (subenvNone env)

subenvOnehot :: SList STy env -> Idx env t -> Subenv env '[t]
subenvOnehot (SCons _ env) IZ = SEYes (subenvNone env)
subenvOnehot (SCons _ env) (IS i) = SENo (subenvOnehot env i)
subenvOnehot SNil i = case i of {}

subenvPlus :: SList STy env
           -> Subenv env env1 -> Subenv env env2
           -> (forall env3. Subenv env env3
                         -> Subenv env3 env1
                         -> Subenv env3 env2
                         -> (Ex exenv (Tup (D2E env1))
                             -> Ex exenv (Tup (D2E env2))
                             -> Ex exenv (Tup (D2E env3)))
                         -> r)
           -> r
subenvPlus SNil SETop SETop k = k SETop SETop SETop (\_ _ -> ENil ext)
subenvPlus (SCons _ env) (SENo sub1) (SENo sub2) k =
  subenvPlus env sub1 sub2 $ \sub3 s31 s32 pl -> k (SENo sub3) s31 s32 pl
subenvPlus (SCons _ env) (SEYes sub1) (SENo sub2) k =
  subenvPlus env sub1 sub2 $ \sub3 s31 s32 pl ->
    k (SEYes sub3) (SEYes s31) (SENo s32) $ \e1 e2 ->
      ELet ext e1 $
        EPair ext (pl (EFst ext (EVar ext (typeOf e1) IZ))
                      (weakenExpr WSink e2))
                  (ESnd ext (EVar ext (typeOf e1) IZ))
subenvPlus (SCons _ env) (SENo sub1) (SEYes sub2) k =
  subenvPlus env sub1 sub2 $ \sub3 s31 s32 pl ->
    k (SEYes sub3) (SENo s31) (SEYes s32) $ \e1 e2 ->
      ELet ext e2 $
        EPair ext (pl (weakenExpr WSink e1)
                      (EFst ext (EVar ext (typeOf e2) IZ)))
                  (ESnd ext (EVar ext (typeOf e2) IZ))
subenvPlus (SCons t env) (SEYes sub1) (SEYes sub2) k =
  subenvPlus env sub1 sub2 $ \sub3 s31 s32 pl ->
    k (SEYes sub3) (SEYes s31) (SEYes s32) $ \e1 e2 ->
      ELet ext e1 $
      ELet ext (weakenExpr WSink e2) $
        EPair ext (pl (EFst ext (EVar ext (typeOf e1) (IS IZ)))
                      (EFst ext (EVar ext (typeOf e2) IZ)))
                  (plus t
                    (ESnd ext (EVar ext (typeOf e1) (IS IZ)))
                    (ESnd ext (EVar ext (typeOf e2) IZ)))

expandSubenvZeros :: SList STy env0 -> Subenv env0 env0Merge -> Ex env (Tup (D2E env0Merge)) -> Ex env (Tup (D2E env0))
expandSubenvZeros _ SETop _ = ENil ext
expandSubenvZeros (SCons t ts) (SEYes sub) e =
  ELet ext e $
    let var = EVar ext (STPair (tTup (d2e (subList ts sub))) (d2 t)) IZ
    in EPair ext (expandSubenvZeros ts sub (EFst ext var)) (ESnd ext var)
expandSubenvZeros (SCons t ts) (SENo sub) e = EPair ext (expandSubenvZeros ts sub e) (zero t)

unscope :: Descr env0 sto
        -> STy a -> Storage s
        -> Subenv (Select (a : env0) (s : sto) "merge") envSub
        -> Ex env (TEVM (D2E (Select (a : env0) (s : sto) "accum")) (Tup (D2E envSub)))
        -> (forall envSub'.
               Subenv (Select env0 sto "merge") envSub'
            -> Ex env (TEVM (D2E (Select env0 sto "accum")) (TPair (Tup (D2E envSub')) (D2 a)))
            -> r)
        -> r
unscope des ty s sub e k = case s of
  SAccum -> k sub (EMScope e)
  SMerge -> case sub of
    SEYes sub' -> k sub' e
    SENo sub' -> k sub' $
                   EMBind e $
                     EMReturn (d2e (select SAccum des)) $
                       EPair ext (EVar ext (tTup (d2e (subList (select SMerge des) sub'))) IZ)
                                 (zero ty)

-- 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 :: SList STy shbinds -> env :> env' -> RetPair env0 sto env shbinds t -> RetPair env0 sto env' shbinds t
weakenRetPair bindslist w (RetPair e1 sub e2) = RetPair (weakenExpr (weakenOver bindslist w) e1) sub e2

weakenRets :: env :> env' -> Rets env0 sto env list -> Rets env0 sto env' list
weakenRets w (Rets binds list) =
  let (binds', _) = weakenBindings weakenExpr w binds
  in Rets binds' (slistMap (weakenRetPair (bindingsBinds binds) w) list)

rebaseRetPair :: forall env b1 b2 env0 sto t f. SList f b1 -> SList f b2 -> RetPair env0 sto (Append b1 env) b2 t -> RetPair env0 sto env (Append b2 b1) t
rebaseRetPair b1 b2 (RetPair p sub d)
  | Refl <- lemAppendAssoc @b2 @b1 @env =
      RetPair p sub (weakenExpr (WCopy (wRaiseAbove b2 (slistMap (\_ -> Const ()) b1))) d)

retConcat :: forall env0 sto list. SList (Ret env0 sto) list -> Rets env0 sto (D1E env0) list
retConcat SNil = Rets BTop SNil
retConcat (SCons (Ret (b :: Bindings Ex (D1E env0) shbinds) p sub d) list)
  | Rets binds1 pairs1 <- retConcat list
  , Rets (binds :: Bindings Ex (Append shbinds (D1E env0)) shbinds2) pairs <- weakenRets (sinkWithBindings b) (Rets binds1 pairs1)
  , Refl <- lemAppendAssoc @shbinds2 @shbinds @(D1E env0)
  = Rets (bconcat b binds)
         (SCons (RetPair (weakenExpr (sinkWithBindings binds) p)
                         sub
                         (weakenExpr (WCopy (sinkWithBindings binds)) d))
                (slistMap (rebaseRetPair (bindingsBinds b) (bindingsBinds binds)) pairs))
-- list ~ a : list'
-- SCons (Ret b p sub d) list :: SList (Ret env0 sto) list
-- Ret b p sub d :: Ret env0 sto a   <- existential shbinds
-- b :: Bindings Ex (D1E env0) shbinds
-- p :: Ex (Append shbinds (D1E env0)) (D1 a)
-- d :: Ex (D2 a : shbinds) (TEVM ...)
--
-- list :: SList (Ret env0 sto) list'
-- retConcat list :: Rets env0 sto (D1E env0) list'   <- existential shbinds1
-- binds1 :: Bindings Ex (D1E env0) shbinds1
-- pairs1 :: SList (RetPair env0 sto (D1E env0) shbinds1) list'
--
-- sinkWithBindings b :: forall e. e :> Append shbinds e
-- Rets binds pairs :: Rets env0 sto (Append shbinds (D1E env0)) list'   <- existential shbinds2
-- binds :: Bindings Ex (Append shbinds (D1E env0)) shbinds2
-- pairs :: SList (RetPair env0 sto (Append shbinds (D1E env0)) shbinds2) list'
--
-- we choose shbindsR ~ Append shbinds2 shbinds
-- result :: Rets env0 sto (D1E env0) list
-- result.1 :: Bindings Ex (D1E env0) shbindsR  ==  Bindings Ex (D1E env0) (Append shbinds2 shbinds)
-- result.2 :: SList (RetPair env0 sto (D1E env0) shbindsR) list
-- result.2.head :: RetPair env0 sto (D1E env0) shbindsR a
-- result.2.tail :: SList (RetPair env0 sto (D1E env0) shbindsR) list'
--                = SList (RetPair env0 sto (D1E env0) (Append shbinds2 shbinds)) list'
--
-- wanted: shbinds1 :> shbindsR

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
d1op OIf e = EOp ext OIf e

-- | Both primal and dual must be duplicable expressions
data D2Op a t = Linear (forall env. Ex env (D2 t) -> Ex env (D2 a))
              | Nonlinear (forall env. Ex env (D1 a) -> Ex env (D2 t) -> Ex env (D2 a))

d2op :: SOp a t -> D2Op a t
d2op op = case op of
  OAdd _ -> Linear $ \d -> EInr ext STNil (EPair ext d d)
  OMul t -> d2opBinArrangeInt t $ Nonlinear $ \e d ->
    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 $ Linear $ \d -> EOp ext (ONeg t) d
  OLt t -> Linear $ \_ -> EInl ext (STPair (d2 (STScal t)) (d2 (STScal t))) (ENil ext)
  OLe t -> Linear $ \_ -> EInl ext (STPair (d2 (STScal t)) (d2 (STScal t))) (ENil ext)
  OEq t -> Linear $ \_ -> EInl ext (STPair (d2 (STScal t)) (d2 (STScal t))) (ENil ext)
  ONot -> Linear $ \_ -> ENil ext
  OIf -> Linear $ \_ -> ENil ext
  where
    d2opUnArrangeInt :: SScalTy a
                     -> (D2s a ~ TScal a => D2Op (TScal a) t)
                     -> D2Op (TScal a) t
    d2opUnArrangeInt ty float = case ty of
      STI32 -> Linear $ \_ -> ENil ext
      STI64 -> Linear $ \_ -> ENil ext
      STF32 -> float
      STF64 -> float
      STBool -> Linear $ \_ -> ENil ext

    d2opBinArrangeInt :: SScalTy a
                      -> (D2s a ~ TScal a => D2Op (TPair (TScal a) (TScal a)) t)
                      -> D2Op (TPair (TScal a) (TScal a)) t
    d2opBinArrangeInt ty float = case ty of
      STI32 -> Linear $ \_ -> EInl ext (STPair STNil STNil) (ENil ext)
      STI64 -> Linear $ \_ -> EInl ext (STPair STNil STNil) (ENil ext)
      STF32 -> float
      STF64 -> float
      STBool -> Linear $ \_ -> EInl ext (STPair STNil STNil) (ENil ext)

type Storage :: Symbol -> Type
data Storage s where
  SAccum :: Storage "accum"  -- ^ in the monad state as a mutable accumulator
  SMerge :: Storage "merge"  -- ^ just return and merge
deriving instance Show (Storage s)

-- | Environment description
data Descr env sto where
  DTop :: Descr '[] '[]
  DPush :: Descr env sto -> (STy t, Storage s) -> Descr (t : env) (s : sto)
deriving instance Show (Descr env sto)

select :: Storage s -> Descr env sto -> SList STy (Select env sto s)
select _ DTop = SNil
select s@SAccum (DPush des (t, SAccum)) = SCons t (select s des)
select s@SMerge (DPush des (_, SAccum)) = select s des
select s@SAccum (DPush des (_, SMerge)) = select s des
select s@SMerge (DPush des (t, SMerge)) = SCons t (select s des)

sD1eEnv :: Descr env sto -> SList (Const ()) (D1E env)
sD1eEnv DTop = SNil
sD1eEnv (DPush d _) = SCons (Const ()) (sD1eEnv d)

freezeRet :: Descr env sto
          -> Ret env sto t
          -> (forall env'. Ex env' (D2 t))  -- the incoming cotangent value
          -> Ex (D1E env) (TPair (D1 t) (TEVM (D2E (Select env sto "accum")) (Tup (D2E (Select env sto "merge")))))
freezeRet descr (Ret e0 e1 sub e2) d =
  let e2' = weakenExpr (WCopy (wRaiseAbove (bindingsBinds e0) (sD1eEnv descr))) e2
  in letBinds e0 $
       EPair ext
         e1
         (ELet ext d
           (EMBind e2'
             (EMReturn (d2e (select SAccum descr))
               (expandSubenvZeros (select SMerge descr) sub (EVar ext (tTup (d2e (subList (select SMerge descr) sub))) IZ)))))

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

{-
drev :: forall env sto t.
        Descr env sto
     -> (forall env' sto' t'. Descr env' sto' -> STy t' -> Some Storage)
     -> Ex env t -> Ret env sto t
drev des policy = \case
  EVar _ t i ->
    case conv2Idx des i of
      Left accumI ->
        Ret BTop
            (EVar ext (d1 t) (conv1Idx i))
            (subenvNone (select SMerge des))
            (EMOne d2acc accumI (EVar ext (d2 t) IZ))

      Right tupI ->
        Ret BTop
            (EVar ext (d1 t) (conv1Idx i))
            (subenvOnehot (select SMerge des) tupI)
            (EMReturn d2acc (EPair ext (ENil ext) (EVar ext (d2 t) IZ)))

  ELet _ rhs body
    | Ret rhs0 rhs1 subRHS rhs2 <- drev des policy rhs
    , Some storage <- policy des (typeOf rhs)
    , Ret body0 body1 subBody body2 <- drev (des `DPush` (typeOf rhs, storage)) policy body ->
    weakenBindings weakenExpr (WCopy (sinkWithBindings rhs0)) body0 $ \body0' wbody0' ->
    unscope des (typeOf rhs) storage subBody body2 $ \subBody' body2' ->
    subenvPlus (select SMerge des) subRHS subBody' $ \subBoth _ _ plus_RHS_Body ->
    let bodyResType = STPair (tTup (d2e (subList (select SMerge des) subBody'))) (d2 (typeOf rhs)) in
    Ret (bconcat (rhs0 `BPush` (d1 (typeOf rhs), rhs1)) body0')
        (weakenExpr wbody0' body1)
        subBoth
        (EMBind
           (weakenExpr (WCopy wbody0') body2')
           (EMBind
             (ELet ext (ESnd ext (EVar ext bodyResType IZ)) $
               weakenExpr (WCopy (wSinks' @[_,_] .> WPop (sinkWithBindings body0'))) rhs2)
             (EMReturn d2acc (plus_RHS_Body
                                (EVar ext (tTup (d2e (subList (select SMerge des) subRHS))) IZ)
                                (EFst ext (EVar ext bodyResType (IS IZ)))))))

  EPair _ a b
    | Rets binds (RetPair a1 subA a2 `SCons` RetPair b1 subB b2 `SCons` SNil)
        <- retConcat $ drev des policy a `SCons` drev des policy b `SCons` SNil
    , let dt = STPair (d2 (typeOf a)) (d2 (typeOf b)) ->
    subenvPlus (select SMerge des) subA subB $ \subBoth _ _ plus_A_B ->
    Ret binds
        (EPair ext a1 b1)
        subBoth
        (ECase ext (EVar ext (STEither STNil (STPair (d2 (typeOf a)) (d2 (typeOf b)))) IZ)
           (EMReturn d2acc (zeroTup (subList (select SMerge des) subBoth)))
           (EMBind (ELet ext (EFst ext (EVar ext dt IZ))
                      (weakenExpr (WCopy (wSinks' @[_,_])) a2)) $
            EMBind (ELet ext (ESnd ext (EVar ext dt (IS IZ)))
                      (weakenExpr (WCopy (wSinks' @[_,_,_])) b2)) $
            EMReturn d2acc
              (plus_A_B
                (EVar ext (tTup (d2e (subList (select SMerge des) subA))) (IS IZ))
                (EVar ext (tTup (d2e (subList (select SMerge des) subB))) IZ))))

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

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

  ENil _ -> Ret BTop (ENil ext) (subenvNone (select SMerge des)) (EMReturn d2acc (ENil ext))

  EInl _ t2 e
    | Ret e0 e1 sub e2 <- drev des policy e ->
    Ret e0
        (EInl ext (d1 t2) e1)
        sub
        (ECase ext (EVar ext (STEither STNil (STEither (d2 (typeOf e)) (d2 t2))) IZ)
           (EMReturn d2acc (zeroTup (subList (select SMerge des) sub)))
           (ECase ext (EVar ext (STEither (d2 (typeOf e)) (d2 t2)) IZ)
              (weakenExpr (WCopy (wSinks' @[_,_])) e2)
              (EError (STEVM d2acc (tTup (d2e (subList (select SMerge des) sub)))) "inl<-dinr")))

  EInr _ t1 e
    | Ret e0 e1 sub e2 <- drev des policy e ->
    Ret e0
        (EInr ext (d1 t1) e1)
        sub
        (ECase ext (EVar ext (STEither STNil (STEither (d2 t1) (d2 (typeOf e)))) IZ)
           (EMReturn d2acc (zeroTup (subList (select SMerge des) sub)))
           (ECase ext (EVar ext (STEither (d2 t1) (d2 (typeOf e))) IZ)
              (EError (STEVM d2acc (tTup (d2e (subList (select SMerge des) sub)))) "inr<-dinl")
              (weakenExpr (WCopy (wSinks' @[_,_])) e2)))

  ECase _ e a b
    | STEither t1 t2 <- typeOf e
    , Ret e0 e1 subE e2 <- drev des policy e
    , Some storageA <- policy des t1
    , Some storageB <- policy des t2
    , Ret a0 a1 subA a2 <- drev (des `DPush` (t1, storageA)) policy a
    , Ret b0 b1 subB b2 <- drev (des `DPush` (t2, storageB)) policy 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' ->
    unscope des t1 storageA subA a2 $ \subA' a2' ->
    unscope des t2 storageB subB b2 $ \subB' b2' ->
    subenvPlus (select SMerge des) subA' subB' $ \subAB sAB_A sAB_B _ ->
    subenvPlus (select SMerge des) subAB subE $ \subOut _ _ plus_AB_E ->
    let tCaseRet = STPair (tTup (d2e (subList (select SMerge des) subAB))) (STEither (d2 t1) (d2 t2)) in
    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))
        subOut
        (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) a2')
                                 (EMReturn d2acc
                                    (EPair ext
                                      (expandSubenvZeros (subList (select SMerge des) subAB) sAB_A $
                                         EFst ext (EVar ext (STPair (tTup (d2e (subList (select SMerge des) subA'))) (d2 t1)) IZ))
                                      (EInl ext (d2 t2)
                                        (ESnd ext (EVar ext (STPair (tTup (d2e (subList (select SMerge des) subA'))) (d2 t1)) IZ))))))
                 (EError (STEVM d2acc tCaseRet) "dcase l/rtape"))
              (ECase ext (ESnd ext (EVar ext tPrimal (IS (IS IZ))))
                 (EError (STEVM d2acc tCaseRet) "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) b2')
                                 (EMReturn d2acc
                                    (EPair ext
                                       (expandSubenvZeros (subList (select SMerge des) subAB) sAB_B $
                                          EFst ext (EVar ext (STPair (tTup (d2e (subList (select SMerge des) subB'))) (d2 t2)) IZ))
                                       (EInr ext (d2 t1)
                                         (ESnd ext (EVar ext (STPair (tTup (d2e (subList (select SMerge des) subB'))) (d2 t2)) IZ))))))))
           (EMBind (ELet ext (EInr ext STNil (ESnd ext (EVar ext tCaseRet IZ))) $
                      weakenExpr (WCopy (wSinks' @[_,_,_,_])) e2) $
              EMReturn d2acc $
                plus_AB_E
                  (EFst ext (EVar ext tCaseRet (IS IZ)))
                  (EVar ext (tTup (d2e (subList (select SMerge des) subE))) IZ)))

  EConst _ t val ->
    Ret BTop
        (EConst ext t val)
        (subenvNone (select SMerge des))
        (EMReturn d2acc (ENil ext))

  EOp _ op e
    | Ret e0 e1 sub e2 <- drev des policy e ->
    case d2op op of
      Linear d2opfun ->
        Ret e0
            (d1op op e1)
            sub
            (ELet ext (d2opfun (EVar ext (d2 (opt2 op)) IZ))
               (weakenExpr (WCopy WSink) e2))
      Nonlinear d2opfun ->
        Ret (e0 `BPush` (d1 (typeOf e), e1))
            (d1op op $ EVar ext (d1 (typeOf e)) IZ)
            sub
            (ELet ext (d2opfun (EVar ext (d1 (typeOf e)) (IS IZ))
                               (EVar ext (d2 (opt2 op)) IZ))
               (weakenExpr (WCopy (wSinks' @[_,_])) e2))

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

  where
    d2acc = d2e (select SAccum des)
-}