content: module Data.Scoped.Telescope where
  
  import Rebound.Classes
  import Rebound.Env (Shiftable (..))
  import Data.Fin (Fin)
  import Data.Nat
  import Data.Type.Equality ((:~:) (..))
  import Data.Type.Nat
  import Data.Vec.Lazy qualified as Vec
  import Rebound.Lib (axiomAssoc, axiomPlusZ)
  import Data.SNat
  
  -- | Unlike 'Scoped.TeleList', this datatype does not nest: it is effectively a
  -- 'List'/'Vec' but with extra scoping inside.
  data Telescope u s n m where
    TNil :: Telescope u s Z m
    TCons :: (u, s (n + m)) -> !(Telescope u s n m) -> Telescope u s (S n) m
  
  tmap :: (forall k. u -> s k -> (u', s' k)) -> Telescope u s n m -> Telescope u' s' n m
  tmap f TNil = TNil
  tmap f (TCons (u, s) xs) = TCons (f u s) (tmap f xs)
  
  empty :: Telescope u s Z m
  empty = TNil
  
  singleton :: (u, s n) -> Telescope u s (S Z) n
  singleton h = TCons h TNil
  
  append :: forall u s nl nr m. Telescope u s nl (nr + m) -> Telescope u s nr m -> (SNat nl, Telescope u s (nl + nr) m)
  append TNil r = (SZ, r)
  append (TCons l ls) r =
    case axiomAssoc @nl @nr @m of
      Refl -> let (k, ls') = append ls r in withSNat k (SS, TCons l ls')
  
  toTelescope :: forall p n u s. (Shiftable s) => Vec.Vec p (u, s n) -> Telescope u s p n
  toTelescope = snd . iter
    where
      iter :: forall p. Vec.Vec p (u, s n) -> (SNat p, Telescope u s p n)
      iter Vec.VNil = (SZ, TNil)
      iter ((Vec.:::) @_ @p' (u, s) xs) =
        let (sp', sc') :: (SNat p', Telescope u s p' n) = iter xs
            s' :: s (p' + n) = shift sp' s
         in (withSNat sp' SS, TCons (u, s') sc')
  
  -- fromTelescope :: forall s u p n. (Shiftable s) => Telescope u s p n -> (SNat p, Vec.Vec p (u, s (p + n)))
  -- fromTelescope = iter SZ
  --   where
  --     iter :: forall u s k n m. (Shiftable s) => SNat k -> Telescope u s n m -> (SNat (k + n), Vec.Vec n (u, s (k + n + m)))
  --     iter sk TNil = case axiomPlusZ @k of Refl -> (sk, Vec.empty)
  --     iter sk (TCons @_ @_ @n' (u, s) sc) =
  --       case axiomSus @k @n' of
  --         Refl ->
  --           let x' :: (u, s (k + n + m)) = case axiomAssoc @k @n' @m of Refl -> (u, shift (addOne sk) s)
  --               (sn', sc') :: (SNat (k + n), Vec.Vec n' (u, s (k + n + m))) = iter (addOne sk) sc
  --            in (sn', x' Vec.::: sc')
  
  --     addOne :: SNat k -> SNat (S k)
  --     addOne k = withSNat k SS
  
  emptyTelescope = TNil
  
  -- nth :: forall s n m u. (Shiftable s) => Fin n -> Telescope u s n m -> (u, s (n + m))
  -- nth i s = snd (fromTelescope s) Vec.! i
  
  instance Sized (Telescope u s n m) where
    type Size (Telescope u s n m) = n
    size :: Telescope u s n m -> SNat n
    size TNil = s0
    size (TCons _ t) = withSNat (size t) SS