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