{-# LANGUAGE DeriveDataTypeable, GeneralizedNewtypeDeriving, Rank2Types #-}
{-# OPTIONS_HADDOCK hide #-}
-----------------------------------------------------------------------------
-- |
-- Module     : Numeric.LinearAlgebra.Matrix.Base
-- Copyright  : Copyright (c) 2010, Patrick Perry <patperry@gmail.com>
-- License    : BSD3
-- Maintainer : Patrick Perry <patperry@gmail.com>
-- Stability  : experimental
--

module Numeric.LinearAlgebra.Matrix.Base (
    Matrix(..),
    dim,

    fromList,
    fromRow,
    fromCol,
    zero,
    constant,

    at,
    unsafeAt,

    indices,
    elems,
    assocs,

    col,
    unsafeCol,
    cols,

    update,
    unsafeUpdate,
    accum,
    unsafeAccum,

    map,
    zipWith,
    unsafeZipWith,

    slice,
    unsafeSlice,

    splitRowsAt,
    dropRows,
    takeRows,

    splitColsAt,
    dropCols,
    takeCols,

    fromVector,
    toVector,

    isContig,
    unsafeFromForeignPtr,
    unsafeToForeignPtr,
    unsafeWith,

    ) where

import Prelude hiding ( read, map, zipWith )
import qualified Prelude as P

import Data.AEq( AEq(..) )
import Data.Typeable( Typeable )
import Foreign( ForeignPtr, Ptr, Storable )
import Text.Printf( printf )

import Numeric.LinearAlgebra.Vector( Vector )
import qualified Numeric.LinearAlgebra.Vector as V

-- | Immutable dense matrices.
data Matrix e =
      Matrix {-# UNPACK #-} !(Vector e) -- matrix data
             {-# UNPACK #-} !Int        -- row dimension
             {-# UNPACK #-} !Int        -- column dimension
             {-# UNPACK #-} !Int        -- leading dimension
  deriving (Typeable)

-- | Get the matrix dimensions (number of rows and number of columns).
dim :: (Storable e) => Matrix e -> (Int,Int)
dim (Matrix _ m n _) = (m,n)
{-# INLINE dim #-}

-- | Indicates if the elements of the matrix are stored contigously
isContig :: (Storable e) => Matrix e -> Bool
isContig (Matrix _ m _ lda) = lda == m || m == 0
{-# INLINE isContig #-}

-- | Create a matrix of the given dimension with elements initialized
-- to the values from the list, in column major order.
fromList :: (Storable e) => (Int,Int) -> [e] -> Matrix e
fromList (m,n) es
    | m < 0 || n < 0 = error $
        printf "fromList (%d,%d): negative dimension" m n
    | otherwise = let
        v = V.fromList (m*n) es
        lda = max 1 m
        in Matrix v m n lda
{-# INLINE fromList #-}

-- | Create a matrix of the given dimension with all elements initialized
-- to the given value
constant :: (Storable e) => (Int,Int) -> e -> Matrix e
constant (m,n) e
    | m < 0 || n < 0 = error $
        printf "constant (%d,%d): negative dimension" m n
    | otherwise = let
        v = V.constant (m*n) e
        lda = max 1 m
        in Matrix v m n lda
{-# INLINE constant #-}

-- | Create a zero of the given dimension with all elements initialized
-- to zero.
zero :: (Storable e, Num e) => (Int,Int) -> Matrix e
zero (m,n)
    | m < 0 || n < 0 = error $
        printf "zero (%d,%d): negative dimension" m n
    | otherwise = let
        v = V.zero (m*n)
        lda = max 1 m
        in Matrix v m n lda
{-# INLINE zero #-}

-- | Returns the element of a matrix at the specified index.
at :: (Storable e) => Matrix e -> (Int,Int) -> e
at a ij@(i,j)
    | i < 0 || i >= m || j < 0 || j >= n = error $
        printf ("at <matrix with dim (%d,%d)> (%d,%d):"
                ++ " invalid index") m n i j
    | otherwise =
        unsafeAt a ij
  where
      (m,n) = dim a
{-# INLINE at #-}

unsafeAt :: (Storable e) => Matrix e -> (Int,Int) -> e
unsafeAt (Matrix v _ _ lda) (i,j) =
    V.unsafeAt v (i + j * lda)
{-# INLINE unsafeAt #-}

-- | Get the indices of the elements in the matrix, in column-major order.
indices :: (Storable e) => Matrix e -> [(Int,Int)]
indices a = [ (i,j) | j <- [ 0..n-1 ], i <- [ 0..m-1 ] ]
  where (m,n) = dim a

-- | Returns a list of the elements of a matrix, in the same order as their
-- indices.
elems :: (Storable e) => Matrix e -> [e]
elems (Matrix v m _ lda)
    | lda == m  = V.elems v
    | otherwise = let
        breakCols [] = []
        breakCols es = let (c,es') = splitAt lda es in c:(breakCols es')

        dropJunk c = take m c

        in concatMap dropJunk $ breakCols (V.elems v)
{-# INLINE elems #-}

-- | Returns the contents of a matrix as a list of associations.
assocs :: (Storable e) => Matrix e -> [((Int,Int),e)]
assocs x = zip (indices x) (elems x)
{-# INLINE assocs #-}

-- | Version of 'update' that doesn't range-check indices.
unsafeUpdate :: (Storable e) => Matrix e -> [((Int,Int),e)] -> Matrix e
unsafeUpdate (Matrix v m n lda) ijes = let
    ies = [ (i + j * lda, e) | ((i,j),e) <- ijes ]
    v' = V.unsafeUpdate v ies
    in Matrix v' m n lda

-- | Create a new matrix by replacing the values at the specified indices.
update :: (Storable e) => Matrix e -> [((Int,Int),e)] -> Matrix e
update (Matrix v m n lda) ijes = let
    ies = [ if i < 0 || i >= m || j < 0 || j >= n
                then error $ printf
                         ("update"
                         ++ " <matrix with dim (%d,%d)>"
                         ++ " [ ..((%d,%d),_).. ]"
                         ++ ": invalid index")
                         m n i j
                else (i + j * lda, e)
          | ((i,j),e) <- ijes
          ]
    v' = V.unsafeUpdate v ies
    in Matrix v' m n lda

-- | Same as 'accum' but does not range-check indices.
unsafeAccum :: (Storable e)
            => (e -> e' -> e)
            -> Matrix e
            -> [((Int,Int), e')]
            -> Matrix e
unsafeAccum f (Matrix v m n lda) ijes = let
    ies = [ (i + j * lda, e) | ((i,j),e) <- ijes ]
    v' = V.unsafeAccum f v ies
    in Matrix v' m n lda

-- | @accum f@ takes a matrix and an association list and accumulates
-- pairs from the list into the matrix with the accumulating function @f@.
accum :: (Storable e)
      => (e -> e' -> e)
      -> Matrix e
      -> [((Int,Int), e')]
      -> Matrix e
accum f (Matrix v m n lda) ijes = let
    ies = [ if i < 0 || i >= m || j < 0 || j >= n
                then error $ printf
                         ("accum"
                         ++ " <matrix with dim (%d,%d)>"
                         ++ " [ ..((%d,%d),_).. ]"
                         ++ ": invalid index")
                         m n i j
                else (i + j * lda, e)
          | ((i,j),e) <- ijes
          ]
    v' = V.unsafeAccum f v ies
    in Matrix v' m n lda

-- | Construct a new matrix by applying a function to every element of
-- a matrix.
map :: (Storable e, Storable e')
    => (e -> e')
    -> Matrix e
    -> Matrix e'
map f a = fromList (dim a) $ P.map f (elems a)
{-# INLINE map #-}

-- | Construct a new matrix by applying a function to every pair of elements
-- of two matrices.  The two matrices must have identical dimensions.
zipWith :: (Storable e, Storable e', Storable f)
        => (e -> e' -> f)
        -> Matrix e
        -> Matrix e'
        -> Matrix f
zipWith f a a'
    | mn /= mn' = error $
        printf ("zipWith"
                ++ " <matrix with dim %s> "
                ++ " <matrix with dim %s>"
                ++ ": dimension mismatch"
                ) (show mn) (show mn')
    | otherwise =
        unsafeZipWith f a a'
  where
    mn  = dim a
    mn' = dim a'
{-# INLINE zipWith #-}

-- | Version of 'zipWith' that does not check if the input matrices
-- have the same dimensions.
unsafeZipWith :: (Storable e, Storable e', Storable f)
              => (e -> e' -> f)
              -> Matrix e
              -> Matrix e'
              -> Matrix f
unsafeZipWith f a a' =
    fromList (dim a') $ P.zipWith f (elems a) (elems a')
{-# INLINE unsafeZipWith #-}

-- | Get the given column of the matrix.
col :: (Storable e) => Matrix e -> Int -> Vector e
col a j
    | j < 0 || j >= n = error $
        printf ("col <matrix with dim (%d,%d)> %d:"
                ++ " index out of range") m n j
    | otherwise =
        unsafeCol a j
  where
    (m,n) = dim a
{-# INLINE col #-}

-- | Version of 'col' that doesn't range-check indices.
unsafeCol :: (Storable e) => Matrix e -> Int -> Vector e
unsafeCol (Matrix v m _ lda) j =
    V.unsafeSlice (j*lda) m v
{-# INLINE unsafeCol #-}

-- | Get a list of the columns of the matrix.
cols :: (Storable e) => Matrix e -> [Vector e]
cols a = P.map (unsafeCol a) [ 0..n-1 ]
  where
    (_,n) = dim a
{-# INLINE cols #-}

-- | @slice (i,j) (m,n) a@ creates a submatrix view of @a@ starting at
-- element @(i,j)@ and having dimensions @(m,n)@.
slice :: (Storable e)
      => (Int,Int)
      -> (Int,Int)
      -> Matrix e
      -> Matrix e
slice (i,j) (m',n') a
    | (i < 0 || m' < 0 || i + m' > m
       || j < 0 || n' < 0 || j + n' > n) = error $
        printf ( "slice"
               ++ " (%d,%d)"
               ++ " (%d,%d)"
               ++ " <matrix with dim (%d,%d)>"
               ++ ": index out of range"
               ) i j m' n' m n
    | otherwise =
        unsafeSlice (i,j) (m',n') a
  where
    (m,n) = dim a
{-# INLINE slice #-}

-- | Version of 'slice' that doesn't range-check indices.
unsafeSlice :: (Storable e)
            => (Int,Int) -> (Int,Int) -> Matrix e -> Matrix e
unsafeSlice (i,j) (m',n') (Matrix v _ _ lda) = let
    o = i + j*lda
    l = if m' == 0
            then 0
            else lda * n'
    v' = V.unsafeSlice o l v
    in Matrix v' m' n' lda
{-# INLINE unsafeSlice #-}

-- | Create a view of a matrix by taking the initial rows.
takeRows :: (Storable e) => Int -> Matrix e -> Matrix e
takeRows i a = slice (0,0) (i,n) a
  where
    (_,n) = dim a

-- | Create a view of a matrix by dropping the initial rows.
dropRows :: (Storable e) => Int -> Matrix e -> Matrix e
dropRows i a = slice (i,0) (m-i,n) a
  where
    (m,n) = dim a

-- | Split a matrix into two blocks and returns views into the blocks.  If
-- @(a1, a2) = splitRowsAt i a@, then
-- @a1 = slice (0,0) (i,n) a@ and
-- @a2 = slice (i,0) (m-i,n) a@, where @(m,n)@ is the dimension of @a@.
splitRowsAt :: (Storable e) => Int -> Matrix e -> (Matrix e, Matrix e)
splitRowsAt i a
    | i < 0 || i > m = error $
        printf ("splitRowsAt %d <matrix with dim (%d,%d)>:"
                ++ " invalid index") i m n
    | otherwise = let
        a1 = unsafeSlice (0,0) (i,n)   a
        a2 = unsafeSlice (i,0) (m-i,n) a
    in (a1,a2)
  where
    (m,n) = dim a
{-# INLINE splitRowsAt #-}

-- | Create a view of a matrix by taking the initial columns.
takeCols :: (Storable e) => Int -> Matrix e -> Matrix e
takeCols j a = slice (0,0) (m,j) a
  where
    (m,_) = dim a

-- | Create a view of a matrix by dropping the initial columns.
dropCols :: (Storable e) => Int -> Matrix e -> Matrix e
dropCols j a = slice (0,j) (m,n-j) a
  where
    (m,n) = dim a

-- | Split a matrix into two blocks and returns views into the blocks.  If
-- @(a1, a2) = splitColsAt j a@, then
-- @a1 = slice (0,0) (m,j) a@ and
-- @a2 = slice (0,j) (m,n-j) a@, where @(m,n)@ is the dimension of @a@.
splitColsAt :: (Storable e) => Int -> Matrix e -> (Matrix e, Matrix e)
splitColsAt j a
    | j < 0 || j > n = error $
        printf ("splitColsAt %d <matrix with dim (%d,%d)>:"
                ++ " invalid index") j m n
    | otherwise = let
        a1 = unsafeSlice (0,0) (m,j)   a
        a2 = unsafeSlice (0,j) (m,n-j) a
    in (a1,a2)
  where
    (m,n) = dim a
{-# INLINE splitColsAt #-}

-- | Convert a matrix to a vector by stacking its columns.
toVector :: (Storable e)
         => Matrix e
         -> Vector e
toVector a@(Matrix v m n _lda)
    | isContig a = v
    | otherwise  = V.fromList (m*n) $ elems a

-- | Cast a vector to a matrix of the given shape.
fromVector :: (Storable e)
           => (Int,Int)
           -> Vector e
           -> Matrix e
fromVector (m,n) v
    | nv /= m * n = error $
        printf ("fromVector"
               ++ " (%d,%d)"
               ++ " <vector with dim %d>"
               ++ ": dimension mismatch"
               ) m n nv
    | otherwise =
        Matrix v m n (max 1 m)
  where
    nv = V.dim v
{-# INLINE fromVector #-}

-- | Cast a vector to a matrix with one column.
fromCol :: (Storable e)
        => Vector e
        -> Matrix e
fromCol v = Matrix v m 1 (max 1 m)
  where
    m = V.dim v

-- | Cast a vector to a matrix with one row.
fromRow :: (Storable e)
        => Vector e
        -> Matrix e
fromRow v = Matrix v 1 n 1
  where
    n = V.dim v

instance (Storable e, Show e) => Show (Matrix e) where
    show x = "fromList " ++ show (dim x) ++ " " ++ show (elems x)
    {-# INLINE show #-}

instance (Storable e, Eq e) => Eq (Matrix e) where
    (==) = compareWith (==)
    {-# INLINE (==) #-}

instance (Storable e, AEq e) => AEq (Matrix e) where
    (===) = compareWith (===)
    {-# INLINE (===) #-}
    (~==) = compareWith (~==)
    {-# INLINE (~==) #-}

compareWith :: (Storable e, Storable e')
            => (e -> e' -> Bool)
            -> Matrix e
            -> Matrix e'
            -> Bool
compareWith cmp a a' =
    dim a == dim a'
    && and (P.zipWith cmp (elems a) (elems a'))
{-# INLINE compareWith #-}

-- | Create a matrix from a 'ForeignPtr' with offset, dimensions, and lda. The
-- data may not be modified through the ForeignPtr afterwards.
unsafeFromForeignPtr :: (Storable e)
                     => ForeignPtr e -- ^ pointer
                     -> Int	         -- ^ offset
                     -> (Int,Int)    -- ^ dimensions
                     -> Int          -- ^ leading dimension (lda)
                     -> Matrix e
unsafeFromForeignPtr p o (m,n) lda = let
    nv = if m == 0 then 0 else n * lda
    v = V.unsafeFromForeignPtr p o nv
    in Matrix v m n lda
{-# INLINE unsafeFromForeignPtr #-}

-- | Yield the underlying 'ForeignPtr' together with the offset to the data
-- the matrix dimensions, and the lda. The data may not be modified through
-- the 'ForeignPtr'.
unsafeToForeignPtr :: (Storable e)
                   => Matrix e
                   -> (ForeignPtr e, Int, (Int,Int), Int)
unsafeToForeignPtr (Matrix v m n lda) = let
    (f,o,_) = V.unsafeToForeignPtr v
    in (f, o, (m,n), lda)
{-# INLINE unsafeToForeignPtr #-}

-- | Execute an 'IO' action with a pointer to the first element in the
-- matrix and the leading dimension (lda).
unsafeWith :: (Storable e) => Matrix e -> (Ptr e -> Int -> IO a) -> IO a
unsafeWith (Matrix v _ _ lda) f =
    V.unsafeWith v $ \p -> f p lda
{-# INLINE unsafeWith #-}


{-# LANGUAGE CPP #-}
{-# LANGUAGE MultiParamTypeClasses #-}
{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE FlexibleInstances #-}
-----------------------------------------------------------------------------
-- |
-- License     :  BSD-style (see the file LICENSE)
-- Maintainer  :  Edward Kmett <ekmett@gmail.com>
-- Stability   :  provisional
-- Portability :  portable
--
-----------------------------------------------------------------------------
module Linear.Algebra
  ( Algebra(..)
  , Coalgebra(..)
  , multRep, unitalRep
  , comultRep, counitalRep
  ) where

import Control.Lens hiding (index)
import Data.Functor.Rep
import Data.Complex
import Data.Void
import Linear.Vector
import Linear.Quaternion
import Linear.Conjugate
import Linear.V0
import Linear.V1
import Linear.V2
import Linear.V3
import Linear.V4

-- | An associative unital algebra over a ring
class Num r => Algebra r m where
  mult :: (m -> m -> r) -> m -> r
  unital :: r -> m -> r

multRep :: (Representable f, Algebra r (Rep f)) => f (f r) -> f r
multRep ffr = tabulate $ mult (index . index ffr)

unitalRep :: (Representable f, Algebra r (Rep f)) => r -> f r
unitalRep = tabulate . unital

instance Num r => Algebra r Void where
  mult _ _ = 0
  unital _ _ = 0

instance Num r => Algebra r (E V0) where
  mult _ _ = 0
  unital _ _ = 0

instance Num r => Algebra r (E V1) where
  mult f _ = f ex ex
  unital r _ = r

instance Num r => Algebra r () where
  mult f () = f () ()
  unital r () = r

instance (Algebra r a, Algebra r b) => Algebra r (a, b) where
  mult f (a,b) = mult (\a1 a2 -> mult (\b1 b2 -> f (a1,b1) (a2,b2)) b) a
  unital r (a,b) = unital r a * unital r b

instance Num r => Algebra r (E Complex) where
  mult f = \ i -> c^.el i where
   c = (f ee ee - f ei ei) :+ (f ee ei + f ei ee)
  unital r i = (r :+ 0)^.el i

instance (Num r, TrivialConjugate r) => Algebra r (E Quaternion) where
  mult f = index $ Quaternion
    (f ee ee - (f ei ei + f ej ej + f ek ek))
    (V3 (f ee ei + f ei ee + f ej ek - f ek ej)
        (f ee ej + f ej ee + f ek ei - f ei ek)
        (f ee ek + f ek ee + f ei ej - f ej ei))
  unital r = index (Quaternion r 0)

-- | A coassociative counital coalgebra over a ring
class Num r => Coalgebra r m where
  comult :: (m -> r) -> m -> m -> r
  counital :: (m -> r) -> r

comultRep :: (Representable f, Coalgebra r (Rep f)) => f r -> f (f r)
comultRep fr = tabulate $ \i -> tabulate $ \j -> comult (index fr) i j

counitalRep :: (Representable f, Coalgebra r (Rep f)) => f r -> r
counitalRep = counital . index

instance Num r => Coalgebra r Void where
  comult _ _ _ = 0
  counital _ = 0

instance Num r => Coalgebra r () where
  comult f () () = f ()
  counital f = f ()

instance Num r => Coalgebra r (E V0) where
  comult _ _ _ = 0
  counital _ = 0

instance Num r => Coalgebra r (E V1) where
  comult f _ _ = f ex
  counital f = f ex

instance Num r => Coalgebra r (E V2) where
  comult f = index . index v where
    v = V2 (V2 (f ex) 0) (V2 0 (f ey))
  counital f = f ex + f ey

instance Num r => Coalgebra r (E V3) where
  comult f = index . index q where
    q = V3 (V3 (f ex) 0 0)
           (V3 0 (f ey) 0)
           (V3 0 0 (f ez))
  counital f = f ex + f ey + f ez

instance Num r => Coalgebra r (E V4) where
  comult f = index . index v where
    v = V4 (V4 (f ex) 0 0 0) (V4 0 (f ey) 0 0) (V4 0 0 (f ez) 0) (V4 0 0 0 (f ew))
  counital f = f ex + f ey + f ez + f ew

instance Num r => Coalgebra r (E Complex) where
  comult f = \i j -> c^.el i.el j where
    c = (f ee :+ 0) :+ (0 :+ f ei)
  counital f = f ee + f ei

instance (Num r, TrivialConjugate r) => Coalgebra r (E Quaternion) where
  comult f = index . index
    (Quaternion (Quaternion (f ee) (V3 0 0 0))
            (V3 (Quaternion 0 (V3 (f ei) 0 0))
                (Quaternion 0 (V3 0 (f ej) 0))
                (Quaternion 0 (V3 0 0 (f ek)))))
  counital f = f ee + f ei + f ej + f ek

instance (Coalgebra r m, Coalgebra r n) => Coalgebra r (m, n) where
  comult f (a1, b1) (a2, b2) = comult (\a -> comult (\b -> f (a, b)) b1 b2) a1 a2
  counital k = counital $ \a -> counital $ \b -> k (a,b)


module Mandelbrot where

import Data.Complex
import qualified Data.Set as Set
import Data.Set (Set)
import Data.Function (on)
import Data.List (partition)
--import Data.Ratio

c = 0 :+ 0
man :: RealFloat a => Complex a -> Complex a -> Complex a
man c z = z*z + c

pixelSize
	:: (Int, Int)
	-> Double
	-> Double
	-> (Int, Int)
pixelSize maxSize w h =
	if (w/h) < (divInt mw mh)
		then (ceiling $ (fromIntegral mh)*w/h , mh)
		else (mw , ceiling $ (fromIntegral mw)*h/w)
	where
		(mw, mh) = maxSize

-- size = (px,py)
points
	:: (Int, Int)
	-> (Double, Double, Double, Double)
	-> ([((Int, Int), (Complex Double, Complex Double))], Int)
points size bounds =
	([
		(
			(x,(y+offset)) ,
			(\z -> (z,z))
			(
				( w * (divInt x (px-1)) + lB)
				:+
				(-h * (divInt y (py-1)) + uB)
			)
		)
		| y <- ys , x <- [0..(px-1)]
	], py) where
		ys = [0..(py-1)]
		rys = reverse ys
		(l, r, u, d) = bounds
		(lB, rB, uB, dB, offset) =
			if (u > 0.0) && (d < 0.0)
				then
					if u > -d
						then (l, r, u, 0, 0)
						else (l, r, 0, d, ((snd size) - py))
				else (l,r,u,d,0)
		(w, h) = (rB - lB, uB - dB)
		(px,py) = pixelSize size w h

divInt = (/) `on` fromIntegral

reduce = partition (\(x,(c,z)) -> magnitude z <= 2)
apply = map (\(x,(c,z)) -> (x,(c, man c z)))


{-|
Module      : PRNG
Description : Class definitions of PRNG api

Copyright   : Thodoris Papakonstantinou, 2017
License     : GPL-3
Maintainer  : mail@tpapak.com
Stability   : experimental
Portability : POSIX

Deterministic (Pure) Pseudo Random Number Generator Class definitions.
 -}

module Data.PRNG
    ( Seed
    , RNG (..)
    , randomBools
    , randomInts
    , randomPermutation
    , sample
    , normalSample
    , truncatedNormalSample
    ) where

import Data.Maybe
import Data.List
import qualified Data.Vector as V
import qualified Data.Set as S
import qualified Data.IntMap as IM
import Data.Number.Erf
import qualified Numeric.SpecFunctions as NS

type Seed = Int
class RNG r where
  getRNG :: Seed -> r
  randomDouble :: r -> (Double, r)
  randomInt :: r -> (Int, r)
  uniformSample :: r -> Int -> [Double]

-- | booleans given the length of the list the probability p and the seed
randomBools :: RNG r => r -> Int -> Double -> [Bool]
randomBools r l p        =
  let rndbls 0 _         = []
      rndbls n g         = b:(rndbls (n-1) g')
        where (d,g')     = randomDouble g
              b          = d > p
  in rndbls l r

-- | booleans given the length of the list the probability p and the seed
randomInts :: RNG r => r -> Int -> [Int]
randomInts r l         =
  let rnds 0 _         = []
      rnds n g         = i:(rnds (n-1) g')
        where (i,g')   = randomInt g
  in rnds l r

-- | Permutates list using a RNG, given that the list has fewer elements than
-- the minimum between  2^64 which is the precision of Double and the period of the prng,the position probability of every element is uniform.
randomPermutation :: RNG r => r -> [a] -> [a]
randomPermutation r ls =
  let randoms = zip (uniformSample r (length ls)) ls
   in map snd $ sortOn fst randoms

-- | Sample from list
sample :: RNG r => r -> Int -> [a] -> [a]
sample r k c
  | k < 0  || k > length c = []
  | otherwise = take k $ randomPermutation r c

-- | List of doubles from nornal distribution given
-- RNG -> μ -> σ -> length of list
normalSample :: RNG r => r -> Double -> Double -> Int -> [Double]
normalSample r μ σ n
  | σ == 0 = replicate n μ
  | otherwise =
    let uds = uniformSample r n
        invx x = σ * x + μ
     in map (invx . invnormcdf) uds

-- | Sample truncated nomral μ -> σ -> left bound -> right bound -> length of list
truncatedNormalSample :: RNG r => r -> Double -> Double -> Double -> Double -> Int -> [Double]
truncatedNormalSample r μ σ f t n
  | (μ < f) || (μ > t) = []
  | σ == 0 = replicate n μ
  | otherwise =
    let uds = uniformSample r n
        sqt = sqrt 2
        invtruncnormcdf y = μ - sqt * σ * inverf (
          - ((-1) + y) * erf((-f + μ)/(sqt * σ)) +
          y * erf((μ - t)/(sqt * σ))
          )
     in map invtruncnormcdf uds

-- | List of doubles from the binomial distribution given
-- RNG -> n -> p -> length of list
binomialSample :: RNG r => r -> Int -> Double -> Int -> [Int]
binomialSample r n p l =
  let uds = uniformSample r l
      invbinomcdf a n p = let sumpr i = foldl (\ac i -> ac + (binomialProb i n p)) 0 [1..i]
                              binomialProb i n p = n `NS.choose` i * p^^i * (1-p)^^(n-i)
                              firstBigger x a
                                | sumpr (x-1) > a = x
                                | x > n = n
                                | otherwise = firstBigger (x+1) a
                              in firstBigger 1 a
 in map (\a -> invbinomcdf a n p) uds


{-# LANGUAGE
    MultiWayIf,
    FlexibleInstances,
    NumDecimals #-}

module Jac where

import Data.Matrix (Matrix(..), fromLists, fromList, toList)
import Data.Complex

import Point
import Data (F)

type FPD = Point -> Double

-----------------------------------------------------------------
-- производная.
diff :: (Floating a, Eq a) => F a -> a -> a
diff f x = (f (x - d/2) - f (x + d/2))/d
  where
    d = if
        | x == 0    ->   1e-9
        | otherwise -> x*1e-5

-- якобиан.
jac :: F Point -> Point -> Matrix Double
jac f p@(Point x' y') = fromLists [
    [diff (\a -> f1 p{x = a}) x', diff (\a -> f1 p{y = a}) y'],
    [diff (\a -> f2 p{x = a}) x', diff (\a -> f2 p{y = a}) y']]
  where
    f1 = x.f ; f2 = y.f

jac' :: F Point -> Point -> Double
jac' f p@(Point x' y') = log $ abs $
    diff (\a -> f1 p{x = a}) x'*diff (\a -> f2 p{y = a}) y'
  - diff (\a -> f1 p{y = a}) y'*diff (\a -> f2 p{x = a}) x'
  where
    f1 = x.f ; f2 = y.f

-- единичный вектор.
e :: Double -> Matrix Double
e f = fromList 2 1 [realPart x, imagPart x]
  where x = cis f

-- см. стр. 102. в книге 1.
formA :: F Point -> Point -> Double -> Double
formA f p k = log.abs.toList $ jac f p * e k
    where abs [x, y] = magnitude $ x :+ y

-----------------------------------------------------------------
-- преобразование вектора в точку.
instance Points (Matrix Double) where
    toPoint = toPoint.toList
    fromPoint = fromList 2 1.fromPoint

-----------------------------------------------------------------
--                   Справочная литература                     --
-----------------------------------------------------------------
-- 1. Г.С. Осипенко, Н.Б.Ампилова:
--      ВВЕДЕНИЕ В СИМВОЛИЧЕСКИЙ АНАЛИЗ ДИНАМИЧЕСКИХ СИСТЕМ


-- monadic computations
-- (contributed by Vivian McPhail)

import Numeric.LinearAlgebra
import Control.Monad.State.Strict
import Control.Monad.Maybe
import Foreign.Storable(Storable)
import System.Random(randomIO)

-------------------------------------------

-- an instance of MonadIO, a monad transformer
type VectorMonadT = StateT Int IO

test1 :: Vector Int -> IO (Vector Int)
test1 = mapVectorM $ \x -> do
    putStr $ (show x) ++ " "
    return (x + 1)

-- we can have an arbitrary monad AND do IO
addInitialM :: Vector Int -> VectorMonadT ()
addInitialM = mapVectorM_ $ \x -> do
    i <- get
    liftIO $ putStr $ (show $ x + i) ++ " "
    put $ x + i

-- sum the values of the even indiced elements
sumEvens :: Vector Int -> Int
sumEvens = foldVectorWithIndex (\x a b -> if x `mod` 2 == 0 then a + b else b) 0

-- sum and print running total of evens
sumEvensAndPrint :: Vector Int -> VectorMonadT ()
sumEvensAndPrint = mapVectorWithIndexM_ $ \ i x -> do
    when (i `mod` 2 == 0) $ do
        v <- get
        put $ v + x
        v' <- get
        liftIO $ putStr $ (show v') ++ " "

indexPlusSum :: Vector Int -> VectorMonadT ()
indexPlusSum v' = do
    let f i x = do
            s <- get
            let inc = x+s
            liftIO $ putStr $ show (i,inc) ++ " "
            put inc
            return inc
    v <- mapVectorWithIndexM f v'
    liftIO $ do
        putStrLn ""
        putStrLn $ show v

-------------------------------------------

-- short circuit
monoStep :: Double -> MaybeT (State Double) ()
monoStep d = do
    dp <- get
    when (d < dp) (fail "negative difference")
    put d
{-# INLINE monoStep #-}

isMonotoneIncreasing :: Vector Double -> Bool
isMonotoneIncreasing v =
    let res = evalState (runMaybeT $ (mapVectorM_ monoStep v)) (v @> 0)
     in case res of
        Nothing -> False
        Just _  -> True

-------------------------------------------

-- | apply a test to successive elements of a vector, evaluates to true iff test passes for all pairs
successive_ :: Storable a => (a -> a -> Bool) -> Vector a -> Bool
successive_ t v = maybe False (\_ -> True) $ evalState (runMaybeT (mapVectorM_ step (subVector 1 (dim v - 1) v))) (v @> 0)
   where step e = do
                  ep <- lift $ get
                  if t e ep
                     then lift $ put e
                     else (fail "successive_ test failed")

-- | operate on successive elements of a vector and return the resulting vector, whose length 1 less than that of the input
successive :: (Storable a, Storable b) => (a -> a -> b) -> Vector a -> Vector b
successive f v = evalState (mapVectorM step (subVector 1 (dim v - 1) v)) (v @> 0)
   where step e = do
                  ep <- get
                  put e
                  return $ f ep e

-------------------------------------------

v :: Vector Int
v = 10 |> [0..]

w = fromList ([1..10]++[10,9..1]) :: Vector Double

main = do
    v' <- test1 v
    putStrLn ""
    putStrLn $ show v'
    evalStateT (addInitialM v) 0
    putStrLn ""
    putStrLn $ show (sumEvens v)
    evalStateT (sumEvensAndPrint v) 0
    putStrLn ""
    evalStateT (indexPlusSum v) 0
    putStrLn "-----------------------"
    mapVectorM_ print v
    print =<< (mapVectorM (const randomIO) v :: IO (Vector Double))
    print =<< (mapVectorM (\a -> fmap (+a) randomIO) (5|>[0,100..1000]) :: IO (Vector Double))
    putStrLn "-----------------------"
    print $ isMonotoneIncreasing w
    print $ isMonotoneIncreasing (subVector 0 7 w)
    print $ successive_ (>) v
    print $ successive_ (>) w
    print $ successive (+) v


{-# LANGUAGE FlexibleInstances #-}
import Data.Complex
import Data.Numeric.Function
import Graphics.Gnuplot.Simple

-- the physical constant are too small and creates rounding errors
l=1.0e-11
h=6.62606957e-34
m=9.10938188e-31
ts = [0.00,0.01e-19..0.60e-19]

e :: Complex Double -> Complex Double
e n = n^2 * h^2 * pi^2 / (2 * m * l^2)

k :: Complex Double -> Complex Double
k n = n*pi/l

w :: Complex Double -> Complex Double
w n = (e n)/h

a :: Complex Double
a = sqrt(2/l)

img :: Complex Double
img = 0 :+ 1

eexp :: Complex Double -> Complex Double
eexp c = (mkComp (2.7182818**(realPart c))) * ((cos (imagPart c)) :+ (sin (imagPart c)))

psi :: Complex Double -> Double -> Complex Double -> Complex Double
psi n t x = a * sin((k n) * x) * (eexp (-img*(w n)*(t:+0)))

nums = map (\x -> x :+ 0) [1..3]

psis = (1/(sqrt . fromIntegral $ (length nums))) * sum $ map psi nums

sans = (^2) . psis

mkComp x = x :+ 0

y :: [[[Double]]]
y = map (\gress -> [map magnitude gress, map realPart gress, map imagPart gress]) funcs
  where xs = map mkComp (linearScale 300 (0, realPart l))
        funcs :: [[Complex Double]]
        funcs = map (\t -> map (sans t) xs) ts

-- | dirty string format
gress :: Int -> String
gress i
  | i < 10 = "00" ++ show i
  | i < 100 = "0" ++ show i
  | otherwise = show i

onePlot :: [[Double]] -> Int -> IO ()
--onePlot ys i = plotLists ([PNG ("/tmp/kvantePNG" ++ gress i ++ ".png")]) ys
onePlot ys i = plotLists ([YRange (-5.0e11,5.0e11), PNG ("/tmp/kvantePNG" ++ gress i ++ ".png")]) ys
--onePlot ys i = plotLists ([YRange (-3,3), PNG ("/tmp/kvantePNG" ++ gress i ++ ".png")]) ys

fun :: (Int, [[Double]]) -> IO ()
fun (i,ys) = onePlot ys i

jazz :: [(Int, [[Double]])]
jazz = zip [1..] y

main = mapM fun jazz


import Numeric.LinearAlgebra
import HVX

n = 4
a = EConst $ (n><n) ([0, 2, 0, 0,
                      -3, 0, 0, 0,
                       0, 0, 1, 0,
                       0, 0, 0, 4] :: [Double])
b = EConst $ (n><n) ([4, 8, -1, 0,
                      0, -2, 1, 0,
                      9, -3, 0, 3,
                      0, 0,  2, 0] :: [Double])
c = EConst $ (n><1) [1..]

x = EVar "x"
y = EVar "y"
constNeg3 = EConst $ scale (-3.0) $ konst 1.0 (n, 1)
constTwos = EConst $ scale 2.0 $ konst 1.0 (n, 1)

ans = subgradMinimize
  (norm 2 (a *~ x) +~ norm 2 (b *~ y) +~ norm 2 (c +~ x +~ neg y))
  [y <=~ constTwos, x <=~ constNeg3]
  (decNonSumStep 10.0) 5000
  [("x", konst 0.0 (n, 1)), ("y", konst 0.0 (n, 1))]

main :: IO ()
main = print ans


{-# LANGUAGE BangPatterns                             #-}
{-# LANGUAGE DataKinds                                #-}
{-# LANGUAGE DeriveGeneric                            #-}
{-# LANGUAGE FlexibleContexts                         #-}
{-# LANGUAGE MultiWayIf                               #-}
{-# LANGUAGE RecordWildCards                          #-}
{-# LANGUAGE ScopedTypeVariables                      #-}
{-# LANGUAGE TypeApplications                         #-}
{-# LANGUAGE TypeFamilies                             #-}
{-# LANGUAGE TypeOperators                            #-}
{-# LANGUAGE ViewPatterns                             #-}
{-# OPTIONS_GHC -fplugin GHC.TypeLits.KnownNat.Solver #-}
{-# OPTIONS_GHC -fplugin GHC.TypeLits.Normalise       #-}

-- |
-- Module      : Numeric.HHT
-- Copyright   : (c) Justin Le 2019
-- License     : BSD3
--
-- Maintainer  : justin@jle.im
-- Stability   : experimental
-- Portability : non-portable
--
-- Hilbert-Huang transform in pure Haskell.
--
-- The main data type is 'HHT', which can be generated using 'hht' or
-- 'hhtEmd'.  See "Numeric.EMD" for information on why this module uses
-- "sized vectors", and how to convert unsized vectors to sized vectors.
--
-- @since 0.1.2.0

module Numeric.HHT (
  -- * Hilbert-Huang Transform
    HHT(..), HHTLine(..)
  , hhtEmd
  , hht
  , ihhtEmd
  , ihht
  -- ** Hilbert-Huang Spectrum
  , hhtSpectrum, hhtSparseSpectrum, hhtDenseSpectrum
  -- ** Properties of spectrum
  , meanMarginal, marginal, instantaneousEnergy, degreeOfStationarity
  , expectedFreq, dominantFreq
  , foldFreq
  -- ** Options
  , EMDOpts(..), defaultEO, BoundaryHandler(..), defaultSifter, SplineEnd(..)
  -- * Hilbert transforms (internal usage)
  , hilbert
  , hilbertIm
  , hilbertPolar
  , hilbertMagFreq
  ) where

import           Control.DeepSeq
import           Data.Complex
import           Data.Finite
import           Data.Fixed
import           Data.Foldable
import           Data.Proxy
import           Data.Semigroup
import           GHC.Generics              (Generic)
import           GHC.TypeNats
import           Numeric.EMD
import           Numeric.HHT.Internal.FFT
import qualified Data.Binary               as Bi
import qualified Data.Map                  as M
import qualified Data.Vector.Generic       as VG
import qualified Data.Vector.Generic.Sized as SVG
import qualified Data.Vector.Sized         as SV
import qualified Math.FFT.Base             as FFT

-- | A Hilbert Trasnform of a given IMF, given as a "skeleton line".
data HHTLine v n a = HHTLine
    { -- | IMF HHT Magnitude as a time series.
      --
      -- It may be useful to "zip" this vector with 'hlFreqs'.  To do this,
      -- use a function like 'SVG.init' or 'SVG.tail' to make these two
      -- vectors contain the same length, or 'weaken'/'shift' to make
      -- indices in 'hlFreqs' usable as indices in 'hlMags'.
      --
      -- Prior to v0.1.9.0, this was a length-n vector, just like
      -- 'hlFreqs'.  To get the same behavior, use 'SVG.init' on this new
      -- field's value.
      hlMags      :: !(SVG.Vector v (n + 1) a)
      -- | IMF HHT instantaneous frequency as a time series (between 0 and 1).
      --
      -- In reality, these frequencies are the frequencies "in between"
      -- each step in 'hlMags'.
    , hlFreqs     :: !(SVG.Vector v n a)
      -- | Initial phase of skeleton line (between -pi and pi)
      --
      -- @since 0.1.9.0
    , hlInitPhase :: !a
    }
  deriving (Show, Eq, Ord, Generic)

-- | @since 0.1.3.0
instance (VG.Vector v a, KnownNat n, Bi.Binary (v a), Bi.Binary a) => Bi.Binary (HHTLine v n a)

-- | @since 0.1.5.0
instance (NFData (v a), NFData a) => NFData (HHTLine v n a)

-- | A Hilbert-Huang Transform.  An @'HHT' v n a@ is a Hilbert-Huang
-- transform of an @n@-item time series of items of type @a@ represented
-- using vector @v@.
--
-- Create using 'hht' or 'hhtEmd'.
data HHT v n a = HHT
    { -- | Skeleton lines corresponding to each IMF
      hhtLines    :: [HHTLine v n a]
      -- | Residual from EMD
      --
      -- @since 0.1.9.0
    , hhtResidual :: SVG.Vector v (n + 1) a
    }
  deriving (Show, Eq, Ord, Generic)

-- | @since 0.1.3.0
instance (VG.Vector v a, KnownNat n, Bi.Binary (v a), Bi.Binary a) => Bi.Binary (HHT v n a)

-- | @since 0.1.5.0
instance (NFData (v a), NFData a) => NFData (HHT v n a)

-- | Directly compute the Hilbert-Huang transform of a given time series.
-- Essentially is a composition of 'hhtEmd' and 'emd'.  See 'hhtEmd' for
-- a more flexible version.
hht :: forall v n a. (VG.Vector v a, VG.Vector v (Complex a), KnownNat n, FFT.FFTWReal a)
    => EMDOpts v (n + 1) a
    -> SVG.Vector v (n + 1) a
    -> HHT v n a
hht eo = hhtEmd . emd eo

-- | Compute the Hilbert-Huang transform from a given Empirical Mode
-- Decomposition.
hhtEmd
    :: forall v n a. (VG.Vector v a, VG.Vector v (Complex a), KnownNat n, FFT.FFTWReal a)
    => EMD v (n + 1) a
    -> HHT v n a
hhtEmd EMD{..} = HHT (map go emdIMFs) emdResidual
  where
    go i = HHTLine m f φ0
      where
        (m, (f, φ0)) = hilbertMagFreq i

-- | Invert a Hilbert-Huang transform back to an Empirical Mode
-- Decomposition
--
-- @since 0.1.9.0
ihhtEmd
    :: (VG.Vector v a, Floating a)
    => HHT v n a
    -> EMD v (n + 1) a
ihhtEmd HHT{..} = EMD (map go hhtLines) hhtResidual
  where
    go HHTLine{..} = SVG.zipWith (\m θ -> m * cos θ) hlMags θs
      where
        θs = SVG.scanl' (+) hlInitPhase ((* (2 * pi)) `SVG.map` hlFreqs)

-- | Construct a time series correpsonding to its hilbert-huang transform.
--
-- @since 0.1.9.0
ihht
    :: (VG.Vector v a, Floating a)
    => HHT v n a
    -> SVG.Vector v (n + 1) a
ihht = iemd . ihhtEmd

-- | Fold and collapse a Hilbert-Huang transform along the frequency axis
-- at each step in time along some monoid.
--
-- @since 0.1.8.0
foldFreq
    :: forall v u n a b c. (VG.Vector v a, VG.Vector u c, KnownNat n, Monoid b)
    => (a -> a -> b)  -- ^ Combining function, taking frequency, then magnitude
    -> (b -> c)       -- ^ Projecting function
    -> HHT v n a
    -> SVG.Vector u n c
foldFreq f g = pullBack
             . foldl' (SV.zipWith (<>)) (SV.replicate mempty)
             . map split
             . hhtLines
  where
    split :: HHTLine v n a -> SV.Vector n b
    split HHTLine{..} = SVG.generate $ \i ->
      f (hlFreqs `SVG.index` i) (hlMags `SVG.index` weaken i)
    {-# INLINE split #-}
    pullBack :: SV.Vector n b -> SVG.Vector u n c
    pullBack v = SVG.generate $ \i -> g (v `SV.index` i)
    {-# INLINE pullBack #-}
{-# INLINE foldFreq #-}

newtype SumMap k a = SumMap { getSumMap :: M.Map k a }

instance (Ord k, Num a) => Semigroup (SumMap k a) where
    SumMap x <> SumMap y = SumMap $ M.unionWith (+) x y

instance (Ord k, Num a) => Monoid (SumMap k a) where
    mempty = SumMap M.empty

-- | Compute the full Hilbert-Huang Transform spectrum.  At each timestep
-- is a sparse map of frequency components and their respective magnitudes.
-- Frequencies not in the map are considered to be zero.
--
-- Takes a "binning" function to allow you to specify how specific you want
-- your frequencies to be.
--
-- See 'hhtSparseSpectrum' for a sparser version, and 'hhtDenseSpectrum'
-- for a denser version.
hhtSpectrum
    :: forall v n a k. (VG.Vector v a, KnownNat n, Ord k, Num a)
    => (a -> k)     -- ^ binning function.  takes rev/tick freq between 0 and 1.
    -> HHT v n a
    -> SV.Vector n (M.Map k a)
hhtSpectrum f = foldFreq (\k -> SumMap . M.singleton (f k)) getSumMap

-- | A sparser vesion of 'hhtSpectrum'.  Compute the full Hilbert-Huang
-- Transform spectrum.  Returns a /sparse/ matrix representing the power at
-- each time step (the @'Finite' n@) and frequency (the @k@).
--
-- Takes a "binning" function to allow you to specify how specific you want
-- your frequencies to be.
--
-- @since 0.1.4.0
hhtSparseSpectrum
    :: forall v n a k. (VG.Vector v a, KnownNat n, Ord k, Num a)
    => (a -> k)     -- ^ binning function.  takes rev/tick freq between 0 and 1.
    -> HHT v n a
    -> M.Map (Finite n, k) a
hhtSparseSpectrum f = M.unionsWith (+) . concatMap go . hhtLines
  where
    go :: HHTLine v n a -> [M.Map (Finite n, k) a]
    go HHTLine{..} = flip fmap (finites @n) $ \i ->
      M.singleton (i, f $ hlFreqs `SVG.index` i) $
        hlMags `SVG.index` weaken i

-- | A denser version of 'hhtSpectrum'.  Compute the full  Hilbert-Huang
-- Transform spectrum, returning a dense matrix (as a vector of vectors)
-- representing the power at each time step and each frequency.
--
-- Takes a "binning" function that maps a frequency to one of @m@ discrete
-- slots, for accumulation in the dense matrix.
--
-- @since 0.1.4.0
hhtDenseSpectrum
    :: forall v n m a. (VG.Vector v a, KnownNat n, KnownNat m, Num a)
    => (a -> Finite m)     -- ^ binning function.  takes rev/tick freq between 0 and 1.
    -> HHT v n a
    -> SV.Vector n (SV.Vector m a)
hhtDenseSpectrum f h = SV.generate $ \i -> SV.generate $ \j ->
    M.findWithDefault 0 (i, j) ss
  where
    ss = hhtSparseSpectrum f h

-- | Compute the marginal spectrum given a Hilbert-Huang Transform. It
-- provides the "total power" over the entire time series for each
-- frequency component.  See 'meanMarginal' for a version that averages
-- over the length of the time series, making it more close in nature to
-- the purpose of a Fourier Transform.
--
-- A binning function is accepted to allow you to specify how specific you
-- want your frequencies to be.
marginal
    :: forall v n a k. (VG.Vector v a, KnownNat n, Ord k, Num a)
    => (a -> k)     -- ^ binning function.  takes rev/tick freq between 0 and 1.
    -> HHT v n a
    -> M.Map k a
marginal f = M.unionsWith (+) . concatMap go . hhtLines
  where
    go :: HHTLine v n a -> [M.Map k a]
    go HHTLine{..} = flip fmap (finites @n) $ \i ->
      M.singleton (f $ hlFreqs `SVG.index` i) (hlMags `SVG.index` weaken i)

-- | Compute the mean marginal spectrum given a Hilbert-Huang Transform. It
-- is similar to a Fourier Transform; it provides the "total power" over
-- the entire time series for each frequency component, averaged over the
-- length of the time series.
--
-- A binning function is accepted to allow you to specify how specific you
-- want your frequencies to be.
--
-- @since 0.1.8.0
meanMarginal
    :: forall v n a k. (VG.Vector v a, KnownNat n, Ord k, Fractional a)
    => (a -> k)     -- ^ binning function.  takes rev/tick freq between 0 and 1.
    -> HHT v n a
    -> M.Map k a
meanMarginal f = fmap (/ n) . marginal f
  where
    n = fromIntegral $ natVal (Proxy @n)

-- | Returns the "expected value" of frequency at each time step,
-- calculated as a weighted average of all contributions at every frequency
-- at that time step.
--
-- @since 0.1.4.0
expectedFreq
    :: forall v n a. (VG.Vector v a, KnownNat n, Fractional a)
    => HHT v n a
    -> SVG.Vector v n a
expectedFreq = foldFreq (\x y -> (Sum (x * y), Sum y)) (\(Sum x, Sum y) -> x / y)

-- | Returns the dominant frequency (frequency with largest magnitude
-- contribution) at each time step.
--
-- @since 0.1.4.0
dominantFreq
    :: forall v n a. (VG.Vector v a, KnownNat n, Ord a)
    => HHT v n a
    -> SVG.Vector v n a
dominantFreq = foldFreq comb proj
  where
    comb :: a -> a -> Maybe (Max (Arg a a))
    comb x y = Just $ Max $ Arg y x
    proj :: Maybe (Max (Arg a a)) -> a
    proj Nothing = errorWithoutStackTrace
      "Numeric.HHT.dominantFreq: HHT was formed with no Intrinsic Mode Functions"
    proj (Just (Max (Arg _ x))) = x

-- | Compute the instantaneous energy of the time series at every step via
-- the Hilbert-Huang Transform.
instantaneousEnergy
    :: forall v n a. (VG.Vector v a, KnownNat n, Num a)
    => HHT v n a
    -> SVG.Vector v n a
instantaneousEnergy = foldFreq (\_ x -> Sum (x * x)) getSum

-- | Degree of stationarity, as a function of frequency.
degreeOfStationarity
    :: forall v n a k. (VG.Vector v a, KnownNat n, Ord k, Fractional a, Eq a)
    => (a -> k)     -- ^ binning function.  takes rev/tick freq between 0 and 1.
    -> HHT v n a
    -> M.Map k a
degreeOfStationarity f h = fmap (/ n)
                         . M.unionsWith (+)
                         . concatMap go
                         . hhtLines
                         $ h
  where
    n = fromIntegral $ natVal (Proxy @n)
    meanMarg = meanMarginal f h
    go :: HHTLine v n a -> [M.Map k a]
    go HHTLine{..} = flip fmap (finites @n) $ \i ->
        let fr = f $ hlFreqs `SVG.index` i
            mm = meanMarg M.! fr
        in  M.singleton fr $
              if mm == 0
                then 0
                else (1 - (hlMags `SVG.index` weaken i / mm)) ^ (2 :: Int)

-- | Given a time series, return a time series of the /magnitude/ of the
-- hilbert transform and the /frequency/ of the hilbert transform, in units
-- of revolutions per tick.  Is only expected to taken in proper/legal
-- IMFs.
--
-- The frequency will always be between 0 and 1, since we can't determine
-- anything faster given the discretization, and we exclude negative values
-- as physically unmeaningful for an IMF.
hilbertMagFreq
    :: forall v n a. (VG.Vector v a, VG.Vector v (Complex a), KnownNat n, FFT.FFTWReal a)
    => SVG.Vector v (n + 1) a
    -> (SVG.Vector v (n + 1) a, (SVG.Vector v n a, a))
hilbertMagFreq v = (hilbertMag, (hilbertFreq, φ0))
  where
    v'           = hilbert v
    hilbertMag   = SVG.map magnitude v'
    hilbertPhase = SVG.map phase v'
    φ0           = SVG.head hilbertPhase
    hilbertFreq  = SVG.map ((`mod'` 1) . (/ (2 * pi))) $ SVG.tail hilbertPhase - SVG.init hilbertPhase

-- | The polar form of 'hilbert': returns the magnitude and phase of the
-- discrete hilbert transform of a series.
--
-- The computation of magnitude is unique, but computing phase gives us
-- some ambiguity.  The interpretation of the hilbert transform for
-- instantaneous frequency is that the original series "spirals" around the
-- complex plane as time progresses, like a helix.  So, we impose
-- a constraint on the phase to uniquely determine it: \(\phi_{t+1}\) is
-- the /minimal valid phase/ such that \(\phi_{t+1} \geq \phi_{t}\).  This
-- enforces the phase to be monotonically increasing at the slowest
-- possible detectable rate.
--
-- @since 0.1.6.0
hilbertPolar
    :: forall v n a. (VG.Vector v a, VG.Vector v (Complex a), KnownNat n, FFT.FFTWReal a)
    => SVG.Vector v (n + 1) a
    -> (SVG.Vector v (n + 1) a, SVG.Vector v (n + 1) a)
hilbertPolar v = (hilbertMag, hilbertPhase)
  where
    hilbertMag :: SVG.Vector v (n + 1) a
    hilbertFreq :: SVG.Vector v n a
    (hilbertMag, (hilbertFreq, φ0)) = hilbertMagFreq v
    hilbertPhase :: SVG.Vector v (n + 1) a
    hilbertPhase = SVG.scanl' (+) φ0 ((* (2 * pi)) `SVG.map` hilbertFreq)

-- | Real part is original series and imaginary part is hilbert transformed
-- series.  Creates a "helical" form of the original series that rotates
-- along the complex plane.
--
-- Note that since /0.1.7.0/, this uses the same algorithm as the matlab
-- implementation <https://www.mathworks.com/help/signal/ref/hilbert.html>
hilbert
    :: forall v n a.
      ( VG.Vector v a
      , VG.Vector v (Complex a)
      , KnownNat n
      , FFT.FFTWReal a
      )
    => SVG.Vector v n a
    -> SVG.Vector v n (Complex a)
hilbert v = ifft u'
  where
    v' = SVG.map (:+ 0) v
    u  = fft v'
    u' = flip SVG.imap u $ \(fromIntegral->i) x ->
      if | i == 0 || i == (n `div` 2) -> x
         | i < (n `div` 2)            -> 2 * x
         | otherwise                  -> 0
    n  = natVal (Proxy @n)

-- | Hilbert transformed series.  Essentially the same series, but
-- phase-shifted 90 degrees.  Is so-named because it is the "imaginary
-- part" of the proper hilbert transform, 'hilbert'.
--
-- Note that since /0.1.7.0/, this uses the same algorithm as the matlab
-- implementation <https://www.mathworks.com/help/signal/ref/hilbert.html>
hilbertIm
    :: forall v n a.
      ( VG.Vector v a
      , VG.Vector v (Complex a)
      , KnownNat n
      , FFT.FFTWReal a
      )
    => SVG.Vector v n a
    -> SVG.Vector v n a
hilbertIm = SVG.map imagPart . hilbert


{-# OPTIONS_GHC -fplugin GHC.TypeLits.KnownNat.Solver #-}
{-# LANGUAGE GADTs #-}
{-# LANGUAGE DataKinds #-}
{-# LANGUAGE TypeOperators #-}
{-# LANGUAGE KindSignatures #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE StandaloneDeriving #-}
{-# LANGUAGE TypeApplications #-}
{-# LANGUAGE TypeSynonymInstances #-}
{-# LANGUAGE FlexibleInstances #-}
module AI.Singularity.Data.Matrix where
import GHC.TypeLits
import qualified Numeric.LinearAlgebra.Data as LD
import qualified Numeric.LinearAlgebra.Devel as LDev
import qualified Numeric.LinearAlgebra as LA
import Control.Lens
import Data.Proxy
import Control.Monad.ST
import AI.Singularity.Utils
import System.Random
import Control.Monad

data GeneralMatrix :: (* -> *) -> Nat -> Nat -> * where
  Mat :: mat Float -> GeneralMatrix mat m n

deriving instance Show (mat Float) => Show (GeneralMatrix mat m n)

type Matrix = GeneralMatrix LD.Matrix

type MutMatrix s = GeneralMatrix (LDev.STMatrix s)

matr :: Lens' (GeneralMatrix mat m n) (mat Float)
{-# INLINE matr #-}
matr f (Mat m) = Mat <$> f m

data GeneralInTrainingMat mat m n =
  GeneralInTrainingMat (GeneralMatrix mat m n) (GeneralMatrix mat m n) (GeneralMatrix mat m n)

type InTrainMat m n = GeneralInTrainingMat LD.Matrix m n

type STInTrainMat s = GeneralInTrainingMat (LDev.STMatrix s)

deriving instance Show (mat Float) => Show (GeneralInTrainingMat mat m n)

trMatr :: Lens' (GeneralInTrainingMat mat m n) (GeneralMatrix mat m n)
{-# INLINE trMatr #-}
trMatr f (GeneralInTrainingMat m g h) = (\x -> GeneralInTrainingMat x g h) <$> f m

trGrad :: Lens' (GeneralInTrainingMat mat m n) (GeneralMatrix mat m n)
{-# INLINE trGrad #-}
trGrad f (GeneralInTrainingMat m g h) = (\x -> GeneralInTrainingMat m x h) <$> f g

hGrad :: Lens' (GeneralInTrainingMat mat m n) (GeneralMatrix mat m n)
{-# INLINE hGrad #-}
hGrad f (GeneralInTrainingMat m g h) = GeneralInTrainingMat m g <$> f h

matSlice,gradSlice :: forall s m n. (KnownNat m, KnownNat n) => STInTrainMat s m n -> LDev.Slice s Float
matSlice  m = LDev.Slice (m^.trMatr.matr) 0 0 (fromIntegral . natVal $ Proxy @ n) (fromIntegral . natVal $ Proxy @ m)
gradSlice m = LDev.Slice (m^.trGrad.matr) 0 0 (fromIntegral . natVal $ Proxy @ n) (fromIntegral . natVal $ Proxy @ m)

thawM :: Matrix m n -> ST s (MutMatrix s m n)
thawM = Mat <$< LDev.thawMatrix . view matr

toSTInTrainM :: (KnownNat m, KnownNat n) => Matrix m n -> ST s (STInTrainMat s m n)
toSTInTrainM m = GeneralInTrainingMat <$> thawM m <*> thawM 0 <*> thawM 0

matHeight :: forall m n mat. KnownNat n => GeneralMatrix mat m n -> Int
matHeight _ = natToInt @ n

matWidth :: forall m n mat. KnownNat m => GeneralMatrix mat m n -> Int
matWidth _ = natToInt @ m

mat :: forall m n. (KnownNat m, KnownNat n) => [Float] -> Matrix m n
mat = Mat . (valm LD.>< valn) . take (valm * valn) . cycle
    where valm = matHeight ( undefined :: Matrix m n)
          valn = matWidth ( undefined :: Matrix m n)

randomMat :: forall m n. (KnownNat m, KnownNat n) => IO (Matrix m n)
randomMat = Mat . (valm LD.>< valn) <$> replicateM (valm * valn) (randomRIO (0.001, 0.1))
    where valm = matHeight ( undefined :: Matrix m n)
          valn = matWidth ( undefined :: Matrix m n)

zeroMat :: forall m n. (KnownNat m, KnownNat n) => Matrix m n
zeroMat = 0

ident :: forall n. (KnownNat n) => Matrix n n
ident = Mat $ LD.ident valn
  where valn = matHeight ( undefined :: Matrix n n)

transposeM :: forall m n. (KnownNat m, KnownNat n) => Matrix m n -> Matrix n m
transposeM = Mat . LA.tr . view matr

instance (KnownNat n, KnownNat m) => Num (Matrix m n) where
  (+) (Mat m1) (Mat m2) = Mat (m1 + m2)
  (*) (Mat m1) (Mat m2) = Mat (m1 * m2)
  negate (Mat m1)       = Mat $ LD.cmap negate m1
  fromInteger i         = mat [fromInteger i]
  abs (Mat m1)          = Mat $ LD.cmap abs m1
  signum (Mat m1)       = Mat $ LD.cmap signum m1

instance (KnownNat n, KnownNat m) => Fractional (Matrix m n) where
  (/) (Mat m1) (Mat m2) = Mat $ m1 / m2
  fromRational          = mat . (:[]) . fromRational

instance (KnownNat n, KnownNat m, n ~ m) => Monoid (Matrix n m) where
  mempty                    = Mat $ LA.ident $ fromInteger $ natVal (Proxy :: Proxy n)
  mappend (Mat m1) (Mat m2) = Mat $ (LA.<>) m1 m2

mmap :: (KnownNat n, KnownNat m) => (Float -> Float) -> Matrix m n -> Matrix m n
mmap f (Mat m) = Mat $ LD.cmap f m

checkNan :: forall m n. Matrix m n -> Matrix m n
checkNan (Mat x) = Mat (LD.cmap nancheck x)


{-# LANGUAGE FlexibleContexts, Rank2Types, ScopedTypeVariables #-}
-- | Student's T-test is for assessing whether two samples have
--   different mean. This module contain several variations of
--   T-test. It's a parametric tests and assumes that samples are
--   normally distributed.
module Statistics.Test.StudentT
    (
      studentTTest
    , welchTTest
    , pairedTTest
    , module Statistics.Test.Types
    ) where

import Statistics.Distribution hiding (mean)
import Statistics.Distribution.StudentT
import Statistics.Sample (mean, varianceUnbiased)
import Statistics.Test.Types
import Statistics.Types    (mkPValue,PValue)
import Statistics.Function (square)
import qualified Data.Vector.Generic  as G
import qualified Data.Vector.Unboxed  as U
import qualified Data.Vector.Storable as S
import qualified Data.Vector          as V

-- | Two-sample Student's t-test. It assumes that both samples are
--   normally distributed and have same variance. Returns @Nothing@ if
--   sample sizes are not sufficient.
studentTTest :: (G.Vector v Double)
             => PositionTest  -- ^ one- or two-tailed test
             -> v Double      -- ^ Sample A
             -> v Double      -- ^ Sample B
             -> Maybe (Test StudentT)
studentTTest test sample1 sample2
  | G.length sample1 < 2 || G.length sample2 < 2 = Nothing
  | otherwise                                    = Just Test
      { testSignificance = significance test t ndf
      , testStatistics   = t
      , testDistribution = studentT ndf
      }
  where
    (t, ndf) = tStatistics True sample1 sample2
{-# INLINABLE  studentTTest #-}
{-# SPECIALIZE studentTTest :: PositionTest -> U.Vector Double -> U.Vector Double -> Maybe (Test StudentT) #-}
{-# SPECIALIZE studentTTest :: PositionTest -> S.Vector Double -> S.Vector Double -> Maybe (Test StudentT) #-}
{-# SPECIALIZE studentTTest :: PositionTest -> V.Vector Double -> V.Vector Double -> Maybe (Test StudentT) #-}

-- | Two-sample Welch's t-test. It assumes that both samples are
--   normally distributed but doesn't assume that they have same
--   variance. Returns @Nothing@ if sample sizes are not sufficient.
welchTTest :: (G.Vector v Double)
           => PositionTest  -- ^ one- or two-tailed test
           -> v Double      -- ^ Sample A
           -> v Double      -- ^ Sample B
           -> Maybe (Test StudentT)
welchTTest test sample1 sample2
  | G.length sample1 < 2 || G.length sample2 < 2 = Nothing
  | otherwise                                    = Just Test
      { testSignificance = significance test t ndf
      , testStatistics   = t
      , testDistribution = studentT ndf
      }
  where
    (t, ndf) = tStatistics False sample1 sample2
{-# INLINABLE  welchTTest #-}
{-# SPECIALIZE welchTTest :: PositionTest -> U.Vector Double -> U.Vector Double -> Maybe (Test StudentT) #-}
{-# SPECIALIZE welchTTest :: PositionTest -> S.Vector Double -> S.Vector Double -> Maybe (Test StudentT) #-}
{-# SPECIALIZE welchTTest :: PositionTest -> V.Vector Double -> V.Vector Double -> Maybe (Test StudentT) #-}

-- | Paired two-sample t-test. Two samples are paired in a
-- within-subject design. Returns @Nothing@ if sample size is not
-- sufficient.
pairedTTest :: forall v. (G.Vector v (Double, Double), G.Vector v Double)
            => PositionTest          -- ^ one- or two-tailed test
            -> v (Double, Double)    -- ^ paired samples
            -> Maybe (Test StudentT)
pairedTTest test sample
  | G.length sample < 2 = Nothing
  | otherwise           = Just Test
      { testSignificance = significance test t ndf
      , testStatistics   = t
      , testDistribution = studentT ndf
      }
  where
    (t, ndf) = tStatisticsPaired sample
{-# INLINABLE  pairedTTest #-}
{-# SPECIALIZE pairedTTest :: PositionTest -> U.Vector (Double,Double) -> Maybe (Test StudentT) #-}
{-# SPECIALIZE pairedTTest :: PositionTest -> V.Vector (Double,Double) -> Maybe (Test StudentT) #-}

-------------------------------------------------------------------------------

significance :: PositionTest    -- ^ one- or two-tailed
             -> Double          -- ^ t statistics
             -> Double          -- ^ degree of freedom
             -> PValue Double   -- ^ p-value
significance test t df =
  case test of
    -- Here we exploit symmetry of T-distribution and calculate small tail
    SamplesDiffer -> mkPValue $ 2 * tailArea (negate (abs t))
    AGreater      -> mkPValue $ tailArea (negate t)
    BGreater      -> mkPValue $ tailArea  t
  where
    tailArea = cumulative (studentT df)

-- Calculate T statistics for two samples
tStatistics :: (G.Vector v Double)
            => Bool               -- variance equality
            -> v Double
            -> v Double
            -> (Double, Double)
{-# INLINE tStatistics #-}
tStatistics varequal sample1 sample2 = (t, ndf)
  where
    -- t-statistics
    t = (m1 - m2) / sqrt (
      if varequal
        then ((n1 - 1) * s1 + (n2 - 1) * s2) / (n1 + n2 - 2) * (1 / n1 + 1 / n2)
        else s1 / n1 + s2 / n2)

    -- degree of freedom
    ndf | varequal  = n1 + n2 - 2
        | otherwise = square (s1 / n1 + s2 / n2)
                    / (square s1 / (square n1 * (n1 - 1)) + square s2 / (square n2 * (n2 - 1)))
    -- statistics of two samples
    n1 = fromIntegral $ G.length sample1
    n2 = fromIntegral $ G.length sample2
    m1 = mean sample1
    m2 = mean sample2
    s1 = varianceUnbiased sample1
    s2 = varianceUnbiased sample2

-- Calculate T-statistics for paired sample
tStatisticsPaired :: (G.Vector v (Double, Double), G.Vector v Double)
                  => v (Double, Double)
                  -> (Double, Double)
{-# INLINE tStatisticsPaired #-}
tStatisticsPaired sample = (t, ndf)
  where
    -- t-statistics
    t = let d    = G.map (uncurry (-)) sample
            sumd = G.sum d
        in sumd / sqrt ((n * G.sum (G.map square d) - square sumd) / ndf)
    -- degree of freedom
    ndf = n - 1
    n   = fromIntegral $ G.length sample


{-# LANGUAGE DataKinds       #-}
{-# LANGUAGE PatternSynonyms #-}
{-# LANGUAGE ViewPatterns    #-}

import           Data.Maybe
import           Numeric.Hamilton
import           Numeric.LinearAlgebra.Static
import qualified Data.Vector.Sized            as V

twoBody :: System 4 2
twoBody = mkSystem masses coordinates potential
  where
    masses :: R 4
    masses = vec4 10 10 1 1
    coordinates
        :: Floating a
        => V.Vector 2 a
        -> V.Vector 4 a
    coordinates (V2 r θ) = V4 (r1 * cos θ) (r1 * sin θ)
                              (r2 * cos θ) (r2 * sin θ)
      where
        r1 =   r *  1 / 11
        r2 = - r * 10 / 11
    potential
        :: Fractional a
        => V.Vector 2 a
        -> a
    potential (V2 r _) = - 10 / r       -- G = 1

pattern V2 :: a -> a -> V.Vector 2 a
pattern V2 x y <- (V.toList->[x,y])
  where
    V2 x y = fromJust (V.fromList [x,y])

pattern V4 :: a -> a -> a -> a -> V.Vector 4 a
pattern V4 x y z a <- (V.toList->[x,y,z,a])
  where
    V4 x y z a = fromJust (V.fromList [x,y,z,a])

config0 :: Config 2
config0 = Cfg (vec2 2   0)  -- initial positions
              (vec2 0 0.5)  -- initial velocities

phase0 :: Phase 2
phase0 = toPhase twoBody config0

evolution :: [Phase 2]
evolution = evolveHam' twoBody phase0 [0,0.1 .. 1]

evolution' :: [Phase 2]
evolution' = iterate (stepHam 0.1 twoBody) phase0

positions :: [R 2]
positions = phsPositions <$> evolution'

positions' :: [R 4]
positions' = underlyingPos twoBody <$> positions

main :: IO ()
main = withRows (take 25 positions) (disp 4)


module Main where

import qualified Numeric.LinearAlgebra.Data

import Optimization.Types
import Optimization.NelderMead

trivialOpt :: Vector Double -> Double
trivialOpt v
  | length v != 3 = error "Vector is not length-3!"
  | otherwise     = ((v ! 0)**2) * (((v ! 1) - 1)**2) * (((v ! 2) - 2) ** 2)

main :: IO ()
main = do
  let
    tc      =  terminateAtOrAfter 0.001 1000
    optimum =  nelderMead tc trivialOpt (fromList [10.0, 12.1, -4.4])

  putStrLn $ show optimum


{-# LANGUAGE DeriveDataTypeable    #-}
{-# LANGUAGE DeriveGeneric         #-}
{-# LANGUAGE MultiParamTypeClasses #-}
{-# LANGUAGE RecordWildCards       #-}

-- |
-- Module    : Statistics.Distribution.Normal
-- Copyright : (c) 2021 Alexander E. Zarebski
-- License   : MIT
--
-- Maintainer  : aezarebski@gmail.com
-- Stability   : experimental
-- Portability : portable
--
-- The lognormal distribution.

module Statistics.Distribution.Lognormal
   (
     LognormalDistribution
   -- * constructors
   , lognormalDistr
   , lognormalDistrE
   ) where

import           Data.Data                       (Data, Typeable)
import           GHC.Generics                    (Generic)
import           Numeric.MathFunctions.Constants (m_sqrt_2, m_sqrt_2_pi)
import           Numeric.SpecFunctions           (erf, invErf)

import qualified Statistics.Distribution         as D

-- | The lognormal distribution.
data LognormalDistribution = LND {
      meanlog :: {-# UNPACK #-} !Double
    , sdlog   :: {-# UNPACK #-} !Double
    } deriving (Eq, Typeable, Data, Generic)

instance D.Distribution LognormalDistribution where
    cumulative      = cumulative

instance D.ContDistr LognormalDistribution where
    logDensity    = logDensity
    quantile      = quantile

instance D.MaybeMean LognormalDistribution where
    maybeMean = Just . D.mean

instance D.Mean LognormalDistribution where
    mean LND {..} = exp (meanlog + 0.5 * sdlog ** 2.0)

instance D.MaybeVariance LognormalDistribution where
    maybeStdDev   = Just . D.stdDev
    maybeVariance = Just . D.variance

instance D.Variance LognormalDistribution where
    variance LND {..} = exp (2.0 * meanlog + sdlog ** 2.0) * (exp (sdlog ** 2.0) - 1)

-- | Create lognormal distribution from parameters.
lognormalDistr :: Double            -- ^ Mean of the logarithm
            -> Double            -- ^ Standard deviation of the logarithm
            -> LognormalDistribution
lognormalDistr m sd = maybe (error $ errMsg m sd) id $ lognormalDistrE m sd

-- | Create normal distribution from parameters.
lognormalDistrE :: Double            -- ^ Mean of logarithm
             -> Double            -- ^ Standard deviation of logarithm
             -> Maybe LognormalDistribution
lognormalDistrE mln sdln
  | sdln > 0    = Just LND { meanlog       = mln
                           , sdlog = sdln
                           }
  | otherwise = Nothing

errMsg :: Double -> Double -> String
errMsg _ sd = "Statistics.Distribution.Lognormal.lognormalDistr: standard deviation must be positive. Got " ++ show sd

logDensity :: LognormalDistribution -> Double -> Double
logDensity LND {..} x = - (log x - meanlog) ** 2.0 / (2 * sdlog ** 2.0) - log (x * sdlog * m_sqrt_2_pi)

cumulative :: LognormalDistribution -> Double -> Double
cumulative LND {..} x = 0.5 * (1 + erf ((log x - meanlog) / (sdlog * m_sqrt_2)))

quantile :: LognormalDistribution -> Double -> Double
quantile LND {..} p
  | p == 0 = 0
  | p == 1 = inf
  | p == 0.5 = exp meanlog
  | p > 0 && p < 1 = exp (meanlog + sdlog * m_sqrt_2 * invErf (2 * p - 1))
  | otherwise = error $ "Statistics.Distribution.Lognormal.quantile: p must be in [0,1] range. Got: "++show p
  where inf = 1/0


-- https://github.com/minoki/my-atcoder-solutions
{-# LANGUAGE BangPatterns          #-}
{-# LANGUAGE DataKinds             #-}
{-# LANGUAGE MultiParamTypeClasses #-}
{-# LANGUAGE NoStarIsType          #-}
{-# LANGUAGE ScopedTypeVariables   #-}
{-# LANGUAGE TypeApplications      #-}
{-# LANGUAGE TypeFamilies          #-}
{-# LANGUAGE TypeOperators         #-}
import           Control.Exception           (assert)
import           Control.Monad
import           Data.Bits
import qualified Data.ByteString.Char8       as BS
import           Data.Char                   (isSpace)
import           Data.Coerce
import           Data.Complex
import           Data.Int                    (Int64)
import           Data.List                   (unfoldr)
import qualified Data.Vector.Generic         as G
import qualified Data.Vector.Unboxing         as U
import qualified Data.Vector.Unboxing.Mutable as UM
import           GHC.TypeNats                (type (+), type (*), KnownNat, Nat, type (^),
                                              natVal)

type R1 = IntMod (5 * 2^25 + 1)
type R2 = IntMod (7 * 2^26 + 1)
type R3 = IntMod (3 * 2^18 + 1)
type R4 = IntMod (7 * 2^20 + 1)
type R5 = IntMod (483 * 2^21 + 1)

order :: (Num a, Eq a) => a -> Int
order !x = go 1 x
  where
    go !n 1 = n
    go !n y = go (n + 1) (x * y)

order' :: (Num a, Eq a) => Int -> a -> Int
order' !m !x = go 1 x
  where
    go !n 1 = n
    go !n y | n > m = m + 1
    go !n y = go (n + 1) (x * y)

findPrimitiveNthRoot :: (Num a, Eq a) => Int -> a
findPrimitiveNthRoot n = head [ x | k <- [1..], let x = fromInteger k, order' n x == n ]

main = do
  -- print (findPrimitiveNthRoot 2)
  -- print (findPrimitiveNthRoot (2^10))
  print (findPrimitiveNthRoot (2^25) :: R1) -- 17 mod 5 * 2^25 + 1
  print (findPrimitiveNthRoot (2^26) :: R2) -- 30 mod 7 * 2^26 + 1
  print (findPrimitiveNthRoot (2^18) :: R3) -- 5 mod 3 * 2^18 + 1
  print (findPrimitiveNthRoot (2^20) :: R4) -- 5 mod 7 * 2^20 + 1
  print (findPrimitiveNthRoot (2^21) :: R5) -- 198 mod 483 * 2^21 + 1

--
-- Fast Fourier Transform (FFT)
--

halve :: G.Vector vec a => vec a -> vec a
halve v = let n = G.length v
          in G.generate (n `quot` 2) $ \j -> v G.! (j * 2)

fft :: forall vec a. (Num a, G.Vector vec a)
    => [vec a] -- ^ For a primitive n-th root of unity @u@, @iterate halve [1,u,u^2 .. u^(n-1)]@
    -> vec a -- ^ a polynomial of length n (= 2^k for some k)
    -> vec a
fft (u:u2) f | n == 1 = f
             | otherwise = let !n2 = n `quot` 2
                               r0, r1', t0, t1' :: vec a
                               r0 = G.generate n2 $ \j -> (f G.! j) + (f G.! (j + n2))
                               r1' = G.generate n2 $ \j -> ((f G.! j) - (f G.! (j + n2))) * u G.! j
                               !t0 = fft u2 r0
                               !t1' = fft u2 r1'
                           in G.generate n $ \j -> if even j then t0 G.! (j `quot` 2) else t1' G.! (j `quot` 2)
  where n = G.length f

mulFFT :: U.Vector Int -> U.Vector Int -> U.Vector Int
mulFFT !f !g = let n' = U.length f + U.length g - 2
                   k = finiteBitSize n' - countLeadingZeros n'
                   !_ = assert (n' < 2^k) ()
                   n = bit k
                   u :: U.Vector (Complex Double)
                   u = U.generate n $ \j -> cis (fromIntegral j * (2 * pi / fromIntegral n))
                   us = iterate halve u
                   f' = U.generate n $ \j -> if j < U.length f then
                                               fromIntegral (f U.! j)
                                             else
                                               0
                   g' = U.generate n $ \j -> if j < U.length g then
                                               fromIntegral (g U.! j)
                                             else
                                               0
                   f'' = fft us f'
                   g'' = fft us g'
                   fg = U.generate n $ \j -> (f'' U.! j) * (g'' U.! j)
                   fg' = fft (map (U.map conjugate) us) fg
               in U.generate n $ \j -> round (realPart (fg' U.! j) / fromIntegral n)

--
-- Univariate polynomial
--

newtype Poly vec a = Poly { coeffAsc :: vec a } deriving Eq

normalizePoly :: (Eq a, Num a, G.Vector vec a) => vec a -> vec a
normalizePoly v | G.null v || G.last v /= 0 = v
                | otherwise = normalizePoly (G.init v)

addPoly :: (Eq a, Num a, G.Vector vec a) => vec a -> vec a -> vec a
addPoly v w = case compare n m of
                LT -> G.generate m $ \i -> if i < n
                                           then v G.! i + w G.! i
                                           else w G.! i
                GT -> G.generate n $ \i -> if i < m
                                           then v G.! i + w G.! i
                                           else v G.! i
                EQ -> normalizePoly $ G.zipWith (+) v w
  where n = G.length v
        m = G.length w

subPoly :: (Eq a, Num a, G.Vector vec a) => vec a -> vec a -> vec a
subPoly v w = case compare n m of
                LT -> G.generate m $ \i -> if i < n
                                           then v G.! i - w G.! i
                                           else negate (w G.! i)
                GT -> G.generate n $ \i -> if i < m
                                           then v G.! i - w G.! i
                                           else v G.! i
                EQ -> normalizePoly $ G.zipWith (-) v w
  where n = G.length v
        m = G.length w

naiveMulPoly :: (Num a, G.Vector vec a) => vec a -> vec a -> vec a
naiveMulPoly v w = G.generate (n + m - 1) $
                   \i -> sum [(v G.! (i-j)) * (w G.! j) | j <- [max (i-n+1) 0..min i (m-1)]]
  where n = G.length v
        m = G.length w

doMulP :: (Eq a, Num a, G.Vector vec a) => Int -> vec a -> vec a -> vec a
doMulP n !v !w | n <= 16 = naiveMulPoly v w
doMulP n !v !w
  | G.null v = v
  | G.null w = w
  | G.length v < n2 = let (w0, w1) = G.splitAt n2 w
                          u0 = doMulP n2 v w0
                          u1 = doMulP n2 v w1
                      in G.generate (G.length v + G.length w - 1)
                         $ \i -> case () of
                                   _ | i < n2     ->  u0 `at` i
                                     | i < n      -> (u0 `at` i) + (u1 `at` (i - n2))
                                     | i < n + n2 -> (u1 `at` (i - n2))
  | G.length w < n2 = let (v0, v1) = G.splitAt n2 v
                          u0 = doMulP n2 v0 w
                          u1 = doMulP n2 v1 w
                      in G.generate (G.length v + G.length w - 1)
                         $ \i -> case () of
                                   _ | i < n2     ->  u0 `at` i
                                     | i < n      -> (u0 `at` i) + (u1 `at` (i - n2))
                                     | i < n + n2 -> (u1 `at` (i - n2))
  | otherwise = let (v0, v1) = G.splitAt n2 v
                    (w0, w1) = G.splitAt n2 w
                    v0_1 = v0 `addPoly` v1
                    w0_1 = w0 `addPoly` w1
                    p = doMulP n2 v0_1 w0_1
                    q = doMulP n2 v0 w0
                    r = doMulP n2 v1 w1
                    -- s = (p `subPoly` q) `subPoly` r -- p - q - r
                    -- q + s*X^n2 + r*X^n
                in G.generate (G.length v + G.length w - 1)
                   $ \i -> case () of
                             _ | i < n2     ->   q `at` i
                               | i < n      -> ((q `at` i) + (p `at` (i - n2))) - ((q `at` (i - n2)) + (r `at` (i - n2)))
                               | i < n + n2 -> ((r `at` (i - n)) + (p `at` (i - n2))) - ((q `at` (i - n2)) + (r `at` (i - n2)))
                               | otherwise  ->   r `at` (i - n)
  where n2 = n `quot` 2
        at :: (Num a, G.Vector vec a) => vec a -> Int -> a
        at v i = if i < G.length v then v G.! i else 0
{-# INLINE doMulP #-}

mulPoly :: (Eq a, Num a, G.Vector vec a) => vec a -> vec a -> vec a
mulPoly !v !w = let k = ceiling ((log (fromIntegral (max n m)) :: Double) / log 2) :: Int
                in doMulP (2^k) v w
  where n = G.length v
        m = G.length w
{-# INLINE mulPoly #-}

zeroPoly :: (G.Vector vec a) => Poly vec a
zeroPoly = Poly G.empty

constPoly :: (Eq a, Num a, G.Vector vec a) => a -> Poly vec a
constPoly 0 = Poly G.empty
constPoly x = Poly (G.singleton x)

scalePoly :: (Eq a, Num a, G.Vector vec a) => a -> Poly vec a -> Poly vec a
scalePoly a (Poly xs)
  | a == 0 = zeroPoly
  | otherwise = Poly $ G.map (* a) xs

valueAtPoly :: (Num a, G.Vector vec a) => Poly vec a -> a -> a
valueAtPoly (Poly xs) t = G.foldr' (\a b -> a + t * b) 0 xs

instance (Eq a, Num a, G.Vector vec a) => Num (Poly vec a) where
  (+) = coerce (addPoly :: vec a -> vec a -> vec a)
  (-) = coerce (subPoly :: vec a -> vec a -> vec a)
  negate (Poly v) = Poly (G.map negate v)
  (*) = coerce (mulPoly :: vec a -> vec a -> vec a)
  fromInteger = constPoly . fromInteger
  abs = undefined; signum = undefined

divModPoly :: (Eq a, Fractional a, G.Vector vec a) => Poly vec a -> Poly vec a -> (Poly vec a, Poly vec a)
divModPoly f g@(Poly w)
  | G.null w = error "divModPoly: divide by zero"
  | degree f < degree g = (zeroPoly, f)
  | otherwise = loop zeroPoly (scalePoly (recip b) f)
  where
    g' = toMonic g
    b = leadingCoefficient g
    -- invariant: f == q * g + scalePoly b p
    loop q p | degree p < degree g = (q, scalePoly b p)
             | otherwise = let q' = Poly (G.drop (degree' g) (coeffAsc p))
                           in loop (q + q') (p - q' * g')

    toMonic :: (Fractional a, G.Vector vec a) => Poly vec a -> Poly vec a
    toMonic f@(Poly xs)
      | G.null xs = zeroPoly
      | otherwise = Poly $ G.map (* recip (leadingCoefficient f)) xs

    leadingCoefficient :: (Num a, G.Vector vec a) => Poly vec a -> a
    leadingCoefficient (Poly xs)
      | G.null xs = 0
      | otherwise = G.last xs

    degree :: G.Vector vec a => Poly vec a -> Maybe Int
    degree (Poly xs) = case G.length xs - 1 of
      -1 -> Nothing
      n  -> Just n

    degree' :: G.Vector vec a => Poly vec a -> Int
    degree' (Poly xs) = case G.length xs of
      0 -> error "degree': zero polynomial"
      n -> n - 1

-- 組立除法
-- second constPoly (divModByDeg1 f t) = divMod f (Poly (G.fromList [-t, 1]))
divModByDeg1 :: (Eq a, Num a, G.Vector vec a) => Poly vec a -> a -> (Poly vec a, a)
divModByDeg1 f t = let w = G.postscanr (\a b -> a + b * t) 0 $ coeffAsc f
                   in (Poly (G.tail w), G.head w)

--
-- Modular Arithmetic
--

newtype IntMod (m :: Nat) = IntMod { unwrapN :: Int64 } deriving (Eq)

instance Show (IntMod m) where
  show (IntMod x) = show x

instance KnownNat m => Num (IntMod m) where
  t@(IntMod x) + IntMod y
    | x + y >= modulus = IntMod (x + y - modulus)
    | otherwise = IntMod (x + y)
    where modulus = fromIntegral (natVal t)
  t@(IntMod x) - IntMod y
    | x >= y = IntMod (x - y)
    | otherwise = IntMod (x - y + modulus)
    where modulus = fromIntegral (natVal t)
  t@(IntMod x) * IntMod y = IntMod ((x * y) `rem` modulus)
    where modulus = fromIntegral (natVal t)
  fromInteger n = let result = IntMod (fromInteger (n `mod` fromIntegral modulus))
                      modulus = natVal result
                  in result
  abs = undefined; signum = undefined

{-# RULES
"^9/Int" forall x. x ^ (9 :: Int) = let u = x; v = u * u * u in v * v * v
"^9/Integer" forall x. x ^ (9 :: Integer) = let u = x; v = u * u * u in v * v * v
 #-}

fromIntegral_Int64_IntMod :: KnownNat m => Int64 -> IntMod m
fromIntegral_Int64_IntMod n = result
  where
    result | 0 <= n && n < modulus = IntMod n
           | otherwise = IntMod (n `mod` modulus)
    modulus = fromIntegral (natVal result)

{-# RULES
"fromIntegral/Int->IntMod" fromIntegral = fromIntegral_Int64_IntMod . (fromIntegral :: Int -> Int64) :: Int -> IntMod (10^9 + 7)
"fromIntegral/Int64->IntMod" fromIntegral = fromIntegral_Int64_IntMod :: Int64 -> IntMod (10^9 + 7)
 #-}

instance U.Unboxable (IntMod m) where
  type Rep (IntMod m) = Int64

exEuclid :: (Eq a, Integral a) => a -> a -> (a, a, a)
exEuclid !f !g = loop 1 0 0 1 f g
  where loop !u0 !u1 !v0 !v1 !f 0 = (f, u0, v0)
        loop !u0 !u1 !v0 !v1 !f g =
          case divMod f g of
            (q,r) -> loop u1 (u0 - q * u1) v1 (v0 - q * v1) g r

instance KnownNat m => Fractional (IntMod m) where
  recip t@(IntMod x) = IntMod $ case exEuclid x modulus of
    (1,a,_)  -> a `mod` modulus
    (-1,a,_) -> (-a) `mod` modulus
    _        -> error "not invertible"
    where modulus = fromIntegral (natVal t)
  fromRational = undefined

recipM :: (Eq a, Integral a, Show a) => a -> a -> a
recipM !x modulo = case exEuclid x modulo of
                     (1,a,_) -> a `mod` modulo
                     (-1,a,_) -> (-a) `mod` modulo
                     (g,a,b) -> error $ show x ++ "^(-1) mod " ++ show modulo ++ " failed: gcd=" ++ show g

-- |
-- >>> crt 3 6 2 7
-- 9
-- >>> crt 2 5 3 9
-- 12
crt :: Int64 -> Int64 -> Int64 -> Int64 -> Int64
crt !a1 !m1 !a2 !m2 = let m1' = recipM m1 m2
                          m2' = recipM m2 m1
                      in (m2 * m2' * a1 + m1 * m1' * a2) `mod` (m1 * m2)


{-|
Module      : Approx
Description : Implement Approx for Double, Floats and structures
Copyright   : (c) 2020 Kishaloy Neogi
License     : MIT
Maintainer  : Kishaloy Neogi
Email       : nkishaloy@yahoo.com

The library is created to allow for a easy-to-use reasonable way of emulating approx in Haskell. The codes are all in /pure/ Haskell. The idea is to have a natural mathematical feel in writing code, with operators which just works when working with Double and Float and their composite types like lists, vectors etc.

The __Approx__ module defines 2 operators __@=~@__ and __@/~@__, which are for checking /nearly equal to/ and /not nearly equal to/ respectively.

Both the operators __=~__ and __/~__ are put under the class __Approx__.

At least one of the operators have to be defined and the other gets automatically defined.

The library already defines the functions for some of the basic / common types.

For types where __Eq__ is defined like __Char, Bool, Int, Day, Text__ the approx is simply replaced with __==__.

For __Float__ and __Double__, the following formula is used,

@
if max ( |x|, |y| ) < epsilon_Zero
then True
else
  if |x - y| / max ( |x|, |y| ) < epsilon_Eq
  then True
  else False
@

The motivation for defining Approx for classes for which Eq is also defined is to allow for composite types where both Eq and Approx would be present. For example, the following code evaluates to __@True@__, even though the tuple is of type @(Int,Double,Text,[Int],[[Char]],[Double])@.

@
((2,5.35,"happ",[1,2],["ret","we"],[6.78,3.5])
    :: (Int,Double,Text,[Int],[[Char]],[Double]))

    =~ (2,5.35,"happ",[1,2],["ret","we"],[6.78,3.5])
@

For UTCTime, the approx operator checks for equality to the nearest minute. The following expression evaluates to __@True@__.

@
(parseTimeM True defaultTimeLocale "%Y-%m-%d %H:%M:%S" "2020-01-15 15:02:15"
    :: Maybe UTCTime)

    =~ parseTimeM True defaultTimeLocale "%Y-%m-%d %H:%M:%S" "2020-01-15 15:01:50"
@

The library also provides approx for Complex and common structures like __List, Boxed and Unboxed Vector, Hashmap, Tuples__ and __Maybe__. For all lists, tuples, hashmaps and vectors, the approximation is checked right down to the elements and the order for lists and vectors are important.

For lists, only finite lists are supported. Any use of infinite lists would cause a runtime error.

There are addtional functions __inRange__, __safeInRange__ and __inTol__, which checks for values within Ranges either /explictily/ defined as in __inRange__ and __safeInRange__ or through tolerances as in __inTol__.

You may see the github repository at <https://github.com/n-kishaloy/approx>

-}

{-# LANGUAGE Strict #-}

module Data.Approx
(
-- *How to use this library
-- |Add @approx@ to build-depends and @import Data.Approx@

-- *Documentation
  Approx (..)
, inRange, safeInRange, inTol

) where

import Data.List (foldl')
import qualified Data.Vector.Unboxed as U
import qualified Data.Vector.Unboxed.Mutable as M
import qualified Data.Vector as V
import Data.Time (Day, UTCTime)
import Data.Time.Clock (diffUTCTime)
import qualified Data.HashMap.Strict as Hm
import Data.Hashable

import Data.Text (Text)
import qualified Data.Complex as Cx
import Data.Complex (Complex ( (:+) ) )

eZerD :: Double
eZerD = 1e-8;   {-# INLINE eZerD  #-}

eEqD :: Double
eEqD   = 1e-7;  {-# INLINE eEqD    #-}

eZerF :: Float
eZerF = 1e-6;   {-# INLINE eZerF   #-}

eEqF :: Float
eEqF   = 1e-5;  {-# INLINE eEqF     #-}

infix 4 =~, /~

-- |The class @Approx@ defines 2 operators __@=~@__ and __@/~@__, which are for checking /nearly equal to/ and /not nearly equal to/ respectively.
class Approx a where
  (=~), (/~) :: a -> a -> Bool

  (=~) x y = not (x /~ y)
  {-# INLINE (=~) #-}
  (/~) x y = not (x =~ y)
  {-# INLINE (/~) #-}

  {-# MINIMAL (=~) | (/~) #-}

instance Approx Day where x =~ y = x == y; {-# INLINE (=~) #-}

instance Approx Char where x =~ y = x == y; {-# INLINE (=~) #-}

instance Approx Bool where x =~ y = x == y; {-# INLINE (=~) #-}

instance Approx Text where x =~ y = x == y; {-# INLINE (=~) #-}

instance Approx Int where x =~ y = x == y; {-# INLINE (=~) #-}

instance Approx Integer where x =~ y = x == y; {-# INLINE (=~) #-}

instance Approx UTCTime where
  x =~ y = (round . (/60.0) . realToFrac $ x `diffUTCTime` y) == 0
  {-# INLINE (=~) #-}

instance Approx a => Approx (Cx.Complex a) where
  (a :+ b) =~ (x :+ y) = (a =~ x) && (b =~ y); {-# INLINE (=~) #-}

instance Approx Float where
  x =~ y = (mx < eZerF) || abs (x-y) / mx < eEqF where mx = max (abs x) (abs y)
  {-# INLINE (=~) #-}

instance Approx Double where
  x =~ y = (mx < eZerD) || abs (x-y) / mx < eEqD where mx = max (abs x) (abs y)
  {-# INLINE (=~) #-}

instance Approx a => Approx (Maybe a) where
  Nothing =~ Nothing  =   True
  Just x  =~ Just y   =   x =~ y
  _       =~ _        =   False
  {-# INLINE (=~) #-}

instance Approx a => Approx [a] where
  x =~ y = (length x == length y) && and (zipWith (=~) x y)

instance (Approx a, Approx b) => Approx (a, b) where
  (x,y) =~ (a,b) =
    (   (x =~ a)
    &&  (y =~ b)
    )
  {-# INLINE (=~) #-}

instance (Approx a, Approx b, Approx c) => Approx (a, b, c) where
  (x,y,z) =~ (a,b,c) =
    (   (x =~ a)
    &&  (y =~ b)
    &&  (z =~ c)
    )
  {-# INLINE (=~) #-}

instance (Approx a,Approx b,Approx c,Approx d) => Approx (a,b,c,d) where
  (x,y,z,u) =~ (a,b,c,d) =
    (   (x =~ a)
    &&  (y =~ b)
    &&  (z =~ c)
    &&  (u =~ d)
    )
  {-# INLINE (=~) #-}

instance (Approx a,Approx b,Approx c,Approx d, Approx e) => Approx (a,b,c,d,e) where
  (x,y,z,u,v) =~ (a,b,c,d,e) =
    (   (x =~ a)
    &&  (y =~ b)
    &&  (z =~ c)
    &&  (u =~ d)
    &&  (v =~ e)
    )
  {-# INLINE (=~) #-}

instance (Approx a,Approx b,Approx c,Approx d, Approx e, Approx f) => Approx (a,b,c,d,e,f) where
  (x,y,z,u,v,w) =~ (a,b,c,d,e,f) =
    (   (x =~ a)
    &&  (y =~ b)
    &&  (z =~ c)
    &&  (u =~ d)
    &&  (v =~ e)
    &&  (w =~ f)
    )
  {-# INLINE (=~) #-}

instance (Approx a,Approx b,Approx c,Approx d, Approx e, Approx f, Approx g) => Approx (a,b,c,d,e,f,g) where
  (x,y,z,u,v,w,p) =~ (a,b,c,d,e,f,g) =
    (   (x =~ a)
    &&  (y =~ b)
    &&  (z =~ c)
    &&  (u =~ d)
    &&  (v =~ e)
    &&  (w =~ f)
    &&  (p =~ g)
    )
  {-# INLINE (=~) #-}

instance (Approx a,Approx b,Approx c,Approx d, Approx e, Approx f, Approx g, Approx h) => Approx (a,b,c,d,e,f,g,h) where
  (x,y,z,u,v,w,p,q) =~ (a,b,c,d,e,f,g,h) =
    (   (x =~ a)
    &&  (y =~ b)
    &&  (z =~ c)
    &&  (u =~ d)
    &&  (v =~ e)
    &&  (w =~ f)
    &&  (p =~ g)
    &&  (q =~ h)
    )
  {-# INLINE (=~) #-}

instance (Approx a,Approx b,Approx c,Approx d, Approx e, Approx f, Approx g, Approx h, Approx i) => Approx (a,b,c,d,e,f,g,h,i) where
  (x,y,z,u,v,w,p,q,r) =~ (a,b,c,d,e,f,g,h,i) =
    (   (x =~ a)
    &&  (y =~ b)
    &&  (z =~ c)
    &&  (u =~ d)
    &&  (v =~ e)
    &&  (w =~ f)
    &&  (p =~ g)
    &&  (q =~ h)
    &&  (r =~ i)
    )
  {-# INLINE (=~) #-}

instance (Approx a,Approx b,Approx c,Approx d, Approx e, Approx f, Approx g, Approx h, Approx i, Approx j) => Approx (a,b,c,d,e,f,g,h,i,j) where
  (x,y,z,u,v,w,p,q,r,s) =~ (a,b,c,d,e,f,g,h,i,j) =
    (   (x =~ a)
    &&  (y =~ b)
    &&  (z =~ c)
    &&  (u =~ d)
    &&  (v =~ e)
    &&  (w =~ f)
    &&  (p =~ g)
    &&  (q =~ h)
    &&  (r =~ i)
    &&  (s =~ j)
    )
  {-# INLINE (=~) #-}

instance (Approx a,Approx b,Approx c,Approx d, Approx e, Approx f, Approx g, Approx h, Approx i, Approx j, Approx k) => Approx (a,b,c,d,e,f,g,h,i,j,k) where
  (x,y,z,u,v,w,p,q,r,s,t) =~ (a,b,c,d,e,f,g,h,i,j,k) =
    (   (x =~ a)
    &&  (y =~ b)
    &&  (z =~ c)
    &&  (u =~ d)
    &&  (v =~ e)
    &&  (w =~ f)
    &&  (p =~ g)
    &&  (q =~ h)
    &&  (r =~ i)
    &&  (s =~ j)
    &&  (t =~ k)
    )
  {-# INLINE (=~) #-}

instance (M.Unbox a, Approx a) => Approx (U.Vector a) where
  x =~ y = (U.length x == U.length y) && U.and (U.zipWith (=~) x y)

instance (Approx a) => Approx (V.Vector a) where
  x =~ y = (V.length x == V.length y) && V.and (V.zipWith (=~) x y)

instance (Eq a, Hashable a, Approx b) => Approx (Hm.HashMap a b) where
  x =~ y = fz x y && fz y x where
    fz p q = and $ (\(k,v) -> Hm.lookup k p =~ Just v) <$> Hm.toList q

infix 4 `inRange`, `safeInRange`, `inTol`

{-|@inRange (u,v) x = check if x is inside the range (u,v)@

Note: The function __assumes @u < v@__. This is done to ensure speed of operations. Use safeInRange, otherwise.
-}
inRange :: Ord a => (a,a) -> a -> Bool
inRange (u,v) x = (u <= x) && (x <= v)
{-# INLINE inRange #-}

{-|@safeInRange (u,v) x = check if x is inside the range (u,v)@

Note: The function works even if u>v. However, it has addtional checks and is more expensive. Use only if you are not sure that u < v for your use-case.
-}
safeInRange :: Ord a => (a,a) -> a -> Bool
safeInRange (u,v) = inRange $ if u<v then (u,v) else (v,u)
{-# INLINE safeInRange #-}

{-|@inTol t p x = inRange (p - t, p + t) x@

The Function checks if __@x@__ is close to __@p@__ within a tolerance band __@t@__. Please ensure __@t@__ is /positive/ or there would be /incorrect/ results.
-}
inTol :: (Num a, Ord a) => a -> a -> a -> Bool
inTol t a = inRange (a - t, a + t)
{-# INLINE inTol #-}


{-# LANGUAGE DataKinds             #-}
{-# LANGUAGE FlexibleContexts      #-}
{-# LANGUAGE FlexibleInstances     #-}
{-# LANGUAGE MagicHash             #-}
{-# LANGUAGE MultiParamTypeClasses #-}
{-# LANGUAGE PolyKinds             #-}
{-# LANGUAGE RecordWildCards       #-}
{-# LANGUAGE ScopedTypeVariables   #-}
{-# LANGUAGE StandaloneDeriving    #-}
{-# LANGUAGE TypeApplications      #-}
{-# LANGUAGE TypeFamilies          #-}
{-# LANGUAGE TypeOperators         #-}
{-# LANGUAGE UnboxedTuples         #-}
{-# LANGUAGE UndecidableInstances  #-}

module Numeric.Matrix.Bidiagonal
  ( BiDiag (..), biDiag, bidiagonalHouseholder
  ) where

import Control.Monad
import Control.Monad.ST
import Data.Kind
import Numeric.Basics
import Numeric.DataFrame.ST
import Numeric.DataFrame.SubSpace
import Numeric.DataFrame.Type
import Numeric.Dimensions
import Numeric.Matrix.Internal
import Numeric.Scalar.Internal
import Numeric.Subroutine.Householder
import Numeric.Vector.Internal

-- | Put two vectors on the main and first upper diagonal.
biDiag :: forall (t :: Type) (n :: Nat) (m :: Nat)
        . (PrimBytes t, Num t)
       => Dims '[n,m]
       -> Vector t (Min n m)
       -> Vector t (Min n m)
       -> Matrix t n m
biDiag (dn@D :* dm@D :* U) a b = runST $ do
    dnm@D <- pure $ minDim dn dm
    rPtr <- thawDataFrame 0
    forM_ [0 .. dimVal dnm - 1] $ \i -> do
      writeDataFrame rPtr (Idx i :* Idx i :* U) $ a ! i
      when (i+1 < dimVal dm) $
        writeDataFrame rPtr (Idx i :* Idx (i+1) :* U) $ b ! i
    unsafeFreezeDataFrame rPtr

-- | Decomposition of a matrix \( A = U B V^\intercal \) such that
--   \( U \) and \( V \) are orthogonal and \( B \) is bidiagonal.
data BiDiag (t :: Type) (n :: Nat)  (m :: Nat)
  = BiDiag
  { bdU     :: Matrix t n n
    -- ^ \( U \) left orthogonal matrix
  , bdUDet  :: Scalar t
    -- ^ A shortcut for evaluating a determinant of \( |U| = \pm 1 \)
  , bdAlpha :: Vector t (Min n m)
    -- ^ Main diagonal of \( B \)
  , bdBeta  :: Vector t (Min n m)
    -- ^ First upper diagonal of \( B \);
    --   its last element equals zero if \( n \geq m \)
  , bdV     :: Matrix t m m
    -- ^ \( B \) left orthogonal matrix
  , bdVDet  :: Scalar t
    -- ^ A shortcut for evaluating a determinant of \( |V| = \pm 1 \)
  }

deriving instance ( Show t, PrimBytes t
                  , KnownDim n, KnownDim m, KnownDim (Min n m))
                  => Show (BiDiag t n m)
deriving instance ( Eq t, PrimBytes t
                  , KnownDim n, KnownDim m, KnownDim (Min n m)
                  , KnownBackend t '[Min n m])
                  => Eq (BiDiag t n m)

{- |
Decompose a matrix \( A = U B V^\intercal \) such that
\( U \) and \( V \) are orthogonal and \( B \) is bidiagonal.

The first returned number
 -}
bidiagonalHouseholder ::
       forall (t :: Type) (n :: Nat) (m :: Nat)
     . (PrimBytes t, Ord t, Epsilon t, KnownDim n, KnownDim m)
    => Matrix t n m
    -> BiDiag t n m
bidiagonalHouseholder a = runST $ do
      D <- pure $ minDim (dim @n) (dim @m)
      tmpNPtr <- newDataFrame
      tmpMPtr <- newDataFrame
      uPtr <- thawDataFrame eye
      bPtr <- thawDataFrame a
      vPtr <- thawDataFrame eye
      (ud, vd) <-
        let f (ud, vd) i = do
              ud' <- householderReflectionInplaceL tmpNPtr uPtr bPtr
                     (Idx (i - 1) :* Idx (i - 1) :* U)
              vd' <- householderReflectionInplaceR tmpMPtr vPtr bPtr
                     (Idx (i - 1) :* Idx i :* U)
              return (ud /= ud', vd /= vd')
        in foldM f (False, False) [1 .. lim - 1]

      udn <- householderReflectionInplaceL tmpNPtr uPtr bPtr
                     (Idx (lim - 1) :* Idx (lim - 1) :* U)
      vdn <- if (m > lim)
             then householderReflectionInplaceR tmpMPtr vPtr bPtr
                     (Idx (lim - 1) :* Idx lim :* U)
             else pure False
      bdU <- unsafeFreezeDataFrame uPtr
      bdV <- unsafeFreezeDataFrame vPtr
      b <- unsafeFreezeDataFrame bPtr
      let bdAlpha = iwgen @t @'[Min n m]
            (\(Idx i :* U) -> index (Idx i :* Idx i :* U) b)
          bdBeta  = iwgen @t @'[Min n m]
            (\(Idx i :* U) -> if i+1 < m then index (Idx i :* Idx (i+1) :* U) b else 0)
          bdUDet = if ud /= udn then -1 else 1
          bdVDet = if vd /= vdn then -1 else 1
      return BiDiag {..}
  where
    n = dimVal' @n
    m = dimVal' @m
    lim = max 1 (min n m)


-----------------------------------------------------------------------------
-- |
-- Module      :  Numeric.Statistics.Histogram
-- Copyright   :  (c) A. V. H. McPhail 2010, 2014
-- License     :  BSD3
--
-- Maintainer  :  haskell.vivian.mcphail <at> gmail <dot> com
-- Stability   :  provisional
-- Portability :  portable
--
-- create histograms from density functions
--
-----------------------------------------------------------------------------

module Numeric.Statistics.Histogram (
                                     cumulativeToHistogram
                                    , gaussianHistogram
                                  ) where

-----------------------------------------------------------------------------

--import qualified Data.Packed.Vector as V
import qualified Data.Vector.Storable as V

import qualified Numeric.GSL.Histogram as H
--import qualified Numeric.GSL.Histogram2D as H2

import qualified Numeric.GSL.Distribution.Continuous as C

--import Numeric.LinearAlgebra.Algorithms
--import Numeric.LinearAlgebra.Interface()

-----------------------------------------------------------------------------

vectorToTuples = toTuples . V.toList
    where toTuples []         = error "need a minimum of two elements"
          toTuples [_]        = error "need a minimum of two elements"
          toTuples [x1,x2]    = [(x1,x2)]
          toTuples (x1:x2:xs) = (x1,x2) : (toTuples (x2:xs))

-----------------------------------------------------------------------------

cumulativeToHistogram :: (Double -> Double)   -- ^ the cumulative distribution function D(x <= X)
                   -> V.Vector Double        -- ^ the bins
                   -> H.Histogram            -- ^ the resulting histogram
cumulativeToHistogram f v = H.addListWeighted (H.emptyRanges v) $ map (\(x1,x2) -> ((x1 + x2) / 2.0,f x2 - f x1)) (vectorToTuples v)

gaussianHistogram :: Double               -- ^ mean
                  -> Double              -- ^ standard deviation
                  -> V.Vector Double     -- ^ the bins
                  -> H.Histogram         -- ^ the resulting histogram
gaussianHistogram u s = cumulativeToHistogram (\x -> C.density_1p C.Gaussian C.Lower s (x-u))

-----------------------------------------------------------------------------


{-# LANGUAGE TypeSynonymInstances #-}
{-# LANGUAGE FlexibleInstances #-}
{-# LANGUAGE ConstraintKinds #-}
module DSLsofMath.Algebra where
import qualified Data.Ratio
import qualified Prelude
import Prelude (Double, Rational, Int, Integer, Bool(..), otherwise,
                Foldable(foldr), (.), const, (==), (<), error)
import Data.Complex

-------------------------------
-- Classes
infixl 6 -
infixl 6 +

infixl 7 *
infixl 7 /

class Additive a where
  zero :: a
  (+) :: a -> a -> a

sum :: (Foldable t, Additive a) => t a -> a
sum = foldr (+) zero

times :: Additive a => Integer -> a -> a
times n0 = if n0 < 0  then  error "Algebra.Classes.times: negative number of times"
                      else  go n0
  where go 0 _ = zero
        go 1 x = x
        go n x = if r == 0 then twoy else x + twoy
          where (m,r) = n `Prelude.divMod` 2
                y = go m x
                twoy = y+y

(-) :: AddGroup a => a -> a -> a
x - y           =  x + negate y

class Additive a => AddGroup a where
  negate :: a -> a

mult :: AddGroup a => Integer -> a -> a
mult n x = if n < 0 then negate (times (negate n) x) else times n x

class Multiplicative a where
  one :: a
  (*) :: a -> a -> a

two :: (Additive a, Multiplicative a) => a
two = one+one

(^+) :: Multiplicative a => a -> Int -> a

x0 ^+ n0 = if n0 < 0  then  error "Algebra.Classes.^: negative exponent"
                      else  go n0 x0
  where go 0 _ = one
        go 1 x = x
        go n x = if r == 0 then y2 else x * y2
          where  (m,r) = n `Prelude.divMod` 2
                 y = go m x
                 y2 = y * y

type Ring a = (AddGroup a, Multiplicative a)

fromInteger :: Ring a => Integer -> a
fromInteger n = mult n one

class Multiplicative a => MulGroup a where
  {-# MINIMAL (recip | (/)) #-}
  recip :: a -> a
  recip x         =  one / x

  (/) :: a -> a -> a
  x / y           =  x * recip y

(^) :: MulGroup a => a -> Int -> a
a ^ b | b < 0      = recip (a ^+ (negate b))
      | otherwise  = (a ^+ b)

type Field a = (Ring a, MulGroup a)

fromRational :: Field a => Data.Ratio.Ratio Integer -> a
fromRational x  =  fromInteger (Data.Ratio.numerator x) / fromInteger (Data.Ratio.denominator x)

class Field a => Algebraic a where
  sqrt :: a -> a

-- normally it should be "Algebraic" instead of "Field" but we're lazy like that.
--   (Also Transcendental is a terrible name; taken from the "numeric prelude".)
class Field a => Transcendental a where
  pi  :: a
  exp :: a -> a
  sin :: a -> a
  cos :: a -> a

cosh, sinh :: Transcendental a => a -> a
cosh x = (exp x + exp (negate x))/two
sinh x = (exp x - exp (negate x))/two

---------------------------------
-- Instances

instance Additive        Int       where  (+)  = (Prelude.+);  zero  = 0
instance Additive        Integer   where  (+)  = (Prelude.+);  zero  = 0
instance Additive        Rational  where  (+)  = (Prelude.+);  zero  = 0
instance Additive        Double    where  (+)  = (Prelude.+);  zero  = 0

instance AddGroup        Int       where  negate = Prelude.negate
instance AddGroup        Integer   where  negate = Prelude.negate
instance AddGroup        Rational  where  negate = Prelude.negate
instance AddGroup        Double    where  negate = Prelude.negate

instance Multiplicative  Int       where  (*)  = (Prelude.*);  one   = 1
instance Multiplicative  Integer   where  (*)  = (Prelude.*);  one   = 1
instance Multiplicative  Rational  where  (*)  = (Prelude.*);  one   = 1
instance Multiplicative  Double    where  (*)  = (Prelude.*);  one   = 1

instance MulGroup        Rational  where  (/)  = (Prelude./);  recip = Prelude.recip
instance MulGroup        Double    where  (/)  = (Prelude./);  recip = Prelude.recip

lift0 ::  a         ->  (x->a)
lift1 :: (a->b)     ->  (x->a) -> (x->b)
lift2 :: (a->b->c)  ->  (x->a) -> (x->b) -> (x->c)
lift0 = const
lift1 = (.)
lift2 op2 f g = \x -> op2 (f x) (g x)

instance Additive a        => Additive        (x -> a) where (+)  = lift2 (+);  zero  = lift0 zero
instance Multiplicative a  => Multiplicative  (x -> a) where (*)  = lift2 (*);  one   = lift0 one

instance AddGroup a   => AddGroup   (x -> a)  where   negate  =  lift1 negate
instance MulGroup a   => MulGroup   (x -> a)  where   recip   =  lift1 recip
instance Algebraic a  => Algebraic  (x -> a)  where   sqrt    =  lift1 sqrt

instance Transcendental a => Transcendental (x -> a) where
   pi = lift0 pi;  sin = lift1 sin;  cos = lift1 cos;  exp = lift1 exp

instance Algebraic       Double where   sqrt = Prelude.sqrt
instance Transcendental  Double where
   pi = Prelude.pi;  sin = Prelude.sin;  cos = Prelude.cos;  exp = Prelude.exp

instance Additive a  => Additive        (Complex a) where  (+) = addC;  zero  = zeroC
instance Ring a      => Multiplicative  (Complex a) where  (*) = mulC;  one   = oneC
instance AddGroup a  => AddGroup        (Complex a) where  negate  = negateC
instance Field a     => MulGroup        (Complex a) where  recip   = recipC

addC :: Additive a => Complex a -> Complex a -> Complex a
addC (x :+ y) (x' :+ y') = (x + x') :+ (y+y')

negateC :: AddGroup a => Complex a -> Complex a
negateC (a :+ b) = negate a :+ negate b

mulC :: Ring a => Complex a -> Complex a -> Complex a
mulC (a :+ b) (a' :+ b') = (a * a' - b * b') :+ (a * b' + b * a')

toC :: Additive a => a -> Complex a
toC x = x :+ zero

zeroC :: Additive a => Complex a
zeroC = toC zero

oneC :: (Additive a, Multiplicative a) => Complex a
oneC = toC one

recipC (a :+ b) = (a / m)  :+ (negate b / m)
  where m = a*a + b*b

instance (Algebraic a, Prelude.RealFloat a) =>  Algebraic (Complex a) where
  sqrt = Prelude.sqrt

instance (Transcendental a) => Transcendental (Complex a) where
  pi = piC;  exp = expC;  sin = sinC;  cos = cosC

piC :: Transcendental a => Complex a
piC = toC pi

expC, sinC, cosC, sinhC, coshC :: Transcendental a => Complex a -> Complex a
expC   (x:+y)  =  expx  * cos y  :+ expx * sin y  where expx = exp x
sinC   (x:+y)  =  sin x * cosh y :+ cos x * sinh y
cosC   (x:+y)  =  cos x * cosh y :+ negate (sin x * sinh y)

sinhC  (x:+y)  =  cos y * sinh x :+ sin y * cosh x
coshC  (x:+y)  =  cos y * cosh x :+ sin y * sinh x

---------------------------------
-- For RebindableSyntax

ifThenElse :: Bool -> p -> p -> p
ifThenElse c a b = if c then a else b

-------------------------------
-- Typesetting aliases.

neg :: AddGroup a => a -> a
neg = negate   -- |neg| is used to typeset unary minus as a shorter dash, closer to its argument

frac :: MulGroup a => a -> a -> a
frac = (/)     -- |frac| is used to typeset a fraction (more compactly than |x / y|)