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

import Control.Monad (foldM, join)
import Data.Kind (Type)
import Data.Int (Int64)
import Data.IORef
import System.IO.Unsafe (unsafePerformIO)

import Array
import AST
import CHAD.Types
import Data
import Interpreter.Rep
import Data.Bifunctor (first, bimap)


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

interpret :: Ex '[] t -> Rep t
interpret e = runAcM (interpret' SNil e)

newtype Value t = Value (Rep t)

interpret' :: forall env t s. SList Value env -> Ex env t -> AcM s (Rep t)
interpret' env = \case
  EVar _ _ i -> case slistIdx env i of Value x -> return x
  ELet _ a b -> do
    x <- interpret' env a
    interpret' (Value x `SCons` env) b
  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
  EBuild1 _ a b -> do
    n <- fromIntegral @Int64 @Int <$> interpret' env a
    arrayGenerateLinM (ShNil `ShCons` n)
                      (\i -> interpret' (Value (fromIntegral @Int @Int64 i) `SCons` env) b)
  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 -> do
    let f = \x y -> interpret' (Value y `SCons` Value x `SCons` env) a
    arr <- interpret' env b
    let sh `ShCons` n = arrayShape arr
    arrayGenerateM sh $ \idx -> foldl1M f [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 $ arrayGenerateM sh $ \idx -> return $ 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
    arrayGenerateM (sh `ShCons` n) (\(idx `IxCons` _) -> return (arrayIndex arr idx))
  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 _ n a b -> 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
  EWith e1 e2 -> do
    initval <- interpret' env e1
    withAccum (typeOf e1) (typeOf e2) initval $ \accum ->
      interpret' (Value accum `SCons` env) e2
  EAccum i e1 e2 e3 -> do
    let STAccum t = typeOf e3
    idx <- interpret' env e1
    val <- interpret' env e2
    accum <- interpret' env e3
    accumAddSparse t i 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'
  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 -> numericIsNum st $ uncurry (==) arg
  ONot -> not arg
  OIf -> if arg then Left () else Right ()

zeroD2 :: STy t -> Rep (D2 t)
zeroD2 typ = case typ of
  STNil -> ()
  STPair _ _ -> Left ()
  STEither _ _ -> Left ()
  STMaybe _ -> Nothing
  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
    (Left (), _) -> b
    (_, Left ()) -> a
    (Right (x1, x2), Right (y1, y2)) -> Right (addD2s t1 x1 y1, addD2s t2 x2 y2)
  STEither t1 t2 -> case (a, b) of
    (Left (), _) -> b
    (_, Left ()) -> a
    (Right (Left x), Right (Left y)) -> Right (Left (addD2s t1 x y))
    (Right (Right x), Right (Right y)) -> Right (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"

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

newAcZero :: STy t -> IO (RepAcSparse t)
newAcZero = \case
  STNil -> return ()
  STPair t1 t2 -> newIORef =<< (,) <$> newAcZero t1 <*> newAcZero t2
  STMaybe _ -> newIORef Nothing
  STArr n _ -> newIORef (emptyArray n)
  STScal sty -> case sty of
    STI32 -> newIORef 0
    STI64 -> newIORef 0
    STF32 -> newIORef 0.0
    STF64 -> newIORef 0.0
    STBool -> error "Accumulator of Bool"
  STAccum{} -> error "Nested accumulators"
  STEither{} -> error "Bare Either in accumulator"

-- | 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 shape: the outermost dimension is at the /outside/ of this list.
data PartialInvShape n m where
  PIShEnd :: PartialInvShape m m
  PIShCons :: Int -> PartialInvShape n m -> PartialInvShape (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

pattern IIxNil :: () => n ~ Z => InvIndex n
pattern IIxNil = InvNil
pattern IIxCons :: () => S n ~ succn => Int -> InvIndex n -> InvIndex succn
pattern IIxCons i ix = InvCons i ix
{-# COMPLETE IIxNil, IIxCons #-}

pattern IShNil :: () => n ~ Z => InvShape n
pattern IShNil = InvNil
pattern IShCons :: () => S n ~ succn => Int -> InvShape n -> InvShape succn
pattern IShCons n sh = InvCons n sh
{-# COMPLETE IShNil, IShCons #-}

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

enumInvShape :: InvShape n -> [InvIndex n]
enumInvShape IShNil = [IIxNil]
enumInvShape (n `IShCons` sh) = [i `IIxCons` ix | i <- [0 .. n - 1], ix <- enumInvShape sh]

enumShape :: Shape n -> [Index n]
enumShape = map uninvert . enumInvShape . invert

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) (IIxCons 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 = IIxCons i (piindexConcat pix ix)

newAcSparse :: STy t -> SNat i -> Rep (AcIdx t i) -> Rep (AcVal t i) -> IO (RepAcSparse t)
newAcSparse typ SZ () val = case typ of
  STNil -> return ()
  STPair{} -> newIORef =<< newAcDense typ SZ () val
  STMaybe t -> newIORef =<< traverse (newAcDense t SZ ()) val
  STArr{} -> newIORef =<< newAcDense typ SZ () val
  STScal{} -> newIORef val
  STAccum{} -> error "Nested accumulators"
  STEither{} -> error "Bare Either in accumulator"
newAcSparse typ (SS dep) idx val = case typ of
  STNil -> return ()
  STPair{} -> newIORef =<< newAcDense typ (SS dep) idx val
  STMaybe t -> newIORef =<< Just <$> newAcDense t dep idx val
  STArr{} -> newIORef =<< newAcDense typ (SS dep) idx val
  STScal{} -> error "Cannot index into scalar"
  STAccum{} -> error "Nested accumulators"
  STEither{} -> error "Bare Either in accumulator"

newAcDense :: STy t -> SNat i -> Rep (AcIdx t i) -> Rep (AcVal t i) -> IO (RepAcDense t)
newAcDense typ SZ () val = case typ of
  STNil -> return ()
  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
  STMaybe t -> traverse (newAcSparse t SZ ()) val
  STArr _ t -> traverse (newAcSparse t SZ ()) val
  STScal{} -> return val
  STAccum{} -> error "Nested accumulators"
newAcDense typ (SS dep) idx val = case typ of
  STNil -> return ()
  STPair{} -> newAcDense typ (SS dep) idx val
  STMaybe t -> Just <$> newAcSparse t dep idx val
  STArr dim (t :: STy t) -> do
    let sh = unTupRepIdx ShNil ShCons dim (fst val)
    go (SS dep) dim idx (snd val) $ \arr position ->
      arrayGenerateM sh (\i -> case uninvert <$> piindexMatch position (invert i) of
                                 Just i' -> return $ arr `arrayIndex` i'
                                 Nothing -> newAcZero t)
    where
      go :: SNat i -> SNat n -> Rep (AcIdx (TArr n t) i) -> Rep (AcValArr n t i) -> (forall m. Array m (RepAcSparse t) -> PartialInvIndex n m -> IO r) -> IO r
      go SZ _ () val' k = arrayMapM (newAcSparse t SZ ()) val' >>= \arr -> k arr PIIxEnd
      go (SS dep') SZ idx' val' k = newAcSparse t dep' idx' val' >>= \arr -> k (arrayUnit arr) PIIxEnd
      go (SS dep') (SS dim') (i, idx') val' k =
        go dep' dim' idx' val' $ \arr pish ->
          k arr (PIIxCons (fromIntegral @Int64 @Int i) pish)
  STScal{} -> error "Cannot index into scalar"
  STAccum{} -> error "Nested accumulators"
  STEither{} -> error "Bare Either in accumulator"

readAcSparse :: STy t -> RepAcSparse t -> IO (Rep t)
readAcSparse typ val = case typ of
  STNil -> return ()
  STPair t1 t2 -> do
    (a, b) <- readIORef val
    (,) <$> readAcSparse t1 a <*> readAcSparse t2 b
  STMaybe t -> traverse (readAcDense t) =<< readIORef val
  STArr _ t -> traverse (readAcSparse t) =<< readIORef val
  STScal{} -> readIORef val
  STAccum{} -> error "Nested accumulators"
  STEither{} -> error "Bare Either in accumulator"

readAcDense :: STy t -> RepAcDense t -> IO (Rep t)
readAcDense typ val = case typ of
  STNil -> return ()
  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
  STMaybe t -> traverse (readAcSparse t) val
  STArr _ t -> traverse (readAcSparse t) val
  STScal{} -> return val
  STAccum{} -> error "Nested accumulators"

accumAddSparse :: STy t -> SNat i -> RepAcSparse t -> Rep (AcIdx t i) -> Rep (AcVal t i) -> AcM s ()
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 <- newAcDense typ 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 <- newAcDense typ (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') IShNil          idx'      val' kj _  = kj dep' IxNil idx' val'  -- ^ Ran out of array dimensions: accumulating into (part of) element
      go (SS dep') (IShCons _ 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
  STNil -> ((), return ())
  STPair t1 t2 ->
    (ref, do accumAddSparse t1 SZ (fst ref) () (fst val)
             accumAddSparse t2 SZ (snd ref) () (snd val))
  STMaybe t ->
    case (ref, val) of
      (_, Nothing) -> (ref, return ())
      (Nothing, Just val') -> _ val'
      (Just ref', Just val') -> _ ref' val'
  STArr _ t -> AcM $ do
    refs <- readIORef ref
    case (shapeSize (arrayShape refs), shapeSize (arrayShape val)) of
      (_, 0) -> return ()
      (0, _) -> do
        newrefarr <- newAcDense typ 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"

accumAddDense 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 <- newAcDense typ (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') IShNil          idx'      val' kj _  = kj dep' IxNil idx' val'  -- ^ Ran out of array dimensions: accumulating into (part of) element
      go (SS dep') (IShCons _ 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"

-- accumAddVal :: forall t i s. STy t -> SNat i -> RepAc t -> Rep (AcIdx t i) -> Rep (AcVal t i) -> (RepAc t, AcM s ())
-- accumAddVal typ SZ ac () val = case typ of
--   STNil -> ((), return ())
--   STPair t1 t2 ->
--     let (ac1', m1) = accumAddVal t1 SZ (fst ac) () (fst val)
--         (ac2', m2) = accumAddVal t2 SZ (snd ac) () (snd val)
--     in ((ac1', ac2'), m1 >> m2)
--   STMaybe t -> case t of
--     STEither t1 t2 -> (ac, accumAddValME t1 t2 ac val)
--     STNil -> def ; STPair{} -> def ; STMaybe{} -> def ; STArr{} -> def ; STScal{} -> def ; STAccum{} -> def
--     where def :: (t ~ TMaybe a, RepAc (TMaybe a) ~ IORef (Maybe (RepAc a))) => (RepAc t, AcM s ())
--           def = (ac, accumAddValM t ac val)
--   STArr n t
--     | shapeSize (arrayShape ac) == 0 -> makeRepAc (STArr n t) val
--   STEither{} -> error "Bare Either in accumulator"
--   _ -> _
-- accumAddVal typ (SS dep) ac idx val = case typ of
--   STNil -> ((), return ())
--   STPair t1 t2 ->
--     case (idx, val) of
--       (Left idx', Left val') -> first (,snd ac) $ accumAddVal t1 dep (fst ac) idx' val'
--       (Right idx', Right val') -> first (fst ac,) $ accumAddVal t2 dep (snd ac) idx' val'
--       _ -> error "Inconsistent idx and val in accumulator add operation"
--   _ -> _

-- accumAddValME :: STy a -> STy b
--               -> IORef (Maybe (Either (RepAc a) (RepAc b)))
--               -> Maybe (Either (Rep a) (Rep b))
--               -> AcM s ()
-- accumAddValME t1 t2 ac val =
--   case val of
--     Nothing -> return ()
--     Just val' ->
--       join $ AcM $ atomicModifyIORef' ac $ \ac' -> case (ac', val') of
--                      (Nothing, _) -> (Nothing, AcM $ initAccumOrTryAgainME t1 t2 ac val' (unAcM $ accumAddValME t1 t2 ac val))
--                      (Just (Left x), Left val'1) -> first (Just . Left) $ accumAddVal t1 SZ x () val'1
--                      (Just (Right y), Right val'2) -> first (Just . Right) $ accumAddVal t2 SZ y () val'2
--                      _ -> error "Inconsistent accumulator and value in add operation on Maybe Either"

-- initAccumOrTryAgainME :: STy a -> STy b
--                       -> IORef (Maybe (Either (RepAc a) (RepAc b)))
--                       -> Either (Rep a) (Rep b)
--                       -> IO ()
--                       -> IO ()
-- initAccumOrTryAgainME t1 t2 ac val onRace = do
--   newContents <- case val of Left x -> Left <$> makeRepAc t1 x
--                              Right y -> Right <$> makeRepAc t2 y
--   join $ atomicModifyIORef' ac (\case Nothing -> (Just newContents, return ())
--                                       value@Just{} -> (value, onRace))

-- accumAddValM :: STy t
--              -> IORef (Maybe (RepAc t))
--              -> Maybe (Rep t)
--              -> AcM s ()
-- accumAddValM t ac val =
--   case val of
--     Nothing -> return ()
--     Just val' ->
--       join $ AcM $ atomicModifyIORef' ac $ \ac' -> case ac' of
--                      Nothing -> (Nothing, AcM $ initAccumOrTryAgainM t ac val' (unAcM $ accumAddValM t ac val))
--                      Just x -> first Just $ accumAddVal t SZ x () val'

-- initAccumOrTryAgainM :: STy t
--                      -> IORef (Maybe (RepAc t))
--                      -> Rep t
--                      -> IO ()
--                      -> IO ()
-- initAccumOrTryAgainM t ac val onRace = do
--   newContents <- makeRepAc t val
--   join $ atomicModifyIORef' ac (\case Nothing -> (Just newContents, return ())
--                                       value@Just{} -> (value, onRace))

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

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)

foldl1M :: Monad m => (a -> a -> m a) -> [a] -> m a
foldl1M _ [] = error "foldl1M: empty list"
foldl1M f (tophead : toptail) = foldM f tophead toptail