{-# LANGUAGE DataKinds #-}
{-# LANGUAGE DerivingStrategies #-}
{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE GADTs #-}
{-# LANGUAGE GeneralizedNewtypeDeriving #-}
{-# LANGUAGE ImplicitParams #-}
{-# LANGUAGE KindSignatures #-}
{-# LANGUAGE LambdaCase #-}
{-# LANGUAGE MultiWayIf #-}
{-# LANGUAGE PatternSynonyms #-}
{-# LANGUAGE RankNTypes #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE TupleSections #-}
{-# LANGUAGE TypeApplications #-}
{-# LANGUAGE TypeOperators #-}
module Interpreter (
  interpret,
  interpretOpen,
  Value(..),
) where

import Control.Monad (foldM, join, when)
import Data.Bifunctor (bimap)
import Data.Bitraversable (bitraverse)
import Data.Char (isSpace)
import Data.Functor.Identity
import Data.Kind (Type)
import Data.Int (Int64)
import Data.IORef
import System.IO (hPutStrLn, stderr)
import System.IO.Unsafe (unsafePerformIO)

import Debug.Trace

import Array
import AST
import AST.Pretty
import CHAD.Types
import Data
import Interpreter.Rep


newtype AcM s a = AcM { unAcM :: IO a }
  deriving newtype (Functor, Applicative, Monad)

runAcM :: (forall s. AcM s a) -> a
runAcM (AcM m) = unsafePerformIO m

acmDebugLog :: String -> AcM s ()
acmDebugLog s = AcM (hPutStrLn stderr s)

interpret :: Ex '[] t -> Rep t
interpret = interpretOpen False SNil

-- | Bool: whether to trace execution with debug prints (very verbose)
interpretOpen :: Bool -> SList Value env -> Ex env t -> Rep t
interpretOpen prints env e =
  runAcM $
    let ?depth = 0
        ?prints = prints
    in interpret' env e

interpret' :: forall env t s. (?prints :: Bool, ?depth :: Int) => SList Value env -> Ex env t -> AcM s (Rep t)
interpret' env e = do
  let dep = ?depth
  let lenlimit = max 20 (100 - dep)
  let trunc s | length s > lenlimit = take (lenlimit - 3) s ++ "..."
              | otherwise           = s
  when ?prints $ acmDebugLog $ replicate dep ' ' ++ "ev: " ++ trunc (ppExpr env e)
  res <- let ?depth = dep + 1 in interpret'Rec env e
  when ?prints $ acmDebugLog $ replicate dep ' ' ++ "<- " ++ showValue 0 (typeOf e) res ""
  return res

interpret'Rec :: forall env t s. (?prints :: Bool, ?depth :: Int) => SList Value env -> Ex env t -> AcM s (Rep t)
interpret'Rec env = \case
  EVar _ _ i -> case slistIdx env i of Value x -> return x
  ELet _ a b -> do
    x <- interpret' env a
    let ?depth = ?depth - 1 in interpret' (Value x `SCons` env) b
  expr | False && trace ("<i> " ++ takeWhile (not . isSpace) (show expr)) False -> undefined
  EPair _ a b -> (,) <$> interpret' env a <*> interpret' env b
  EFst _ e -> fst <$> interpret' env e
  ESnd _ e -> snd <$> interpret' env e
  ENil _ -> return ()
  EInl _ _ e -> Left <$> interpret' env e
  EInr _ _ e -> Right <$> interpret' env e
  ECase _ e a b -> interpret' env e >>= \case
                     Left x -> interpret' (Value x `SCons` env) a
                     Right y -> interpret' (Value y `SCons` env) b
  ENothing _ _ -> return Nothing
  EJust _ e -> Just <$> interpret' env e
  EMaybe _ a b e -> maybe (interpret' env a) (\x -> interpret' (Value x `SCons` env) b) =<< interpret' env e
  EConstArr _ _ _ v -> return v
  EBuild _ dim a b -> do
    sh <- unTupRepIdx ShNil ShCons dim <$> interpret' env a
    arrayGenerateM sh (\idx -> interpret' (Value (tupRepIdx ixUncons dim idx) `SCons` env) b)
  EFold1Inner _ a b c -> do
    let f = \x y -> interpret' (Value y `SCons` Value x `SCons` env) a
    x0 <- interpret' env b
    arr <- interpret' env c
    let sh `ShCons` n = arrayShape arr
    arrayGenerateM sh $ \idx -> foldM f x0 [arrayIndex arr (idx `IxCons` i) | i <- [0 .. n - 1]]
  ESum1Inner _ e -> do
    arr <- interpret' env e
    let STArr _ (STScal t) = typeOf e
        sh `ShCons` n = arrayShape arr
    numericIsNum t $ return $ arrayGenerate sh $ \idx -> sum [arrayIndex arr (idx `IxCons` i) | i <- [0 .. n - 1]]
  EUnit _ e -> arrayGenerateLinM ShNil (\_ -> interpret' env e)
  EReplicate1Inner _ a b -> do
    n <- fromIntegral @Int64 @Int <$> interpret' env a
    arr <- interpret' env b
    let sh = arrayShape arr
    return $ arrayGenerate (sh `ShCons` n) (\(idx `IxCons` _) -> arrayIndex arr idx)
  EMaximum1Inner _ e -> do
    arr <- interpret' env e
    let STArr _ (STScal t) = typeOf e
        sh `ShCons` n = arrayShape arr
    numericIsNum t $ return $
      arrayGenerate sh (\idx -> maximum [arrayIndex arr (idx `IxCons` i) | i <- [0 .. n-1]])
  EMinimum1Inner _ e -> do
    arr <- interpret' env e
    let STArr _ (STScal t) = typeOf e
        sh `ShCons` n = arrayShape arr
    numericIsNum t $ return $
      arrayGenerate sh (\idx -> minimum [arrayIndex arr (idx `IxCons` i) | i <- [0 .. n-1]])
  EConst _ _ v -> return v
  EIdx0 _ e -> (`arrayIndexLinear` 0) <$> interpret' env e
  EIdx1 _ a b -> arrayIndex1 <$> interpret' env a <*> (fromIntegral @Int64 @Int <$> interpret' env b)
  EIdx _ a b
    | STArr n _ <- typeOf a
    -> arrayIndex <$> interpret' env a <*> (unTupRepIdx IxNil IxCons n <$> interpret' env b)
  EShape _ e | STArr n _ <- typeOf e -> tupRepIdx shUncons n . arrayShape <$> interpret' env e
  EOp _ op e -> interpretOp op <$> interpret' env e
  ECustom _ _ _ _ pr _ _ e1 e2 -> do
    e1' <- interpret' env e1
    e2' <- interpret' env e2
    interpret' (Value e2' `SCons` Value e1' `SCons` SNil) pr
  EWith _ t e1 e2 -> do
    initval <- interpret' env e1
    withAccum t (typeOf e2) initval $ \accum ->
      interpret' (Value accum `SCons` env) e2
  EAccum _ t p e1 e2 e3 -> do
    idx <- interpret' env e1
    val <- interpret' env e2
    accum <- interpret' env e3
    accumAddSparse t p accum idx val
  EZero _ t -> do
    return $ zeroD2 t
  EPlus _ t a b -> do
    a' <- interpret' env a
    b' <- interpret' env b
    return $ addD2s t a' b'
  EOneHot _ t p a b -> do
    a' <- interpret' env a
    b' <- interpret' env b
    return $ onehotD2 p t a' b'
  EError _ _ s -> error $ "Interpreter: Program threw error: " ++ s

interpretOp :: SOp a t -> Rep a -> Rep t
interpretOp op arg = case op of
  OAdd st -> numericIsNum st $ uncurry (+) arg
  OMul st -> numericIsNum st $ uncurry (*) arg
  ONeg st -> numericIsNum st $ negate arg
  OLt st -> numericIsNum st $ uncurry (<) arg
  OLe st -> numericIsNum st $ uncurry (<=) arg
  OEq st -> styIsEq st $ uncurry (==) arg
  ONot -> not arg
  OAnd -> uncurry (&&) arg
  OOr -> uncurry (||) arg
  OIf -> if arg then Left () else Right ()
  ORound64 -> round arg
  OToFl64 -> fromIntegral arg
  ORecip st -> floatingIsFractional st $ recip arg
  OExp st -> floatingIsFractional st $ exp arg
  OLog st -> floatingIsFractional st $ log arg
  OIDiv st -> integralIsIntegral st $ uncurry quot arg
  where
    styIsEq :: SScalTy t -> (Eq (Rep (TScal t)) => r) -> r
    styIsEq STI32 = id
    styIsEq STI64 = id
    styIsEq STF32 = id
    styIsEq STF64 = id
    styIsEq STBool = id

zeroD2 :: STy t -> Rep (D2 t)
zeroD2 typ = case typ of
  STNil -> ()
  STPair _ _ -> Nothing
  STEither _ _ -> Nothing
  STMaybe _ -> Nothing
  STArr SZ t -> arrayUnit (zeroD2 t)
  STArr n _ -> emptyArray n
  STScal sty -> case sty of
                  STI32 -> ()
                  STI64 -> ()
                  STF32 -> 0.0
                  STF64 -> 0.0
                  STBool -> ()
  STAccum{} -> error "Zero of Accum"

addD2s :: STy t -> Rep (D2 t) -> Rep (D2 t) -> Rep (D2 t)
addD2s typ a b = case typ of
  STNil -> ()
  STPair t1 t2 -> case (a, b) of
    (Nothing, _) -> b
    (_, Nothing) -> a
    (Just (x1, x2), Just (y1, y2)) -> Just (addD2s t1 x1 y1, addD2s t2 x2 y2)
  STEither t1 t2 -> case (a, b) of
    (Nothing, _) -> b
    (_, Nothing) -> a
    (Just (Left x), Just (Left y)) -> Just (Left (addD2s t1 x y))
    (Just (Right x), Just (Right y)) -> Just (Right (addD2s t2 x y))
    _ -> error "Plus of inconsistent Eithers"
  STMaybe t -> case (a, b) of
    (Nothing, _) -> b
    (_, Nothing) -> a
    (Just x, Just y) -> Just (addD2s t x y)
  STArr _ t ->
    let sh1 = arrayShape a
        sh2 = arrayShape b
    in if | shapeSize sh1 == 0 -> b
          | shapeSize sh2 == 0 -> a
          | sh1 == sh2 -> arrayGenerateLin sh1 (\i -> addD2s t (arrayIndexLinear a i) (arrayIndexLinear b i))
          | otherwise -> error "Plus of inconsistently shaped arrays"
  STScal sty -> case sty of
    STI32 -> ()
    STI64 -> ()
    STF32 -> a + b
    STF64 -> a + b
    STBool -> ()
  STAccum{} -> error "Plus of Accum"

onehotD2 :: SAcPrj p a b -> STy a -> Rep (AcIdx p a) -> Rep (D2 b) -> Rep (D2 a)
onehotD2 SAPHere _ _ val = val
onehotD2 (SAPFst prj) (STPair a b) idx val = Just (onehotD2 prj a idx val, zeroD2 b)
onehotD2 (SAPSnd prj) (STPair a b) idx val = Just (zeroD2 a, onehotD2 prj b idx val)
onehotD2 (SAPLeft prj) (STEither a _) idx val = Just (Left (onehotD2 prj a idx val))
onehotD2 (SAPRight prj) (STEither _ b) idx val = Just (Right (onehotD2 prj b idx val))
onehotD2 (SAPJust prj) (STMaybe a) idx val = Just (onehotD2 prj a idx val)
onehotD2 (SAPArrIdx prj _) (STArr n a) idx val =
  runIdentity $ onehotArray (\idx' -> Identity (onehotD2 prj a idx' val)) (Identity (zeroD2 a)) n prj idx

withAccum :: STy t -> STy a -> Rep (D2 t) -> (RepAc t -> AcM s (Rep a)) -> AcM s (Rep a, Rep (D2 t))
withAccum t _ initval f = AcM $ do
  accum <- newAcSparse t SAPHere () initval
  out <- case f accum of AcM m -> m
  val <- readAcSparse t accum
  return (out, val)

newAcZero :: STy t -> IO (RepAc t)
newAcZero = \case
  STNil -> return ()
  STPair{} -> newIORef Nothing
  STEither{} -> newIORef Nothing
  STMaybe _ -> newIORef Nothing
  STArr n _ -> newIORef (emptyArray n)
  STScal sty -> case sty of
    STI32 -> return ()
    STI64 -> return ()
    STF32 -> newIORef 0.0
    STF64 -> newIORef 0.0
    STBool -> return ()
  STAccum{} -> error "Nested accumulators"

-- | Inverted index: the outermost index is at the /outside/ of this list.
data PartialInvIndex n m where
  PIIxEnd :: PartialInvIndex m m
  PIIxCons :: Int -> PartialInvIndex n m -> PartialInvIndex (S n) m

-- | Inverted shapey thing: the outermost dimension is at the /outside/ of this list.
data Inverted (f :: Nat -> Type) n where
  InvNil :: Inverted f Z
  InvCons :: Int -> Inverted f n -> Inverted f (S n)

type InvShape = Inverted Shape
type InvIndex = Inverted Index

class Shapey f where
  shapeyNil :: f Z
  shapeyCons :: f n -> Int -> f (S n)
  shapeyCase :: f n -> (n ~ Z => r) -> (forall m. n ~ S m => f m -> Int -> r) -> r
instance Shapey Index where
  shapeyNil = IxNil
  shapeyCons = IxCons
  shapeyCase IxNil k0 _ = k0
  shapeyCase (IxCons idx i) _ k1 = k1 idx i
instance Shapey Shape where
  shapeyNil = ShNil
  shapeyCons = ShCons
  shapeyCase ShNil k0 _ = k0
  shapeyCase (ShCons sh n) _ k1 = k1 sh n

invert :: forall f n. Shapey f => f n -> Inverted f n
invert | Refl <- lemPlusZero @n = flip go InvNil
  where
    go :: forall n' m. f n' -> Inverted f m -> Inverted f (n' + m)
    go sh ish = shapeyCase sh
                  ish
                  (\sh' n -> case lemPlusSuccRight @n' @m of Refl -> go sh' (InvCons n ish))

uninvert :: forall f n. Shapey f => Inverted f n -> f n
uninvert = go shapeyNil
  where
    go :: forall n' m. f n' -> Inverted f m -> f (n' + m)
    go sh InvNil | Refl <- lemPlusZero @n' = sh
    go sh (InvCons n (ish :: Inverted f predm)) | Refl <- lemPlusSuccRight @n' @predm = go (shapeyCons sh n) ish

piindexMatch :: PartialInvIndex n m -> InvIndex n -> Maybe (InvIndex m)
piindexMatch PIIxEnd ix = Just ix
piindexMatch (PIIxCons i pix) (InvCons i' ix)
  | i == i' = piindexMatch pix ix
  | otherwise = Nothing

piindexConcat :: PartialInvIndex n m -> InvIndex m -> InvIndex n
piindexConcat PIIxEnd ix = ix
piindexConcat (PIIxCons i pix) ix = InvCons i (piindexConcat pix ix)

newAcSparse :: STy a -> SAcPrj p a b -> Rep (AcIdx p a) -> Rep (D2 b) -> IO (RepAc a)
newAcSparse typ prj idx val = case (typ, prj) of
  (STNil, SAPHere) -> return ()
  (STPair t1 t2, SAPHere) -> newIORef =<< traverse (bitraverse (newAcSparse t1 SAPHere ()) (newAcSparse t2 SAPHere ())) val
  (STEither t1 t2, SAPHere) -> newIORef =<< traverse (bitraverse (newAcSparse t1 SAPHere ()) (newAcSparse t2 SAPHere ())) val
  (STMaybe t1, SAPHere) -> newIORef =<< traverse (newAcSparse t1 SAPHere ()) val
  (STArr _ t1, SAPHere) -> newIORef =<< traverse (newAcSparse t1 SAPHere ()) val
  (STScal sty, SAPHere) -> case sty of
    STI32 -> return ()
    STI64 -> return ()
    STF32 -> newIORef val
    STF64 -> newIORef val
    STBool -> return ()

  (STPair t1 t2, SAPFst prj') ->
    newIORef . Just =<< (,) <$> newAcSparse t1 prj' idx val <*> newAcZero t2
  (STPair t1 t2, SAPSnd prj') ->
    newIORef . Just =<< (,) <$> newAcZero t1 <*> newAcSparse t2 prj' idx val

  (STEither t1 _, SAPLeft prj') -> newIORef . Just . Left =<< newAcSparse t1 prj' idx val
  (STEither _ t2, SAPRight prj') -> newIORef . Just . Right =<< newAcSparse t2 prj' idx val

  (STMaybe t1, SAPJust prj') -> newIORef . Just =<< newAcSparse t1 prj' idx val

  (STArr n t, SAPArrIdx prj' _) -> newIORef =<< newAcArray n t prj' idx val

  (STAccum{}, _) -> error "Accumulators not allowed in source program"

newAcArray :: SNat n -> STy a -> SAcPrj p a b -> Rep (AcIdx (APArrIdx p) (TArr n a)) -> Rep (D2 b) -> IO (Array n (RepAc a))
newAcArray n t prj idx val = onehotArray (\idx' -> newAcSparse t prj idx' val) (newAcZero t) n prj idx

onehotArray :: Monad m
            => (Rep (AcIdx p a) -> m v)  -- ^ the "one"
            -> m v  -- ^ the "zero"
            -> SNat n -> SAcPrj p a b -> Rep (AcIdx (APArrIdx p) (TArr n a)) -> m (Array n v)
onehotArray mkone mkzero n _ ((arrindex', arrsh'), idx) =
  let arrindex = unTupRepIdx IxNil IxCons n arrindex'
      arrsh = unTupRepIdx ShNil ShCons n arrsh'
  in arrayGenerateM arrsh (\i -> if i == arrindex then mkone idx else mkzero)

-- newAcDense :: STy a -> SAcPrj p a b -> Rep (AcIdx p a) -> Rep (D2 b) -> IO (RepAcDense (D2 a))
-- newAcDense typ SZ () val = case typ of
--   STPair t1 t2 -> (,) <$> newAcSparse t1 SZ () (fst val) <*> newAcSparse t2 SZ () (snd val)
--   STEither t1 t2 -> case val of
--     Left x -> Left <$> newAcSparse t1 SZ () x
--     Right y -> Right <$> newAcSparse t2 SZ () y
--   _ -> error "newAcDense: invalid dense type"
-- newAcDense typ (SS dep) idx val = case typ of
--   STPair t1 t2 ->
--     case (idx, val) of
--       (Left idx', Left val') -> (,) <$> newAcSparse t1 dep idx' val' <*> newAcZero t2
--       (Right idx', Right val') -> (,) <$> newAcZero t1 <*> newAcSparse t2 dep idx' val'
--       _ -> error "Index/value mismatch in newAc pair"
--   STEither t1 t2 ->
--     case (idx, val) of
--       (Left idx', Left val') -> Left <$> newAcSparse t1 dep idx' val'
--       (Right idx', Right val') -> Right <$> newAcSparse t2 dep idx' val'
--       _ -> error "Index/value mismatch in newAc either"
--   _ -> error "newAcDense: invalid dense type"

readAcSparse :: STy t -> RepAc t -> IO (Rep (D2 t))
readAcSparse typ val = case typ of
  STNil -> return ()
  STPair t1 t2 -> traverse (bitraverse (readAcSparse t1) (readAcSparse t2)) =<< readIORef val
  STEither t1 t2 -> traverse (bitraverse (readAcSparse t1) (readAcSparse t2)) =<< readIORef val
  STMaybe t -> traverse (readAcSparse t) =<< readIORef val
  STArr _ t -> traverse (readAcSparse t) =<< readIORef val
  STScal sty -> case sty of
    STI32 -> return ()
    STI64 -> return ()
    STF32 -> readIORef val
    STF64 -> readIORef val
    STBool -> return ()
  STAccum{} -> error "Nested accumulators"

-- readAcDense :: STy t -> RepAcDense t -> IO (Rep t)
-- readAcDense typ val = case typ of
--   STPair t1 t2 -> (,) <$> readAcSparse t1 (fst val) <*> readAcSparse t2 (snd val)
--   STEither t1 t2 -> case val of
--     Left x -> Left <$> readAcSparse t1 x
--     Right y -> Right <$> readAcSparse t2 y
--   _ -> error "readAcDense: invalid dense type"

accumAddSparse :: STy a -> SAcPrj p a b -> RepAc a -> Rep (AcIdx p a) -> Rep (D2 b) -> AcM s ()
accumAddSparse typ prj ref idx val = case (typ, prj) of
  (STNil, SAPHere) -> return ()

  (STPair t1 t2, SAPHere) ->
    case val of
      Nothing -> return ()
      Just (val1, val2) ->
        AcM $ realiseMaybeSparse ref ((,) <$> newAcSparse t1 SAPHere () val1
                                          <*> newAcSparse t2 SAPHere () val2)
                                     (\(ac1, ac2) -> do unAcM $ accumAddSparse t1 SAPHere ac1 () val1
                                                        unAcM $ accumAddSparse t2 SAPHere ac2 () val2)
  (STPair t1 t2, SAPFst prj') ->
    AcM $ realiseMaybeSparse ref ((,) <$> newAcSparse t1 prj' idx val <*> newAcZero t2)
                                 (\(ac1, _) -> do unAcM $ accumAddSparse t1 prj' ac1 idx val)
  (STPair t1 t2, SAPSnd prj') ->
    AcM $ realiseMaybeSparse ref ((,) <$> newAcZero t1 <*> newAcSparse t2 prj' idx val)
                                 (\(_, ac2) -> do unAcM $ accumAddSparse t2 prj' ac2 idx val)

  (STEither t1 t2, SAPHere) -> _ ref val
  (STEither t1 _, SAPLeft prj') -> _ ref idx val
  (STEither _ t2, SAPRight prj') -> _ ref idx val

  (STMaybe t1, SAPHere) -> _ ref val
  (STMaybe t1, SAPJust prj') -> _ ref idx val

  (STArr _ t1, SAPHere) -> _ ref val
  (STArr n t, SAPArrIdx prj' _) -> _ ref idx val

  (STScal sty, SAPHere) -> AcM $ case sty of
    STI32 -> return ()
    STI64 -> return ()
    STF32 -> atomicModifyIORef' ref (\x -> (x + val, ()))
    STF64 -> atomicModifyIORef' ref (\x -> (x + val, ()))
    STBool -> return ()

  (STAccum{}, _) -> error "Accumulators not allowed in source program"

realiseMaybeSparse :: IORef (Maybe a) -> IO a -> (a -> IO ()) -> IO ()
realiseMaybeSparse ref makeval modifyval =
  -- Try modifying what's already in ref. The 'join' makes the snd
  -- of the function's return value a _continuation_ that is run after
  -- the critical section ends.
  join $ atomicModifyIORef' ref $ \ac -> case ac of
           -- Oops, ref's contents was still sparse. Have to initialise
           -- it first, then try again.
           Nothing -> (ac, do val <- makeval
                              join $ atomicModifyIORef' ref $ \ac' -> case ac' of
                                       Nothing -> (Just val, return ())
                                       Just val' -> (ac', modifyval val'))
           -- Yep, ref already had a value in there, so we can just add
           -- val' to it recursively.
           Just val -> (ac, modifyval val)

{-
accumAddSparse typ SZ ref () val = case typ of
  STNil -> return ()
  STPair t1 t2 -> AcM $ do
    (r1, r2) <- readIORef ref
    unAcM $ accumAddSparse t1 SZ r1 () (fst val)
    unAcM $ accumAddSparse t2 SZ r2 () (snd val)
  STMaybe t ->
    case val of
      Nothing -> return ()
      Just val' ->
        -- Try adding val' to what's already in ref. The 'join' makes the snd
        -- of the function's return value a _continuation_ that is run after
        -- the critical section ends.
        AcM $ join $ atomicModifyIORef' ref $ \ac -> case ac of
                       -- Oops, ref's contents was still sparse. Have to initialise
                       -- it first, then try again.
                       Nothing -> (ac, do newac <- newAcDense t SZ () val'
                                          join $ atomicModifyIORef' ref $ \ac2 -> case ac2 of
                                                   Nothing -> (Just newac, return ())
                                                   Just ac2' -> bimap Just unAcM (accumAddDense t SZ ac2' () val'))
                       -- Yep, ref already had a value in there, so we can just add
                       -- val' to it recursively.
                       Just ac' -> bimap Just unAcM (accumAddDense t SZ ac' () val')
  STArr _ t -> AcM $ do
    refs <- readIORef ref
    case (shapeSize (arrayShape refs), shapeSize (arrayShape val)) of
      (_, 0) -> return ()
      (0, _) -> do
        newrefarr <- traverse (newAcSparse t SZ ()) val
        join $ atomicModifyIORef' ref $ \refarr ->
                 if shapeSize (arrayShape refarr) == 0
                   then (newrefarr, return ())
                   else -- someone was faster than us in initialising the reference!
                        (refarr, unAcM $ accumAddSparse typ SZ ref () val)  -- just try again from the start (dropping newrefarr for the GC to clean up)
      _ | arrayShape refs == arrayShape val ->
            sequence_ [unAcM $ accumAddSparse t SZ (arrayIndexLinear refs i) () (arrayIndexLinear val i)
                      | i <- [0 .. shapeSize (arrayShape val) - 1]]
        | otherwise -> error "Array shape mismatch in accum add"
  STScal sty -> AcM $ case sty of
    STI32 -> atomicModifyIORef' ref (\x -> (x + val, ()))
    STI64 -> atomicModifyIORef' ref (\x -> (x + val, ()))
    STF32 -> atomicModifyIORef' ref (\x -> (x + val, ()))
    STF64 -> atomicModifyIORef' ref (\x -> (x + val, ()))
    STBool -> error "Accumulator of Bool"
  STAccum{} -> error "Nested accumulators"
  STEither{} -> error "Bare Either in accumulator"

accumAddSparse typ (SS dep) ref idx val = case typ of
  STNil -> return ()
  STPair t1 t2 -> AcM $ do
    (ref1, ref2) <- readIORef ref
    case (idx, val) of
      (Left idx', Left val') -> unAcM $ accumAddSparse t1 dep ref1 idx' val'
      (Right idx', Right val') -> unAcM $ accumAddSparse t2 dep ref2 idx' val'
      _ -> error "Index/value mismatch in pair accumulator add"
  STMaybe t ->
    AcM $ join $ atomicModifyIORef' ref $ \case
                   -- Oops, ref's contents was still sparse. Have to initialise
                   -- it first, then try again.
                   Nothing -> (Nothing, do newac <- newAcDense t dep idx val
                                           join $ atomicModifyIORef' ref $ \ac2 -> case ac2 of
                                                    Nothing -> (Just newac, return ())
                                                    Just ac2' -> bimap Just unAcM (accumAddDense t dep ac2' idx val))
                   -- Yep, ref already had a value in there, so we can just add
                   -- val' to it recursively.
                   Just ac -> bimap Just unAcM (accumAddDense t dep ac idx val)
  STArr dim (t :: STy t) -> AcM $ do
    refs <- readIORef ref
    if shapeSize (arrayShape refs) == 0
      then do newrefarr <- newAcArray dim t (SS dep) idx val
              join $ atomicModifyIORef' ref $ \refarr ->
                       if shapeSize (arrayShape refarr) == 0
                         then (newrefarr, return ())
                         else -- someone was faster than us in initialising the reference!
                              (refarr, unAcM $ accumAddSparse typ (SS dep) ref idx val)  -- just try again from the start (dropping newrefarr for the GC to clean up)
      else do let sh = unTupRepIdx ShNil ShCons dim (fst val)
              go (SS dep) (invert sh) idx (snd val)
                (\j index idxj valj -> unAcM $ accumAddSparse t j (refs `arrayIndex` index) idxj valj)
                (\piix subsh val' -> unAcM $ sequence_
                                       [accumAddSparse t SZ (refs `arrayIndex` uninvert (piindexConcat piix (invert subix)))
                                                       () (val' `arrayIndex` subix)
                                       | subix <- enumShape subsh])
    where
      go :: SNat i -> InvShape n -> Rep (AcIdx (TArr n t) i) -> Rep (AcValArr n t i)
         -> (forall j. SNat j -> Index n -> Rep (AcIdx t j) -> Rep (AcVal t j) -> r)  -- ^ Indexing into element of the array
         -> (forall m. PartialInvIndex n m -> Shape m -> Rep (TArr m t) -> r)  -- ^ Accumulating onto a subarray
         -> r
      go SZ        ish             ()        val' _  k0 = k0 PIIxEnd (uninvert ish) val'  -- ^ Ran out of AcIdx: accumulating onto subarray
      go (SS dep') InvNil          idx'      val' kj _  = kj dep' IxNil idx' val'  -- ^ Ran out of array dimensions: accumulating into (part of) element
      go (SS dep') (InvCons _ ish) (i, idx') val' kj k0 =
        go dep' ish idx' val'
          (\j index idxj valj -> kj j (IxCons index (fromIntegral @Int64 @Int i)) idxj valj)
          (\pidxm shm valm -> k0 (PIIxCons (fromIntegral @Int64 @Int i) pidxm) shm valm)
  STScal{} -> error "Cannot index into scalar"
  STAccum{} -> error "Nested accumulators"
  STEither{} -> error "Bare Either in accumulator"

accumAddDense :: forall t i s. STy t -> SNat i -> RepAcDense t -> Rep (AcIdx t i) -> Rep (AcVal t i) -> (RepAcDense t, AcM s ())
accumAddDense typ SZ ref () val = case typ of
  STPair t1 t2 ->
    (ref, do accumAddSparse t1 SZ (fst ref) () (fst val)
             accumAddSparse t2 SZ (snd ref) () (snd val))
  STEither t1 t2 ->
    case (ref, val) of
      (Left ref', Left val') -> (ref, accumAddSparse t1 SZ ref' () val')
      (Right ref', Right val') -> (ref, accumAddSparse t2 SZ ref' () val')
      _ -> error "Mismatched Either in accumAddDense either"
  _ -> error "accumAddDense: invalid dense type"

accumAddDense typ (SS dep) ref idx val = case typ of
  STPair t1 t2 ->
    case (idx, val) of
      (Left idx', Left val') -> (ref, accumAddSparse t1 dep (fst ref) idx' val')
      (Right idx', Right val') -> (ref, accumAddSparse t2 dep (snd ref) idx' val')
      _ -> error "Mismatched Either in accumAddDense pair"
  STEither t1 t2 ->
    case (ref, idx, val) of
      (Left ref', Left idx', Left val') -> (Left ref', accumAddSparse t1 dep ref' idx' val')
      (Right ref', Right idx', Right val') -> (Right ref', accumAddSparse t2 dep ref' idx' val')
      _ -> error "Mismatched Either in accumAddDense either"
  _ -> error "accumAddDense: invalid dense type"
-}


numericIsNum :: ScalIsNumeric st ~ True => SScalTy st -> ((Num (ScalRep st), Ord (ScalRep st)) => r) -> r
numericIsNum STI32 = id
numericIsNum STI64 = id
numericIsNum STF32 = id
numericIsNum STF64 = id

floatingIsFractional :: ScalIsFloating st ~ True => SScalTy st -> ((Floating (ScalRep st), Ord (ScalRep st), ScalIsNumeric st ~ True, ScalIsFloating st ~ True) => r) -> r
floatingIsFractional STF32 = id
floatingIsFractional STF64 = id

integralIsIntegral :: ScalIsIntegral st ~ True => SScalTy st -> ((Integral (ScalRep st), Ord (ScalRep st), ScalIsNumeric st ~ True, ScalIsIntegral st ~ True) => r) -> r
integralIsIntegral STI32 = id
integralIsIntegral STI64 = id

unTupRepIdx :: f Z -> (forall m. f m -> Int -> f (S m))
            -> SNat n -> Rep (Tup (Replicate n TIx)) -> f n
unTupRepIdx nil _    SZ _ = nil
unTupRepIdx nil cons (SS n) (idx, i) = unTupRepIdx nil cons n idx `cons` fromIntegral @Int64 @Int i

tupRepIdx :: (forall m. f (S m) -> (f m, Int))
          -> SNat n -> f n -> Rep (Tup (Replicate n TIx))
tupRepIdx _      SZ _ = ()
tupRepIdx uncons (SS n) tup =
  let (tup', i) = uncons tup
  in ((,) $! tupRepIdx uncons n tup') $! fromIntegral @Int @Int64 i

ixUncons :: Index (S n) -> (Index n, Int)
ixUncons (IxCons idx i) = (idx, i)

shUncons :: Shape (S n) -> (Shape n, Int)
shUncons (ShCons idx i) = (idx, i)