{-# LANGUAGE DeriveFunctor #-}
{-# LANGUAGE FlexibleInstances #-}
{-# LANGUAGE GADTs #-}

module Zwirn.Core.Tree where

{-
    Tree.hs - a structure for parallel signals
    Copyright (C) 2025, Martin Gius

    This library is free software: you can redistribute it and/or modify
    it under the terms of the GNU General Public License as published by
    the Free Software Foundation, either version 3 of the License, or
    (at your option) any later version.

    This library is distributed in the hope that it will be useful,
    but WITHOUT ANY WARRANTY; without even the implied warranty of
    MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
    GNU General Public License for more details.

    You should have received a copy of the GNU General Public License
    along with this library.  If not, see <http://www.gnu.org/licenses/>.
-}

import Data.Bifunctor
import Data.Fixed (mod')
import Zwirn.Core.Types

nth :: (RealFrac r) => r -> [a] -> a
nth :: forall r a. RealFrac r => r -> [a] -> a
nth = r -> [a] -> a
forall r a. RealFrac r => r -> [a] -> a
wrapAt

wrapAt :: (RealFrac r) => r -> [a] -> a
wrapAt :: forall r a. RealFrac r => r -> [a] -> a
wrapAt r
t [a]
ls = [a]
ls [a] -> Int -> a
forall a. HasCallStack => [a] -> Int -> a
!! Int
phase
  where
    l :: Integer
l = Int -> Integer
forall a b. (Integral a, Num b) => a -> b
fromIntegral (Int -> Integer) -> Int -> Integer
forall a b. (a -> b) -> a -> b
$ [a] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [a]
ls
    phase :: Int
phase = Int -> Int -> Int
forall a. Integral a => a -> a -> a
mod (r -> Int
forall b. Integral b => r -> b
forall a b. (RealFrac a, Integral b) => a -> b
floor r
t) ([a] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [a]
ls)

frac :: (Real r) => r -> r
frac :: forall r. Real r => r -> r
frac r
d = r -> r -> r
forall a. Real a => a -> a -> a
mod' r
d r
1

data Tree a
  = Leaf a
  | Branch [Tree a]
  deriving (Int -> Tree a -> ShowS
[Tree a] -> ShowS
Tree a -> String
(Int -> Tree a -> ShowS)
-> (Tree a -> String) -> ([Tree a] -> ShowS) -> Show (Tree a)
forall a. Show a => Int -> Tree a -> ShowS
forall a. Show a => [Tree a] -> ShowS
forall a. Show a => Tree a -> String
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
$cshowsPrec :: forall a. Show a => Int -> Tree a -> ShowS
showsPrec :: Int -> Tree a -> ShowS
$cshow :: forall a. Show a => Tree a -> String
show :: Tree a -> String
$cshowList :: forall a. Show a => [Tree a] -> ShowS
showList :: [Tree a] -> ShowS
Show, Tree a -> Tree a -> Bool
(Tree a -> Tree a -> Bool)
-> (Tree a -> Tree a -> Bool) -> Eq (Tree a)
forall a. Eq a => Tree a -> Tree a -> Bool
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
$c== :: forall a. Eq a => Tree a -> Tree a -> Bool
== :: Tree a -> Tree a -> Bool
$c/= :: forall a. Eq a => Tree a -> Tree a -> Bool
/= :: Tree a -> Tree a -> Bool
Eq, (forall a b. (a -> b) -> Tree a -> Tree b)
-> (forall a b. a -> Tree b -> Tree a) -> Functor Tree
forall a b. a -> Tree b -> Tree a
forall a b. (a -> b) -> Tree a -> Tree b
forall (f :: * -> *).
(forall a b. (a -> b) -> f a -> f b)
-> (forall a b. a -> f b -> f a) -> Functor f
$cfmap :: forall a b. (a -> b) -> Tree a -> Tree b
fmap :: forall a b. (a -> b) -> Tree a -> Tree b
$c<$ :: forall a b. a -> Tree b -> Tree a
<$ :: forall a b. a -> Tree b -> Tree a
Functor)

instance ToList Tree where
  toList :: forall a. Tree a -> [a]
toList (Leaf a
a) = [a
a]
  toList (Branch [Tree a]
as) = (Tree a -> [a]) -> [Tree a] -> [a]
forall (t :: * -> *) a b. Foldable t => (a -> [b]) -> t a -> [b]
concatMap Tree a -> [a]
forall a. Tree a -> [a]
forall (k :: * -> *) a. ToList k => k a -> [a]
toList [Tree a]
as

(!!!) :: (Num b, RealFrac b) => [a] -> b -> a
!!! :: forall b a. (Num b, RealFrac b) => [a] -> b -> a
(!!!) [a]
as b
r = b -> [a] -> a
forall r a. RealFrac r => r -> [a] -> a
nth b
r [a]
as

empty :: Tree a
empty :: forall a. Tree a
empty = [Tree a] -> Tree a
forall a. [Tree a] -> Tree a
Branch []

isEmpty :: Tree a -> Bool
isEmpty :: forall a. Tree a -> Bool
isEmpty (Leaf a
_) = Bool
False
isEmpty (Branch [Tree a]
ts) = (Tree a -> Bool) -> [Tree a] -> Bool
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
all Tree a -> Bool
forall a. Tree a -> Bool
isEmpty [Tree a]
ts

singleton :: a -> Tree a
singleton :: forall a. a -> Tree a
singleton = a -> Tree a
forall a. a -> Tree a
Leaf

fromList :: [a] -> Tree a
fromList :: forall a. [a] -> Tree a
fromList [a]
as = [Tree a] -> Tree a
forall a. [Tree a] -> Tree a
Branch ([Tree a] -> Tree a) -> [Tree a] -> Tree a
forall a b. (a -> b) -> a -> b
$ (a -> Tree a) -> [a] -> [Tree a]
forall a b. (a -> b) -> [a] -> [b]
map a -> Tree a
forall a. a -> Tree a
singleton [a]
as

look :: Int -> Tree a -> Tree a
look :: forall a. Int -> Tree a -> Tree a
look Int
_ (Leaf a
x) = a -> Tree a
forall a. a -> Tree a
Leaf a
x
look Int
i (Branch [Tree a]
xs) = [Tree a]
xs [Tree a] -> Double -> Tree a
forall b a. (Num b, RealFrac b) => [a] -> b -> a
!!! Int -> Double
forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
i

look' :: Int -> Tree a -> Tree a
look' :: forall a. Int -> Tree a -> Tree a
look' Int
0 (Leaf a
x) = a -> Tree a
forall a. a -> Tree a
Leaf a
x
look' Int
_ (Leaf a
x) = Tree a
forall a. Tree a
empty
look' Int
i (Branch [Tree a]
xs) = if [Tree a] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [Tree a]
xs Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
> Int
i Bool -> Bool -> Bool
&& Int
i Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= Int
0 then [Tree a]
xs [Tree a] -> Int -> Tree a
forall a. HasCallStack => [a] -> Int -> a
!! Int
i else Tree a
forall a. Tree a
empty

lookup :: [Int] -> Tree a -> Tree a
lookup :: forall a. [Int] -> Tree a -> Tree a
lookup [Int]
is Tree a
x = (Tree a -> Int -> Tree a) -> Tree a -> [Int] -> Tree a
forall b a. (b -> a -> b) -> b -> [a] -> b
forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl ((Int -> Tree a -> Tree a) -> Tree a -> Int -> Tree a
forall a b c. (a -> b -> c) -> b -> a -> c
flip Int -> Tree a -> Tree a
forall a. Int -> Tree a -> Tree a
look) Tree a
x [Int]
is

concatMapTree :: (a -> Tree b) -> Tree a -> Tree b
concatMapTree :: forall a b. (a -> Tree b) -> Tree a -> Tree b
concatMapTree a -> Tree b
f Tree a
x = Tree (Tree b) -> Tree b
forall a. Tree (Tree a) -> Tree a
forall (m :: * -> *) a. MultiMonad m => m (m a) -> m a
squeezeJoin (Tree (Tree b) -> Tree b) -> Tree (Tree b) -> Tree b
forall a b. (a -> b) -> a -> b
$ (a -> Tree b) -> Tree a -> Tree (Tree b)
forall a b. (a -> b) -> Tree a -> Tree b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap a -> Tree b
f Tree a
x

topLength :: Tree a -> Int
topLength :: forall a. Tree a -> Int
topLength (Leaf a
_) = Int
1
topLength (Branch [Tree a]
xs) = [Tree a] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [Tree a]
xs

push :: Tree a -> Tree a -> Tree a
push :: forall a. Tree a -> Tree a -> Tree a
push Tree a
x l :: Tree a
l@(Leaf a
_) = [Tree a] -> Tree a
forall a. [Tree a] -> Tree a
Branch [Tree a
x, Tree a
l]
push Tree a
x (Branch [Tree a]
xs) = [Tree a] -> Tree a
forall a. [Tree a] -> Tree a
Branch (Tree a
x Tree a -> [Tree a] -> [Tree a]
forall a. a -> [a] -> [a]
: [Tree a]
xs)

pop :: Tree a -> Tree a
pop :: forall a. Tree a -> Tree a
pop (Leaf a
_) = Tree a
forall a. Tree a
empty
pop (Branch []) = Tree a
forall a. Tree a
empty
pop (Branch (Tree a
x : [Tree a]
xs)) = [Tree a] -> Tree a
forall a. [Tree a] -> Tree a
Branch [Tree a]
xs

insertT :: Int -> Tree a -> Tree a -> Tree a
insertT :: forall a. Int -> Tree a -> Tree a -> Tree a
insertT Int
0 Tree a
x (Leaf a
y) = [Tree a] -> Tree a
forall a. [Tree a] -> Tree a
Branch [Tree a
x, a -> Tree a
forall a. a -> Tree a
Leaf a
y]
insertT Int
_ Tree a
x (Leaf a
y) = [Tree a] -> Tree a
forall a. [Tree a] -> Tree a
Branch [a -> Tree a
forall a. a -> Tree a
Leaf a
y, Tree a
x]
insertT Int
i Tree a
x (Branch [Tree a]
ys) = [Tree a] -> Tree a
forall a. [Tree a] -> Tree a
Branch ([Tree a]
ys1 [Tree a] -> [Tree a] -> [Tree a]
forall a. [a] -> [a] -> [a]
++ [Tree a
x] [Tree a] -> [Tree a] -> [Tree a]
forall a. [a] -> [a] -> [a]
++ [Tree a]
ys2)
  where
    ([Tree a]
ys1, [Tree a]
ys2) = Int -> [Tree a] -> ([Tree a], [Tree a])
forall a. Int -> [a] -> ([a], [a])
splitAt Int
i [Tree a]
ys

removeT :: Int -> Tree a -> Tree a
removeT :: forall a. Int -> Tree a -> Tree a
removeT Int
_ (Leaf a
x) = Tree a
forall a. Tree a
empty
removeT Int
i (Branch [Tree a]
xs) = case Int -> [Tree a] -> ([Tree a], [Tree a])
forall a. Int -> [a] -> ([a], [a])
splitAt Int
i [Tree a]
xs of
  ([Tree a]
xs1, []) -> [Tree a] -> Tree a
forall a. [Tree a] -> Tree a
Branch [Tree a]
xs1
  ([Tree a]
xs1, Tree a
_ : [Tree a]
xs2) -> [Tree a] -> Tree a
forall a. [Tree a] -> Tree a
Branch ([Tree a] -> Tree a) -> [Tree a] -> Tree a
forall a b. (a -> b) -> a -> b
$ [Tree a]
xs1 [Tree a] -> [Tree a] -> [Tree a]
forall a. [a] -> [a] -> [a]
++ [Tree a]
xs2

-------------------------------------------------------
------------------- APPLICATIVE STUFF -----------------
-------------------------------------------------------

instance MultiApplicative [] where
  liftA2Left :: forall a b c. (a -> b -> c) -> [a] -> [b] -> [c]
liftA2Left a -> b -> c
f [] [b]
_ = []
  liftA2Left a -> b -> c
f [a]
_ [] = []
  liftA2Left a -> b -> c
f [a]
as [b]
bs = (Int -> c) -> [Int] -> [c]
forall a b. (a -> b) -> [a] -> [b]
map (\Int
i -> a -> b -> c
f ([a]
as [a] -> Int -> a
forall a. HasCallStack => [a] -> Int -> a
!! Int
i) ([b]
bs [b] -> Int -> b
forall a. HasCallStack => [a] -> Int -> a
!! Double -> Int
forall b. Integral b => Double -> b
forall a b. (RealFrac a, Integral b) => a -> b
floor ((Int -> Double
forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
i Double -> Double -> Double
forall a. Fractional a => a -> a -> a
/ Int -> Double
forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
n) Double -> Double -> Double
forall a. Num a => a -> a -> a
* Int -> Double
forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
m))) [Int
0 .. Int
n Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1]
    where
      n :: Int
n = [a] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [a]
as
      m :: Int
m = [b] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [b]
bs
  liftA2Right :: forall a b c. (a -> b -> c) -> [a] -> [b] -> [c]
liftA2Right a -> b -> c
f [a]
as [b]
bs = (b -> a -> c) -> [b] -> [a] -> [c]
forall a b c. (a -> b -> c) -> [a] -> [b] -> [c]
forall (f :: * -> *) a b c.
MultiApplicative f =>
(a -> b -> c) -> f a -> f b -> f c
liftA2Left ((a -> b -> c) -> b -> a -> c
forall a b c. (a -> b -> c) -> b -> a -> c
flip a -> b -> c
f) [b]
bs [a]
as

instance Applicative Tree where
  pure :: forall a. a -> Tree a
pure = a -> Tree a
forall a. a -> Tree a
Leaf
  liftA2 :: forall a b c. (a -> b -> c) -> Tree a -> Tree b -> Tree c
liftA2 a -> b -> c
f (Leaf a
x) (Leaf b
y) = c -> Tree c
forall a. a -> Tree a
Leaf (c -> Tree c) -> c -> Tree c
forall a b. (a -> b) -> a -> b
$ a -> b -> c
f a
x b
y
  liftA2 a -> b -> c
f l :: Tree a
l@(Leaf a
_) (Branch [Tree b]
ys) = [Tree c] -> Tree c
forall a. [Tree a] -> Tree a
Branch ([Tree c] -> Tree c) -> [Tree c] -> Tree c
forall a b. (a -> b) -> a -> b
$ (Tree b -> Tree c) -> [Tree b] -> [Tree c]
forall a b. (a -> b) -> [a] -> [b]
map ((a -> b -> c) -> Tree a -> Tree b -> Tree c
forall a b c. (a -> b -> c) -> Tree a -> Tree b -> Tree c
forall (f :: * -> *) a b c.
Applicative f =>
(a -> b -> c) -> f a -> f b -> f c
liftA2 a -> b -> c
f Tree a
l) [Tree b]
ys
  liftA2 a -> b -> c
f (Branch [Tree a]
xs) l :: Tree b
l@(Leaf b
_) = [Tree c] -> Tree c
forall a. [Tree a] -> Tree a
Branch ([Tree c] -> Tree c) -> [Tree c] -> Tree c
forall a b. (a -> b) -> a -> b
$ (Tree a -> Tree c) -> [Tree a] -> [Tree c]
forall a b. (a -> b) -> [a] -> [b]
map (\Tree a
x -> (a -> b -> c) -> Tree a -> Tree b -> Tree c
forall a b c. (a -> b -> c) -> Tree a -> Tree b -> Tree c
forall (f :: * -> *) a b c.
Applicative f =>
(a -> b -> c) -> f a -> f b -> f c
liftA2 a -> b -> c
f Tree a
x Tree b
l) [Tree a]
xs
  liftA2 a -> b -> c
f (Branch [Tree a]
xs) (Branch [Tree b]
ys) = [Tree c] -> Tree c
forall a. [Tree a] -> Tree a
Branch ([Tree c] -> Tree c) -> [Tree c] -> Tree c
forall a b. (a -> b) -> a -> b
$ (Tree a -> Tree b -> Tree c) -> [Tree a] -> [Tree b] -> [Tree c]
forall a b c. (a -> b -> c) -> [a] -> [b] -> [c]
lift2Both ((a -> b -> c) -> Tree a -> Tree b -> Tree c
forall a b c. (a -> b -> c) -> Tree a -> Tree b -> Tree c
forall (f :: * -> *) a b c.
Applicative f =>
(a -> b -> c) -> f a -> f b -> f c
liftA2 a -> b -> c
f) [Tree a]
xs [Tree b]
ys

instance MultiApplicative Tree where
  liftA2Left :: forall a b c. (a -> b -> c) -> Tree a -> Tree b -> Tree c
liftA2Left a -> b -> c
f (Leaf a
x) (Leaf b
y) = c -> Tree c
forall a. a -> Tree a
Leaf (c -> Tree c) -> c -> Tree c
forall a b. (a -> b) -> a -> b
$ a -> b -> c
f a
x b
y
  liftA2Left a -> b -> c
f l :: Tree a
l@(Leaf a
_) (Branch [Tree b]
ys) = [Tree c] -> Tree c
forall a. [Tree a] -> Tree a
Branch ([Tree c] -> Tree c) -> [Tree c] -> Tree c
forall a b. (a -> b) -> a -> b
$ (Tree b -> Tree c) -> [Tree b] -> [Tree c]
forall a b. (a -> b) -> [a] -> [b]
map ((a -> b -> c) -> Tree a -> Tree b -> Tree c
forall a b c. (a -> b -> c) -> Tree a -> Tree b -> Tree c
forall (f :: * -> *) a b c.
MultiApplicative f =>
(a -> b -> c) -> f a -> f b -> f c
liftA2Left a -> b -> c
f Tree a
l) [Tree b]
ys
  liftA2Left a -> b -> c
f (Branch [Tree a]
xs) l :: Tree b
l@(Leaf b
_) = [Tree c] -> Tree c
forall a. [Tree a] -> Tree a
Branch ([Tree c] -> Tree c) -> [Tree c] -> Tree c
forall a b. (a -> b) -> a -> b
$ (Tree a -> Tree c) -> [Tree a] -> [Tree c]
forall a b. (a -> b) -> [a] -> [b]
map (\Tree a
x -> (a -> b -> c) -> Tree a -> Tree b -> Tree c
forall a b c. (a -> b -> c) -> Tree a -> Tree b -> Tree c
forall (f :: * -> *) a b c.
MultiApplicative f =>
(a -> b -> c) -> f a -> f b -> f c
liftA2Left a -> b -> c
f Tree a
x Tree b
l) [Tree a]
xs
  liftA2Left a -> b -> c
f (Branch [Tree a]
xs) (Branch [Tree b]
ys) = [Tree c] -> Tree c
forall a. [Tree a] -> Tree a
Branch ([Tree c] -> Tree c) -> [Tree c] -> Tree c
forall a b. (a -> b) -> a -> b
$ (Tree a -> Tree b -> Tree c) -> [Tree a] -> [Tree b] -> [Tree c]
forall a b c. (a -> b -> c) -> [a] -> [b] -> [c]
forall (f :: * -> *) a b c.
MultiApplicative f =>
(a -> b -> c) -> f a -> f b -> f c
liftA2Left ((a -> b -> c) -> Tree a -> Tree b -> Tree c
forall a b c. (a -> b -> c) -> Tree a -> Tree b -> Tree c
forall (f :: * -> *) a b c.
MultiApplicative f =>
(a -> b -> c) -> f a -> f b -> f c
liftA2Left a -> b -> c
f) [Tree a]
xs [Tree b]
ys
  liftA2Right :: forall a b c. (a -> b -> c) -> Tree a -> Tree b -> Tree c
liftA2Right a -> b -> c
f (Leaf a
x) (Leaf b
y) = c -> Tree c
forall a. a -> Tree a
Leaf (c -> Tree c) -> c -> Tree c
forall a b. (a -> b) -> a -> b
$ a -> b -> c
f a
x b
y
  liftA2Right a -> b -> c
f l :: Tree a
l@(Leaf a
_) (Branch [Tree b]
ys) = [Tree c] -> Tree c
forall a. [Tree a] -> Tree a
Branch ([Tree c] -> Tree c) -> [Tree c] -> Tree c
forall a b. (a -> b) -> a -> b
$ (Tree b -> Tree c) -> [Tree b] -> [Tree c]
forall a b. (a -> b) -> [a] -> [b]
map ((a -> b -> c) -> Tree a -> Tree b -> Tree c
forall a b c. (a -> b -> c) -> Tree a -> Tree b -> Tree c
forall (f :: * -> *) a b c.
MultiApplicative f =>
(a -> b -> c) -> f a -> f b -> f c
liftA2Right a -> b -> c
f Tree a
l) [Tree b]
ys
  liftA2Right a -> b -> c
f (Branch [Tree a]
xs) l :: Tree b
l@(Leaf b
_) = [Tree c] -> Tree c
forall a. [Tree a] -> Tree a
Branch ([Tree c] -> Tree c) -> [Tree c] -> Tree c
forall a b. (a -> b) -> a -> b
$ (Tree a -> Tree c) -> [Tree a] -> [Tree c]
forall a b. (a -> b) -> [a] -> [b]
map (\Tree a
x -> (a -> b -> c) -> Tree a -> Tree b -> Tree c
forall a b c. (a -> b -> c) -> Tree a -> Tree b -> Tree c
forall (f :: * -> *) a b c.
MultiApplicative f =>
(a -> b -> c) -> f a -> f b -> f c
liftA2Right a -> b -> c
f Tree a
x Tree b
l) [Tree a]
xs
  liftA2Right a -> b -> c
f (Branch [Tree a]
xs) (Branch [Tree b]
ys) = [Tree c] -> Tree c
forall a. [Tree a] -> Tree a
Branch ([Tree c] -> Tree c) -> [Tree c] -> Tree c
forall a b. (a -> b) -> a -> b
$ (Tree a -> Tree b -> Tree c) -> [Tree a] -> [Tree b] -> [Tree c]
forall a b c. (a -> b -> c) -> [a] -> [b] -> [c]
forall (f :: * -> *) a b c.
MultiApplicative f =>
(a -> b -> c) -> f a -> f b -> f c
liftA2Right ((a -> b -> c) -> Tree a -> Tree b -> Tree c
forall a b c. (a -> b -> c) -> Tree a -> Tree b -> Tree c
forall (f :: * -> *) a b c.
MultiApplicative f =>
(a -> b -> c) -> f a -> f b -> f c
liftA2Right a -> b -> c
f) [Tree a]
xs [Tree b]
ys

lift2Both :: (a -> b -> c) -> [a] -> [b] -> [c]
lift2Both :: forall a b c. (a -> b -> c) -> [a] -> [b] -> [c]
lift2Both a -> b -> c
f [a]
as [b]
bs =
  if Int
n Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
m
    then (a -> b -> c) -> [a] -> [b] -> [c]
forall a b c. (a -> b -> c) -> [a] -> [b] -> [c]
forall (f :: * -> *) a b c.
MultiApplicative f =>
(a -> b -> c) -> f a -> f b -> f c
liftA2Right a -> b -> c
f [a]
as [b]
bs
    else (a -> b -> c) -> [a] -> [b] -> [c]
forall a b c. (a -> b -> c) -> [a] -> [b] -> [c]
forall (f :: * -> *) a b c.
MultiApplicative f =>
(a -> b -> c) -> f a -> f b -> f c
liftA2Left a -> b -> c
f [a]
as [b]
bs
  where
    n :: Int
n = [a] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [a]
as
    m :: Int
m = [b] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [b]
bs

--------------------------------------------------
------------------- MONAD STUFF ------------------
--------------------------------------------------

instance MultiMonad [] where
  innerJoin :: forall a. [[a]] -> [a]
innerJoin = [[a]] -> [a]
forall (t :: * -> *) a. Foldable t => t [a] -> [a]
concat
  outerJoin :: forall a. [[a]] -> [a]
outerJoin = [[a]] -> [a]
forall (t :: * -> *) a. Foldable t => t [a] -> [a]
concat
  squeezeJoin :: forall a. [[a]] -> [a]
squeezeJoin = [[a]] -> [a]
forall (t :: * -> *) a. Foldable t => t [a] -> [a]
concat

instance Monad Tree where
  >>= :: forall a b. Tree a -> (a -> Tree b) -> Tree b
(>>=) Tree a
x a -> Tree b
f = Tree (Tree b) -> Tree b
forall a. Tree (Tree a) -> Tree a
forall (m :: * -> *) a. MultiMonad m => m (m a) -> m a
innerJoin (Tree (Tree b) -> Tree b) -> Tree (Tree b) -> Tree b
forall a b. (a -> b) -> a -> b
$ a -> Tree b
f (a -> Tree b) -> Tree a -> Tree (Tree b)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Tree a
x

instance MultiMonad Tree where
  innerJoin :: forall a. Tree (Tree a) -> Tree a
innerJoin Tree (Tree a)
t = Tree (Tree a) -> Tree a
forall a. Tree (Tree a) -> Tree a
forall (m :: * -> *) a. MultiMonad m => m (m a) -> m a
squeezeJoin (Tree (Tree a) -> Tree a) -> Tree (Tree a) -> Tree a
forall a b. (a -> b) -> a -> b
$ (([(Int, Int)], Tree a) -> Tree a)
-> Tree ([(Int, Int)], Tree a) -> Tree (Tree a)
forall a b. (a -> b) -> Tree a -> Tree b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap ([(Int, Int)], Tree a) -> Tree a
forall {a}. ([(Int, Int)], Tree a) -> Tree a
select Tree ([(Int, Int)], Tree a)
indx
    where
      indx :: Tree ([(Int, Int)], Tree a)
indx = Tree (Tree a) -> Tree ([(Int, Int)], Tree a)
forall a. Tree a -> Tree ([(Int, Int)], a)
indexTree Tree (Tree a)
t
      select :: ([(Int, Int)], Tree a) -> Tree a
select ([(Int, Int)]
is, Tree a
x) = [(Int, Int)] -> Tree a -> Tree a
forall a. [(Int, Int)] -> Tree a -> Tree a
reduceNested [(Int, Int)]
is Tree a
x
  outerJoin :: forall a. Tree (Tree a) -> Tree a
outerJoin = Tree (Tree a) -> Tree a
forall a. Tree (Tree a) -> Tree a
forall (m :: * -> *) a. MultiMonad m => m (m a) -> m a
innerJoin
  squeezeJoin :: forall a. Tree (Tree a) -> Tree a
squeezeJoin (Leaf Tree a
x) = Tree a
x
  squeezeJoin (Branch [Tree (Tree a)]
xs) = [Tree a] -> Tree a
forall a. [Tree a] -> Tree a
Branch ([Tree a] -> Tree a) -> [Tree a] -> Tree a
forall a b. (a -> b) -> a -> b
$ (Tree (Tree a) -> Tree a) -> [Tree (Tree a)] -> [Tree a]
forall a b. (a -> b) -> [a] -> [b]
map Tree (Tree a) -> Tree a
forall a. Tree (Tree a) -> Tree a
forall (m :: * -> *) a. MultiMonad m => m (m a) -> m a
squeezeJoin [Tree (Tree a)]
xs

indexTree :: Tree a -> Tree ([(Int, Int)], a)
indexTree :: forall a. Tree a -> Tree ([(Int, Int)], a)
indexTree (Leaf a
x) = ([(Int, Int)], a) -> Tree ([(Int, Int)], a)
forall a. a -> Tree a
Leaf ([], a
x)
indexTree (Branch [Tree a]
bs) = [Tree ([(Int, Int)], a)] -> Tree ([(Int, Int)], a)
forall a. [Tree a] -> Tree a
Branch ([Tree ([(Int, Int)], a)] -> Tree ([(Int, Int)], a))
-> [Tree ([(Int, Int)], a)] -> Tree ([(Int, Int)], a)
forall a b. (a -> b) -> a -> b
$ (Int -> Tree ([(Int, Int)], a))
-> [Int] -> [Tree ([(Int, Int)], a)]
forall a b. (a -> b) -> [a] -> [b]
map (\Int
i -> ([(Int, Int)] -> [(Int, Int)])
-> ([(Int, Int)], a) -> ([(Int, Int)], a)
forall a b c. (a -> b) -> (a, c) -> (b, c)
forall (p :: * -> * -> *) a b c.
Bifunctor p =>
(a -> b) -> p a c -> p b c
first ((Int
i, [Tree a] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [Tree a]
bs) (Int, Int) -> [(Int, Int)] -> [(Int, Int)]
forall a. a -> [a] -> [a]
:) (([(Int, Int)], a) -> ([(Int, Int)], a))
-> Tree ([(Int, Int)], a) -> Tree ([(Int, Int)], a)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Tree a -> Tree ([(Int, Int)], a)
forall a. Tree a -> Tree ([(Int, Int)], a)
indexTree ([Tree a]
bs [Tree a] -> Int -> Tree a
forall a. HasCallStack => [a] -> Int -> a
!! Int
i)) [Int
0 .. [Tree a] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [Tree a]
bs Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1]

reduce :: (Int, Int) -> Tree a -> [Tree a]
reduce :: forall a. (Int, Int) -> Tree a -> [Tree a]
reduce (Int
i, Int
_) (Leaf a
x) = [a -> Tree a
forall a. a -> Tree a
Leaf a
x]
reduce (Int
i, Int
n) (Branch [Tree a]
xs) = (Int -> Tree a) -> [Int] -> [Tree a]
forall a b. (a -> b) -> [a] -> [b]
map (\Int
l -> Double -> [Tree a] -> Tree a
forall r a. RealFrac r => r -> [a] -> a
nth (Int -> Double
forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
l) [Tree a]
xs) [Int]
ind
  where
    m :: Int
m = [Tree a] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [Tree a]
xs
    ind :: [Int]
ind = [Int
j | Int
j <- [Int
0 .. Int
m Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1], Int -> Double
forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
i Double -> Double -> Double
forall a. Fractional a => a -> a -> a
/ Int -> Double
forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
n Double -> Double -> Bool
forall a. Ord a => a -> a -> Bool
<= Int -> Double
forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
j Double -> Double -> Double
forall a. Fractional a => a -> a -> a
/ Int -> Double
forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
m, Int -> Double
forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
j Double -> Double -> Double
forall a. Fractional a => a -> a -> a
/ Int -> Double
forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
m Double -> Double -> Bool
forall a. Ord a => a -> a -> Bool
< Int -> Double
forall a b. (Integral a, Num b) => a -> b
fromIntegral (Int
i Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1) Double -> Double -> Double
forall a. Fractional a => a -> a -> a
/ Int -> Double
forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
n]

reduceNested :: [(Int, Int)] -> Tree a -> Tree a
reduceNested :: forall a. [(Int, Int)] -> Tree a -> Tree a
reduceNested [] Tree a
x = Tree a
x
reduceNested ((Int, Int)
i : [(Int, Int)]
is) Tree a
x = [Tree a] -> Tree a
forall a. [Tree a] -> Tree a
Branch ([Tree a] -> Tree a) -> [Tree a] -> Tree a
forall a b. (a -> b) -> a -> b
$ (Tree a -> Tree a) -> [Tree a] -> [Tree a]
forall a b. (a -> b) -> [a] -> [b]
map ([(Int, Int)] -> Tree a -> Tree a
forall a. [(Int, Int)] -> Tree a -> Tree a
reduceNested [(Int, Int)]
is) ((Int, Int) -> Tree a -> [Tree a]
forall a. (Int, Int) -> Tree a -> [Tree a]
reduce (Int, Int)
i Tree a
x)