{-# language AllowAmbiguousTypes, TypeFamilies, DataKinds, UndecidableInstances #-}
module M where
import GHC.TypeLits
import Unsafe.Coerce (unsafeCoerce)
import Data.Type.Equality (gcastWith, (:~:) (Refl))
import Type.Reflection
import Data.Proxy
import           Data.Kind (Type)

class (Typeable s) => Shape (s :: [Nat]) where
  shapeP :: Proxy s -> [Int]
  sizeP  :: Proxy s -> Int

instance Shape '[] where

instance forall n s . (Shape s, KnownNat n) => Shape (n ': s) where

data Array (sh :: [Nat])

type family (++) (xs :: [Nat]) (ys :: [Nat]) :: [Nat] where
  (++) '[] ys = ys
  (++) (x ': xs) ys = x ': (xs ++ ys)

type family Take (n :: Nat) (xs :: [Nat]) :: [Nat] where
    Take 0 xs = '[]
    Take n (x ': xs) = x ': Take (n-1) xs

type family Drop (n :: Nat) (xs :: [Nat]) :: [Nat] where
    Drop 0 xs = xs
    Drop n (x ': xs) = Drop (n-1) xs
    
data ShapedList (sh :: [Nat]) i where
  ZSH :: ShapedList '[] i
  (:$:) :: (KnownNat n, Shape sh)
        => i -> ShapedList sh i -> ShapedList (n ': sh) i

newtype AstVarId = AstVarId Int

type AstIndexS r sh = ShapedList sh Int

data AstShaped :: Type -> [Nat] -> Type where

build1VOccurenceUnknownS
  :: forall k sh r. (KnownNat k, Shape sh)
  => (AstVarId, AstShaped r sh) -> AstShaped r (k ': sh)
build1VOccurenceUnknownS = undefined

astIndexStepS
  :: forall sh1 sh2 r.
     (Shape sh1, Shape sh2, Shape (sh1 ++ sh2))
  => AstShaped r (sh1 ++ sh2) -> AstIndexS r sh1
  -> AstShaped r sh2
astIndexStepS = undefined

build1VIndexS'
  :: forall k p sh r.
     ( KnownNat k, Shape sh  -- !!!! Removing this also causes an unreadable error message
     , Shape (Drop p (Take p sh ++ Drop p sh)) )
  => (AstVarId, AstShaped r sh, AstIndexS r (Take p sh))
  -> AstShaped r (k ': Drop p sh)
build1VIndexS' (var, v0, ix@(_ :$: _)) =
  -- Commenting this out causes the error with the unreadable error message
  -- from which one can't really guess what axiom needs to be added:
  -- gcastWith (unsafeCoerce Refl :: sh :~: Take p sh ++ Drop p sh) $
     case astIndexStepS v0 ix of
       v ->
         build1VOccurenceUnknownS (var, v)