Tree.hs

{-# 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 = wrapAt

wrapAt :: (RealFrac r) => r -> [a] -> a
wrapAt t ls = ls !! phase
  where
    l = fromIntegral $ length ls
    phase = mod (floor t) (length ls)

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

data Tree a
  = Leaf a
  | Branch [Tree a]
  deriving (Show, Eq, Functor)

instance ToList Tree where
  toList (Leaf a) = [a]
  toList (Branch as) = concatMap toList as

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

empty :: Tree a
empty = Branch []

isEmpty :: Tree a -> Bool
isEmpty (Leaf _) = False
isEmpty (Branch ts) = all isEmpty ts

singleton :: a -> Tree a
singleton = Leaf

fromList :: [a] -> Tree a
fromList as = Branch $ map singleton as

look :: Int -> Tree a -> Tree a
look _ (Leaf x) = Leaf x
look i (Branch xs) = xs !!! fromIntegral i

look' :: Int -> Tree a -> Tree a
look' 0 (Leaf x) = Leaf x
look' _ (Leaf x) = empty
look' i (Branch xs) = if length xs > i && i >= 0 then xs !! i else empty

lookup :: [Int] -> Tree a -> Tree a
lookup is x = foldl (flip look) x is

concatMapTree :: (a -> Tree b) -> Tree a -> Tree b
concatMapTree f x = squeezeJoin $ fmap f x

topLength :: Tree a -> Int
topLength (Leaf _) = 1
topLength (Branch xs) = length xs

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

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

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

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

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

instance MultiApplicative [] where
  liftA2Left f [] _ = []
  liftA2Left f _ [] = []
  liftA2Left f as bs = map (\i -> f (as !! i) (bs !! floor ((fromIntegral i / fromIntegral n) * fromIntegral m))) [0 .. n - 1]
    where
      n = length as
      m = length bs
  liftA2Right f as bs = liftA2Left (flip f) bs as

instance Applicative Tree where
  pure = Leaf
  liftA2 f (Leaf x) (Leaf y) = Leaf $ f x y
  liftA2 f l@(Leaf _) (Branch ys) = Branch $ map (liftA2 f l) ys
  liftA2 f (Branch xs) l@(Leaf _) = Branch $ map (\x -> liftA2 f x l) xs
  liftA2 f (Branch xs) (Branch ys) = Branch $ lift2Both (liftA2 f) xs ys

instance MultiApplicative Tree where
  liftA2Left f (Leaf x) (Leaf y) = Leaf $ f x y
  liftA2Left f l@(Leaf _) (Branch ys) = Branch $ map (liftA2Left f l) ys
  liftA2Left f (Branch xs) l@(Leaf _) = Branch $ map (\x -> liftA2Left f x l) xs
  liftA2Left f (Branch xs) (Branch ys) = Branch $ liftA2Left (liftA2Left f) xs ys
  liftA2Right f (Leaf x) (Leaf y) = Leaf $ f x y
  liftA2Right f l@(Leaf _) (Branch ys) = Branch $ map (liftA2Right f l) ys
  liftA2Right f (Branch xs) l@(Leaf _) = Branch $ map (\x -> liftA2Right f x l) xs
  liftA2Right f (Branch xs) (Branch ys) = Branch $ liftA2Right (liftA2Right f) xs ys

lift2Both :: (a -> b -> c) -> [a] -> [b] -> [c]
lift2Both f as bs =
  if n < m
    then liftA2Right f as bs
    else liftA2Left f as bs
  where
    n = length as
    m = length bs

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

instance MultiMonad [] where
  innerJoin = concat
  outerJoin = concat
  squeezeJoin = concat

instance Monad Tree where
  (>>=) x f = innerJoin $ f <$> x

instance MultiMonad Tree where
  innerJoin t = squeezeJoin $ fmap select indx
    where
      indx = indexTree t
      select (is, x) = reduceNested is x
  outerJoin = innerJoin
  squeezeJoin (Leaf x) = x
  squeezeJoin (Branch xs) = Branch $ map squeezeJoin xs

indexTree :: Tree a -> Tree ([(Int, Int)], a)
indexTree (Leaf x) = Leaf ([], x)
indexTree (Branch bs) = Branch $ map (\i -> first ((i, length bs) :) <$> indexTree (bs !! i)) [0 .. length bs - 1]

reduce :: (Int, Int) -> Tree a -> [Tree a]
reduce (i, _) (Leaf x) = [Leaf x]
reduce (i, n) (Branch xs) = map (\l -> nth (fromIntegral l) xs) ind
  where
    m = length xs
    ind = [j | j <- [0 .. m - 1], fromIntegral i / fromIntegral n <= fromIntegral j / fromIntegral m, fromIntegral j / fromIntegral m < fromIntegral (i + 1) / fromIntegral n]

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