Tree.hs

{-# LANGUAGE DeriveFunctor #-}
{-# LANGUAGE FlexibleInstances #-}
{-# LANGUAGE GADTs #-}
{-# LANGUAGE TypeApplications #-}
{-# OPTIONS_GHC -Wno-orphans #-}
{-# OPTIONS_GHC -Wno-type-defaults #-}

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.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
    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

(!!!) :: (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 _ (Branch []) = empty
look i (Branch xs) = xs !!! fromIntegral i

look' :: Int -> Tree a -> Tree a
look' 0 (Leaf x) = Leaf x
look' _ (Leaf _) = 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 [l, x]
push x (Branch xs) = Branch (xs ++ [x])

concat :: Tree a -> Tree a -> Tree a
concat l1@(Leaf _) l2@(Leaf _) = Branch [l1, l2]
concat l@(Leaf _) (Branch xs) = Branch (l : xs)
concat (Branch xs) l@(Leaf _) = Branch (xs ++ [l])
concat (Branch xs) (Branch ys) = Branch (xs ++ ys)

pop :: Tree a -> Tree a
pop (Leaf _) = empty
pop (Branch []) = empty
pop (Branch xs) = Branch (take (length xs - 1) 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 _) = empty
removeT i (Branch xs) = case splitAt i xs of
  (xs1, []) -> Branch xs1
  (xs1, _ : xs2) -> Branch $ xs1 ++ xs2

numberedTree :: Tree a -> Tree (Int, a)
numberedTree = snd . numberedTreeRec 0
  where
    numberedTreeRec :: Int -> Tree a -> (Int, Tree (Int, a))
    numberedTreeRec n (Leaf x) = (n + 1, Leaf (n, x))
    numberedTreeRec n (Branch []) = (n, Branch [])
    numberedTreeRec n (Branch (x : xs)) = (fst $ last fs, Branch $ map snd fs)
      where
        l = numberedTreeRec n x
        fs = foldl' folder [l] xs
        folder is t = is ++ [numberedTreeRec m t]
          where
            m = fst $ last is

numberedTreeBool :: Tree Bool -> Tree (Int, Bool)
numberedTreeBool = snd . numberedTreeRec 0
  where
    numberedTreeRec :: Int -> Tree Bool -> (Int, Tree (Int, Bool))
    numberedTreeRec n (Leaf True) = (n + 1, Leaf (n, True))
    numberedTreeRec n (Leaf False) = (n, Leaf (n, False))
    numberedTreeRec n (Branch []) = (n, Branch [])
    numberedTreeRec n (Branch (x : xs)) = (fst $ last fs, Branch $ map snd fs)
      where
        l = numberedTreeRec n x
        fs = foldl' folder [l] xs
        folder is t = is ++ [numberedTreeRec m t]
          where
            m = fst $ last is

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

instance MultiApplicative [] where
  liftA2Left _ [] _ = []
  liftA2Left _ _ [] = []
  liftA2Left f as bs = map (\i -> f (as !! i) (bs !! floor @Double ((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 = Prelude.concat
  outerJoin = Prelude.concat
  squeezeJoin = Prelude.concat

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

instance MultiMonad Tree where
  innerJoin t = squeezeJoin $ mapWithPos reduceNested t
  outerJoin = innerJoin
  squeezeJoin (Leaf x) = x
  squeezeJoin (Branch xs) = Branch $ map squeezeJoin xs

reduce :: (Int, Int) -> Tree a -> [Tree a]
reduce _ (Leaf x) = [Leaf x]
reduce (i, n) (Branch xs) =
  if count <= 0
    then []
    else take count $ drop (start `mod` m) (cycle xs)
  where
    m = length xs
    start = floor (fromIntegral (m * i) / fromIntegral n)
    end = ceiling (fromIntegral (m * (i + 1)) / fromIntegral n)
    count = end - start

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

mapWithPos :: ([(Int, Int)] -> a -> b) -> Tree a -> Tree b
mapWithPos f = go []
  where
    go path (Leaf x) = Leaf (f (reverse path) x)
    go path (Branch ts) = Branch [go ((i, length ts) : path) t | (i, t) <- zip [0 ..] ts]