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