Map.hs

{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE OverloadedStrings #-}

module Zwirn.Core.Map where

{-
    Map.hs - lifting functions on maps to signals, some adapted
    from https://github.com/tidalcycles/Tidal/blob/dev/src/Sound/Tidal/Control.hs
    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.Map (Map)
import qualified Data.Map as Map
import Data.Maybe (fromMaybe)
import Data.String (IsString)
import Zwirn.Core.Core
import Zwirn.Core.Modulate (fastcat, slow)
import Zwirn.Core.Structure (run)
import Zwirn.Core.Time (Time)
import Zwirn.Core.Types
import Prelude hiding ((*>))

-- | create a singleton map with specific key
singleton :: (MultiApplicative m) => ZwirnT m st i k -> ZwirnT m st i a -> ZwirnT m st i (Map k a)
singleton = liftA2Right Map.singleton

union :: (Applicative m, Ord k) => ZwirnT m st i (Map k a) -> ZwirnT m st i (Map k a) -> ZwirnT m st i (Map k a)
union = liftA2 Map.union

-- | lookup a value via key
lookup :: (HasSilence m, MultiMonad m, Ord k) => ZwirnT m st i k -> ZwirnT m st i (Map k a) -> ZwirnT m st i a
lookup tz xz = outerJoin $ liftA2Right (\t x -> fromLookup $ Map.lookup t x) tz xz
  where
    fromLookup (Just x) = pure x
    fromLookup _ = silence

insert :: (Applicative m, Ord k) => ZwirnT m st i k -> ZwirnT m st i a -> ZwirnT m st i (Map k a) -> ZwirnT m st i (Map k a)
insert k a m = Map.insert <$> k <*> a <*> m

-- | apply a function to a specific key, if key is absent, return the original map
fix :: (HasSilence m, MultiMonad m, Ord k) => ZwirnT m st i k -> ZwirnT m st i (ZwirnT m st i a -> ZwirnT m st i a) -> ZwirnT m st i (Map k a) -> ZwirnT m st i (Map k a)
fix kz fz mz = outerJoin $ fromLookup <$> lookupMaybe kz mz
  where
    fromLookup (Just x) = insert kz (squeezeApply fz (pure x)) mz
    fromLookup Nothing = mz
    lookupMaybe = liftA2Right Map.lookup

chop :: (Fractional a, MultiMonad m, HasSilence m, Ord k, IsString k) => ZwirnT m st i Int -> ZwirnT m st i (Map k a) -> ZwirnT m st i (Map k a)
chop nz = squeezeMap (quickslice nz (run nz))

quickslice :: (Fractional a, MultiMonad m, Ord k, IsString k) => ZwirnT m st i Int -> ZwirnT m st i Int -> ZwirnT m st i (Map k a) -> ZwirnT m st i (Map k a)
quickslice nz iz = squeezeMap (slice nz iz)

loopAt :: (Fractional a, IsString a, HasSilence m, Monad m, Ord k, IsString k) => ZwirnT m st i Time -> ZwirnT m st i (Map k a) -> ZwirnT m st i (Map k a)
loopAt zt zx = _loopAt <$> zt <$$> zx
  where
    _loopAt 0 _ = silence
    _loopAt t x = Map.alter a "speed" . Map.insert "unit" "c" <$> slow (pure t) x
      where
        a (Just s) = Just (s / realToFrac t)
        a Nothing = Just (1 / realToFrac t)

slice :: (Fractional a, MultiApplicative m, Ord k, IsString k) => ZwirnT m st i Int -> ZwirnT m st i Int -> ZwirnT m st i (Map k a) -> ZwirnT m st i (Map k a)
slice nz iz zm = _slice <$> nz *> iz <*> zm
  where
    _slice n i m = Map.unions [Map.singleton "begin" newb, Map.singleton "end" newe, m]
      where
        b = fromMaybe 0 $ Map.lookup "begin" m
        e = fromMaybe 1 $ Map.lookup "end" m
        newrange x = e * x + (1 - x) * b
        newb = newrange $ div' i n
        newe = newrange $ div' i n + div' 1 n
        div' num den = fromIntegral (num `mod` den) / fromIntegral den

striateBy :: (Fractional a, Monad m, HasSilence m, Ord k, IsString k) => ZwirnT m st i Int -> ZwirnT m st i a -> ZwirnT m st i (Map k a) -> ZwirnT m st i (Map k a)
striateBy i f x = _striateBy <$> i <*> f <$$> x
  where
    _striateBy n f mz = fastcat $ map (offset . fromIntegral) [0 .. n - 1]
      where
        offset i = mergePlayRange (slot * i, (slot * i) + f) <$> mz
        slot = (1 - f) / fromIntegral (n - 1)

striate :: (Fractional a, Monad m, HasSilence m, Ord k, IsString k) => ZwirnT m st i Int -> ZwirnT m st i (Map k a) -> ZwirnT m st i (Map k a)
striate i x = _striate <$> i <$$> x
  where
    _striate n z = fastcat $ map offset [0 .. n - 1]
      where
        offset i = mergePlayRange (fromIntegral i / fromIntegral n, fromIntegral (i + 1) / fromIntegral n) <$> z

mergePlayRange :: (Fractional a, Ord k, IsString k) => (a, a) -> Map k a -> Map k a
mergePlayRange (b, e) cm = Map.insert "begin" ((b * d') + b') $ Map.insert "end" ((e * d') + b') cm
  where
    b' = fromMaybe 0 $ Map.lookup "begin" cm
    e' = fromMaybe 1 $ Map.lookup "end" cm
    d' = e' - b'