
   0 | module Data.Container.Base.Morphism.Instances
   1 | 
   2 | import Data.Fin
   3 | import Data.Fin.Split
   4 | import Data.Vect
   5 | import Data.List.Elem
   6 | import Data.List.Quantifiers
   7 | 
   8 | import Data.Container.Base.Object.Definition
   9 | import Data.Container.Base.Morphism.Definition
  10 | import Data.Container.Base.Extension.Definition
  11 | import Data.Container.Base.Properties.Definitions
  12 | import Data.Container.Base.Product.Definitions
  13 | 
  14 | import Data.Container.Base.Object.Instances
  15 | 
  16 | import Data.Container.Base.Quantifiers
  17 | import Data.Container.Base.TreeUtils
  18 | 
  19 | import Control.Monad.Distribution
  20 | import Control.Monad.Sample.Definition
  21 | 
  22 | import Data.Num
  23 | import Data.Layout
  24 | import Misc
  25 | 
  26 | ||| If we model the idea of a container (S !> P) as a box
  27 | |||  ┌──────┐
  28 | |||  │ s:S  │
  29 | |||  ├──────┤
  30 | |||  │  Ps  │
  31 | |||  └──────┘
  32 | ||| then `pushDown` is interpreted as pushing down the container,
  33 | ||| pruning anything that goes out of the box, and using `Unit` for
  34 | ||| anything new that appears:
  35 | |||  ┌──────┐
  36 | |||  │ Unit │
  37 | |||  ├──────┤
  38 | |||  │ s:S  │
  39 | |||  └──────┘
  40 | |||     Ps
  41 | public export
  42 | pushDown : Cont -> Cont
  43 | pushDown c = Const2 Unit c.Shp
  44 | 
  45 | public export
  46 | pushIntoContinuation : {0 d, p, l : Cont} ->
  47 |   d >< p =%> l ->
  48 |   p =%> (pushDown d) >@ l
  49 | pushIntoContinuation f = !% \p => (() <| \d => f.fwd (d, p) **
  50 |   \(d ** l') => snd $ f.bwd (d, p) l')
  51 | 
  52 | 
  53 | namespace CategoricalProduct
  54 |   public export
  55 |   terminal : c =%> UnitCont
  56 |   terminal = !% \_ => (() ** absurd)
  57 | 
  58 | 
  59 | namespace HancockTensorProduct
  60 |   public export
  61 |   leftUnit : Scalar >< c =%> c
  62 |   leftUnit = !% \((), s) => (s ** \p => ((), p))
  63 |   
  64 |   public export
  65 |   rightUnit : c >< Scalar =%> c
  66 |   rightUnit = !% \(x, ()) => (x ** \x' => (x', ()))
  67 | 
  68 |   public export
  69 |   leftUnitInv : c =%> Scalar >< c
  70 |   leftUnitInv = !% \x => (((), x) ** \((), x') => x')
  71 | 
  72 |   public export
  73 |   rightUnitInv : c =%> c >< Scalar
  74 |   rightUnitInv = !% \x => ((x, ()) ** \(x', ()) => x')
  75 | 
  76 |   public export
  77 |   assocL : (a >< b) >< c =%> a >< (b >< c)
  78 |   assocL = !% \((a, b), c) => ((a, (b, c)) ** \(a', (b', c')) => ((a', b'), c'))
  79 | 
  80 |   public export
  81 |   assocR : a >< (b >< c) =%> (a >< b) >< c
  82 |   assocR = !% \(a, (b, c)) => (((a, b), c) ** \((a', b'), c') => (a', (b', c')))
  83 | 
  84 |   public export
  85 |   swap : a >< b =%> b >< a
  86 |   swap = !% \(a, b) => ((b, a) ** \(b', a') => (a', b'))
  87 | 
  88 |   public export
  89 |   swapMiddle : (c1 >< c2) >< (c3 >< c4) =%> (c1 >< c3) >< (c2 >< c4)
  90 |   swapMiddle = assocL {c=_ >< _}
  91 |            %>> (id >< assocR)
  92 |            %>> (id >< swap >< id)
  93 |            %>> (id >< assocL)
  94 |            %>> assocR {c=_ >< _}
  95 | 
  96 | namespace CompositionProduct
  97 |   public export
  98 |   leftUnit : Scalar >@ c =%> c
  99 |   leftUnit = !% \(() <| cShp) => (cShp () ** \c' => (() ** c'))
 100 | 
 101 |   public export
 102 |   rightUnit : c >@ Scalar =%> c
 103 |   rightUnit = !% \(s <| _) => (s ** \cp => (cp ** ()))
 104 | 
 105 |   public export
 106 |   leftUnitInv : c =%> Scalar >@ c
 107 |   leftUnitInv = !% \s => (() <| (\_ => s) ** \(() ** c') => c')
 108 |   
 109 |   public export
 110 |   rightUnitInv : c =%> c >@ Scalar
 111 |   rightUnitInv = !% \s => (s <| const () ** fst)
 112 | 
 113 |   public export
 114 |   assocL : (a >@ b) >@ c =%> a >@ (b >@ c)
 115 |   assocL = !% \((aShp <| f) <| g) =>
 116 |     (aShp <| \aPos => f aPos <| \bPos => g (aPos  ** bPos) ** \(aPos ** bPos ** cPos) => ((aPos ** bPos) ** cPos))
 117 | 
 118 |   public export
 119 |   assocR : a >@ (b >@ c) =%> (a >@ b) >@ c
 120 |   assocR = !% \(aShp <| f) => ((aShp <| shapeExt . f) <| \(aPos ** bPos) =>
 121 |     index (f aPos) bPos **
 122 |       \((aPos ** bPos) ** cPos) => (aPos ** (bPos ** cPos)))
 123 | 
 124 | namespace Coproduct
 125 |   public export
 126 |   elim : c >+< c =%> c
 127 |   elim = !% \case
 128 |     Left x => (x ** id)
 129 |     Right y => (y ** id)
 130 | 
 131 |   public export
 132 |   initial : Empty =%> c
 133 |   initial = !% absurd
 134 | 
 135 |   public export
 136 |   cojoin : c =%> z ->
 137 |     d =%> z ->
 138 |     c >+< d =%> z
 139 |   cojoin f g = (f >+< g) %>> elim
 140 | 
 141 |   public export
 142 |   direct : (a.Shp -> Bool) ->
 143 |     a =%> a >+< a
 144 |   direct p = !% \x => case p x of
 145 |     False => (Left x ** id)
 146 |     True => (Right x ** id)
 147 | 
 148 | 
 149 | 
 150 | namespace CartesianClosure
 151 |   ||| The following is the proof that for any container `c` there is an
 152 |   ||| isomorphism in `Cont` between `c` and `CartesianClosure UnitCont c`
 153 |   ||| This holds in any monoidal closed category: `X ≅ [I, X]`
 154 |   namespace StateIsomorphismProof
 155 |     stateToCartClosureFw : c =%> (CartesianClosure UnitCont c)
 156 |     stateToCartClosureFw = !% \cShp => (!% \() => (cShp ** \_ => Nothing)
 157 |                                        ** \(() ** cPos ** ItIsNothing) => cPos)
 158 | 
 159 |     stateToCartClosureBw : CartesianClosure UnitCont c =%> c
 160 |     stateToCartClosureBw = !% \l => (l.fwd () ** \cPos =>
 161 |       (() ** cPos ** maybeVoidIsNothing (l.bwd () cPos)))
 162 | 
 163 | 
 164 | ||| For a overview of this interaction from the categorical perspective, see
 165 | ||| the Poly book (https://arxiv.org/abs/2312.00990) (Section 6.3.4)
 166 | namespace CompositionTensorInteraction
 167 |   ||| Interaction between composition and tensor product
 168 |   ||| Swaps the operations, and middle two containers
 169 |   ||| Not an isomorphism!
 170 |   public export
 171 |   duoidal : (c >@ d) >< (e >@ f) =%> (c >< e) >@ (d >< f)
 172 |   duoidal = !% \((sc <| idxC), (se <| idxE)) =>
 173 |     ((sc, se) <| \(cp, ep) => (idxC cp, idxE ep) **
 174 |       \((cp, ep) ** (dp, fp)) => ((cp ** dp), (ep ** fp)))
 175 |   
 176 |   ||| Tensor product embeds into composition
 177 |   ||| A special case of `duoidal`
 178 |   public export
 179 |   tensorToComp : c >< f =%> c >@ f
 180 |   tensorToComp =   (rightUnitInv >< leftUnitInv)
 181 |                %>> duoidal {d=Scalar,e=Scalar}
 182 |                %>> (rightUnit >@ leftUnit)
 183 | 
 184 |   ||| Going the other way is impossible without any constraints 
 185 |   ||| Two possibilities on constraints (this, and `compToTensor2`)
 186 |   public export 
 187 |   compToTensor : IsNaperian d =>
 188 |     c >@ d =%> c >< d
 189 |   compToTensor @{(MkIsNaperian dPos)} = !% \(cShp <| content) =>
 190 |     ((cShp,()) ** \(cPos, dPos) => (cPos ** dPos))
 191 |   
 192 |   public export
 193 |   compToTensor2 : IsFlat c =>
 194 |     c >@ d =%> c >< d
 195 |   compToTensor2 @{(ItIsFlat cShp)} = !% \(cShp <| dShp) =>
 196 |     ((cShp, dShp ()) ** \((), dPos') => (() ** dPos'))
 197 |   
 198 |   ||| Specific distributive law we need
 199 |   public export
 200 |   distribute : c >< e =%> s ->
 201 |     c >< (e >@ g) =%> s >@ g
 202 |   distribute f = (rightUnitInv >< id {c=e >@ g})
 203 |                %>> duoidal {d = Scalar}
 204 |                %>> (f >@ leftUnit)
 205 | 
 206 | 
 207 | namespace State
 208 |   ||| "State" as defined in https://arxiv.org/abs/2403.13001 and open games 
 209 |   |||
 210 |   |||       ┌─────────────┐
 211 |   |||       │             ├──► (x : c.Shp)
 212 |   |||       │    State    │
 213 |   |||       │             ├◄── c.Pos x
 214 |   |||       └─────────────┘
 215 |   public export
 216 |   State : Cont -> Type
 217 |   State c = Scalar =%> c
 218 | 
 219 |   ||| Given a shape of any container, state can be defined
 220 |   public export
 221 |   toState : (x : c.Shp) -> State c
 222 |   toState x = !% \() => (x ** \_ => ())
 223 | 
 224 |   public export
 225 |   fromState : State c -> c.Shp
 226 |   fromState f = f.fwd ()
 227 | 
 228 |   public export
 229 |   mapState : State c ->
 230 |     c =%> d ->
 231 |     State d
 232 |   mapState s f = s %>> f
 233 | 
 234 | namespace Costate
 235 |   ||| "Costate" as defined in https://arxiv.org/abs/2403.13001 and open games 
 236 |   |||
 237 |   |||                  ┌─────────────┐
 238 |   |||  (x : c.Shp)  ──►┤             │
 239 |   |||                  │   Costate   │
 240 |   |||     c.Pos x   ◄──┤             │
 241 |   |||                  └─────────────┘
 242 |   public export
 243 |   Costate : Cont -> Type
 244 |   Costate c = c =%> Scalar
 245 | 
 246 |   public export
 247 |   toCostate : ((x : c.Shp) -> c.Pos x) -> Costate c
 248 |   toCostate s = !% \x => (() ** \() => s x)
 249 | 
 250 |   public export
 251 |   fromCostate : Costate c -> (x : c.Shp) -> c.Pos x
 252 |   fromCostate f x = f.bwd x ()
 253 | 
 254 |   public export
 255 |   mapCostate : Costate d ->
 256 |     c =%> d ->
 257 |     Costate c
 258 |   mapCostate s f = f %>> s
 259 | 
 260 |   public export
 261 |   pairCostate : Costate c -> Costate d -> Costate (c >< d)
 262 |   pairCostate f g = (f >< g) %>> leftUnit {c=Scalar} %>> id
 263 |   
 264 | public export
 265 | fromNapCostateToState : Costate (Nap c.Shp) -> State c
 266 | fromNapCostateToState f = toState (f.bwd () ())
 267 | 
 268 | public export
 269 | fromStateToNapCostate : State c -> Costate (Nap c.Shp)
 270 | fromStateToNapCostate f = toCostate f.fwd
 271 | 
 272 | 
 273 | 
 274 | ||| Wraps a dependent lens `c =%> d`
 275 | ||| into one of type `c >@ Scalar =%> d >@ Scalar`
 276 | ||| Needed because `c >@ Scalar` isn't automatically reduced to `c`
 277 | public export
 278 | wrapIntoVector : c =%> d ->
 279 |   Tensor [c] =%> Tensor [d]
 280 | wrapIntoVector f = rightUnit %>> f %>> rightUnitInv
 281 | 
 282 | public export
 283 | wrapIntoMatrix : c >@ c' =%> d >@ d' ->
 284 |   Tensor [c, c'] =%> Tensor [d, d']
 285 | wrapIntoMatrix f =   (id >@ rightUnit)
 286 |                  %>> f
 287 |                  %>> (id >@ rightUnitInv)
 288 | 
 289 | ||| Wraps a dependent lens `c =%> d`
 290 | ||| into one of type `c >< Scalar =%> d >< Scalar`
 291 | ||| Needed because `c >< Scalar` isn't automatically reduced to `c`
 292 | public export
 293 | wrapIntoVectorHancock : c =%> d ->
 294 |   HancockTensor [c] =%> HancockTensor [d]
 295 | wrapIntoVectorHancock f = rightUnit %>> f %>> rightUnitInv
 296 | 
 297 | namespace CubicalHelpers
 298 |   ||| Helper function allowing `shape` in `cubicalShape` to have zero annotation
 299 |   public export
 300 |   cubicalShapeHelper : All IsCubical shape -> List Nat
 301 |   cubicalShapeHelper [] = []
 302 |   cubicalShapeHelper (ic :: ics) = dimHelper ic :: cubicalShapeHelper ics
 303 |     
 304 |   ||| Given a list of cubical containers, return the list of their dimensions
 305 |   public export
 306 |   cubicalShape : (0 shape : List Cont) -> All IsCubical shape => List Nat
 307 |   cubicalShape _ @{ac} = cubicalShapeHelper ac
 308 |     
 309 |   ||| Size of a list of cubical containers is the product of their dimensions
 310 |   public export
 311 |   size : (0 shape : List Cont) -> All IsCubical shape => Nat
 312 |   size shape = prod (cubicalShape shape)
 313 | 
 314 | ||| Layout-aware dependent lens flattening a cubical tensor
 315 | public export
 316 | flattenCubical : {shape : List Cont} ->
 317 |   (ac : All IsCubical shape) =>
 318 |   LayoutOrder ->
 319 |   Tensor shape =%> Vect (size shape)
 320 | flattenCubical {shape = [], ac=[]} _ = !% \() => (() ** \FZ => ())
 321 | flattenCubical {shape = (_ :: ss), ac=(MkIsCubical n :: as)} lo
 322 |   = !% \(() <| t) => (() ** \idx =>
 323 |       let (!% recBackward) = flattenCubical {shape = ss} lo
 324 |           (i, rest) = splitFinProd lo idx
 325 |           (_ ** backRec) = recBackward (t i)
 326 |       in (i ** backRec rest))
 327 | 
 328 | ||| Layout-aware dependent lens unflattening a tensor
 329 | public export
 330 | unflattenCubical : {shape : List Cont} ->
 331 |   (ac : All IsCubical shape) =>
 332 |   LayoutOrder ->
 333 |   Vect (size shape) =%> Tensor shape
 334 | unflattenCubical {shape = [], ac=[]} lo = !% \() => (() ** \() => FZ)
 335 | unflattenCubical {shape = (_ :: ss), ac=((MkIsCubical n) :: as)} lo =
 336 |   let (!% f) = unflattenCubical {shape = ss} lo
 337 |       (innerShape ** innerBack) = f ()
 338 |   in !% \() => ((() <| \_ => innerShape) ** (\(cp ** restPos) =>
 339 |     indexFinProd lo cp (innerBack restPos)))
 340 | 
 341 | ||| This is simply a rewrite!
 342 | public export
 343 | recastFlattenedTensor : {oldShape, newShape : List Cont} ->
 344 |   (oldAc : All IsCubical oldShape) => (newAc : All IsCubical newShape) =>
 345 |   {auto prf : size oldShape = size newShape} ->
 346 |   Vect (size oldShape) =%> Vect (size newShape)
 347 | recastFlattenedTensor = !% \() => (() ** \i => rewrite prf in i)
 348 | 
 349 | ||| Reshapes a cubical tensor by first flattening it to a linear representation,
 350 | ||| casting the type to the new shape, and then unflattening it back
 351 | ||| Is generic over layout order
 352 | public export
 353 | reshape : {oldShape, newShape : List Cont} ->
 354 |   (oldAc : All IsCubical oldShape) => (newAc : All IsCubical newShape) =>
 355 |   LayoutOrder ->
 356 |   {auto prf : size oldShape = size newShape} ->
 357 |   Tensor oldShape =%> Tensor newShape
 358 | reshape lo = flattenCubical lo
 359 |          %>> recastFlattenedTensor
 360 |          %>> unflattenCubical lo
 361 | 
 362 | 
 363 | namespace Transpose
 364 |   public export
 365 |   transposeLens : IsNaperian c => IsNaperian d => c >@ d =%> d >@ c
 366 |   transposeLens @{MkIsNaperian _} @{MkIsNaperian _} = !% \(() <| _) =>
 367 |     (() <| (\_ => ()) ** \(dInd ** cInd) => (cInd ** dInd))
 368 | 
 369 |   public export
 370 |   transpose : IsNaperian c => IsNaperian d =>
 371 |     Tensor [c, d] =%> Tensor [d, c]
 372 |   transpose @{MkIsNaperian _} @{MkIsNaperian _} = wrapIntoMatrix transposeLens
 373 | 
 374 |   -- ||| experiment, does this work?
 375 |   -- public export
 376 |   -- transposeMiddle : IsNaperian c => IsNaperian e =>
 377 |   --   Tensor [c, e, d] =%> 
 378 |   
 379 | 
 380 |   --||| Transpose a given element to the front of the shape
 381 |   --public export
 382 |   --transposeToFront : (shape : List Cont) ->
 383 |   --  (c : Cont) ->
 384 |   --  (elem : Elem c shape) =>
 385 |   --  All IsNaperian (dropAfterElem shape elem) =>
 386 |   --  Tensor shape =%> Tensor (c :: dropElem shape elem)
 387 |   --transposeToFront (_ :: xs) c @{Here} @{allNap} = ?transposeToFront_rhs_0
 388 |   --transposeToFront (y :: xs) c @{(There x)} @{allNap} = ?transposeToFront_rhs_1
 389 |   
 390 | ||| Functionality for transforming a tensor into a hancock tensor
 391 | namespace TransformIntoHancockTensor
 392 |   public export
 393 |   hancockTensorNaperianShape : {shape : List Cont} ->
 394 |     (allNap : All IsNaperian shape) =>
 395 |     (HancockTensor shape).Shp
 396 |   hancockTensorNaperianShape {shape = []} = ()
 397 |   hancockTensorNaperianShape {allNap = ((MkIsNaperian _) :: _)}
 398 |     = ((), hancockTensorNaperianShape)
 399 |   
 400 |   ||| Helper to compute the unique shape of Tensor when all containers are Naperian
 401 |   public export
 402 |   tensorNaperianShape : {shape : List Cont} ->
 403 |     (allNap : All IsNaperian shape) =>
 404 |     (Tensor shape).Shp
 405 |   tensorNaperianShape {shape = []} = ()
 406 |   tensorNaperianShape {shape = (_ :: ss), allNap = ((MkIsNaperian _) :: ns)}
 407 |     = () <| \_ => tensorNaperianShape {shape = ss} @{ns}
 408 |   
 409 |   ||| Analogous to `naperianPosEq` but for the HancockTensor structure
 410 |   ||| We can't use `naperianPosEq` directly because the shape of the resulting
 411 |   ||| container is not Unit, it is only isomorphic to it
 412 |   public export
 413 |   hancockTensorPosEq : {shape : List Cont} ->
 414 |     (allNap : All IsNaperian shape) =>
 415 |     {0 x, y : (HancockTensor shape).Shp} ->
 416 |     (HancockTensor shape).Pos x = (HancockTensor shape).Pos y
 417 |   hancockTensorPosEq {allNap = []} = Refl
 418 |   hancockTensorPosEq {allNap = ((MkIsNaperian _) :: _)} = cong2 Pair
 419 |     (naperianPosEq @{MkIsNaperian _} {x=()} {y=()})
 420 |     hancockTensorPosEq
 421 |   
 422 |   ||| Tensor shape is isomorphic to HancockTensor shape when all containers in
 423 |   ||| the shape are Naperian. This is one arrow of that isomorphism
 424 |   public export
 425 |   transformToHancock : {shape : List Cont} ->
 426 |     All IsNaperian shape =>
 427 |     Tensor shape =%> HancockTensor shape
 428 |   transformToHancock {shape = []} = id
 429 |   transformToHancock {shape = (_ :: _)} @{((MkIsNaperian _) :: _)}
 430 |     = !% \(() <| content) => (((), hancockTensorNaperianShape) **
 431 |        \(p, restPos) =>
 432 |          let (_ ** recBack) = (%!) transformToHancock (content p)
 433 |          in (p ** recBack $ replace {p = id} hancockTensorPosEq restPos))
 434 | 
 435 |   public export
 436 |   transformFromHancock : {shape : List Cont} ->
 437 |     All IsNaperian shape =>
 438 |     HancockTensor shape =%> Tensor shape
 439 |   transformFromHancock {shape = []} = id
 440 |   transformFromHancock {shape = (Nap s :: ss)} @{((MkIsNaperian s) :: _)}
 441 |     = !% \((), hShp) =>
 442 |         let (tShp ** recBack) = (%!) transformFromHancock hShp
 443 |         in (() <| (\_ => tShp) ** \(p ** restPos) => (p, recBack restPos))
 444 | 
 445 |     
 446 | 
 447 |   -- ||| Technically this is Unit, but hard to prove
 448 |   -- public export
 449 |   -- foldOverNaperianShapeComp : {shape : List Cont} ->
 450 |   --   (allNap : All IsNaperian shape) =>
 451 |   --   (Tensor shape).Shp
 452 |   -- foldOverNaperianShapeComp {shape = []} = ()
 453 |   -- foldOverNaperianShapeComp {allNap = ((MkIsNaperian pos) :: ns)}
 454 |   --   = () <| \_ => foldOverNaperianShapeComp
 455 |   -- 
 456 |   -- public export
 457 |   -- naperianHancockShape : {shape : List Cont} ->
 458 |   --   (allNap : All IsNaperian shape) =>
 459 |   --   (HancockTensor shape).Shp = Unit
 460 |   -- naperianHancockShape = believe_me ()
 461 |   -- 
 462 |   -- public export
 463 |   -- foldOverNaperianShapeHancock : {shape : List Cont} ->
 464 |   --   (allNap : All IsNaperian shape) =>
 465 |   --   (HancockTensor shape).Shp
 466 |   -- foldOverNaperianShapeHancock {shape = []} = ()
 467 |   -- foldOverNaperianShapeHancock {allNap = ((MkIsNaperian _) :: _)}
 468 |   --   = ((), foldOverNaperianShapeHancock)
 469 | 
 470 | 
 471 | -- public export
 472 | -- tensorIsNaperianShape : {shape : List Cont} ->
 473 | --   (allNap : All IsNaperian shape) =>
 474 | --   IsNaperian (Tensor shape)
 475 | -- tensorIsNaperianShape {shape = []} = MkIsNaperian ()
 476 | -- tensorIsNaperianShape {shape = (_ :: ss), allNap = ((MkIsNaperian pos) :: ns)}
 477 | --   = let tg = tensorIsNaperianShape {shape = ss} 
 478 | --     in ?tensorIsNaperianShape_rhs_1
 479 | --     --in rewrite naperianShpEq @{tg}
 480 | --     --in (rewrite (EmptyExtEq {c=(Nap pos)})
 481 | --     --in let tg = MkIsNaperian in ?tensorIsNaperianShape_rhs_2)
 482 | 
 483 | -- public export
 484 | -- transformToHancock : {shape : List Cont} ->
 485 | --   All IsNaperian shape =>
 486 | --   Tensor shape =%> HancockTensor shape
 487 | -- transformToHancock {shape = []} = id
 488 | -- transformToHancock {shape = (_ :: ss)} @{((MkIsNaperian pos) :: ns)}
 489 | --   = let f = (%!) (transformToHancock {shape = ss} @{ns})
 490 | --         (_ ** h) = f (foldOverNaperianShapeComp {shape=ss})
 491 | --     in !% \(() <| content) => (((), foldOverNaperianShapeHancock) **
 492 | --       \(p, fld) => (p ** ?hhh))
 493 | --       -- (((), rewrite -- foldOverNaperianShapeHancock {shape=ss} @{ns} in ()) **
 494 | --     --   \(p, fld) => (p ** ?bnn))
 495 | 
 496 | -- need to organise this
 497 | namespace BinTree
 498 |   public export
 499 |   inorderBackward : (b : BinTreeShape) ->
 500 |     Fin (numNodesAndLeaves b) ->
 501 |     BinTreePos b
 502 |   inorderBackward LeafS FZ = AtLeaf
 503 |   inorderBackward (NodeS lt rt) n with (strengthenN {m=numNodesAndLeaves lt} n)
 504 |      _ | Left p = GoLeft (inorderBackward lt p)
 505 |      _ | Right FZ = AtNode
 506 |      _ | Right (FS g) = GoRight (inorderBackward rt g)
 507 | 
 508 | 
 509 |   public export
 510 |   inorder : BinTree =%> List
 511 |   inorder = !% \b => (numNodesAndLeaves b ** inorderBackward b)
 512 | 
 513 | namespace BinTreeNode
 514 |   public export
 515 |   inorderBackward : (b : BinTreeShape) ->
 516 |     Fin (numNodes b) ->
 517 |     BinTreePosNode b
 518 |   inorderBackward (NodeS lt rt) n with (strengthenN {m=numNodes lt} n)
 519 |     _ | Left p = GoLeft (inorderBackward lt p)
 520 |     _ | Right FZ = AtNode
 521 |     _ | Right (FS g) = GoRight (inorderBackward rt g)
 522 | 
 523 |   ||| Traverses a binary tree container in order, producing a list container
 524 |   public export
 525 |   inorder : BinTreeNode =%> List
 526 |   inorder = !% \b => (numNodes b ** inorderBackward b)
 527 | 
 528 |   -- Need to do some rewriting for preorder
 529 |   public export
 530 |   preorderBinTreeNode : (b : BinTreeShape) ->
 531 |     Fin (numNodes b) -> BinTreePosNode b
 532 |   preorderBinTreeNode (NodeS lt rt) x = ?preorderBinTreeNode_rhs_1
 533 |   --preorderBinTreeNode (NodeS lt rt) n with (strengthenN {m=numNodes lt} n)
 534 |   --  _ | Left p = ?whl
 535 |   --  _ | Right FZ = ?whn
 536 |   --  _ | Right (FS g) = ?whr
 537 | 
 538 | namespace BinTreeLeaf
 539 |   public export
 540 |   inorderBackward : (b : BinTreeShape) ->
 541 |     Fin (numLeaves b) ->
 542 |     BinTreePosLeaf b
 543 |   inorderBackward LeafS 0 = AtLeaf
 544 |   inorderBackward (NodeS lt rt) i with (strengthenN {m=numLeaves lt} i)
 545 |     _ | (Left indLeft) = GoLeft (inorderBackward lt indLeft)
 546 |     _ | (Right indRight) = GoRight (inorderBackward rt indRight)
 547 | 
 548 |   public export
 549 |   inorder : BinTreeLeaf =%> List
 550 |   inorder = !% \b => (numLeaves b ** inorderBackward b)
 551 | 
 552 | -- public export
 553 | -- traverseLeaf : (x : BinTreeShape) -> FinBinTreeLeaf x -> Fin (numLeaves x)
 554 | -- traverseLeaf LeafS Done = FZ
 555 | -- traverseLeaf (NodeS lt rt) (GoLeft x) = weakenN (numLeaves rt) (traverseLeaf lt x)
 556 | -- traverseLeaf (NodeS lt rt) (GoRight x) = shift (numLeaves lt) (traverseLeaf rt x)
 557 | -- 
 558 | 
 559 | public export
 560 | vectToList : {n : Nat} -> Vect n =%> List
 561 | vectToList = !% \() => (n ** id)
 562 | 
 563 | public export
 564 | maybeToList : Maybe =%> List
 565 | maybeToList = !% \b => case b of 
 566 |   False => (0 ** absurd)
 567 |   True => (1 ** \_ => ())
 568 | 
 569 | public export
 570 | Sample : MonadSample m => {n : Nat} -> IsSucc n =>
 571 |   (m <!> Dist n) =%> Scalar
 572 | Sample = toCostate sample
 573 | 
 574 | -- TODO here maybe need to uncomment during merge?
 575 | -- public export
 576 | -- selectShape : {cs : Vect k Cont} ->
 577 | --   (shapes : All Shp cs) -> (i : Fin k) -> Any Shp cs
 578 | -- selectShape (s :: _) FZ = Here s
 579 | -- selectShape (_ :: ss) (FS j) = There (selectShape ss j)
 580 | -- 
 581 | -- ||| Extract the position from an AnyPos at a given index
 582 | -- public export
 583 | -- extractPos : {n : Nat} -> {xs : Vect n Cont} ->
 584 | --   {shapes : All Shp xs} ->
 585 | --   (i : Fin n) ->
 586 | --   AnyShpPos (selectShape shapes i) ->
 587 | --   AnyPos shapes
 588 | -- extractPos {shapes = (_ :: _)} FZ (Here x) = Here x
 589 | -- extractPos {shapes = (_ :: _)} (FS j) (There rest)
 590 | --   = There $ extractPos j rest
 591 | -- 
 592 | -- public export
 593 | -- SampleAndChoose : {n : Nat} -> {xs : Vect n Cont} ->
 594 | --   ConvexComb xs =%> (Sample n >@ Any xs)
 595 | -- SampleAndChoose = !% \(d, shapes) =>
 596 | --   (d <| selectShape shapes ** \(i ** grad) => (0, [extractPos i grad]))
 597 | 
 598 | -- SampleAndChooseWithDist = !% \(d, shapes) =>
 599 | --   (d <| electShape shapes ** \(i ** grad) => (0, [(i ** extractPos i grad)]))
 600 | 
 601 | -- public export
 602 | -- GetDist : {n : Nat} -> {xs : Vect n Cont} ->
 603 | --   ConvexComb xs =%> Simplex n
 604 | -- GetDist = !% \(d, shapes) => (d ** \d' => (d', ?GetDist_rhs))
 605 | 
 606 | public export
 607 | handleEffect : Monad m =>
 608 |   (handler : (m <!> effect) =%> Scalar) ->
 609 |   (program : a =%> effect) ->
 610 |   m <!> a =%> Scalar
 611 | handleEffect handler program = !% \x =>
 612 |   let (ef ** nn) = (%! program) x
 613 |       (() ** rest) = (%! handler) ef
 614 |   in (() ** \() => do 
 615 |     e <- rest ()
 616 |     pure (nn e))