{-# 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 #-}
module CHAD where

import Data.Bifunctor (first, second)
import Data.Kind (Type)
import GHC.TypeLits (Symbol)

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

-- | 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 (D2E (Select env sto "merge")) (D2 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"

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

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

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

data Ret env sto t =
  forall env'.
    Ret (Bindings Ex (D1E env) env')
        (Ex env' (D1 t))
        (Ex (D2 t : env') (TEVM (D2E (Select env sto "accum")) (Tup (D2E (Select env sto "merge")))))
deriving instance Show (Ret env sto t)

data RetPair env0 sto env t =
    RetPair (Ex env (D1 t))
            (Ex (D2 t : env) (TEVM (D2E (Select env0 sto "accum")) (Tup (D2E (Select env0 sto "merge")))))
  deriving (Show)

data Rets env0 sto env list =
  forall env'.
    Rets (Bindings Ex env env')
         (SList (RetPair env0 sto env') list)
deriving instance Show (Rets env0 sto 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 sto env t -> RetPair env0 sto env' t
weakenRetPair w (RetPair e1 e2) = RetPair (weakenExpr w e1) (weakenExpr (WCopy w) e2)

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

retConcat :: forall env sto list. SList (Ret env sto) list -> Rets env sto (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
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)

freezeRet :: 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 (Ret e0 e1 e2) d = letBinds e0 $ EPair ext e1 (ELet ext d e2)

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)

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

drev :: Descr env sto -> Ex env t -> Ret env sto t
drev des = \case
  EVar _ t i ->
    case conv2Idx des i of
      Left accumI ->
        Ret BTop
            (EVar ext (d1 t) (conv1Idx i))
            (EMBind
              (EMOne d2mon accumI (EVar ext (d2 t) IZ))
              (EMReturn d2mon (zeroTup (select SMerge des))))
      Right tupI ->
        _

  ELet _ rhs body
    | Ret rhs0 rhs1 rhs2 <- drev des rhs
    , Ret body0 body1 body2 <- drev (DPush des (typeOf rhs) SMerge) 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 des a `SCons` drev des 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 d2mon (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 des 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 des 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 d2mon (ENil ext))

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

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

  ECase _ e a b
    | STEither t1 t2 <- typeOf e
    , Ret e0 e1 e2 <- drev des e
    , Ret a0 a1 a2 <- drev (DPush des t1 SMerge) a
    , Ret b0 b1 b2 <- drev (DPush des t2 SMerge) 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 d2mon
                                    (EInr ext STNil (EInl ext (d2 t2)
                                       (ESnd ext (EVar ext (STPair STNil (d2 t1)) IZ))))))
                 (EError (STEVM d2mon (STEither STNil (STEither (d2 t1) (d2 t2)))) "dcase l/rtape"))
              (ECase ext (ESnd ext (EVar ext tPrimal (IS (IS IZ))))
                 (EError (STEVM d2mon (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 d2mon
                                    (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 d2mon (ENil ext))

  EOp _ op e
    | Ret e0 e1 e2 <- drev des e ->
    case d2op op of
      Linear d2opfun ->
        Ret e0
            (d1op op e1)
            (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)
            (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
    d2mon = d2e (select SAccum des)