-- | This module contains functions to sample sound and to save it in a file.
-- The @LambdaSound.Play@ module exports a function to play sounds directly.
module LambdaSound.Sampling (sampleSound, sampleSoundRaw, saveWav, saveRawPCM) where

import Codec.Audio.Wave
import Data.ByteString qualified as B
import Data.Massiv.Array qualified as M
import LambdaSound.Sound
import LambdaSound.Sound.ComputeSound (sampleComputeSound)
import LambdaSound.Sound.Types (makeSamplingInfo)
import Data.Vector.Storable.ByteString (vectorToByteString)

-- | Samples a sound with the given frequency (usually 44100 is good) without post-processing
sampleSoundRaw :: Hz -> Sound T Pulse -> IO (M.Vector M.S Pulse)
sampleSoundRaw :: Hz -> Sound 'T Pulse -> IO (Vector S Pulse)
sampleSoundRaw Hz
hz (TimedSound Duration
duration ComputeSound Pulse
msc) = do
  let sr :: SamplingInfo
sr = Hz -> Duration -> SamplingInfo
makeSamplingInfo Hz
hz Duration
duration
  SamplingInfo -> ComputeSound Pulse -> IO (Vector S Pulse)
sampleComputeSound SamplingInfo
sr ComputeSound Pulse
msc

-- | Samples a sound with the given frequency (usually 44100 is good) with post-processing
--
-- This is recommended over 'sampleSoundRaw' if you are unsure
sampleSound :: Hz -> Sound T Pulse -> IO (M.Vector M.S Pulse)
sampleSound :: Hz -> Sound 'T Pulse -> IO (Vector S Pulse)
sampleSound Hz
hz Sound 'T Pulse
sound =
  Array D Int Pulse -> Vector S Pulse
