{-# LANGUAGE FlexibleInstances, FlexibleContexts #-}
module Test.Gen where
import Test.Utils
import Test.QuickCheck
import Control.Monad
import qualified Numeric.LinearAlgebra as L
import Numeric.LinearAlgebra.Devel
import qualified Data.Vector.Storable as V

squared_real_matrices :: Int -> Gen (L.Matrix Float)
squared_real_matrices k = do
  vs <- sequence (replicate (k*k) arbitrary) :: Gen [Float]
  return $ L.reshape k $ V.fromList vs

small_matrices :: Gen (L.Matrix Float)
small_matrices = do
  n <- choose (2,10)
  squared_real_matrices n

pair :: Gen a -> Gen b -> Gen (a,b)
pair = liftM2 (,)


module LibSpec where

import           Debug.Trace
import           Lib                        (sigmoid)
import           Numeric.LinearAlgebra      (( #> ))
import           Numeric.LinearAlgebra.Data (Matrix, Vector, cond, tr, vector,
                                             (><))
import           Test.Hspec                 hiding (Spec)
import           Test.Hspec                 as Hspec
import           Test.HUnit                 (Assertion, assertBool)

spec :: Hspec.Spec
spec = do
  describe "sigmoid" $
    it "Should be correct." $ do
      (sigmoid ((3 >< 2) [1.0 ..]) :: Matrix Double) `shouldBe`
        (((3 >< 2)
            [ 0.7310585786300049
            , 0.8807970779778823
            , 0.9525741268224334
            , 0.9820137900379085
            , 0.9933071490757153
            , 0.9975273768433653
            ]) :: Matrix Double)
  describe "cond behaviour" $
    it "Should be correct." $ do
      (traceShowId (cond ((3 >< 2) [1.0 ..] :: Matrix Double) 3 0 1 0)) `shouldBe`
        ((3 >< 2) [0, 0, 1, 0, 0, 0] :: Matrix Double)
  describe "m.v behaviour" $
    it "Should be correct." $ do
      let x = (3 >< 2) [1.0 ..] :: Matrix Double
          theta = vector [10.0, 5.0] :: Vector Double
      (x #> theta) `shouldBe` (vector [20.0, 50.0, 80.0] :: Vector Double)


{-# 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 #-}


module Main where

import Criterion.Main

import Control.Monad
import Data.IORef
import Data.Random
import Data.Random.Distribution.Uniform
import Data.Vector (Vector, fromList)
import System.Random.MWC (GenIO, uniformVector, createSystemRandom, create)
import qualified System.Random.MWC as MWC
import Statistics.Sample (mean)

import Data.Random.Source.PureMT (newPureMT)
import Data.Random.Source.DevRandom

k = 10000

randomListFu :: RVar (Vector Double)
randomListFu = fromList <$> replicateM k stdUniform

randomListMWC :: GenIO -> IO (Vector Double)
randomListMWC gen = fromList <$> replicateM k (MWC.uniform gen)

main :: IO ()
main = do
  mwc <- createSystemRandom
  puremt <- newPureMT >>= newIORef

  defaultMain
    [ bgroup "MWC vs randomu-fu abstraction"
      [ bench "random-fu"
        $ nfIO (mean <$> sampleFrom mwc randomListFu)
      , bench "mw'"
        $ nfIO (mean <$> randomListMWC mwc)
      ]
    , bgroup "Different sources"
      [ bench "pureMT"
        $ nfIO (mean <$> sampleFrom puremt randomListFu)
      , bench "std"
        $ nfIO (mean <$> sample randomListFu)
      , bench "urandom"
        $ nfIO (mean <$> sampleFrom DevURandom randomListFu)
      ]
    ]


{-# 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)


{-# LANGUAGE AllowAmbiguousTypes #-}
{-# LANGUAGE InstanceSigs #-}
{-# LANGUAGE MultiParamTypeClasses #-}
{-# LANGUAGE TypeFamilies #-}

module Interpreter where

import Data.Array as A
import Data.Complex as DC
import qualified Data.IntMap.Strict as IM
import Data.Map (Map, fromList)
import qualified Data.Map as Map
import Expression
import Utils

type family Output d rc where
    Output Scalar R = Double
    Output Scalar C = DC.Complex Double
    Output One R = Array Int Double
    Output One C = Array Int (DC.Complex Double)

-- |
--
data ValMaps =
    ValMaps
        { vm0 :: Map String Double
        , vm1 :: Map String (Array Int Double)
        }
    deriving (Eq, Show, Ord)

subs :: [(String, Double)] -> [(String, Array Int Double)] -> ValMaps
subs vm0 vm1 = ValMaps (fromList vm0) (fromList vm1)

-- |
--
class Evaluable d rc where
    eval :: ValMaps -> Expression d rc -> Output d rc

-- |
--
instance Evaluable Scalar R where
    eval :: ValMaps -> Expression Scalar R -> Double
    eval valMap e@(Expression n mp) =
        case IM.lookup n mp of
            Just ([], Var name) ->
                case Map.lookup name $ vm0 valMap of
                    Just val -> val
                    _ -> error "no value associated with the variable"
            Just ([], Sum Real [node1, node2]) ->
                let subExp1 = Expression node1 mp :: Expression Scalar R
                    subExp2 = Expression node2 mp :: Expression Scalar R
                 in eval valMap subExp1 + eval valMap subExp2
            Just ([], Mul Real [node1, node2]) ->
                let subExp1 = Expression node1 mp :: Expression Scalar R
                    subExp2 = Expression node2 mp :: Expression Scalar R
                 in eval valMap subExp1 * eval valMap subExp2
            Just ([], Scale Real [node1, node2]) ->
                let subExp1 = Expression node1 mp :: Expression Scalar R
                    subExp2 = Expression node2 mp :: Expression Scalar R
                 in eval valMap subExp1 * eval valMap subExp2
            Just ([], InnerProd Real [node1, node2]) ->
                case IM.lookup node1 mp of
                    Just ([], _) ->
                        let subExp1 = Expression node1 mp :: Expression Scalar R -- shape is [], so must be Scalar R
                            subExp2 = Expression node2 mp :: Expression Scalar R -- shape is [], so must be Scalar R
                         in eval valMap subExp1 * eval valMap subExp2
                    Just ([size], _) ->
                        let subExp1 = Expression node1 mp :: Expression One R -- shape is [size], so must be One R
                            subExp2 = Expression node2 mp :: Expression One R -- shape is [size], so must be One R
                            lst1 = A.elems $ eval valMap subExp1
                            lst2 = A.elems $ eval valMap subExp2
                         in sum $ zipWith (*) lst1 lst2
            _ -> error "expression structure Scalar R is wrong"

-- |
--
instance Evaluable Scalar C where
    eval :: ValMaps -> Expression Scalar C -> DC.Complex Double
    eval valMap e@(Expression n mp) =
        case IM.lookup n mp of
            Just ([], Sum Complex [node1, node2]) ->
                let subExp1 = Expression node1 mp :: Expression Scalar C
                    subExp2 = Expression node2 mp :: Expression Scalar C
                 in eval valMap subExp1 + eval valMap subExp2
            Just ([], Mul Complex [node1, node2]) ->
                let subExp1 = Expression node1 mp :: Expression Scalar C
                    subExp2 = Expression node2 mp :: Expression Scalar C
                 in eval valMap subExp1 * eval valMap subExp2
            Just ([], Scale Complex [node1, node2]) ->
                let subExp2 = Expression node2 mp :: Expression Scalar C
                    scale =
                        case nodeNumType . retrieveNode mp $ node1 of
                            Real ->
                                fromReal . eval valMap $
                                (Expression node1 mp :: Expression Scalar R)
                            Complex ->
                                eval
                                    valMap
                                    (Expression node1 mp :: Expression Scalar C)
                 in scale * eval valMap subExp2
            Just ([], RealImg [node1, node2]) ->
                let subExp1 = Expression node1 mp :: Expression Scalar R
                    subExp2 = Expression node2 mp :: Expression Scalar R
                 in eval valMap subExp1 :+ eval valMap subExp2
            Just ([], InnerProd Complex [node1, node2]) ->
                case IM.lookup node1 mp of
                    Just ([], _) ->
                        let subExp1 = Expression node1 mp :: Expression Scalar C -- shape is [], so must be Scalar C
                            subExp2 = Expression node2 mp :: Expression Scalar C -- shape is [], so must be Scalar C
                         in eval valMap subExp1 * eval valMap subExp2

                    Just ([size], _) ->
                        let subExp1 = Expression node1 mp :: Expression One C -- shape is [size], so must be One C
                            subExp2 = Expression node2 mp :: Expression One C -- shape is [size], so must be One C
                            lst1 = A.elems $ eval valMap subExp1
                            lst2 = A.elems $ eval valMap subExp2
                         in sum $ zipWith (*) lst1 lst2
            _ -> error "expression structure Scalar C is wrong"

-- |
--
instance Evaluable One R where
    eval :: ValMaps -> Expression One R -> Array Int Double
    eval valMap e@(Expression n mp) =
        case IM.lookup n mp of
            Just ([size], Var name) ->
                case Map.lookup name $ vm1 valMap of
                    Just val -> val
                    _ -> error "no value associated with the variable"
            Just ([size], Sum Real [node1, node2]) ->
                let subExp1 = Expression node1 mp :: Expression One R
                    subExp2 = Expression node2 mp :: Expression One R
                    lst1 = A.elems $ eval valMap subExp1
                    lst2 = A.elems $ eval valMap subExp2
                    lstRes = zipWith (+) lst1 lst2
                 in A.listArray (0, size - 1) lstRes
            Just ([size], Mul Real [node1, node2]) ->
                let subExp1 = Expression node1 mp :: Expression One R
                    subExp2 = Expression node2 mp :: Expression One R
                    lst1 = A.elems $ eval valMap subExp1
                    lst2 = A.elems $ eval valMap subExp2
                    lstRes = zipWith (*) lst1 lst2
                 in A.listArray (0, size - 1) lstRes
            Just ([size], Scale Real [node1, node2]) ->
                let subExp1 = Expression node1 mp :: Expression Scalar R
                    subExp2 = Expression node2 mp :: Expression One R
                    scale = eval valMap subExp1
                 in fmap (* scale) $ eval valMap subExp2
            _ -> error "expression structure One R is wrong"

-- |
--
instance Evaluable One C where
    eval :: ValMaps -> Expression One C -> Array Int (DC.Complex Double)
    eval valMap e@(Expression n mp) =
        case IM.lookup n mp of
            Just ([size], Sum Complex [node1, node2]) ->
                let subExp1 = Expression node1 mp :: Expression One C
                    subExp2 = Expression node2 mp :: Expression One C
                    lst1 = A.elems $ eval valMap subExp1
                    lst2 = A.elems $ eval valMap subExp2
                    lstRes = zipWith (+) lst1 lst2
                 in A.listArray (0, size - 1) lstRes
            Just ([size], Mul Complex [node1, node2]) ->
                let subExp1 = Expression node1 mp :: Expression One C
                    subExp2 = Expression node2 mp :: Expression One C
                    lst1 = A.elems $ eval valMap subExp1
                    lst2 = A.elems $ eval valMap subExp2
                    lstRes = zipWith (*) lst1 lst2
                 in A.listArray (0, size - 1) lstRes
            Just ([size], Scale Complex [node1, node2]) ->
                let subExp2 = Expression node2 mp :: Expression One C
                    lst = A.elems $ eval valMap subExp2
                    scale =
                        case nodeNumType . retrieveNode mp $ node1 of
                            Real ->
                                fromReal . eval valMap $
                                (Expression node1 mp :: Expression Scalar R)
                            Complex ->
                                eval
                                    valMap
                                    (Expression node1 mp :: Expression Scalar C)
                 in A.listArray (0, size - 1) $ map (* scale) lst
            Just ([size], RealImg [node1, node2]) ->
                let subExp1 = Expression node1 mp :: Expression One R
                    subExp2 = Expression node2 mp :: Expression One R
                    lst1 = A.elems $ eval valMap subExp1
                    lst2 = A.elems $ eval valMap subExp2
                    lstRes = zipWith (:+) lst1 lst2
                 in A.listArray (0, size - 1) lstRes
            _ -> error "expression structure One C is wrong"


module Phi.Window
  ( Window, mkWindow
  , length
  , push
  , phi
  ) where

import           Data.Vector.Unboxed            (Vector)
import qualified Data.Vector.Unboxed            as V
import           Prelude                        hiding (length)
import           Statistics.Distribution        (complCumulative)
import           Statistics.Distribution.Normal (normalFromSample)

data Window =
  Window
  { sample :: !(Vector Double)
  , size   :: !Int
  }

mkWindow :: Int -> Window
mkWindow = Window V.empty

push :: Double -> Window -> Window
push d window = window {sample = sample'}
  where sample' =
          V.take (size window)
                 (V.cons d (sample window))

phi :: Window -> Double -> Double
phi w = negate . logBase 10 . pLater w
  where pLater :: Window -> Double -> Double
        pLater = complCumulative . normalFromSample . sample

length :: Window -> Int
length = V.length . sample


{-# LANGUAGE RecordWildCards #-}
{-# LANGUAGE TypeOperators   #-}
module GenerateInParallel where

import Data.Complex
import Control.Monad
import Codec.Picture
import Data.Array.Repa
import Data.Word

import Mand

generateInParallel = do
  mapM_ (\ x-> (computeUnboxedP $ repaMandArray x (-2) (-2) (0.0025) 1600 1600) >>= (savePngImage ("./repatestcol" Prelude.++ show x Prelude.++ ".png") . ImageRGB8 . toImage)  ) [2..3]

  --img <- computeUnboxedP $ repaMandArray (-2) (-2) (0.0005) 8000 8000
  --(savePngImage "./repatest1.png" . ImageRGB8 . toImage) img

repaMandArray :: Int -> Double -> Double -> Double -> Int -> Int -> Array D DIM2 Word16
repaMandArray pow start_re start_im step width height = fromFunction (Z :. width :. height) (arrayToMand pow start_re start_im step)

arrayToMand :: Int -> Double -> Double -> Double ->  Z :. Int :. Int -> Word16
arrayToMand pow start_re start_im step (Z :. x :. y) =  fromIntegral $ general (mand_pow_iteration pow) Nothing ((start_re + (fromIntegral x) * step ) :+(start_im + step * (fromIntegral y) ) ) 32767

toImage :: Array U DIM2 Word16 -> Image PixelRGB8
toImage a = generateImage gen width height
  where
    Z :. width :. height = extent a
    gen x y =
      let d = a ! (Z :. x :. y)
      in PixelRGB8 (fromIntegral $ 2 * mod d 128) (fromIntegral $ 2* mod (div d 128) 128) ( fromIntegral $ 2 * div d 16384)


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 Strategy.RandomWalkMetropolis where

import Math.Probably.Sampler
import Math.Probably.MCMC
import Numeric.LinearAlgebra
import Control.Monad
import Data.Maybe

rwm = VStrategy rwmTrans rwmIni

rwmTrans posterior xi (sigma,(i, iaccept)) mpi = do
   let pi = fromMaybe (posterior xi) mpi
   xstar <- fmap fromList $ mapM (\x -> normal x sigma) $ toList xi
   let pstar =  posterior xstar
   let ratio = exp $ pstar - pi
   u <- unit
   let accept = u < ratio

   --diminishing adaptation:
   let sigmaNext = if accept then sigma*(min 1.4 $ 1+5.0/i)^3
                             else sigma*(max 0.7143 $ 1-5.0/i)

   return $ if accept
               then ((xstar, (sigmaNext, (i+1, iaccept+1))), Just pstar)
               else ((xi, (sigmaNext, (i+1, iaccept))), Just pi)

rwmIni _ = (0.1,(1.0, 0.0))


{-|
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. Г.С. Осипенко, Н.Б.Ампилова:
--      ВВЕДЕНИЕ В СИМВОЛИЧЕСКИЙ АНАЛИЗ ДИНАМИЧЕСКИХ СИСТЕМ


module Fourier where

import           Data.Complex
import           Data.Complex.Integrate

coefficient :: (RealFloat a, Integral b) => (a -> Complex a) -> b -> Complex a
coefficient = coefficientWithPrecision 1000

coefficientWithPrecision
  :: (RealFloat a, Integral b) => Integer -> (a -> Complex a) -> b -> Complex a
coefficientWithPrecision prec f n = integrate g prec 0 1
 where
  term t = f t * cis ((-2) * pi * fromIntegral n * t)
  -- "integrate" requires a function with a complex input, but "term" takes a real input
  g = term . realPart


{-|
Module: MachineLearning.SoftmaxClassifier
Description: Softmax Classifier.
Copyright: (c) Alexander Ignatyev, 2017-2018.
License: BSD-3
Stability: experimental
Portability: POSIX

Softmax Classifier (Multiclass Logistic Regression).
-}

module MachineLearning.SoftmaxClassifier
(
  module MachineLearning.Model
  , module MachineLearning.Classification.MultiClass
  , SoftmaxClassifier(..)
)

where

import Prelude hiding ((<>))
import qualified Numeric.LinearAlgebra as LA
import Numeric.LinearAlgebra((<>), (<.>), (|||))
import qualified Data.Vector.Storable as V

import qualified MachineLearning as ML
import MachineLearning.Types (R, Vector, Matrix)
import MachineLearning.Utils (reduceByRows, sumByRows)
import MachineLearning.Model
import MachineLearning.Classification.MultiClass

-- | Softmax Classifier, takes number of classes.
data SoftmaxClassifier = Softmax Int

instance Classifier SoftmaxClassifier where
  cscore (Softmax _) x theta = scores - reduceByRows V.maximum scores
    where scores = x <> (LA.tr theta)

  chypothesis m x theta = V.fromList predictions
    where scores = cscore m x theta
          scores' = LA.toRows scores
          predictions = map (fromIntegral . LA.maxIndex) scores'

  ccost m lambda x y theta =
    let nSamples = fromIntegral $ LA.rows x
        scores = cscore m x theta
        sum_probs = sumByRows $ exp scores
        loss = LA.sumElements $ (log sum_probs) - remap scores y
        reg = ccostReg lambda theta
    in (loss + reg) / nSamples

  cgradient m lambda x y theta =
    let nSamples = fromIntegral $ LA.rows x
        ys = processOutput m y
        scores = cscore m x theta
        sum_probs = sumByRows $ exp scores
        probs = (exp scores) / sum_probs
        probs' = probs - ys
        dw =  (LA.tr probs') <> x
        reg = cgradientReg lambda theta
    in (dw + reg)/ nSamples

  cnumClasses (Softmax nLabels) = nLabels

remap :: Matrix -> Vector -> Matrix
remap m v = LA.remap cols rows m
  where cols = LA.asColumn $ V.fromList [0..(fromIntegral $ LA.rows m)-1]
        rows = LA.toInt $ LA.asColumn v


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


module CNC.DeclarativeTests where

import CNC.Declarative
import CNC.Geometry(eps)

import Data.Complex
import Control.Monad

fig1 = do [p1, p2] <- declarePoints 2
          xSize p1 p2 1

points1 = [0 :+ 0, 1 :+ 0]

fig2 = do [p1, p2, p3] <- declarePoints 3
          len p1 p2 1
          xAngle p1 p2 (pi/2)
          xAngle p2 p3 (-pi/2)
          len p2 p3 1

points2 = [0 :+ 0, 0 :+ 1.0, 0 :+ 0]

fig3 = do
  [p1,p2,p3,p4] <- declarePoints 4
  xSize p1 p2 1
  xAngle p2 p3 ( pi / 3)
  len p2 p3 2
  len p4 p3 2
  xAngle p3 p4 (-pi / 3)
--  xSize p3 p4 2

points3 = [0, 1, 1 + mkPolar 2 (pi/3), 3] --1 + mkPolar 2 (- 2*pi/3) + 2]

tests = [ (fig1,points1),
          (fig2, points2),
          (fig3, points3)
        ]

declarative_tests = do mapM_ run' tests
  where run' (fig, points) = do
          let solved = figure fig
          print solved
          when (any ((> 0.1) . magnitude) $ zipWith (-) (unPath solved) points) $ fail $ "test failed, should be " ++ show points


{-# 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


import Data.List(transpose)
import Numeric.LinearAlgebra.Algorithms

data Matrix a = Matrix { width :: Int, height :: Int, array :: [a] } deriving (Show, Eq)

mkMatrix :: (Num a) => Int -> Int -> [a] -> Matrix a
mkMatrix w h a = Matrix w h (take (w*h) (a ++ repeat 0))

--add :: (Num a, Num b, Num c) => Matrix a -> Matrix b -> Maybe (Matrix c)
add (Matrix w h a) (Matrix w' h' a')
  | w == w' && h == h' = Just $ Matrix w h (zipWith (+) a a')
  | otherwise = Nothing

multiply (Matrix w h a) (Matrix w' h' a')
  | w == h' = undefined --normal
  | otherwise = Nothing

groupN n xs@(_:_) = (take n xs) : groupN n (drop n xs)
groupN n [] = []

transpose' :: [[a]] -> [[a]]
transpose' xs@(x:_)
  | null x = []
  | otherwise = heads : transpose' tails
    where heads = fmap head xs
          tails = fmap tail xs


module Loss where

import Numeric.LinearAlgebra

getLoss :: String -> Matrix R -> Matrix R -> Matrix R
getLoss lossType yMat tMat =
  case lossType of
    "mse" -> mse yMat tMat
    otherwise -> yMat -- Warning: Never Used Necessery!

mse :: Matrix R -> Matrix R -> Matrix R
mse yMat tMat =
  let dyMat = yMat - tMat
      r = rows yMat
      c = cols yMat
      halfMat = (r >< c) [0.5 ..]::(Matrix R) in
  halfMat * dyMat * dyMat


{-# 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


{-# LANGUAGE TypeApplications #-}
{-# LANGUAGE TypeFamilies #-}
{-# LANGUAGE FlexibleContexts #-}
module Ch3 where

import Text.Printf
import Numeric.LinearAlgebra
import qualified Data.ByteString.Char8 as BS
import qualified Data.Yaml as Yaml

import qualified MNIST as MNIST
import NeuralNetwork

type ActivationFunction = [Double] -> [Double]
type Layer = (Matrix R, Vector R, ActivationFunction)
type Network = [Layer]

-- ch3_4
network :: Network
network = [ ( (2><3) [0.1, 0.3, 0.5, 0.2, 0.4, 0.6]
            , vector [0.1, 0.2, 0.3]
            , map sigmoidFunction
            )
          , ( (3><2) [0.1, 0.4, 0.2, 0.5, 0.3, 0.6]
            , vector [0.1, 0.2]
            , map sigmoidFunction
            )
          , ( (2><2) [0.1, 0.3, 0.2, 0.4]
            , vector [0.1, 0.2]
            , id
            )
          ]

forward :: Network -> Vector R -> Vector R
forward nw x = foldl propagate x nw
  where propagate x (w, b, f) = fromList . f . toList $ x <# w + b

x = vector [1.0, 0.5]
y = forward network x

-- ch3_6
readMatrix :: FilePath -> IO (Matrix R)
readMatrix file = do
  s <- BS.readFile file
  let Just m = Yaml.decode s :: Maybe [[Double]]
  return $ fromLists m

readVector :: FilePath -> IO (Vector R)
readVector file = do
  s <- BS.readFile file
  let Just v = Yaml.decode s :: Maybe [Double]
  return $ fromList v

readNetwork :: IO Network
readNetwork = do
  w1 <- readMatrix' "W1"
  b1 <- readVector' "b1"
  w2 <- readMatrix' "W2"
  b2 <- readVector' "b2"
  w3 <- readMatrix' "W3"
  b3 <- readVector' "b3"
  return [ (w1, b1, map sigmoidFunction)
         , (w2, b2, map sigmoidFunction)
         , (w3, b3, (toList . softmax . vector))
         ]
  where readMatrix' f = readMatrix $ "data/mnist/" ++ f ++ ".yml"
        readVector' f = readVector $ "data/mnist/" ++ f ++ ".yml"

predict :: Network -> [Double] -> [Double]
predict network = toList . (forward network) . vector

displayPrediction :: [Double] -> IO ()
displayPrediction xs = do
  mapM_ print' $ zip xs [0..]
  where print' (x, i) = do
          printf "%d %5.2f%% " (i :: Int) (x * 100)
          putStrLn $ replicate (floor (x * 20)) '*'

readImages :: IO [MNIST.Image]
readImages = do
  (MNIST.Images _ _ _ imgs) <- MNIST.getImages MNIST.Test
  return imgs

run :: [MNIST.Image] -> IO ()
run imgs = do
  predictor <- predict <$> readNetwork
  mapM_ (predict' predictor) imgs
  where predict' predictor img = do
          MNIST.displayImage img
          displayPrediction $ predictor $ concat img

main :: IO ()
main = run =<< readImages


{-# 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


module Statistics.Information.Continuous.Entropy where

import Data.Matrix
import Numeric.SpecFunctions
import Statistics.Information.Utils.KNN
import Statistics.Information.Utils.Random
import System.Random

entropy :: RandomGen g => g -> Int -> Int -> Matrix Double -> Double
entropy g k base xs | k < nrows xs =
  let noised = toLists $ noisy g xs
      distances = knnDistances chebyshevDistance noised k in
    (const + (fromIntegral d) * mean (map log distances)) /
    (log $ fromIntegral base)
  where
    d = ncols xs
    n = nrows xs
    const = digamma (fromIntegral n) - digamma (fromIntegral k) +
            (fromIntegral d) * log 2
    mean xs = sum xs / (fromIntegral $ length xs)

centropy :: RandomGen g => g -> Int -> Int -> Matrix Double -> Matrix Double ->
            Double
centropy g k base xs ys = joint - hy where
  joint = entropy g1 k base (xs <|> ys)
  hy = entropy g2 k base ys
  (g1, g2) = split g


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