text:
{-# LANGUAGE BangPatterns, TypeFamilies, TypeOperators, FlexibleInstances, FlexibleContexts, GADTs #-}
{-# LANGUAGE MultiParamTypeClasses, UndecidableInstances #-}
module Data.NeuralNetwork.Backend.HMatrix.Layers where

import Numeric.LinearAlgebra hiding (R, C)
import Numeric.LinearAlgebra.Devel
import qualified Data.Vector as V
import qualified Data.Vector.Storable as SV
import qualified Data.Vector.Storable.Mutable as SVM
import System.Random.MWC
import System.Random.MWC.Distributions
import Control.Monad.ST
import Control.Monad (liftM2, forM_, when)
import GHC.Float
import Data.STRef
import Data.Functor.Identity
import Control.DeepSeq
import Data.NeuralNetwork

type R = Float

-- We parameterise the activation layer T, where the parameter indicates how
-- elements are contained:
--   SinglC :. Vector, SinglC :. Matrix, MultiC :. Vector, MultiC :.  Matrix
-- SinglC means the input has only one channel, while
-- MultiC means the input has more than one.
--
-- type function composition
data (f :: * -> *) :. (g :: * -> *) :: * -> *
-- type function: Identity
data SinglC :: * -> *
data MultiC :: * -> *

-- Tags for each form of layer
data F
data C
data A
data M
data T (c :: * -> *)
data S a b

data RunLayer :: * -> * where
  -- Densely connected layer
  -- input:   vector of size m
  -- output:  vector of size n
  -- weights: matrix of size m x n
  -- biases:  vector of size n
  Full :: !(Matrix R) -> !(Vector R) -> RunLayer F
  -- convolutional layer
  -- input:  channels of 2D floats, of the same size (a x b), # of input channels:  m
  -- output: channels of 2D floats, of the same size (c x d), # of output channels: n
  --         where c = a + 2*padding + 1 - s
  --               d = b + 2*padding + 1 - t
  -- feature:  matrix of (s x t), # of features: m x n
  -- padding:  number of 0s padded at each side of channel
  -- biases:   bias for each output, # of biases: n
  Conv  :: !(V.Vector (V.Vector (Matrix R))) -> !(V.Vector R) -> Int -> RunLayer C
  -- Reshape from channels of matrix to a single vector
  -- input:  m channels of 2D matrices
  --         assuming that all matrices are of the same size a x b
  -- output: 1D vector of the concatenation of all input channels
  --         its size: m x a x b
  As1D  :: RunLayer A
  -- max pooling layer
  -- input:  channels of 2D floats, of the same size (a x b), # of input channels:  m
  --         assuming that a and b are both multiple of stride
  -- output: channels of 2D floats, of the same size (c x d), # of output channels: m
  --         where c = a / stride
  --               d = b / stride
  MaxP :: Int -> RunLayer M
  -- Activator
  -- the input can be either a 1D vector, 2D matrix, or channels of either.
  Activation :: (R->R, R->R) -> RunLayer (T c)
  -- stacking two components a and b
  -- the output of a should matches the input of b
  Stack :: !(RunLayer a) -> !(RunLayer b) -> RunLayer (S a b)

instance Component (RunLayer F) where
    type Run (RunLayer F) = Identity
    type Inp (RunLayer F) = Vector R
    type Out (RunLayer F) = Vector R
    -- trace is (input, weighted-sum)
    newtype Trace (RunLayer F) = DTrace (Vector R, Vector R)
    forwardT (Full w b) !inp =
        let !bv = (inp <# w) `add` b
        in return $ DTrace (inp,bv)
    output (DTrace (_,!a)) = a
    backward l (DTrace (!iv,!bv)) !odelta rate =
        let Full w b = l
            !d = scale (negate rate) odelta
            !m = iv `outer` d
            -- back-propagated error at input
            !idelta = w #> odelta
            -- update to weights
            ---- for what reason, could this expression: w `add` (iv `outer` d)
            ---- entails a huge space leak? especially, neither 'seq' nor
            ---- 'deepseq' helps a bit. The only workaround is to expand the
            ---- add function, and call SV.force on the result vector, which
            ---- explcitly copy and drop reference to orignal computed result.
            !w'= w `add` m
            -- !w'= let (r,c) = size w
            --          dat1 = flatten (tr' w)
            --          dat2 = flatten (tr' m)
            --      in matrixFromVector ColumnMajor r c $ SV.force $ dat1 `add` dat2
            !b'= b `add` d
        in return $ (Full w' b', idelta)

instance Component (RunLayer C) where
    type Run (RunLayer C) = Identity
    type Inp (RunLayer C) = V.Vector (Matrix R)
    type Out (RunLayer C) = V.Vector (Matrix R)
    -- trace is (input, convoluted output)
    newtype Trace (RunLayer C) = CTrace (Inp (RunLayer C), V.Vector (Matrix R))
    forwardT (Conv fs bs p) !inp =
        let !ov = V.zipWith feature
                            (tr fs) -- feature matrix indexed majorly by each output
                            bs      -- biases by each output
        in return $ CTrace (inp,ov)
      where
        !osize = let (x,y) = size (V.head inp)
                     (u,v) = size (V.head $ V.head fs)
                 in (x+2*p-u+1, y+2*p-v+1)
        -- transpose the features matrix
        tr :: V.Vector (V.Vector a) -> V.Vector (V.Vector a)
        tr uv = let n   = V.length (V.head uv)
                    !vu = V.map (\i -> V.map (V.! i) uv) $ V.enumFromN 0 n
                in vu
        feature :: V.Vector (Matrix R) -> R -> Matrix R
        feature f b = V.foldl1' add (V.zipWith (layerCorr2 p) f inp) `add` konst b osize
    output (CTrace (_,a)) = a
    backward l (CTrace (!iv,!av)) !odelta rate =
      let Conv fs bs p = l
          -- update to the feature matrix
          m :: V.Vector (V.Vector (Matrix R))
          !m = V.zipWith (\flts chn ->
                            -- chn:  a single input channel
                            -- flts: all features used for chn
                            V.zipWith (\f d ->
                              let upd = scale (negate rate) (layerCorr2 p chn d)
                              in f `add` upd
                            ) flts odelta
                          ) fs iv
          -- update to the biases
          b :: V.Vector R
          !b = V.zipWith (\b d -> b + (negate rate) * sumElements d) bs odelta
          -- back-propagated error at input
          idelta :: V.Vector (Matrix R)
          !idelta = V.map (\f -> V.foldl1' add $ V.zipWith (layerConv2 p) f odelta) fs
      in --trace ("CL:" ++ show odelta)
         return $ (Conv m b p, idelta)

instance Component (RunLayer A) where
  type Run (RunLayer A) = Identity
  type Inp (RunLayer A) = V.Vector (Matrix R)
  type Out (RunLayer A) = Vector R
  -- trace keeps information of (m, axb, b, output)
  newtype Trace (RunLayer A) = ReshapeTrace (Int, Int, Int, Vector R)
  forwardT _ !inp =
    let !b = V.length inp
        (!r,!c) = size (V.head inp)
        !o = V.foldr' (\x y -> flatten x SV.++ y) SV.empty inp
    in return $ ReshapeTrace (b, r*c, c, o)
  output (ReshapeTrace (_,_,_,a)) = a
  backward a (ReshapeTrace (b,n,c,_)) !odelta _ =
    let !idelta = V.fromList $ map (reshape c) $ takesV (replicate b n) odelta
    in return $ (a, idelta)

instance Component (RunLayer M) where
  type Run (RunLayer M) = Identity
  type Inp (RunLayer M) = V.Vector (Matrix R)
  type Out (RunLayer M) = V.Vector (Matrix R)
  -- trace is (dimension of pools, index of max in each pool, pooled matrix)
  -- for each channel.
  newtype Trace (RunLayer M) = PTrace (V.Vector (IndexOf Matrix, Vector Int, Matrix R))
  -- forward is to divide the input matrix in stride x stride sub matrices,
  -- and then find the max element in each sub matrices.
  forwardT (MaxP stride) !inp = return $ PTrace $ V.map mk inp
    where
      mk inp = let (!i,!v) = pool stride inp in (size v, i, v)
  output (PTrace a) = V.map (\(_,_,!o) ->o) a
  -- use the saved index-of-max in each pool to propagate the error.
  backward l@(MaxP stride) (PTrace t) odelta _ =
      let !idelta = V.zipWith gen t odelta in return $ (l, idelta)
    where
      gen (!si,!iv,_) od = unpool stride iv od

instance (Component (RunLayer a),
          Component (RunLayer b),
          Run (RunLayer a) ~ Identity,
          Run (RunLayer b) ~ Identity,
          Out (RunLayer a) ~ Inp (RunLayer b)
         ) => Component (RunLayer (S a b)) where
    type Run (RunLayer (S a b)) = Identity
    type Inp (RunLayer (S a b)) = Inp (RunLayer a)
    type Out (RunLayer (S a b)) = Out (RunLayer b)
    newtype Trace (RunLayer (S a b)) = TTrace (Trace (RunLayer b), Trace (RunLayer a))
    forwardT (Stack a b) !i = do
        !tra <- forwardT a i
        !trb <- forwardT b (output tra)
        return $ TTrace (trb, tra)
    output (TTrace !a) = output (fst a)
    backward (Stack a b) (TTrace (!trb,!tra)) !odelta rate = do
        (b', !odelta') <- backward b trb odelta  rate
        (a', !idelta ) <- backward a tra odelta' rate
        return (Stack a' b', idelta)

instance (Container c R, Pairwise c R) => Component (RunLayer (T (MultiC :. c))) where
    type Run (RunLayer (T (MultiC :. c))) = Identity
    type Inp (RunLayer (T (MultiC :. c))) = V.Vector (c R)
    type Out (RunLayer (T (MultiC :. c))) = V.Vector (c R)
    newtype Trace (RunLayer (T (MultiC :. c))) = TTraceM (V.Vector (Trace (RunLayer (T (SinglC :. c)))))
    forwardT (Activation ac) !inp =
      TTraceM <$> V.mapM (forwardT (Activation ac)) inp
    output (TTraceM a) = V.map output a
    backward a@(Activation ac) (TTraceM ts) !odelta r = do
      idelta <- V.zipWithM (\t d -> snd <$> backward (Activation ac) t d r) ts odelta
      return (a, idelta)

instance (Container c R, Pairwise c R) => Component (RunLayer (T (SinglC :. c))) where
    type Run (RunLayer (T (SinglC :. c))) = Identity
    type Inp (RunLayer (T (SinglC :. c))) = c R
    type Out (RunLayer (T (SinglC :. c))) = c R
    newtype Trace (RunLayer (T (SinglC :. c))) = TTraceS (c R, c R)
    forwardT (Activation (af,_)) !inp = return $ TTraceS (inp, cmap af inp)
    output (TTraceS (_,!a)) = a
    backward a@(Activation (_,ag)) (TTraceS (!iv,_)) !odelta _ = return $ (a, odelta `hadamard` cmap ag iv)

newFLayer :: Int                -- number of input values
          -> Int                -- number of neurons (output values)
          -> IO (RunLayer F)    -- new layer
newFLayer m n =
    withSystemRandom . asGenIO $ \gen -> do
        -- we build the weights in column major because in the back-propagation
        -- algo, the computed update to weights is in column major. So it is
        -- good for performance to keep the matrix always in column major.
        w <- buildMatrix (normal 0 0.01 gen) ColumnMajor (m,n)
        b <- return $ konst 1 n
        return $ Full w b

newCLayer :: Int                -- number of input channels
          -> Int                -- number of output channels
          -> Int                -- size of each feature
          -> Int                -- size of padding
          -> IO (RunLayer C)    -- new layer
newCLayer inpsize outsize sfilter npadding =
  withSystemRandom . asGenIO $ \gen -> do
      fs <- V.replicateM inpsize $ V.replicateM outsize $
              buildMatrix (truncNormal 0 0.1 gen) RowMajor (sfilter, sfilter)
      bs <- return $ V.replicate outsize 0.1
      return $ Conv fs bs npadding
  where
    truncNormal m s g = do
      x <- standard g
      if x >= 2.0 || x <= -2.0
        then truncNormal m s g
        else return $! m + s * x

buildMatrix g order (nr, nc) = do
  vals <- SV.replicateM (nr*nc) (double2Float <$> g)
  return $ matrixFromVector order nr nc vals

layerCorr2 :: Int -> Matrix R -> Matrix R -> Matrix R
layerCorr2 p k m = corr2 k padded
  where
    padded = zeroPadded p m
    (w,_)  = size k

layerConv2 :: Int -> Matrix R -> Matrix R -> Matrix R
layerConv2 p k m = corr2 (fliprl . flipud $ k) padded
  where
    padded = zeroPadded p m

-- max pool, picking out the maximum element
-- in each stride x stride sub-matrices.
-- assuming that the original matrix row and column size are
-- both multiple of stride
pool :: Int -> Matrix Float -> (Vector Int, Matrix Float)
pool 1 mat = let (r,c) = size mat in (SV.replicate (r*c) 0, mat)
-- pool 2 mat | orderOf mat == RowMajor = c_max_pool2_f mat
pool stride mat = runST $ do
  ori <- unsafeThawMatrix mat
  mxv <- newUndefinedMatrix RowMajor r' c'
  mxi <- newUndefinedVector (r'*c')
  forM_ [0..r'-1] $ \i -> do
    forM_ [0..c'-1] $ \j -> do
      (n,v) <- unsafeMaxIndEle ori (i*stride) (j*stride) stride stride
      unsafeWriteVector mxi (i*c'+j) n
      unsafeWriteMatrix mxv i j v
  a <- unsafeFreezeVector mxi
  b <- unsafeFreezeMatrix mxv
  return (a,b)
  where
    (r,c) = size mat
    r'    = r `div` stride
    c'    = c `div` stride
    unsafeMaxIndEle mm x y r c = do
      mp <- newSTRef 0
      mv <- newSTRef (-10000.0)
      forM_ [0..r-1] $ \ i -> do
        forM_ [0..c-1] $ \ j -> do
          v1 <- unsafeReadMatrix mm (x+i) (y+j)
          v0 <- readSTRef mv
          when (v1 > v0) $ do
            writeSTRef mv v1
            writeSTRef mp (i*2+j)
      p <- readSTRef mp
      v <- readSTRef mv
      return (p, v)

-- the reverse of max pool.
-- assuming idx and mat are of the same size
unpool :: Int -> Vector Int -> Matrix Float -> Matrix Float
unpool stride idx mat = runSTMatrix $ do
  mat' <- newMatrix' 0 r' c'
  forM_ [0..r-1] $ \i -> do
    forM_ [0..c-1] $ \j -> do
      let pos     = idx SV.! (i*c+j)
      let (oi,oj) = pos `divMod` 2
      let val     = mat `atIndex` (i,j)
      unsafeWriteMatrix mat' (i*stride+oi) (j*stride+oj) val
  return mat'
  where
    (r,c) = size mat
    (r',c') = (r*stride, c*stride)

-- a slightly faster way to pading the matrix
-- camparing to fromBlocks provided by hmatrix.
zeroPadded :: Int -> Matrix Float -> Matrix Float
zeroPadded p mat = runSTMatrix $ do
  mat' <- newMatrix' 0 r' c'
  setMatrix mat' p p mat
  return mat'
  where
    (r,c) = size mat
    (r',c') = (r+2*p,c+2*p)

-- a slightly faster version of newMatrix, which based
-- directly on lower level Vector.Storable creation.
newMatrix' :: SVM.Storable t => t -> Int -> Int -> ST s (STMatrix s t)
newMatrix' v r c = do
  vec <- SVM.replicate (r*c) v
  vec <- SV.unsafeFreeze vec
  unsafeThawMatrix $ reshape c vec

-- The hmatrix does has a SIMD enabled pairwise product named 'mul' in the
-- 'Container' class, but unfortunately it is not exported at all.
-- We define a slow hadamard product here.
class Container c e => Pairwise c e where
  hadamard :: c e -> c e -> c e
instance Pairwise Vector R where
  hadamard = zipVectorWith (*)
instance Pairwise Matrix R where
  hadamard a b = reshape (cols a) $ hadamard (flatten a) (flatten b)

max_stars_repo_name: pierric/neural-network

max_issues_repo_name: pierric/neural-network

max_forks_repo_name: pierric/neural-network


text:
{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE ScopedTypeVariables #-}

{-

A lot of error handling needs to be added to the code. For example
0 < hopsize < windowsize < fftsize needs to be reinforced.
Code currently just throws exceptions. Maybe change to something like;
http://www.mega-nerd.com/erikd/Blog/CodeHacking/Haskell/what_do_you_mean.html

-}

module STFT (

    stft
  , readWav
  , writeWav
  , stftSynth
  , hammingC

  ) where

import System.IO (writeFile)
import Data.WAVE
import Window
import qualified Data.Vector.Generic as V
import Numeric.FFT.Vector.Invertible
import Data.Vector
import qualified Data.Complex as C
import Data.Bifunctor (first)
import Data.Vector.Split (divvy)

epsilon = 2.2204460492503131e-16

-- This code from the Linear.Epsilon library
-- by Edward Kmett. There were some dependency conflicts
-- so I just copied it in. Will fix.
class Num a => Epsilon a where
  -- | Determine if a quantity is near zero.
  nearZero :: a -> Bool

-- | @'abs' a '<=' 1e-6@
instance Epsilon Float where
  nearZero a = abs a <= 1e-6

-- | @'abs' a '<=' 1e-12@
instance Epsilon Double where
  nearZero a = abs a <= 1e-12

type MagSpect = Vector Double -- Magnitude Spectrum

type PhaseSpect = Vector Double -- Phase Spectrum

type CSignal = Vector (C.Complex Double) -- Complex Signal

type Signal = Vector Double

type FFTsz = Int -- Should be Greater than 0 and a power of 2

type Hopsz = Int -- Should be about 1/4 the fft sz

type Path = String

type Winsz = Int -- Should be an odd number and smaller than FFTsz

-- Phase unwrapping algorithm.
-- Converting to and from list for pattern matching is a little cumbersome.
-- See this post for possible improvement;
-- http://stackoverflow.com/questions/36993937/haskell-pattern-matching-on-vectors
unwrap :: PhaseSpect -> PhaseSpect
unwrap xs = fromList $ diff (toList xs) 0 where
  diff (x:y:ys) accm
    | (y+accm) - x > pi    = x: diff (y+accm-(2*pi):ys) (accm-(2*pi))
    | (y+accm) - x < (-pi) = x: diff (y+accm+(2*pi):ys) (accm+(2*pi))
    | otherwise             = x: diff (y+accm:ys) accm
  diff [x] _     = [x]
  diff []  _     = []

-- Takes a complex signal and fftsz does zero phase windows
zeroPhaseWindow :: CSignal -> FFTsz -> CSignal
zeroPhaseWindow xs fftsz = let win = V.length xs
                               hM1 = floor $ fromIntegral (win + 1) / 2
                               hM2 = floor $ (fromIntegral win) / 2
                               zs = V.replicate (fftsz-hM1-hM2) (c 0)
                            in V.concat [(V.slice hM2 hM1 xs), zs,
                                         (V.slice 0 hM2 xs)]

-- Normalize signal to 0 Db.
normTo0Db :: Vector (Double,Double) -> Vector (Double, Double)
normTo0Db xs = V.map (first ((-) (fst $ V.maximumBy compare xs))) xs

-- Next power of 2 greater than n.
po2gtn :: Int -> Int
po2gtn n = 2^(ceiling $ logBase 2 (fromIntegral n))

isPo2 :: (Ord a, Fractional a) => a -> Bool
isPo2 x
  | x > 2     = isPo2 (x / 2)
  | x == 2    = True
  | otherwise = False

-- Calculates the magnitude spectrum in Db
magSpect :: CSignal -> MagSpect
magSpect vec = V.map (dB.(near0).(C.magnitude)) vec where
  dB x = 20 * (logBase 10 x)
  near0 x = if nearZero x then epsilon else x

-- Calculates the phase spectrum
phaseSpect :: CSignal -> PhaseSpect
phaseSpect vec = unwrap $ V.map ((C.phase).to0) vec where
  to0 x = go (C.realPart x) (C.imagPart x)
  go x y | nearZero x && nearZero y = (0 C.:+ 0)
         | nearZero x               = (0 C.:+ y)
         | nearZero y               = (x C.:+ 0)
         | otherwise                = (x C.:+ y)

-- Takes a windowed signal and FFT size returns a vector of tuples
-- (Magnitude, Phase). Magnitude in Db, phase unwrapped, both positve
-- half of the spectrum.
dftAnal :: FFTsz -> CSignal -> (MagSpect, PhaseSpect)
dftAnal fftsz winSig = (mag, phase) where
    inputVec = zeroPhaseWindow winSig fftsz
    hN = fftsz `div` 2 + 1
    trsf :: Vector (C.Complex Double)
    trsf = V.slice 0 hN (run dft inputVec)
    mag = magSpect trsf
    phase = phaseSpect trsf

-- Takes an input signal, a window, an fftsize and a hop size and
-- computes the short time Fourier transform, returning a list of
-- (vector of magnitude, vector of phase)
stft :: Signal -> CWindow ->
  FFTsz -> Hopsz -> [(MagSpect, PhaseSpect)]
stft dsig win fftsz hopsz = let sig = V.map c dsig -- Convert to complex
                                winsz = V.length win
                                wdiv = V.sum win
                                w = V.map (/wdiv) win -- Normalize window function
                                halfWin = floor $ (fromIntegral winsz) / 2
                                zs = V.replicate halfWin (c 0)
                                sigzs = V.concat [zs, sig, zs]
                                splitSig = divvy winsz hopsz sigzs
                             in fmap ((dftAnal fftsz).(V.zipWith (*) w)) splitSig

-- Takes a vector of tuples (Magnitude, Phase) and a window size
-- and returns the original signal
dftSynth :: Winsz -> (MagSpect, PhaseSpect) -> Signal
dftSynth win (magVec, phaseVec) = let vec = V.zipWith (\x y -> (x, y)) magVec phaseVec
                                      hN = V.length vec
                                      n = (hN-1)*2
                                      hM1 = floor $ fromIntegral (win+1) / 2
                                      hM2 = floor $ fromIntegral win / 2
                                      posFreqs = V.map f vec where
                                        f (mag, phase) = (10**(mag/20) C.:+ 0) *
                                                         (exp ((phase C.:+ 0) *
                                                         (0 C.:+ 1)))
                                      negFreqs = ((V.map f) . V.reverse) (g vec) where
                                        g = (V.init . V.tail)
                                        f (mag, phase) = (10**(mag/20) C.:+ 0) *
                                                         (exp ((phase C.:+ 0) *
                                                         (0 C.:+ (-1))))
                                      x = (run idft (posFreqs V.++ negFreqs))
                                  in V.map C.realPart $ V.concat [(V.slice (V.length x - hM2) hM2 x), (V.take hM1 x)]

stftSynth :: [(MagSpect, PhaseSpect)] -> Hopsz -> Winsz -> Signal
stftSynth magphase hopsz winsz = let hM1 = floor $ fromIntegral (winsz+1) / 2
                                     hM2 = floor $ fromIntegral winsz / 2
                                     signalFrames = fmap ((V.map ((fromIntegral hopsz)*)).(dftSynth winsz)) magphase
                                     signalTuples = fmap (V.splitAt hM1) signalFrames
                                     overlapAdd (x1, x2) (y1, y2) = (x1 V.++ (V.zipWith (+) x2 y1), y2)
                                  in V.drop hM1 $ fst (Prelude.foldl overlapAdd
                                        (Prelude.head signalTuples) (Prelude.tail signalTuples))

-- Takes a vector of doubles, the sample frequency, and a name
-- and writes the audio file. Written in 32 bit.
writeWav :: Signal -> Int -> String -> IO ()
writeWav vec sf name = let samples = fmap ((:[]) . doubleToSample) (V.toList vec)
                           header = WAVEHeader 1 sf 32 Nothing
                       in putWAVEFile name (WAVE header samples)

-- Takes a Path and returns IO (sampling frequency, Vector signal).
readWav :: Path -> IO (Int, Signal)
readWav path = do
  audio <- getWAVEFile "singing-female.wav"
  let header = waveHeader audio
      samples = waveSamples audio
      channels = waveNumChannels header
      sampRate = waveFrameRate header
  case channels of
    1 -> let sig = fromList $ fmap (sampleToDouble.(Prelude.head)) samples
          in return $ (sampRate, sig)
    _ -> error "Should be mono."

ext: hs

max_stars_repo_name: davlum/haskell-stft

max_issues_repo_name: davlum/haskell-stft

max_forks_repo_name: davlum/haskell-stft


text:
{-# LANGUAGE DataKinds #-}
{-# LANGUAGE TypeFamilies #-}
{-# LANGUAGE GADTs #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE TypeSynonymInstances #-}
{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE FlexibleInstances #-}
{-# LANGUAGE IncoherentInstances #-}

module Test.LinearAlgebra (linearAlgebra) where

import Test.Tasty
import Data.Matrix.Static.LinearAlgebra
import qualified Data.Matrix.Static.Dense as D
import qualified Data.Matrix.Static.Generic as G
import qualified Data.Matrix.Static.Sparse as S
import Data.Singletons hiding ((@@))
import Data.Singletons.Prelude (Min)
import Data.Complex
import Data.Vector.Storable (Storable)
import qualified Data.Vector.Storable as V
import GHC.TypeNats (KnownNat)
import Control.Monad.IO.Class (liftIO)
import Test.Tasty.QuickCheck

import Test.Utils

linearAlgebra :: TestTree
linearAlgebra = testGroup "Linear algebra"
    [ basicTest
    , svdTest
    , eigenTest
    ]

basicTest :: TestTree
basicTest = testGroup "Basic"
    [ testProperty "Row sum (Dense)" drs
    , testProperty "Column sum (Dense)" dcs
    , testProperty "Row sum (Sparse)" srs
    , testProperty "Column sum (Sparse)" scs ]
  where
    drs :: Matrix 50 30 Double -> Bool
    drs m = D.toList (rowSum m) == map V.sum (G.toRows m)
    dcs :: Matrix 50 30 Double -> Bool
    dcs m = D.toList (colSum m) == map V.sum (G.toColumns m)
    srs :: SparseMatrix 50 30 Double -> Bool
    srs m = D.toList (rowSum m) == map V.sum (G.toRows m)
    scs :: SparseMatrix 50 30 Double -> Bool
    scs m = D.toList (colSum m) == map V.sum (G.toColumns m)

svdTest :: TestTree
svdTest = testGroup "SVD"
    [ testProperty "SVD (Float)" svd1
    , testProperty "SVD (Double)" svd2
    ]
  where
    svd1 :: Matrix 50 30 Float -> Bool
    svd1 m = m' ~= m
      where
        m' = u @@ S.diag d @@ D.transpose v
        (u,d,v) = svd m
    svd2 :: Matrix 50 30 Double -> Bool
    svd2 m = m' ~= m
      where
        m' = u @@ S.diag d @@ D.transpose v
        (u,d,v) = svd m

eigenTest :: TestTree
eigenTest = testGroup "Eigendecomposition"
    [ testProperty "Full" eigen1
    , testProperty "Partial dense" eigen2
    , testProperty "Partial symmetric dense" eigen3
    , testProperty "Partial symmetric sparse" eigen4
    ]
  where
    eigen1 :: Matrix 100 100 Double -> Bool
    eigen1 m = (m' @@ v) ~= (v @@ S.diag d)
      where
        m' = D.map (\x -> mkPolar x 0) m
        (d, v)= eig m
    eigen2 :: Matrix 100 100 Double -> Bool
    eigen2 m = m' @@ v ~= v @@ S.diag d
      where
        m' = D.map (\x -> mkPolar x 0) m
        (d, v)= eigS (sing :: Sing 8) m
    eigen3 :: Matrix 100 100 Double -> Bool
    eigen3 raw = m @@ v ~= v @@ S.diag d
      where
        m = raw %+% D.transpose raw
        (d, v)= eigSH (sing :: Sing 99) m
    eigen4 :: SparseMatrix 100 100 Double -> Bool
    eigen4 raw = m @@ v ~= v @@ S.diag d
      where
        m = raw %+% D.transpose raw
        (d, v)= eigSH (sing :: Sing 99) m

max_stars_repo_name: kaizhang/matrix-sized

max_issues_repo_name: kaizhang/matrix-sized

max_forks_repo_name: kaizhang/matrix-sized


text:
{-# LANGUAGE ScopedTypeVariables #-}

module Types
    ( Planet(..)
    , mass
    , diameter
    , pullOn
    , distance
    , distanceScale
    , alterSpeed
    , newPoint
    )
where

import           Numeric.LinearAlgebra

data Planet = Earth  | Jupiter deriving (Eq, Show)

type Point = (Float, Float)

mass :: Planet -> Int
mass Earth   = 5972
mass Jupiter = 1898000

distanceScale :: Float
distanceScale = 1000000

diameter :: Planet -> Integer
diameter Earth   = 12756000 -- m
diameter Jupiter = 142984000

distance :: Point -> Point -> Float
distance (aX, aY) (bX, bY) =
    let height  = abs (bY - aY)
        width   = abs (bX - aX)
        height2 = height ^ 2
        width2  = width ^ 2
    in  sqrt (height2 + width2) * distanceScale * 100 -- planets actual size so brag about distance

pullOn :: Planet -> Point -> Planet -> Point -> Vector Float -- a accelaration pointing out from a
pullOn on a from b =
    let
        gConstant = 15 * 10 ^ 12
        massA     = mass on
        massB     = mass from
        dist      = distance b a ^ 2
        forceScalar =
            (-gConstant * fromIntegral massA * fromIntegral massB) / dist
        accScalar = forceScalar / fromIntegral massA
    in
        scale accScalar (vectorize b a)

vectorize :: Point -> Point -> Vector Float
vectorize (ax, ay) (bx, by) = fromList [bx - ax, by - ay]

alterSpeed :: Vector Float -> Vector Float -> Vector Float
alterSpeed current pullForce = add current pullForce -- todo

newPoint :: Point -> Vector Float -> Point
newPoint (ax, ay) v = case toList v of
    [bx, by] -> (ax + bx, ay + by)
    x ->
        error
            $  "cannot create new point from speed vector with /= 2 elements "
            ++ show x

max_stars_repo_name: mariatsji/planets

max_issues_repo_name: mariatsji/planets

max_forks_repo_name: mariatsji/planets


text:
-- nonlinear least-squares fitting

import Numeric.GSL.Fitting
import Numeric.LinearAlgebra

xs = map return [0 .. 39]
sigma = 0.1
ys = map return $ toList $ fromList (map (head . expModel [5,0.1,1]) xs)
            + scalar sigma * (randomVector 0 Gaussian 40)

dat :: [([Double],([Double],Double))]

dat = zip xs (zip ys (repeat sigma))

expModel [a,lambda,b] [t] = [a * exp (-lambda * t) + b]

expModelDer [a,lambda,b] [t] = [[exp (-lambda * t), -t * a * exp(-lambda*t) , 1]]

(sol,path) = fitModelScaled 1E-4 1E-4 20 (expModel, expModelDer) dat [1,0,0]

main = do
    print dat
    print path
    print sol

max_stars_repo_name: curiousleo/liquidhaskell

max_issues_repo_name: curiousleo/liquidhaskell

max_forks_repo_name: curiousleo/liquidhaskell


text:
-- Fingerprint
{-# LANGUAGE DeriveFunctor #-}
{-# LANGUAGE Strict #-}
{-# LANGUAGE StrictData #-}

module Fingerprint (
  Fingerprint
, ColorVector
, fp2colvec
, makeMap
, findSimilarImage
, colvec2id
, colvec2status
, isValid
) where

import qualified Data.KdTree.Static as KT
import           Data.List (transpose)
import           Data.List.Split (chunksOf)
import           Numeric (readHex)
import           Numeric.LinearAlgebra

type Fingerprint = String
type Channel = [Double]

type Color = [Double]  -- (r,g,b) => [r,g,b]
data ColorVector = ColorVector
  { pixel  :: [Color]
  , imgid  :: Int
  , fp     :: Fingerprint
  , status :: CvStatus
  } deriving (Eq, Show)
type ColorVectorMap = KT.KdTree Double ColorVector

data CvStatus = FILED | PEND | DISCARDED | DUPLICATED | INFERIOR deriving (Eq, Show)

-- constants
--   status

-- public functions

makeMap :: [ColorVector] -> ColorVectorMap
makeMap cvs = KT.buildWithDist colvec2points colvecDistance cvs

{-|
  IN:
    r1  : radius for KdTree search
    r2  : difference for each pixel
-}

findSimilarImage :: Double -> Double -> Double -> ColorVectorMap -> ColorVector -> [ColorVector]
findSimilarImage r1 r2 rt cvmap cv = ss2
  where
    ss1 = filter (\s -> colvec2id cv < colvec2id s) $ KT.inRadius cvmap r1 cv
    ss2 = filter (\s -> isNearImage r2 rt (colvec2fp cv) (colvec2fp s)) ss1

{-|

>>> let fp1 = [1.0, 2.0, 1.0, 2.0]
>>> let fp2 = [1.0, 2.0, 1.0, 2.0]
>>> isNearImage 0.0 1.0 fp1 fp2
True
>>> let fp1 = [1.0, 2.0, 1.0, 2.0]
>>> let fp2 = [1.1, 2.0, 1.0, 2.0]
>>> let ds = map abs $ zipWith (-) fp1 fp2
>>> length ds
4
>>> let ds' = filter (<= 0.0) $ ds
>>> length ds'
3
>>> (fromIntegral (length ds') / fromIntegral (length ds))
0.75
>>> isNearImage 0.0 0.5 fp1 fp2
True

-}

isNearImage :: Double -> Double -> [Double] -> [Double] -> Bool
isNearImage rad rate fp1 fp2 = (fromIntegral (length ds') / fromIntegral (length ds) >= rate)
  where
    ds = map abs $ zipWith (-) fp1 fp2
    ds' = filter (<= rad) ds

{-|
  IN:
      oreso - original resolution of x (ex. 8)
      vreso - ColorVector resolution of x (ex. 2)
      imgid - ID of image
      fp    - fingerprint of image (ex. 8 x 8 x 3ch)

>>> let fp = "000000010101020202030303010101020202030303040404020202030303040404050505030303040404050505060606"
>>> fp2colvec 4 2 1 fp "filed"
ColorVector {pixel = [[1.0,1.0,1.0],[3.0,3.0,3.0],[3.0,3.0,3.0],[5.0,5.0,5.0]], imgid = 1, fp = "000000010101020202030303010101020202030303040404020202030303040404050505030303040404050505060606", status = FILED}

-}

fp2colvec :: Int -> Int -> Int -> Fingerprint -> String -> ColorVector
fp2colvec oreso vreso imgid fp st = ColorVector pixel imgid fp st'
  where
    chs = map (matrix oreso) (fp2channels fp)
    d = oreso `div` vreso
    pixel = transpose $ map (summary d vreso) chs
    st' = case st of
      "filed"      -> FILED
      "pending"    -> PEND
      "deleted"    -> DISCARDED
      "duplicated" -> DUPLICATED
      "inferior"   -> INFERIOR
      _            -> DISCARDED

colvec2id :: ColorVector -> Int
colvec2id (ColorVector _ i _ _) = i

colvec2fp :: ColorVector -> [Double]
colvec2fp (ColorVector _ _ fp _) = fp2double fp

colvec2status :: ColorVector -> CvStatus
colvec2status (ColorVector _ _ _ st) = st

isValid :: ColorVector -> Bool
isValid (ColorVector _ _ _ st) = (st == FILED || st == PEND)

-- internal functions

{-|

>>> let cv1 = ColorVector [[1.0, 2.0, 3.0], [4.0, 5.0, 6.0], [7.0, 8.0, 9.0], [10.0, 11.0, 12.0]] 1 "001122334455" FILED
>>> colvec2points cv1
[1.0,2.0,3.0,4.0,5.0,6.0,7.0,8.0,9.0,10.0,11.0,12.0]

-}

colvec2points :: ColorVector -> [Double]
colvec2points (ColorVector points _ _ _) = concat points

{-|

>>> let cv1 = ColorVector [[1.0, 2.0, 3.0], [4.0, 5.0, 6.0], [7.0, 8.0, 9.0], [10.0, 11.0, 12.0]] 1 "001122334455" FILED
>>> let cv2 = ColorVector [[1.0, 1.0, 4.0], [3.0, 6.5, 6.0], [8.2, 8.0, 7.8], [10.0, 12.0, 11.0]] 2 "101122334455" FILED
>>> colvecDistance cv1 cv2
2.5

-}

colvecDistance :: ColorVector -> ColorVector -> Double
colvecDistance cv1 cv2 = maximum ls
  where
    p1s = colvec2points cv1
    p2s = colvec2points cv2
    ls = map sum $ chunksOf 3 $ zipWith (\x y -> abs (x - y)) p1s p2s

{-|

>>> let m = matrix 4 [0,1,2,3,1,2,3,4,2,3,4,5,3,4,5,6]
>>> summary 2 2 m
[1.0,3.0,3.0,5.0]
>>> let m2 = matrix 4 [6,5,4,3,5,4,3,2,4,3,2,1,3,2,1,0]
>>> summary 2 2 m2
[5.0,3.0,3.0,1.0]

-}

summary :: Int -> Int -> Matrix R -> [Double]
summary d vreso m = map (\x -> x / fromIntegral (d*d)) col
  where
    ele = map (*d) [0..(vreso-1)]
    offset = [(x, y)| x <- ele, y <- ele]
    col = map (\o -> sum.toList.flatten $ subMatrix o (d, d) m) offset

{-|

>>> fp2channels "000102030405060708090a0b"
[[0.0,3.0,6.0,9.0],[1.0,4.0,7.0,10.0],[2.0,5.0,8.0,11.0]]

-}

fp2channels :: Fingerprint -> [Channel]
fp2channels fp = splitRgb ds ([], [], [])
  where
    ds = fp2double fp
    splitRgb :: [Double] -> (Channel, Channel, Channel) -> [Channel]
    splitRgb [] (r, g, b) = [reverse r, reverse g, reverse b]
    splitRgb (r:g:b:xs) (rs, gs, bs) = splitRgb xs ((r:rs), (g:gs), (b:bs))

{-|

>>> fp2double "000102030405060708090a0b0c0d0e0f"
[0.0,1.0,2.0,3.0,4.0,5.0,6.0,7.0,8.0,9.0,10.0,11.0,12.0,13.0,14.0,15.0]
>>> fp2double "00102030405060708090a0b0c0d0e0f0"
[0.0,16.0,32.0,48.0,64.0,80.0,96.0,112.0,128.0,144.0,160.0,176.0,192.0,208.0,224.0,240.0]

-}

fp2double :: Fingerprint -> [Double]
fp2double fp = map hex2double $ chunksOf 2 fp

hex2double :: String -> Double
hex2double hex = fromIntegral h
  where
    [(h, _)] = (readHex hex)

max_stars_repo_name: eiji-a/epconv

max_issues_repo_name: eiji-a/epconv

max_forks_repo_name: eiji-a/epconv


text:
{-# LANGUAGE RankNTypes, FlexibleContexts, TypeFamilies #-}
module Quipp.Util where

import Control.Monad (liftM)
import Control.Monad.State.Lazy (runState)
import Control.Monad.Identity
import Debug.Trace
import Data.Function (on)
import Data.List (groupBy, sortBy)
import Data.Map (Map)
import qualified Data.Map as Map
import Data.Random (MonadRandom, RandomSource, RVarT, RVar, StdRandom(StdRandom), sampleState, runRVar, runRVarT, runRVarTWith, stdUniform)
import qualified Data.Packed.Matrix as Mat
import Numeric.LinearAlgebra.Algorithms (linearSolve, pinv)
import Numeric.AD (diff, Mode, Scalar)
import System.Random (StdGen, mkStdGen)

fst3 (a, _, _) = a
snd3 (_, b, _) = b
thd3 (_, _, c) = c

insertAll :: Ord k => [(k, v)] -> Map k v -> Map k v
insertAll kvs m = foldr (\(k, v) m -> Map.insert k v m) m kvs

infixr 9 .:
(f .: g) x y = f (g x y)

traced label x = trace (label ++ " " ++ show x) x

takeEvery :: Int -> [a] -> [a]
takeEvery _ [] = []
takeEvery n (x:xs) = x : takeEvery n (drop (n-1) xs)

sampleRVar :: MonadRandom m => RVar a -> m a
sampleRVar v = runRVar v StdRandom

sampleRVarT v = runRVarT v StdRandom

sampleRVarTWith :: RandomSource m StdRandom => (forall t. n t -> m t) -> RVarT n a -> m a
sampleRVarTWith f v = runRVarTWith f v StdRandom

infinity :: Double
infinity = read "Infinity"

negInfinity :: Double
negInfinity = read "-Infinity"

funPow :: Int -> (a -> a) -> a -> a
funPow n f x = iterate f x !! n

zipWithSameLength f [] [] = []
zipWithSameLength f (x:xs) (y:ys) = f x y : zipWithSameLength f xs ys
zipWithSameLength _ _ _ = error "zipWithSameLength: different lengths"

zipSameLength = zipWithSameLength (,)

iterateM :: Monad m => Int -> (a -> m a) -> a -> m [a]
iterateM 0 _ x = return [x]
iterateM n f x = liftM (x:) (f x >>= iterateM (n-1) f)

stateInfList :: (s -> (a, s)) -> s -> [a]
stateInfList f s =
  let (a, s') = f s in a : stateInfList f s'

iterateRVar :: (a -> RVar a) -> a -> RVar [a]
iterateRVar f x = do
  seed <- stdUniform
  return $ stateInfList (\(y, gen) -> let (y', gen') = sampleState (f y) gen in (y', (y', gen'))) (x, mkStdGen seed)

groupAnywhereBy :: Ord b => (a -> b) -> [a] -> [[a]]
groupAnywhereBy f = groupBy ((==) `on` f) . sortBy (compare `on` f)

logSumExp :: RealFloat a => [a] -> a
logSumExp lps = mx + log (sum [exp (lp - mx) | lp <- lps])
  where mx = maximum lps

-- logSumExp :: RealFloat a => [a] -> a
-- logSumExp lps = log $ sum $ map exp lps

logProbsToProbs :: [Double] -> [Double]
logProbsToProbs lps = [exp (lp - lse) | lp <- lps]
  where lse = logSumExp lps

realToDouble :: Real s => s -> Double
realToDouble = fromRational . toRational

type Matrix a = [[a]]

diagonalEntries :: Matrix a -> [a]
diagonalEntries ((x:xs):rows) = x : diagonalEntries (map tail rows)

outerProduct :: Num a => [a] -> [a] -> Matrix a
outerProduct as bs = [[a*b | b <- bs] | a <- as]

splitListIntoBlocks :: Int -> [a] -> [[a]]
splitListIntoBlocks n lst
  | blocksize * n /= length lst = undefined
  | blocksize == 0 = replicate n []
  | otherwise = go blocksize lst
  where blocksize = length lst `div` n
        go _ [] = []
        go k lst = take k lst : go k (drop k lst)

scaleVec :: Num a => a -> [a] -> [a]
scaleVec x = map (x *)

transpose xs = if maximum (map length xs) == 0 then [] else map head xs : transpose (map tail xs)

dotProduct :: Num a => [a] -> [a] -> a
dotProduct x y = sum (zipWith (*) x y)

squareMagnitude :: Num a => [a] -> a
squareMagnitude x = dotProduct x x

matMulByVector :: Num a => Matrix a -> [a] -> [a]
matMulByVector m v = map (dotProduct v) m

matMul :: Num a => Matrix a -> Matrix a -> Matrix a
matMul m1 m2 = map (matMulByVector m1) (transpose m2)

matInv :: Matrix Double -> Matrix Double
matInv = Mat.toLists . pinv . Mat.fromLists

linSolve :: Matrix Double -> [Double] -> [Double]
linSolve mat d =
  matMulByVector (Mat.toLists $ pinv $ Mat.fromLists mat) d

diagEntries :: Matrix a -> [a]
diagEntries m
  | length m == 0 || length m == length (head m) =
    zipWith (!!) m [0..]
  | otherwise = error $ "Cannot get diagonal entries of non-square matrix " ++ show (length m, length (head m))

fromDouble :: Fractional a => Double -> a
fromDouble = fromRational . toRational

toDouble :: Real a => a -> Double
toDouble = fromRational . toRational

{- f(x) = ax^2 + bx + c
 - f'(x) = 2ax + b
 - f''(x) = 2a
 - a = f''(x) / 2
 - b = f'(x) - 2ax
 - c = f(x) - bx - ax^2
 -}
quadApproximation :: (forall a. (Mode a, RealFloat a) => a -> a) -> Double -> (Double, Double, Double)
quadApproximation f x =
  let deriv = diff f x
      deriv2 = diff (diff f :: (forall a. (Mode a, RealFloat a) => a -> a)) x
      a = deriv2 / 2
      b = deriv - 2 * a * x
      c = f x - b * x - a * x * x
  in (c, b, a)

mean xs = sum xs / fromIntegral (length xs)

variance xs = sum [(x-m)^2 | x <- xs] / fromIntegral (length xs) where m = mean xs

covariance xys = sum [(x - ux) * (y - uy) | (x, y) <- xys] / fromIntegral (length xys)
  where ux = mean (map fst xys)
        uy = mean (map snd xys)


text:
module Shader ( sampleBRDF, probability, reflectance
              , sampleVec
              , sampleCol
              , sampleVal
              ) where

import Types
import Numeric.Vector
import Numeric.Scalar (fromScalar, scalar)
import Control.Monad.State
import VecUtil (unpack3, unpack2)
import qualified SceneRandom as R
import VecUtil (vec3Of)

-- UTIL STUFF

fmod :: (RealFrac a) => a -> a -> a
fmod n d = n - (fromIntegral $ truncate $ n / d) * d

-- END UTIL STUFF

-- BRDF STUFF

sampleBRDF :: BRDF -> State SceneContext (Maybe Ray)
sampleBRDF BRDFEmpty = return Nothing
sampleBRDF (Emission _ _) = return Nothing
sampleBRDF (Diffuse _) = do
  ctx <- get
  let normal = rh_getNormal $ ss_getHit ctx
  let pos = (rh_getPos (ss_getHit ctx)) + (normal * (vec3Of 0.01))
  dir <- R.randVecHemisphere normal
  return $ Just $ Ray pos dir

probability :: BRDF -> Vec3d -> State SceneContext Double
probability BRDFEmpty _ = return 0
probability (Emission _ _) _ = return 0.15915494309
probability (Diffuse _) _ = return 0.15915494309 -- 1 / (2 * PI)

reflectance :: BRDF -> State SceneContext Vec3d
reflectance BRDFEmpty = return $ vec3 0 0 0

reflectance (Emission col str) = do
  col' <- sampleCol col
  str' <- sampleVal str
  return $ col' * (fromScalar (scalar str'))

reflectance (Diffuse col) = do
  col' <- sampleCol col
  return $ col'

reflectance (Glossy col roughness) = do
  col' <- sampleCol col
  return col'

-- END BRDF STUFF

-- VECTOR STUFF

sampleVec :: SVector -> State SceneContext Vec3d

sampleVec UV = do
  ss <- get
  let (u, v) = unpack2 $ rh_getTexCoord $ ss_getHit ss
  return $ vec3 u v 0

sampleVec Position = do
  ss <- get
  return $ rh_getPos $ ss_getHit ss

sampleVec Normal = do
  ss <- get
  return $ rh_getNormal $ ss_getHit ss

sampleVec (VecConst x y z) = return $ vec3 x y z

sampleVec (VecMath op l r) = do
  l' <- sampleVec l
  r' <- sampleVec r
  case op of
    VAdd -> return $ l' + r'
    VSub -> return $ l' - r'
    VMul -> return $ l' * r'
    VDiv -> return $ l' / r'
    VMod -> let (x,y,z) = unpack3 l'; (i,j,k) = unpack3 r' in return $ vec3 (fmod x i) (fmod y j) (fmod z k)
    VAbs -> return $ abs l'
    -- VDot -> return $ l' · r'
    -- VCross -> return $ l' × r'

sampleVec (CombineXYZ x y z) = do
  x' <- sampleVal x
  y' <- sampleVal y
  z' <- sampleVal z
  return $ vec3 x' y' z'

-- END VECTOR STUFF

-- COLOR STUFF

sampleCol :: SColor -> State SceneContext Vec3d

sampleCol ColVertex = do
  ss <- get
  return $ rh_getColor $ ss_getHit ss

sampleCol (ColVector v) = sampleVec v

sampleCol (ColConst r g b) = return $ vec3 r g b

sampleCol (CombineRGB vr vg vb) = do
  r <- sampleVal vr
  g <- sampleVal vg
  b <- sampleVal vb
  return $ vec3 r g b

-- END COLOR STUFF

-- VALUE STUFF
sampleVal :: SValue -> State SceneContext Double

sampleVal (ValConst x) = return x

sampleVal RayLength = do
  ss <- get
  return $ rh_getDistance $ ss_getHit ss

sampleVal (ValMath op l r) = do
  l' <- sampleVal l
  r' <- sampleVal r
  case op of
    Add -> return $ l' + r'
    Sub -> return $ l' - r'
    Mul -> return $ l' * r'
    Div -> return $ l' / r'
    Abs -> return $ abs l'
    Mod -> return $ fmod l' r'
    LessThan -> return $ if l' < r' then 1 else 0
    GreaterThan -> return $ if l' > r' then 1 else 0

sampleVal (SeparateX v) = sampleVec v >>= (return . unpack3) >>= \(x,_,_) -> return x
sampleVal (SeparateY v) = sampleVec v >>= (return . unpack3) >>= \(_,y,_) -> return y
sampleVal (SeparateZ v) = sampleVec v >>= (return . unpack3) >>= \(_,_,z) -> return z

sampleVal (SeparateR v) = sampleCol v >>= (return . unpack3) >>= \(r,_,_) -> return r
sampleVal (SeparateG v) = sampleCol v >>= (return . unpack3) >>= \(_,g,_) -> return g
sampleVal (SeparateB v) = sampleCol v >>= (return . unpack3) >>= \(_,_,b) -> return b

-- END VALUE STUFF


text:
{-# LANGUAGE DataKinds #-}
{-# LANGUAGE TypeOperators #-}
{-# LANGUAGE TypeApplications #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE ViewPatterns #-}
{-# LANGUAGE KindSignatures #-}
{-# LANGUAGE TypeFamilies #-}
{-# LANGUAGE RankNTypes #-}
{-# LANGUAGE GADTs #-}

{-# OPTIONS_GHC -fplugin GHC.TypeLits.KnownNat.Solver #-}
{-# OPTIONS_GHC -fplugin GHC.TypeLits.Normalise #-}

{-|
Module      : Numeric.LinearProgramming.LP.Static
Description : Solve linear programs using static vectors/matrices.
Copyright   : (c) Ryan Orendorff, 2020
License     : BSD3
Stability   : experimental
-}
module Numeric.LinearProgramming.LP.Static (
  affineScaling
)
where

import           GHC.TypeLits
import           Numeric.LinearAlgebra.Static
import qualified Numeric.LinearAlgebra.Static as S
import qualified Numeric.LinearAlgebra as LA
import Numeric.UnitaryRational

-- | https://en.wikipedia.org/wiki/Karmarkar%27s_algorithm
affineScaling :: (KnownNat m, KnownNat n, KnownNat p, KnownNat q) => L m n -> R m -> R n -> R n -> UnitaryRational p q -> [Maybe (R n)]
affineScaling a b c x0 gamma = iterate (>>= go) (Just x0)
    where
        gamma' = unitaryToFractional gamma

        min_vh :: (KnownNat n) => R n -> R n -> Double
        min_vh v h = minimum . map fst . filter ((<0) . snd) $ vals_with_h
            where
                vals = zipWithVector (\vi hi -> -vi/hi) v h
                vals_with_h = zip (LA.toList $ extract vals) (LA.toList $ extract h)

        go xk = let
            vk = b - a #> xk
            dm_negsq = diag $ dvmap (**(-2)) vk
            hx = inv ((tr a) S.<> dm_negsq S.<> a) #> c
            hv = ((konst (-1)) * a #> hx)
            alpha = gamma' * min_vh vk hv
            xkp1 = xk + (konst alpha) * hx
            in
                if ((>=0) . maximum . LA.toList . extract $ hv)
                then Nothing
                else Just xkp1

-- Example
-- -------
--
--
-- p = vector [0.0, 0.1, 0.2, 0.3, 0.4, 0.5, 0.6, 0.7, 0.8, 0.9, 1.0] :: R 11
--
-- -- min c^Tx
-- c = vec2 1 1
--
-- -- s.t Ax <= b
-- a = ((konst 2) * (col p)) ||| (1 :: L 11 1)
-- b = p**2 + 1
--
-- lpSol = affineScaling a b c (konst 0) (UnitaryRational :: UnitaryRational 50 100)
--
-- mapM_ (disp 10 . fromJust) . take 20 $ lpSol

max_stars_repo_name: ryanorendorff/convex

max_issues_repo_name: ryanorendorff/convex

max_forks_repo_name: ryanorendorff/convex


text:
{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE TypeFamilies     #-}
-- |
-- Module      : Main
-- Description : Test runner
-- Copyright   : (c) Tom Westerhout, 2017
-- License     : BSD3
-- Maintainer  : t.westerhout@student.ru.nl
-- Stability   : experimental

module Main where

import qualified System.Random.MWC as MWC
import           Control.Lens
import           Control.Monad.Reader
import           Control.Monad.State
import qualified Data.List            as L
import Data.Complex
import qualified Data.Vector.Storable as V
import qualified Numeric.LinearAlgebra as LA
import Data.Vector.Storable((!))

import PSO.Random
import PSO.Energy
import PSO.Swarm

class TestFunction f where
  function :: f -> V.Vector Float -> Float
  initBounds :: f -> (Float, Float)
  dim :: f -> Int

data RosenbrockFn = RosenbrockFn

instance TestFunction RosenbrockFn where
  function _ x =
    let p1 = (100.0 *) . V.sum . V.map (^^2)
                $ V.zipWith (-) (V.tail x) (V.map (^^2) $ V.init x)
        p2 = V.sum . V.map (^^2) . V.map (1.0 -) $ V.init x
    in p1 + p2
  initBounds _ = (15.0, 30.0)
  dim _ = 30

unpack :: (Float, Float) -> Int -> (V.Vector Float, V.Vector Float)
unpack (x, y) d = (f x, f y)
  where f = V.fromList . replicate d

wpg :: (Float, Float, Float)
wpg = (0.7298 :: Float, 1.49618, 1.49618)

cmUpdater :: (Num χ, RandomScalable m χ, VectorSpace χ Float)
        => PhaseUpdater m (SwarmGuide χ r) (BeeGuide χ r) (CMState χ) r
cmUpdater = standardUpdater wpg

qmUpdater :: (RandomScalable m χ, Randomisable m Float, χ ~ V.Vector Float)
        => PhaseUpdater m (SwarmGuide χ r) (BeeGuide χ r) (QMState χ) r
qmUpdater = (PhaseUpdater $ deltaWellUpdater (2 * log 2 * 0.7 :: Float))

cmRunND :: IO ()
cmRunND = do
  let xs = optimiseND
            (mkCMState (unpack (initBounds RosenbrockFn) (dim RosenbrockFn)))
            cmUpdater
            (function RosenbrockFn)
            20
            (\s -> (s^.guide.val) < 1.0
                   || (s^.guide.iteration == 10000))
  gen <- mkMWCGen (Just 123)
  swarms <- runReaderT xs gen
  writeEnergies2TSV "Function.dat" (view val) swarms
  let swarm = last swarms
  -- mapM_ (print . (!!0) . (view bees)) $ swarms
  putStrLn ""
  putStrLn $ "[+] Best[f] = " ++ show (swarm ^. guide . val)

qmRunND :: IO ()
qmRunND = do
  let xs = optimiseND
            (mkQMState (unpack (initBounds RosenbrockFn) (dim RosenbrockFn)))
            qmUpdater
            (function RosenbrockFn)
            20
            (\s -> (s^.guide.val) < 1.0
                   || (s^.guide.iteration == 5000))
  gen <- mkMWCGen (Just 123)
  swarms <- runReaderT xs gen
  writeEnergies2TSV "Function.dat" (view val) swarms
  let swarm = last swarms
  -- mapM_ (print . (!!0) . (view bees)) $ swarms
  putStrLn ""
  putStrLn $ "[+] Best[f] = " ++ show (swarm ^. guide . val)

main :: IO ()
main = qmRunND

max_stars_repo_name: twesterhout/tcm-swarm

max_issues_repo_name: twesterhout/tcm-swarm

max_forks_repo_name: twesterhout/tcm-swarm


text:
{-# language BangPatterns #-}
{-# language GeneralizedNewtypeDeriving #-}
{-# language MultiParamTypeClasses #-}
{-# language RankNTypes #-}
{-# language RebindableSyntax #-}
{-# language TypeFamilies #-}
module W40K.Core.Prob
  ( QQ
  , Event(..)
  , Prob
  , events
  , traceLength
  , traceEvents
  , traceNumEvents
  , fmapProb
  , fmapProbMonotone
  , fmapProbMonotone'
  , (|>>=|), (|=<<|)
  , uniformly
  , probTrue
  , probFalse
  , probOf
  , given
  , bernoulli
  , binomial
  , foldlProbs'
  , foldrProbs
  , foldProbs
  , sumProbs
  , foldAssocIID
  , foldIID
  , sumIID
  , addImpossibleEvents
  , cdf
  , ccdf
  , summary
  , summaryInt
  , mean
  , variance
  , stDev
  , maxEvent
  ) where

import Prelude hiding (Functor(..), Applicative(..), Monad(..))

import Control.DeepSeq (NFData(..))

import Data.List (foldl', sort)
import Data.MemoTrie (memo)

import Numeric.SpecFunctions (choose)

import Debug.Trace (trace)

import W40K.Core.SortedList (SortedList, SortedListItem(..))
import qualified W40K.Core.SortedList as SortedList

import W40K.Core.ConstrMonad as ConstrMonad
import W40K.Core.Util

type QQ = Double

data Event a = Event !a {-# unpack #-} !QQ
    deriving (Eq, Ord)

mapEvent :: (a -> b) -> Event a -> Event b
mapEvent f (Event a p) = Event (f a) p
{-# inline mapEvent #-}

instance NFData a => NFData (Event a) where
    rnf (Event a _) = rnf a

instance Ord a => SortedListItem (Event a) where
    type ItemKey (Event a) = a
    itemKey (Event a _) = a
    combineItems (Event _ p) (Event a p') = Event a (p + p')

instance Show a => Show (Event a) where
    show (Event a p) = show a ++ ": " ++ show (realToFrac p :: Float)

newtype Prob a = Prob { density :: SortedList (Event a) }
    deriving (Eq, Ord, Show, NFData)

events :: Prob a -> [Event a]
events (Prob es) = SortedList.toAscList es
{-# inline events #-}

traceLength :: [a] -> [a]
traceLength as = trace (show (length as)) as

traceEvents :: Show a => Prob a -> Prob a
traceEvents df@(Prob _) =
    trace (show (events df)) df

traceNumEvents :: Prob a -> Prob a
traceNumEvents df@(Prob _) =
    trace (show (length (events df))) df

fmapProb :: Ord b => (a -> b) -> Prob a -> Prob b
fmapProb f (Prob evts) = Prob (SortedList.map (mapEvent f) evts)
{-# inline fmapProb #-}

{-# rules "fmap/fmapProb" fmap = fmapProb #-}
{-# specialize fmapProb :: (a -> Int) -> Prob a -> Prob Int #-}

-- only for monotone f
fmapProbMonotone :: Ord b => (a -> b) -> Prob a -> Prob b
fmapProbMonotone f (Prob evts) = Prob (SortedList.mapMonotone (mapEvent f) evts)
{-# inline fmapProbMonotone #-}

-- only for *strictly* monotone f
fmapProbMonotone' :: (a -> b) -> Prob a -> Prob b
fmapProbMonotone' f (Prob evts) = Prob (SortedList.mapMonotone' (mapEvent f) evts)
{-# inline fmapProbMonotone' #-}

bindProbWithStrat :: Ord b => (forall c. [c] -> [c]) -> Prob a -> (a -> Prob b) -> Prob b
bindProbWithStrat evalList (Prob evts) f =
    case SortedList.toAscList evts of
      [Event a _] -> f a
      evtList     -> Prob $ SortedList.concat $ evalList $ totalProb [Event (f a) p | Event a p <- evtList]
  where
    totalProb :: [Event (Prob a)] -> [SortedList (Event a)]
    totalProb ess = [SortedList.fromAscList [Event b (p*p') | Event b p' <- events y] | Event y p <- ess]

{-# inlinable bindProbWithStrat #-}

instance ConstrMonad Ord Prob where
    return a = Prob (SortedList.singleton (Event a 1))
    {-# inline return #-}
    (>>=) = bindProbWithStrat seqItems
    {-# inline (>>=) #-}
    {-# specialize (>>=) :: Prob a -> (a -> Prob Int) -> Prob Int #-}

(|>>=|) :: (Ord a, Ord b) => Prob a -> (a -> Prob b) -> Prob b
(|>>=|) = bindProbWithStrat parItems
{-# inline (|>>=|) #-}

(|=<<|) :: (Ord a, Ord b) => (a -> Prob b) -> Prob a -> Prob b
f |=<<| ma = ma |>>=| f
{-# inline (|=<<|) #-}

uniformly :: Ord a => [a] -> Prob a
uniformly as = Prob $ SortedList.fromAscList [Event a (1/n) | a <- sort as]
  where
    n = fromIntegral (length as)

probTrue :: Prob Bool -> QQ
probTrue prob =
    case events prob of
        [Event False _, Event True p] -> p
        [Event True p]                -> p
        [Event False q]               -> 1 - q
        []                            -> error "empty Prob!"
        _                             -> error ("unnormalized Prob!" ++ show prob)

probFalse :: Prob Bool -> QQ
probFalse df = 1 - probTrue df

probOf :: (a -> Bool) -> Prob a -> QQ
probOf test df = sum [p | Event a p <- events df, test a]

given :: (a -> Bool) -> Prob a -> Prob a
given hyp prob =
    Prob (SortedList.fromAscList [Event a (p / probHyp) | Event a p <- events prob, hyp a])
  where
    probHyp = probOf hyp prob

bernoulli :: QQ -> Prob Bool
bernoulli p = Prob $ SortedList.fromAscList [Event False (1-p), Event True p]

binomial :: Int -> QQ -> Prob Int
binomial 0 _ = return 0
binomial n p = binomialFlip p n

binomialFlip :: QQ -> Int -> Prob Int
binomialFlip p
  | p == 1/2  = binomial12
  | p == 1/3  = binomial13
  | p == 1/6  = binomial16
  | p == 2/3  = binomial23
  | otherwise = binomialMemo p

binomialMemo :: QQ -> Int -> Prob Int
binomialMemo p = memo $ \n ->
    case p of
      0 -> return 0
      1 -> return n
      _ -> Prob $ SortedList.fromAscList [Event k (binomProbOf n p k) | k <- [0..n]]
  where
    binomProbOf :: Int -> QQ -> Int -> QQ
    binomProbOf n p k
      | k < 0 || k > n = 0
      | n == 0         = 1
      | otherwise      = realToFrac (choose n k) * p^k * (1-p)^(n-k)

binomial23 :: Int -> Prob Int
binomial23 = binomialMemo (2/3)

binomial12 :: Int -> Prob Int
binomial12 = binomialMemo (1/2)

binomial13 :: Int -> Prob Int
binomial13 = binomialMemo (1/3)

binomial16 :: Int -> Prob Int
binomial16 = binomialMemo (1/6)

foldlProbs' :: (Ord a, Ord b) => (b -> a -> b) -> Prob b -> [Prob a] -> Prob b
foldlProbs' = foldl' . liftA2

foldrProbs :: (Ord a, Ord b) => (a -> b -> b) -> Prob b -> [Prob a] -> Prob b
foldrProbs = foldr . liftA2

foldProbs :: (Ord m, Monoid m) => [Prob m] -> Prob m
foldProbs = foldlProbs' mappend (return mempty)

sumProbs :: (Ord a, Num a) => [Prob a] -> Prob a
sumProbs []     = return 0
sumProbs [p]    = p
sumProbs (p:ps) = foldlProbs' (+) p ps
{-# specialize sumProbs :: [Prob Int] -> Prob Int #-}

foldAssocIID :: (Ord a) => (a -> a -> a) -> a -> Int -> Prob a -> Prob a
foldAssocIID _ z 0 _ = return z
foldAssocIID f z n p
  | n < 0     = error "foldAssocIID: n must be >= 0"
  | otherwise = noCheck f z n p
  where
    noCheck _ _ 1 p = p
    noCheck f z n p
      | n `mod` 2 == 0 = twiceHalf
      | otherwise      = liftA2 f p twiceHalf
      where
        !twiceHalf = liftA2 f recHalf recHalf
        !recHalf   = foldAssocIID f z (n`div`2) p

foldIID :: (Ord a, Monoid a) => Int -> Prob a -> Prob a
foldIID = foldAssocIID mappend mempty

-- TODO: Option to use Central Limit Theorem for sufficiently large n
sumIID :: (Ord a, Num a) => Int -> Prob a -> Prob a
sumIID = foldAssocIID (+) 0

addImpossibleEvents :: (Ord a, Enum a) => Prob a -> Prob a
addImpossibleEvents prob =
    let es = events prob
        as = [a | Event a _ <- es]
    in
        Prob $ SortedList.fromAscList $ merge es (zipWith Event [minimum as .. maximum as] (repeat 0))
  where
    merge es []  = es
    merge [] es' = es'
    merge ees@(e@(Event a _):es)
          ees'@(e'@(Event a' _):es')
      | a == a' = e : merge es es'
      | a <  a' = e : merge es ees'
      | a >  a' = e' : merge ees es'

cdf :: Ord a => Prob a -> [Event a]
cdf = scanl1 sumEvents . events
  where
    sumEvents (Event _ p) (Event a p') = Event a (p + p')

ccdf :: Ord a => Prob a -> [Event a]
ccdf = scanr1 sumEvents . events
  where
    sumEvents (Event a p) (Event _ p') = Event a (p + p')

summary :: Prob QQ -> IO ()
summary p =
  let mu    = "µ=" ++ show (realToFrac (mean p) :: Double)
      sigma = "σ=" ++ show (sqrt $ realToFrac (variance p) :: Double)
  in
      putStrLn $ mu ++ ", " ++ sigma

summaryInt :: Prob Int -> IO ()
summaryInt = summary . fmap fromIntegral

mean :: Prob QQ -> QQ
mean df = sum [k * p | Event k p <- events df]

variance :: Prob QQ -> QQ
variance df = sum [squared k * p | Event k p <- events df] - squared (mean df)
  where
    squared x = x*x

stDev :: Prob QQ -> QQ
stDev df = sqrt (variance df)

maxEvent :: Ord a => Prob a -> a
maxEvent df = maximum [a | Event a _ <- events df]


text:
{-# LANGUAGE FlexibleContexts #-}
module FourierPinwheel.Filtering where

import           Data.Array.Repa            as R
import           Data.Complex
import           Data.List                  as L
import           Data.Vector.Generic        as VG
import           Data.Vector.Storable       as VS
import           Data.Vector.Unboxed        as VU
import           DFT.Plan
import           Filter.Utils
import           Numeric.GSL.Special.Bessel
import           Utils.List

{-# INLINE idealLowPassFilter #-}
idealLowPassFilter ::
     (VG.Vector vector (Complex Double))
  => Double
  -> Double
  -> Int
  -> Int
  -> vector (Complex Double)
idealLowPassFilter radius period numRhoFreqs numRFreqs =
  let centerRhoFreq = div numRhoFreqs 2
      centerRFreq = div numRFreqs 2
      sinc x =
        if x == 0
          then 1
          else sin (pi * x) / (pi * x)
   in VG.convert .
      toUnboxed . computeS . fromFunction (Z :. numRFreqs :. numRhoFreqs) $ \(Z :. i' :. j') ->
        let i = i' - centerRFreq
            j = j' - centerRhoFreq
         in -- if i == (j)
            --   then
          (2 * radius) / period *
                   sinc (fromIntegral j  / period * (2 * radius)) :+
                   0
              -- else 0

{-# INLINE applyFilterRho #-}
applyFilterRho ::
     DFTPlan
  -> VU.Vector (Complex Double)
  -> R.Array U DIM4 (Complex Double)
  -> IO (R.Array U DIM4 (Complex Double))
applyFilterRho plan filterF arr' = do
  let (Z :. numRFreqs :. numThetaFreqs :. numRhoFreqs :. numPhiFreqs) =
        extent arr'
      arr =
        computeUnboxedS $
        R.backpermute
          (Z :. numRFreqs :. numRhoFreqs :. numThetaFreqs :. numPhiFreqs)
          (\(Z :. r :. rho :. theta :. phi) -> (Z :. r :. theta :. rho :. phi))
          arr'
  arrF <-
    fmap (fromUnboxed (extent arr) . VG.convert) .
    dftExecute
      plan
      (DFTPlanID
         DFT1DG
         [numRFreqs, numRhoFreqs, numThetaFreqs, numPhiFreqs]
         [0, 1]) .
    VG.convert . toUnboxed $
    arr
  let convolvedArrF =
        computeUnboxedS .
        R.traverse2
          arrF
          (fromUnboxed (Z :. numRFreqs :. numRhoFreqs) filterF)
          const $ \fArr fFilter idx@(Z :. r :. rho :. theta :. phi) ->
          fArr idx * fFilter (Z :. r :. rho)
  fmap
    (computeS .
     R.backpermute
       (Z :. numRFreqs :. numThetaFreqs :. numRhoFreqs :. numPhiFreqs)
       (\(Z :. r :. theta :. rho :. phi) -> (Z :. r :. rho :. theta :. phi)) .
     fromUnboxed (extent arr) . VG.convert) .
    dftExecute
      plan
      (DFTPlanID
         IDFT1DG
         [numRFreqs, numRhoFreqs, numThetaFreqs, numPhiFreqs]
         [0, 1]) .
    VG.convert . toUnboxed $
    convolvedArrF

hollowCoefficients ::
     DFTPlan
  -> Double
  -> Double
  -> R.Array U DIM4 (Complex Double)
  -> IO (R.Array U DIM4 (Complex Double))
hollowCoefficients plan radius period coefficients = do
  let (Z :. numRFreqs :. _ :. numRhoFreqs :. _) = extent coefficients
      lpf = idealLowPassFilter (log radius) (log period) numRhoFreqs numRFreqs
      filter =
        VG.convert .
        toUnboxed .
        computeS . makeFilter2D . fromUnboxed (Z :. numRFreqs :. numRhoFreqs) $
        lpf
  filterF <-
    VG.convert <$>
    dftExecute plan (DFTPlanID DFT1DG [numRFreqs, numRhoFreqs] [0, 1]) filter
  -- computeS . R.zipWith (-) coefficients <$>
  applyFilterRho plan filterF coefficients

max_stars_repo_name: XinhuaZhang/Stochastic-Completion-Field

max_issues_repo_name: XinhuaZhang/Stochastic-Completion-Field

max_forks_repo_name: XinhuaZhang/Stochastic-Completion-Field


text:
{-# LANGUAGE RankNTypes #-}

module Val
  ( -- * Val and support types
    Val (..), Symbol, Port(..)
  , Number (..), RealNumber (..)
  , LispErr (..), isTerminationError
  , Primitive (..)
  , Macro (..)
  , Builtin, InTail, Arity (..)

    -- * Utilities for 'Val's
  , truthy, isConstant
  , makePrimitive
  , showArity, usePluralForArity, testArity
  , symbolName, symbolAsString

    -- * Re-export Text type
  , T.Text, T.pack, T.unpack

    -- * Basic handling and construction of lists
  , getList, getListOrError, testCircularList
  , isListSH, requireList, makeMutableList, makeImmutableList
  , consSSS, makeImproperMutableList, makeImproperImmutableList

    -- * Freezing lists for evaluation
  , FrozenList (..), freezeList, lintFrozenList

    -- * Direct handling of pairs
  , carPR, cdrPR, carPS, cdrPS
  , setCarSSS, setCdrSSS, pairSH

    -- * Vectors
  , vectorSH
  , makeVector, makeByteVector

    -- * Test for immutable data
  , isImmutablePair, isImmutable

    -- * Making and working with Numbers
  , makeBignum, makeFlonum, makeRatnum
  , approxToRatnum, realInexact
  , getExactInteger, isExactZero
  , numberSH, complexSH, realSH, rationalSH, integerSH
  , implicitRealConversion, implicitNumericConversion

    -- * Showing 'Val' in IO
  , showValIO
  ) where

import Types
import EvaluationMonad
import Control.Monad.Trans.Maybe (MaybeT(..))
import Data.Bits (shiftL, bit)
import Data.Complex (Complex((:+)), magnitude)
import Data.IORef (readIORef) -- for show in IO
import Data.Foldable (foldrM)
import Data.Ratio ((%), numerator, denominator)
import qualified Data.Text as T
import Data.Word (Word8)

import qualified Data.Vector as V
import qualified Data.Vector.Unboxed as U
import Data.Maybe (fromMaybe)
import System.Mem.StableName (makeStableName)

import GHC.Float (Floating(..))

makePrimitive :: Primitive -> Val
makePrimitive (Prim name arity func) = Primitive arity func name

-- | Is this 'Val' truthy?
truthy :: Val -> Bool
truthy (Bool False) = False
truthy _ = True

-- | Is this 'Val' an atomic constant?
isConstant :: Val -> Bool
isConstant Number{} = True
isConstant Bool{}   = True
isConstant Char{}   = True
isConstant Nil{}    = True
isConstant _        = False

-------------------------------------------------------------------------------
--
-- Pure 'Val' instances. You probably don't want to use these, but they are
-- ocassionally useful in the test suite etc. Be careful!
--
-------------------------------------------------------------------------------

-- | 'show' a 'Val'. This almost certainly isn't what you want;
-- it's a good escape hatch if you need a pure way to display a 'Val'
-- but it doesn't display unicode characters in strings properly,
-- can't display anything mutable, and loops on cyclic immutable data.
showVal :: Val -> ShowS
showVal (Symbol s) = showString $ T.unpack $ symbolName s
showVal (Number n) = shows n
showVal (String _s) = showString "<can't show mutable string>"
showVal (IString s) = shows s
showVal (Char c)   = showString "#\\" . showString (case c of
    ' '  -> "space"
    '\t' -> "tab"
    '\n' -> "newline"
    '\r' -> "carriage-return"
    _    -> [c])
showVal (Bool True) = showString "#t"
showVal (Bool False) = showString "#f"
showVal Port{} = showString "#<port>"
showVal Undefined = showString "#<undefined>"
showVal (Primitive arity _ name) = shows fakeClosure
  where
    fakeClosure = Closure fakeArgs fakeVarargs [] [] (Just name)
    (fakeArgs, fakeVarargs) = case arity of
      Exactly n     -> (replicate n "_", Nothing)
      Between lo hi -> (fakeBetween lo hi, Nothing)
      AtLeast n     -> (replicate n "_", Just "_")
    fakeBetween lo hi =
      let reqs = replicate lo "_"
          numOpts = hi - lo
          opts = replicate numOpts "[_]"
      in reqs ++ opts

showVal Continuation{} = showString "#<cont>"
showVal (Closure [] (Just varargs) _body _env name) =
      showString "#<procedure "
    . showString (symbolAsString $ fromMaybe "lambda" name)
    . showChar ' ' . showString (symbolAsString varargs)
    . showChar '>'
showVal (Closure args varargs _body _env name) =
      showString "#<procedure "
      . displayName . showChar ' '
      . showParen True (displayArgs . displayVarargs)
      . showChar '>'
  where
    displayName = showString $ symbolAsString $ fromMaybe "lambda" name
    displayArgs = showString $ unwords $ map symbolAsString args
    displayVarargs = case varargs of
        Nothing  -> id
        Just arg -> showString " . " . showString (symbolAsString arg)
showVal (PrimMacro _ _ name) = showString $ "#<macro "
                                            ++ symbolAsString name
                                            ++ ">"
showVal Pair{} = showString "<can't show mutable pair>"
showVal p@IPair{} = showParen True $ showDList $ fromList p
  where
    fromList :: Val -> ([Val], Val)
    fromList (IPair c d) =
      first (c:) $ fromList d
    fromList obj = ([], obj)

    first f (a, b) = (f a, b)

    showDList (vs, v) = unwordsList vs . case v of
      Nil -> id
      obj -> showString " . " . shows obj

showVal Vector{} = showString "<can't show mutable vector>"
showVal (IVector v) = showChar '#' . showParen True (unwordsList (V.toList v))
showVal (IByteVector v) = showString "#u8" . showParen True (showBytes (U.toList v))
  where showBytes [] = id
        showBytes [b] = shows b
        showBytes (b:bs) = shows b . showChar ' ' . showBytes bs
showVal Nil = showString "()"
showVal EOF = showString "#<eof>"
showVal (Error msg irritants) =
  showString "#<error: "
  . prefix (T.unpack msg)
  . intercalateS " " (map shows irritants)
  . showChar '>'
  where prefix "" = showString ""
        prefix s  = showString s . showChar ' '
showVal (Exception err) = showString $ "#<exception: " ++ showErr err ++ ">"

showVal (MultipleValues vs) =
    showString "#<values: "
  . unwordsList vs
  . showChar '>'

-- | This instance can't show mutable things.
instance Show Val where
    showsPrec _ v s = showVal v s

-- | Doesn't display exactness indicators and always uses base 10.
instance Show Number where
  show (Real r) = show r
  show (Complex (r :+ i)) = show r ++ sign ++ show i ++ "i"
    where sign = if i < 0 then "" else "+"

showErr :: LispErr -> String
showErr (UnboundVar message varname)
  = message ++ " unbound symbol: " ++ symbolAsString varname
showErr (EvaluateDuringInit name)
  = symbolAsString name ++ " referred to itself during initialization"
showErr (SetImmutable tyname) = "can't set immutable " ++ tyname
showErr (BadSpecialForm form) = "bad form: " ++ show form
showErr (NotFunction message func)    = message ++ ": " ++ show func
showErr (NumArgs expected found)      = "expected " ++ showArity expected
    ++ " arg" ++ (if usePluralForArity expected then "s" else "")
    ++ ", found values " ++ show found
showErr (TypeMismatch expected found) = "invalid type: expected " ++ expected
    ++ ", found " ++ show found
showErr CircularList                  = "circular list"
showErr EmptyBody                     = "attempt to define function with no body"
showErr (Parser parseErr)             = parseErr
showErr (Condition _ obj)             = show obj
showErr (Default message)             = message
showErr Quit                          = "quit invoked"

-- | This instance can't show internal values with mutable pairs and will
-- likely be removed soon instead of maintaining it. If you want to turn an
-- error into a string, you should use 'showErrIO' in Primitives.WriteLib.
instance Show LispErr where show = ("Error: " ++ ) . showErr

intercalateS :: String -> [ShowS] -> ShowS
intercalateS sep = go
  where go []     = id
        go [s]    = s
        go (s:ss) = s . showString sep . go ss

unwordsList :: [Val] -> ShowS
unwordsList = intercalateS " " . map shows

showArity :: Arity -> String
showArity (Exactly n) = show n
showArity (AtLeast n) = "at least " ++ show n
showArity (Between i j) = unwords ["between", show i, "and", show j]

usePluralForArity :: Arity -> Bool
usePluralForArity (Exactly 1) = False
usePluralForArity (AtLeast 1) = False
usePluralForArity _ = True

-- | Test an 'Arity' against the length of a list, reporting an error
-- if the arity is not satisfied.
--
-- More efficient than using 'length'.
testArity :: Arity -> [Val] -> EM ()
testArity (Exactly n0) vs0
  | n0 < 0 = panic "negative exact arity"
  | otherwise = check $ cmpLength n0 vs0
  where
    check EQ = pure ()
    check _  = throwError $ NumArgs (Exactly n0) vs0
testArity (AtLeast n0) vs0 = check $ cmpLength n0 vs0
  where
    check GT = throwError $ NumArgs (AtLeast n0) vs0
    check _  = pure ()
testArity a@(Between lo hi) vs0 = check (cmpLength lo vs0) (cmpLength hi vs0)
  where
    check GT _ = raise
    check _ LT = raise
    check _ _  = pure ()
    raise = throwError $ NumArgs a vs0

cmpLength :: Int -> [a] -> Ordering
cmpLength 0 [] = EQ
cmpLength _ [] = GT
cmpLength 0 (_:_) = LT
cmpLength n (_:vs) = cmpLength (n-1) vs

{- Note: [List Operation Efficiency]
There's some low-hanging fruit here for optimizing list operations.
Primarily, most operations end up walking the structure more than once.
'testCircularList' walks the whole list (in the event that the list is
not circular) since it terminates when it sees a cdr that isn't a Pair.
But then 'flattenFiniteList' also walks the whole list, to flatten it,
and is always called separately from 'testCircularList'. Notably,
'isListSH' currently tests, then flattens to check the ending. But then
'getList' calls 'isListSH', then flattens to get the list! So 'getList'
ends up walking the whole structure 3 times.

I think the most frequently-called function here should end up being
'freezeList'
-}

-- | Flatten a Scheme list into a Haskell list.
--
-- If the Scheme value is not a list (in the sense of 'list?')
-- then 'Nothing' is returned. Otherwise, 'Just' the contents of the
-- list are returned.
getList :: Val -> EM (Maybe [Val])
getList v = do
  b <- isListSH v
  if b
    then Just . fst <$> flattenFiniteList v
    else return Nothing

-- | 'getList', but if the object is not a list, a type error is raised
-- instead of returning 'Nothing'.
getListOrError :: Val -> EM [Val]
getListOrError v = do
  mlst <- getList v
  case mlst of
    Nothing  -> throwError $ TypeMismatch "list" v
    Just lst -> return lst

-- | Represents a (shallow copy of a) Scheme list at the Haskell level.
-- Objects inside the list can be mutated as normal but the car/cdr
-- Refs that would be needed to mutate the shape of the list itself
-- are not present. This is useful for implementing Scheme syntax;
-- for example the first thing `eval` does is freeze its input.
--
-- Invariants:
--
--  (1) The Haskell lists in FrozenLists are always finite (unchecked)
--  (2) FNotList never contains 'Nil', a 'Pair', or an 'IPair'.
--      See Note: [Freezing Nil].
data FrozenList
  = FList [Val]
  | FDottedList [Val] Val
  | FNotList Val

lintFrozenList :: FrozenList -> EM ()
lintFrozenList fl = lintAssert "FNotList contains Nil or a pair" $ pure $
  case fl of
    FNotList Nil     -> False
    FNotList Pair{}  -> False
    FNotList IPair{} -> False
    _ -> True

{- Note: [Freezing Nil]
There are multiple possible representations of any frozen list that represents
a list. True lists can be represented with either FList or FDottedList
or FNotList (where the 'Val's in the last two cases are pairs). The empty list,
'Nil', can be represented as both 'FList []' and 'FNotList Nil'.

Invariant (2) of 'FrozenList' resolve the ambiguities. All lists should be
represented as frozen lists, as expected, meaning that lists can't appear in
'FNotList'. This is perfectly sane.

However, 'Nil' is a bit of a corner case. 'Nil' is the only list that is
self-evaluating. In most cases where we freeze a list, we care more about the
fact that 'Nil' is a list than the fact that it's self-evaluating.
In fact, the only place where we _do_ care that it's self evaluating is in
'eval'! So we always freeze 'Nil' as 'FList []', and we check the special
case in 'eval'.
-}

-- | Get an immutable Haskell representation of a Scheme list.
-- This is useful for core operations like `eval`, which can evaluate
-- a Scheme list by freezing it and then evaluating the frozen list instead.
-- Attempting to freeze a circular list raises a CircularList error.
--
-- If the object is not a list, it is returned unchanged in an
-- 'FNotList'.
--
-- The 'FrozenList' is linted before it is returned, so it is not necessary
-- to immediately lint it again.
freezeList :: Val -> EM FrozenList
freezeList v = do
  requireFiniteList v
  (lst, dot) <- flattenFiniteList v
  let flist = case (lst, dot) of
        ([], Nil) -> FList []
        ([], obj) -> FNotList obj
        (xs, Nil) -> FList xs
        (xs, obj) -> FDottedList xs obj
  lintFrozenList flist
  return flist

-- | Require that a Scheme value is not a circular list.
-- Circular data structures that are not lists may or may not be detected.
-- If the value is a circular list, a CircularList error is raised.
requireFiniteList :: Val -> EM ()
requireFiniteList v = whenM (testCircularList v) (throwError CircularList)

-- | Test if a Scheme value is a circular list.
testCircularList :: Val -> EM Bool
testCircularList v = next v >>= \case
    Nothing   -> pure False
    Just node -> lift2 go (pure $ Just node) (next node)
  where
    -- Floyd's Algorithm
    go :: Maybe Val -> Maybe Val -> EM Bool
    go tortoise hare
      | Just t <- tortoise
      , Just h <- hare
      = ifM ((==) <$> liftIO (makeStableName t) <*> liftIO (makeStableName h))
            (return True)
            (lift2 go (next t) (next2 h))
      | otherwise = return False

    -- takes a val and returns the cdr if it is a pair, else Nothing
    next p@Pair{}  = Just <$> cdrPS p
    next p@IPair{} = Just <$> cdrPS p
    next _ = pure Nothing

    next2 = runMaybeT . (next' >=> next')
      where next' = MaybeT . next

    lift2 f x y = do
      xv <- x
      yv <- y
      f xv yv

-- | Flatten a _finite_ Scheme list into an init/improper tail pair.
-- A result of (init, tl) means that "(init . tl)" is an external
-- representation of the input object.
--
-- See 'freezeList', which wraps this function and further converts
-- the pair to a proper Scheme object.
flattenFiniteList :: Val -> EM ([Val], Val)
flattenFiniteList p@Pair{} = do
  hd <- carPS p
  ~(tl, dot) <- cdrPS p >>= flattenFiniteList
  return (hd:tl, dot)
flattenFiniteList p@IPair{} = do
  car <- carPS p
  cdr <- cdrPS p
  ~(tl, dot) <- flattenFiniteList cdr
  return (car:tl, dot)

flattenFiniteList other = return ([], other)

-- | Make a freshly-allocated mutable Scheme list containing
-- the given values.
makeMutableList :: [Val] -> EM Val
makeMutableList = foldrM consSSS Nil

-- | Make a freshly-allocated, mutable, improper Scheme list containing
-- the given values.
makeImproperMutableList :: [Val] -- ^ the "proper" part of the list
                        -> Val   -- ^ the improper tail
                        -> EM Val
makeImproperMutableList prop improp = foldrM consSSS improp prop

-- | Make a freshly-allocated _immutable_ Scheme list containing
-- the given values.
--
-- Unless you're working on the parser, you want 'makeMutableList'
-- instead of this.
makeImmutableList :: [Val] -> Val
makeImmutableList = foldr immutableCons Nil

-- | Make a freshly allocated, immutable, improper Scheme list.
--
-- You probably don't actually want this function.
-- See 'makeImproperMutableList' and 'makeImmutableList'.
makeImproperImmutableList :: [Val] -> Val -> Val
makeImproperImmutableList prop improp = foldr immutableCons improp prop

-- | Primitive list? operation.
--
-- See page 42 of the r7rs report:
-- "By definition, all lists have finite length
-- and are terminated by the empty list."
isListSH :: Val -> EM Bool
isListSH v = callCC $ \exit -> do
  whenM (testCircularList v) $ exit False
  (_, tl) <- flattenFiniteList v
  case tl of
    Nil -> return True
    _   -> return False

-- | Primitive cons operation.
consSSS :: Val -> Val -> EM Val
consSSS car cdr = Pair <$> newRef car <*> newRef cdr

-- | Primitive cons operation for immutable lists.
--
-- In an ideal world, this function would perform linting to check that
-- the car and cdr are immutable objects. Unfortunately, that would require
-- access to the repl 'Opts' in the parser (the only place this function
-- should be used) which are currently (12/10/2020) not available. As the
-- language grows, they will probably become available eventually, at which
-- point this unfortunate issue should be patched.
--
-- Unless you're working on the parser, you really shouldn't use this function!
immutableCons :: Val -> Val -> Val
immutableCons = IPair

carPR :: Val -> Ref Val
carPR (Pair c _d) = c
carPR _ = panic "carPR: not a mutable pair"

cdrPR :: Val -> Ref Val
cdrPR (Pair _c d) = d
cdrPR _ = panic "cdrPR: not a mutable pair"

carPS :: MonadIO m => Val -> m Val
carPS p@Pair{} = liftIO $ readIORef $ carPR p
carPS (IPair car _) = pure car
carPS _ = panic "carPS: not a pair"
{-# SPECIALIZE carPS :: Val -> EM Val #-}

cdrPS :: MonadIO m => Val -> m Val
cdrPS p@Pair{} = liftIO $ readIORef $ cdrPR p
cdrPS (IPair _ cdr) = pure cdr
cdrPS _ = panic "cdrPS: not a pair"
{-# SPECIALIZE cdrPS :: Val -> EM Val #-}

setCarSSS :: Val -> Val -> EM Val
setCarSSS (Pair c _d) v = do
  writeRef c v
  return v
setCarSSS IPair{} _ = throwError $ SetImmutable "pair"
setCarSSS v _ = throwError $ TypeMismatch "pair" v

setCdrSSS :: Val -> Val -> EM Val
setCdrSSS (Pair _c d) v = do
  writeRef d v
  return v
setCdrSSS IPair{} _ = throwError $ SetImmutable "pair"
setCdrSSS v _ = throwError $ TypeMismatch "pair" v

-- | Require that the given Val is a list. Otherwise, raises a 'TypeError'.
requireList :: Val -> EM ()
requireList v =
  whenM (not <$> isListSH v) $ throwError $ TypeMismatch "list" v

ifM :: Monad m => m Bool -> m a -> m a -> m a
ifM b t f = b >>= \b' -> if b' then t else f

whenM :: Monad m => m Bool -> m () -> m ()
whenM test thing = ifM test thing (pure ())

-- | Test if a given Scheme object is an immutable pair.
isImmutablePair :: Val -> Bool
isImmutablePair IPair{} = True
isImmutablePair _ = False

-- | Test if a given Scheme object is immutable.
--
-- Specifically, we mean that an object is "immutable" if it can be output
-- by the parser. That means the type is not a reference, or its constructor's
-- name starts with I.
--
-- Most objects qualify; 'Pair', 'Vector', 'U8Vector', and 'String' do
-- not.
isImmutable :: Val -> Bool
isImmutable Pair{}   = False
isImmutable Vector{} = False
isImmutable String{} = False
isImmutable _ = True

makeBignum :: Integer -> Val
makeBignum = Number . Real . Bignum

makeFlonum :: Double -> Val
makeFlonum = Number . Real . Flonum

makeRatnum :: Integer -> Integer -> Val
makeRatnum x y = Number $ Real $ Ratnum (x % y)

approxToRatnum :: Double -> Val
approxToRatnum = Number . Real . Ratnum . toRational

getExactInteger :: Val -> Maybe Integer
getExactInteger v = case v of
  Number (Real (Bignum i)) -> return i
  Number (Real (Flonum _)) -> Nothing
  Number (Real (Ratnum r)) -> case denominator r of
    1 -> return $ numerator r
    _ -> Nothing
  Number (Complex (r :+ i)) -> case i of
    Bignum 0 -> getExactInteger (Number (Real r))
    Flonum _ -> Nothing
    Ratnum z -> if z == 0 then getExactInteger (Number (Real r)) else Nothing
    _notNplusZeroI -> Nothing
  _other -> Nothing

mkNumericTypecheck :: (Number -> Bool) -> Val -> Bool
mkNumericTypecheck p = f
  where f (Number n) = p n
        f _notNumber = False
{-# INLINE mkNumericTypecheck #-}

numberSH, complexSH, realSH, rationalSH, integerSH :: Val -> Bool
numberSH = mkNumericTypecheck $ const True
complexSH = numberSH -- see note on p35 of r7rs paper
realSH = mkNumericTypecheck $ \case
  Real{} -> True
  Complex (_r :+ i)
    | i == 0    -> True
    | otherwise -> False
rationalSH = mkNumericTypecheck test
  where
    test (Real (Flonum f))
      | isInfinite f || isNaN f = False
    test Real{} = True -- truncated real numbers are always rational!
    test (Complex (r :+ i))
      | i == 0 = test $ Real r
      | otherwise = False
integerSH = mkNumericTypecheck test
  where
    test (Real x) = x == round x -- (==)@RealNumber is Scheme =
    test (Complex (r :+ i))
      | i == 0 = test $ Real r
      | otherwise = False

implicitNumericConversion :: Number -> Number -> (Number, Number)
implicitNumericConversion (Real x) (Real y) = (Real x, Real y)
implicitNumericConversion x y = (asComplex x, asComplex y)
  where
    asComplex :: Number -> Number
    asComplex (Real r) = Complex (r :+ 0)
    asComplex (Complex c) = Complex c

-- | Perform an implicit Scheme conversion of two numbers.
--
-- The "order" of numeric types is bignum < ratnum < flonum.
-- The lesser argument will be casted upwards to match the greater one.
-- The returned Numbers are guaranteed to use the same Number constructor.
implicitRealConversion :: RealNumber -> RealNumber -> (RealNumber, RealNumber)
implicitRealConversion = work
  where
    work x@Flonum{} y = (x, Flonum $ asFlonum y)
    work x y@Flonum{} = (Flonum $ asFlonum x, y)
    work x@Ratnum{} y = (x, asRatnum y)
    work x y@Ratnum{} = (asRatnum x, y)
    work x y = (x, y) -- both Bignum

    asRatnum (Ratnum r) = Ratnum r
    asRatnum (Bignum n) = Ratnum (n % 1)
    asRatnum Flonum{} = error "implicitRealConversion.asRatnum: impossible flonum"

pairSH :: Val -> Bool
pairSH Pair{}  = True
pairSH IPair{} = True
pairSH _ = False

vectorSH :: Val -> Bool
vectorSH Vector{} = True
vectorSH IVector{} = True
vectorSH _ = False

makeVector :: [Val] -> Val
makeVector = IVector . V.fromList

makeByteVector :: [Word8] -> Val
makeByteVector = IByteVector . U.fromList

-- | Show a 'Val' in IO. This function loops forever
-- on cyclic data; it's morally equivalent to write-simple.
-- However we don't use it to _implement_ write-simple
-- because we can 't do the linting that we want to do
-- since we aren't in EM. This is mostly just a temporary
-- (12/10/2020) thing so that we can look at values until
-- write-simple/write-shared are implemented (in 'MonadIO',
-- preferably).
--
-- TODO: it should be good to get rid of this now.
showValIO :: Val -> IO String
showValIO = fmap ($ "") . showsValIO

-- | ShowS version of 'showValIO'.
showsValIO :: Val -> IO ShowS
showsValIO p@Pair{}  = showParen True <$> showsListIO p
showsValIO p@IPair{} = showParen True <$> showsListIO p
showsValIO s = pure $ shows s

showsListIO :: Val -> IO ShowS
-- only accepts pairs
showsListIO (Pair c d) = do
  car <- readIORef c
  showsCar <- showsValIO car
  (showsCar .) <$> (readIORef d >>= showsListIO1)
-- we're in IO, so have to forego linting here.
showsListIO (IPair car cdr) = do
    showsCar <- showsValIO car
    (showsCar .) <$> showsListIO1 cdr
showsListIO _notPair = error "showsListIO: not a list"

-- | 'showsListIO', but puts a space before showing
-- the next object. If the next object is 'Nil' (i.e., we're done)
-- then no space is printed.
--
-- This is suitable for printing objects in lists other than the first.
showsListIO1 :: Val -> IO ShowS
showsListIO1 p@Pair{}  = (showString " " .) <$> showsListIO p
showsListIO1 p@IPair{} = (showString " " .) <$> showsListIO p
showsListIO1 Nil = pure id
showsListIO1 dot = (showString " . " .) <$> showsValIO dot

-------------------------------------------------------------------------------
-- Numeric Class Instances for Number
-------------------------------------------------------------------------------

type Binop a = a -> a -> a
type Monop a = a -> a

{- HLINT ignore "Redundant lambda" -}
mkRealNumberBinop :: (forall a. Num a => Binop a)
                  -> Binop RealNumber
mkRealNumberBinop op = \x y -> case implicitRealConversion x y of
  (Bignum a, Bignum b) -> Bignum (a `op` b)
  (Ratnum a, Ratnum b) -> Ratnum (a `op` b)
  (Flonum a, Flonum b) -> Flonum (a `op` b)
  _ -> error "mkRealNumberBinop: implicitRealConversion failure"
{-# INLINE mkRealNumberBinop #-}

{- HLINT ignore "Redundant lambda" -}
mkRealNumberMonop :: (forall a. Num a => Monop a)
                  -> Monop RealNumber
mkRealNumberMonop op = \case
  Bignum a -> Bignum (op a)
  Ratnum a -> Ratnum (op a)
  Flonum a -> Flonum (op a)
{-# INLINE mkRealNumberMonop #-}

{- HLINT ignore "Redundant lambda" -}
mkNumberBinop :: (forall a. Num a => Binop a)
              -> Binop Number
mkNumberBinop op = \x y -> case implicitNumericConversion x y of
  (Real a, Real b) -> Real (a `op` b)
  (Complex w, Complex z) -> Complex (w `op` z)
  _ -> panic " mkNumberBinop: implicitNumericConversion failure"
{-# INLINE mkNumberBinop #-}

{- HLINT ignore "Redundant lambda" -}
mkNumberMonop :: (forall a. Num a => Monop a)
              -> Monop Number
mkNumberMonop op = \case
  Real n -> Real $ op n
  Complex c -> Complex $ op c
{-# INLINE mkNumberMonop #-}

realNumberPlus, realNumberSub, realNumberTimes :: Binop RealNumber
realNumberPlus  = mkRealNumberBinop (+)
realNumberSub   = mkRealNumberBinop (-)
realNumberTimes x y
  | isExactZero (Real x) || isExactZero (Real y) = Bignum 0
  | otherwise = mkRealNumberBinop (*) x y

realNumberNegate, realNumberAbs, realNumberSignum :: Monop RealNumber
realNumberNegate = mkRealNumberMonop negate
realNumberAbs    = mkRealNumberMonop abs
realNumberSignum = mkRealNumberMonop signum

numberTimes :: Binop Number
numberTimes x y
  | isExactZero x || isExactZero y = 0
  | otherwise = mkNumberBinop (*) x y

isExactZero :: Number -> Bool
isExactZero (Real (Bignum 0)) = True
isExactZero (Real (Ratnum 0)) = True
isExactZero (Complex (r :+ i)) = isExactZero (Real r) && isExactZero (Real i)
isExactZero _          = False
{-# INLINABLE isExactZero #-}

instance Num RealNumber where
  (+) = realNumberPlus
  (-) = realNumberSub
  (*) = realNumberTimes
  negate = realNumberNegate
  abs    = realNumberAbs
  signum = realNumberSignum
  fromInteger = Bignum

instance Num Number where
  (+) = mkNumberBinop (+)
  (-) = mkNumberBinop (-)
  (*) = numberTimes -- checks for exact zeroes
  negate = mkNumberMonop negate
  abs    = mkNumberMonop abs
  signum = mkNumberMonop signum
  fromInteger = Real . Bignum

asFlonum :: RealNumber -> Double
asFlonum (Bignum i) = fromInteger i
asFlonum (Ratnum r) = fromRational r
asFlonum (Flonum f) = f

realInexact :: RealNumber -> RealNumber
realInexact = Flonum . asFlonum

flonumPredicate :: (Double -> Bool) -> RealNumber -> Bool
flonumPredicate p = applyPred
  where applyPred (Flonum f) = p f
        applyPred _notFlonum = False
{-# INLINE flonumPredicate #-}

-- | Implementation of 'scaleFloat' for Number. Does the right thing
-- without converting where possible. Bignums are converted to Ratnums
-- if the power is negative.
--
-- realNumberScale p n returns 0 if p is negative and n is 0. This is the
-- same behavior as scaleFloat for Float/Double and gives correct
-- results for Complex divison.
realNumberScale :: Int -> RealNumber -> RealNumber
realNumberScale p (Flonum f) = Flonum (scaleFloat p f)
realNumberScale p bigOrRat
  | p < 0 = case bigOrRat of
    Bignum 0 -> Bignum 0
    Ratnum 0 -> Bignum 0 -- simplify. No cost to doing this because a
                         -- rational would need to be deconstructed to do
                         -- the next step anyway.
    _notZero -> bigOrRat * recip (Bignum (bit $ negate p))
  | otherwise = case bigOrRat of
    Bignum i -> Bignum (i `shiftL` p)
    Ratnum r -> Ratnum (r * (2 ^ p))
    Flonum{} -> error "realNumberScale: impossible flonum"

instance RealFloat RealNumber where
  floatRadix  = floatRadix  . asFlonum
  floatDigits = floatDigits . asFlonum
  floatRange  = floatRange  . asFlonum
  decodeFloat = decodeFloat . asFlonum
  encodeFloat s e = Flonum (encodeFloat s e)
  exponent (Flonum f) = exponent f
  exponent _bigOrRat  = 0
  significand = Flonum . significand . asFlonum
  scaleFloat = realNumberScale
  isNaN = flonumPredicate isNaN
  isInfinite = flonumPredicate isInfinite
  isDenormalized = flonumPredicate isDenormalized
  isNegativeZero = flonumPredicate isNegativeZero
  -- Bignums are not IEEE floating-point numbers!
  isIEEE = flonumPredicate isIEEE
  atan2 x y = Flonum (atan2 (asFlonum x) (asFlonum y))

mkRoundingMode :: Integral b
               => (forall a. RealFrac a => a -> b)
               -> RealNumber -> b
mkRoundingMode mode = applyMode
  where applyMode (Bignum i) = fromInteger i -- no rounding to do
        applyMode (Ratnum r) = mode r
        applyMode (Flonum f) = mode f
{-# INLINE mkRoundingMode #-}

instance RealFrac RealNumber where
  properFraction (Bignum a) = (fromInteger a, 0)
  properFraction (Ratnum r) = let (n, f) = properFraction r
                              in  (n, Ratnum f)
  properFraction (Flonum f) = let (n, f') = properFraction f
                              in  (n, Flonum f')

  truncate = mkRoundingMode truncate
  round    = mkRoundingMode round
  ceiling  = mkRoundingMode ceiling
  floor    = mkRoundingMode floor

instance Real RealNumber where
  toRational (Bignum i) = i % 1
  toRational (Ratnum r) = r
  toRational (Flonum f) = toRational f

floatingOp :: Monop Double -> Monop RealNumber
floatingOp f n = Flonum (f (asFlonum n))
{-# INLINE floatingOp #-}

instance Floating RealNumber where
  pi = Flonum pi
  exp = floatingOp exp
  log = floatingOp log
  sqrt = floatingOp sqrt
  sin = floatingOp sin
  cos = floatingOp cos
  tan = floatingOp tan
  asin = floatingOp asin
  acos = floatingOp acos
  atan = floatingOp atan
  sinh = floatingOp sinh
  cosh = floatingOp cosh
  tanh = floatingOp tanh
  asinh = floatingOp asinh
  acosh = floatingOp acosh
  atanh = floatingOp atanh
  log1p = floatingOp log1p
  expm1 = floatingOp expm1
  log1pexp = floatingOp log1pexp
  log1mexp = floatingOp log1mexp

  x ** Bignum y = x ^^ y
  x ** y = Flonum (asFlonum x ** asFlonum y)
  logBase x y = Flonum (logBase (asFlonum x) (asFlonum y))

realPart :: Number -> RealNumber
realPart (Real r) = r
realPart (Complex (x:+_)) = x
inf, nan :: RealNumber
inf = Flonum $ 1/0
nan = Flonum $ 0/0

numberIsInfinite :: Number -> Bool
numberIsInfinite (Real r) = isInfinite r
numberIsInfinite (Complex (x:+y)) = isInfinite x || isInfinite y

-- | This works for sin, cos, tan, sinh, cosh, and tanh,
-- which all have real results wherever they are defined
-- on the real line.
--
-- It does not work for inverse trig functions, hence the name.
forwardTrigFunction :: (forall a. Floating a => Monop a)
                    -> Monop Number
forwardTrigFunction f = \case
  Real r    -> Real (f r)
  Complex z -> Complex (f z)

instance Floating Number where
  pi = Real pi

  exp (Real r) = Real (exp r)
  exp (Complex (x:+y)) = Complex $ expx * cos y :+ expx * sin y
    where expx = exp x

  log (Real r)
    | r < 0 || isNegativeZero r = Complex $ log (abs r) :+ pi
    | otherwise = Real $ log r
  log (Complex z@(x:+y)) = Complex $ log (magnitude z) :+ atan2 y x

  x ** y
    | isExactZero y = Real 1
    | isExactZero x = case compare (realPart y) 0 of
      GT -> Real 0
      LT -> Real inf
      EQ -> Real nan
    | numberIsInfinite x = case compare (realPart y) 0 of
      GT -> Real inf
      LT -> Real 0
      EQ -> Real nan
      -- note that exp and log give real results when applied
      -- to real arguments, when possible.
  Real x ** Real (Bignum y) = Real (x ** Bignum y)
  Real x ** Real y | x >= 0 = Real (x ** y)
  x ** y = exp (log x * y)

  sqrt z | isExactZero z = 0
  sqrt (Real r)
    | r >= 0    = Real $ sqrt r
    | otherwise = Complex $ 0 :+ sqrt (-r)
  sqrt (Complex z) = Complex $ sqrt z

  sin = forwardTrigFunction sin
  cos = forwardTrigFunction cos
  tan = forwardTrigFunction tan

  sinh = forwardTrigFunction sinh
  cosh = forwardTrigFunction cosh
  tanh = forwardTrigFunction tanh

  -- redefine these functions to use our definition of log which behaves
  -- correctly at -0. Otherwise they are copies of Data.Complex.
  asin (Real r)
    | r >= -1 && r <= 1 = Real $ asin r
    | otherwise = asin (Complex (r :+ 0))
  asin z@(Complex (x:+y)) = Complex $ y' :+ (-x')
    where Complex (x':+y') = log (iz + sqrt (1 - z * z))
          -- take care to use log @Number and not @(Complex RealNumber)
          iz = Complex ((-y) :+ x)

  acos (Real r)
    | r >= -1 && r <= 1 = Real $ acos r
    | otherwise = acos (Complex (r :+ 0))
  acos z = Complex $ y'' :+ (-x'')
    where Complex (x'':+y'') = log (z + Complex ((-y'):+x'))
          Complex (x':+y')   = sqrt (1 - z * z)

  atan (Real r) = Real $ atan r
  atan z@(Complex (x:+y)) = Complex $ y':+(-x')
    where Complex (x':+y') = log (z' / sqrt (1 + z * z))
          z' = Complex ((1-y) :+ x)

  asinh z = log (z + sqrt (1 + z*z))
  -- see comment in Data.Complex:
  -- essentially, don't use the formula with a division because it is
  -- undefined when z = -1.
  acosh z = log (z + sqrt (z+1) * sqrt (z-1))
  atanh z = 0.5 * log ((1+z) / (1-z))

  log1p (Real r)
    | abs r < 0.5 = Real (log1p r)
    | otherwise = Real (log (1 + r))
  log1p x@(Complex (a :+ b))
    | abs a < 0.5 && abs b < 0.5 =
      let u = 2*a + a*a + b*b
      in Complex $ log1p (u/(1 + sqrt(u+1))) :+ atan2 (1+a) b
    | otherwise = log (1 + x)

  expm1 (Real r) = Real (expm1 r)
  expm1 x@(Complex (a :+ b))
    | a*a + b*b < 1 =
      let u = expm1 a
          v = sin (b/2)
          w = -2*v*v
      in Complex $ (u*w + u + w) :+ (u+1)*sin b
    | otherwise = exp x - 1

-- distinct from numberDiv
realNumberDivide :: Binop RealNumber
realNumberDivide x y = case implicitRealConversion x y of
  (Bignum a, Bignum b) -> Ratnum (a % b)
  (Ratnum a, Ratnum b) -> Ratnum (a / b)
  (Flonum a, Flonum b) -> Flonum (a / b)
  _ -> error "realNumberDivide: implicitRealConversion failure"

realNumberRecip :: Monop RealNumber
realNumberRecip (Bignum i) = Ratnum (1 % i)
realNumberRecip (Ratnum r) = Ratnum (recip r)
realNumberRecip (Flonum f) = Flonum (recip f)

numberDivide :: Binop Number
numberDivide x y = case implicitNumericConversion x y of
  (Real a, Real b) -> Real (a / b)
  (Complex a, Complex b) -> Complex (a / b)
  _ -> error "numberDivide: implicitNumericConversion failure"

numberRecip :: Monop Number
numberRecip (Real r)    = Real (recip r)
numberRecip (Complex z) = Complex (recip z)

instance Fractional RealNumber where
  (/) = realNumberDivide
  recip = realNumberRecip
  fromRational = Ratnum

instance Fractional Number where
  (/) = numberDivide
  recip = numberRecip
  fromRational = Real . Ratnum

realNumberCompare :: RealNumber -> RealNumber -> Ordering
realNumberCompare = (<=>) where
  Flonum f <=> Flonum g = f `compare` g
  Flonum f <=> bigOrRat
    | isInfinite f = f `compare` 0
    | otherwise    = r1 `compare` r2
    -- note that we make the fp number exact instead of making the
    -- exact number fp. This is critical. An example from the R7RS
    -- paper is: let big be (exp 2 1000). Assume big is exact and
    -- that inexact numbers are IEEE doubles [they are]. Then
    -- (= (inexact big) (+ big 1)) should be #f, but if inexact
    -- conversion is used, it will be #t.
    where r1 = toRational f
          r2 = case bigOrRat of
            Bignum i -> i % 1
            Ratnum r -> r
            Flonum{} -> error "realNumberCompare.<=>: impossible flonum"
  bigOrRat <=> n@Flonum{} = case n <=> bigOrRat of
    LT -> GT
    EQ -> EQ
    GT -> LT
  Bignum i <=> Bignum j = i `compare` j
  Bignum i <=> Ratnum r = (i % 1) `compare` r
  Ratnum r <=> Bignum i = r `compare` (i % 1)
  Ratnum r <=> Ratnum s = r `compare` s

-- | Number is not really an Integral type, since it's possible it has
-- a floating-point backing. Here's what we do:
-- 1) toInteger truncates
-- 2) quotRem uses Number division unless both arguments are Bignums.

-- This instance exists mostly so that the rounding modes from RealFrac
-- can be used with RealNumbers.
instance Integral RealNumber where
  toInteger (Bignum i) = i
  toInteger (Ratnum r) = truncate r
  toInteger (Flonum f) = truncate f

  quotRem (Bignum x) (Bignum y) = let (q, r) = quotRem x y
                                  in (Bignum q, Bignum r)
  quotRem x y = (x / y, 0)

instance Ord RealNumber where
  compare = realNumberCompare

-- | 'Enum' essentially does not have any of the laws you might expect
-- it to have; for types that are not also Bounded, it is lawless.
-- This is a useful-ish instance which generally behaves as expected
-- and which is needed to give the Integral instance.
instance Enum RealNumber where
  succ (Bignum i) = Bignum (succ i)
  succ (Ratnum r) = Ratnum (succ r)
  succ (Flonum f) = Flonum (succ f)
  pred (Bignum i) = Bignum (pred i)
  pred (Ratnum r) = Ratnum (pred r)
  pred (Flonum f) = Flonum (pred f)
  toEnum i = Bignum (toEnum i)
  fromEnum (Bignum i) = fromEnum i
  fromEnum (Ratnum r) = fromEnum r
  fromEnum (Flonum f) = fromEnum f
  -- The rest of the functions are omitted because they are not used
  -- and it is not clear what the 'correct' implementations are for
  -- Scheme Numbers.

max_stars_repo_name: JKTKops/Thcheme

max_issues_repo_name: JKTKops/Thcheme

max_forks_repo_name: JKTKops/Thcheme