forall r ix e r'.
(Manifest r e, Load r' ix e) =>
Array r' ix e -> Array r ix e
M.compute (Array D Int Pulse -> Vector S Pulse)
-> (Vector S Pulse -> Array D Int Pulse)
-> Vector S Pulse
-> Vector S Pulse
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Vector S Pulse -> Array D Int Pulse
forall r. Source r Pulse => Vector r Pulse -> Array D Int Pulse
postProcess (Vector S Pulse -> Vector S Pulse)
-> IO (Vector S Pulse) -> IO (Vector S Pulse)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Hz -> Sound 'T Pulse -> IO (Vector S Pulse)
sampleSoundRaw Hz
hz Sound 'T Pulse
sound

postProcess :: (M.Source r Pulse) => M.Vector r Pulse -> M.Vector M.D Pulse
postProcess :: forall r. Source r Pulse => Vector r Pulse -> Array D Int Pulse
postProcess = Vector r Pulse -> Array D Int Pulse
forall r. Source r Pulse => Vector r Pulse -> Array D Int Pulse
compressDynamically

-- | Save the sound as raw floats
saveRawPCM :: FilePath -> M.Vector M.S Pulse -> IO ()
saveRawPCM :: FilePath -> Vector S Pulse -> IO ()
saveRawPCM FilePath
filePath Vector S Pulse
floats =
  FilePath -> ByteString -> IO ()
B.writeFile FilePath
filePath (ByteString -> IO ()) -> ByteString -> IO ()
forall a b. (a -> b) -> a -> b
$ Vector Pulse -> ByteString
forall a. Storable a => Vector a -> ByteString
vectorToByteString (Vector Pulse -> ByteString) -> Vector Pulse -> ByteString
forall a b. (a -> b) -> a -> b
$ Vector S Pulse -> Vector Pulse
forall ix e. Index ix => Array S ix e -> Vector e
M.toStorableVector Vector S Pulse
floats

-- | Apply dynamic compression on a vector of samples such that
-- they are constrained within (-1, 1). Quieter sounds are boosted
-- while louder sounds are reduced.
-- This is very important if you use the parallel combinator since
-- parallel sounds are awful without post processing.
compressDynamically :: (M.Source r Pulse) => M.Vector r Pulse -> M.Vector M.D Pulse
compressDynamically :: forall r. Source r Pulse => Vector r Pulse -> Array D Int Pulse
compressDynamically Vector r Pulse
signal = (Pulse -> Pulse) -> Vector r Pulse -> Array D Int Pulse
forall ix r e' e.
(Index ix, Source r e') =>
(e' -> e) -> Array r ix e' -> Array D ix e
M.map (Pulse -> Pulse
scaleToMax (Pulse -> Pulse) -> (Pulse -> Pulse) -> Pulse -> Pulse
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Pulse -> Pulse
sigmoid) Vector r Pulse
signal
  where
    scaleToMax :: Pulse -> Pulse
scaleToMax Pulse
x = (Pulse
1 Pulse -> Pulse -> Pulse
forall a. Fractional a => a -> a -> a
/ Pulse -> Pulse
sigmoid Pulse
maxPulse) Pulse -> Pulse -> Pulse
forall a. Num a => a -> a -> a
* Pulse
x
    sigmoid :: Pulse -> Pulse
sigmoid Pulse
x = Pulse
2 Pulse -> Pulse -> Pulse
forall a. Fractional a => a -> a -> a
/ (Pulse
1 Pulse -> Pulse -> Pulse
forall a. Num a => a -> a -> a
+ Pulse -> Pulse
forall a. Floating a => a -> a
exp (Pulse
g Pulse -> Pulse -> Pulse
forall a. Num a => a -> a -> a
* (-Pulse
x))) Pulse -> Pulse -> Pulse
forall a. Num a => a -> a -> a
- Pulse
1
    g :: Pulse
g = Pulse -> Pulse -> Pulse
forall a. Floating a => a -> a -> a
logBase (Pulse
2 Pulse -> Pulse -> Pulse
forall a. Num a => a -> a -> a
- Pulse
factor) Pulse
factor Pulse -> Pulse -> Pulse
forall a. Fractional a => a -> a -> a
/ (-Pulse
maxPulse)
    maxPulse :: Pulse
maxPulse = Array D Int Pulse -> Pulse
forall r ix e.
(HasCallStack, Shape r ix, Source r e, Ord e) =>
Array r ix e -> e
M.maximum' (Array D Int Pulse -> Pulse) -> Array D Int Pulse -> Pulse
forall a b. (a -> b) -> a -> b
$ (Pulse -> Pulse) -> Vector r Pulse -> Array D Int Pulse
forall ix r e' e.
(Index ix, Source r e') =>
(e' -> e) -> Array r ix e' -> Array D ix e
M.map Pulse -> Pulse
forall a. Num a => a -> a
abs Vector r Pulse
signal
    factor :: Pulse
factor = Pulse
0.8

-- | Save a sound to a wave file with the given sampling frequency
saveWav :: FilePath -> Int -> M.Vector M.S Pulse -> IO ()
saveWav :: FilePath -> Int -> Vector S Pulse -> IO ()
saveWav FilePath
filepath Int
sampleRate Vector S Pulse
floats = do
  let floatsLength :: Int
floatsLength = Sz Int -> Int
forall ix. Sz ix -> ix
M.unSz (Sz Int -> Int) -> Sz Int -> Int
forall a b. (a -> b) -> a -> b
$ Vector S Pulse -> Sz Int
forall r ix e. Size r => Array r ix e -> Sz ix
forall ix e. Array S ix e -> Sz ix
M.size Vector S Pulse
floats
      wave :: Wave
wave =
        Wave
          { waveFileFormat :: WaveFormat
waveFileFormat = WaveFormat
WaveVanilla,
            waveSampleRate :: Word32
waveSampleRate = Int -> Word32
forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
sampleRate,
            waveSampleFormat :: SampleFormat
waveSampleFormat = SampleFormat
SampleFormatIeeeFloat32Bit,
            waveChannelMask :: Set SpeakerPosition
waveChannelMask = Set SpeakerPosition
speakerMono,
            waveDataOffset :: Word32
waveDataOffset = Word32
0,
            waveDataSize :: Word64
waveDataSize = Int -> Word64
forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
floatsLength Word64 -> Word64 -> Word64
forall a. Num a => a -> a -> a
* Word64
4,
            waveSamplesTotal :: Word64
waveSamplesTotal = Int -> Word64
forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
floatsLength,
            waveOtherChunks :: [(ByteString, ByteString)]
waveOtherChunks = []
          }
  FilePath -> Wave -> (Handle -> IO ()) -> IO ()
forall (m :: * -> *).
MonadIO m =>
FilePath -> Wave -> (Handle -> IO ()) -> m ()
writeWaveFile FilePath
filepath Wave
wave ((Handle -> IO ()) -> IO ()) -> (Handle -> IO ()) -> IO ()
forall a b. (a -> b) -> a -> b
$ \Handle
handle ->
    Handle -> ByteString -> IO ()
B.hPut Handle
handle (ByteString -> IO ()) -> ByteString -> IO ()
forall a b. (a -> b) -> a -> b
$ Vector Pulse -> ByteString
forall a. Storable a => Vector a -> ByteString
vectorToByteString (Vector Pulse -> ByteString) -> Vector Pulse -> ByteString
forall a b. (a -> b) -> a -> b
$ Vector S Pulse -> Vector Pulse
forall ix e. Index ix => Array S ix e -> Vector e
M.toStorableVector Vector S Pulse
floats