{-# LANGUAGE FlexibleContexts, FlexibleInstances #-}
{-# LANGUAGE CPP #-}
{-# LANGUAGE MultiParamTypeClasses #-}
{-# LANGUAGE UndecidableInstances #-}
{-# LANGUAGE TypeFamilies #-}

-----------------------------------------------------------------------------
{- |
Module      :  Numeric.LinearAlgebra.Algorithms
Copyright   :  (c) Alberto Ruiz 2006-9
License     :  GPL-style

Maintainer  :  Alberto Ruiz (aruiz at um dot es)
Stability   :  provisional
Portability :  uses ffi

High level generic interface to common matrix computations.

Specific functions for particular base types can also be explicitly
imported from "Numeric.LinearAlgebra.LAPACK".

-}
-----------------------------------------------------------------------------

module Numeric.LinearAlgebra.Algorithms (
-- * Supported types
    Field(),
-- * Linear Systems
    linearSolve,
    luSolve,
    cholSolve,
    linearSolveLS,
    linearSolveSVD,
    inv, pinv,
    det, invlndet,
    rank, rcond,
-- * Matrix factorizations
-- ** Singular value decomposition
    svd,
    fullSVD,
    thinSVD,
    compactSVD,
    singularValues,
    leftSV, rightSV,
-- ** Eigensystems
    eig, eigSH, eigSH',
    eigenvalues, eigenvaluesSH, eigenvaluesSH',
    geigSH',
-- ** QR
    qr, rq,
-- ** Cholesky
    chol, cholSH, mbCholSH,
-- ** Hessenberg
    hess,
-- ** Schur
    schur,
-- ** LU
    lu, luPacked,
-- * Matrix functions
    expm,
    sqrtm,
    matFunc,
-- * Nullspace
    nullspacePrec,
    nullVector,
    nullspaceSVD,
    orth,
-- * Norms
    Normed(..), NormType(..),
    relativeError,
-- * Misc
    eps, peps, i,
-- * Util
    haussholder,
    unpackQR, unpackHess,
    pinvTol,
    ranksv
) where

import Data.Packed.Internal hiding ((//))
import Data.Packed.Matrix
import Numeric.LinearAlgebra.LAPACK as LAPACK
import Data.List(foldl1')
import Data.Array
import Numeric.ContainerBoot

{- | Class used to define generic linear algebra computations for both real and complex matrices. Only double precision is supported in this version (we can
transform single precision objects using 'single' and 'double').

-}
class (Product t,
       Convert t,
       Container Vector t,
       Container Matrix t,
       Normed Matrix t,
       Normed Vector t) => Field t where
    svd'         :: Matrix t -> (Matrix t, Vector Double, Matrix t)
    thinSVD'     :: Matrix t -> (Matrix t, Vector Double, Matrix t)
    sv'          :: Matrix t -> Vector Double
    luPacked'    :: Matrix t -> (Matrix t, [Int])
    luSolve'     :: (Matrix t, [Int]) -> Matrix t -> Matrix t
    linearSolve' :: Matrix t -> Matrix t -> Matrix t
    cholSolve'   :: Matrix t -> Matrix t -> Matrix t
    linearSolveSVD' :: Matrix t -> Matrix t -> Matrix t
    linearSolveLS'  :: Matrix t -> Matrix t -> Matrix t
    eig'         :: Matrix t -> (Vector (Complex Double), Matrix (Complex Double))
    eigSH''      :: Matrix t -> (Vector Double, Matrix t)
    eigOnly      :: Matrix t -> Vector (Complex Double)
    eigOnlySH    :: Matrix t -> Vector Double
    cholSH'      :: Matrix t -> Matrix t
    mbCholSH'    :: Matrix t -> Maybe (Matrix t)
    qr'          :: Matrix t -> (Matrix t, Matrix t)
    hess'        :: Matrix t -> (Matrix t, Matrix t)
    schur'       :: Matrix t -> (Matrix t, Matrix t)

instance Field Double where
    svd' = svdRd
    thinSVD' = thinSVDRd
    sv' = svR
    luPacked' = luR
    luSolve' (l_u,perm) = lusR l_u perm
    linearSolve' = linearSolveR                 -- (luSolve . luPacked) ??
    cholSolve' = cholSolveR
    linearSolveLS' = linearSolveLSR
    linearSolveSVD' = linearSolveSVDR Nothing
    eig' = eigR
    eigSH'' = eigS
    eigOnly = eigOnlyR
    eigOnlySH = eigOnlyS
    cholSH' = cholS
    mbCholSH' = mbCholS
    qr' = unpackQR . qrR
    hess' = unpackHess hessR
    schur' = schurR

instance Field (Complex Double) where
#ifdef NOZGESDD
    svd' = svdC
    thinSVD' = thinSVDC
#else
    svd' = svdCd
    thinSVD' = thinSVDCd
#endif
    sv' = svC
    luPacked' = luC
    luSolve' (l_u,perm) = lusC l_u perm
    linearSolve' = linearSolveC
    cholSolve' = cholSolveC
    linearSolveLS' = linearSolveLSC
    linearSolveSVD' = linearSolveSVDC Nothing
    eig' = eigC
    eigOnly = eigOnlyC
    eigSH'' = eigH
    eigOnlySH = eigOnlyH
    cholSH' = cholH
    mbCholSH' = mbCholH
    qr' = unpackQR . qrC
    hess' = unpackHess hessC
    schur' = schurC

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

square m = rows m == cols m

vertical m = rows m >= cols m

exactHermitian m = m `equal` ctrans m

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

-- | Full singular value decomposition.
svd :: Field t => Matrix t -> (Matrix t, Vector Double, Matrix t)
svd = {-# SCC "svd" #-} svd'

-- | A version of 'svd' which returns only the @min (rows m) (cols m)@ singular vectors of @m@.
--
-- If @(u,s,v) = thinSVD m@ then @m == u \<> diag s \<> trans v@.
thinSVD :: Field t => Matrix t -> (Matrix t, Vector Double, Matrix t)
thinSVD = {-# SCC "thinSVD" #-} thinSVD'

-- | Singular values only.
singularValues :: Field t => Matrix t -> Vector Double
singularValues = {-# SCC "singularValues" #-} sv'

-- | A version of 'svd' which returns an appropriate diagonal matrix with the singular values.
--
-- If @(u,d,v) = fullSVD m@ then @m == u \<> d \<> trans v@.
fullSVD :: Field t => Matrix t -> (Matrix t, Matrix Double, Matrix t)
fullSVD m = (u,d,v) where
    (u,s,v) = svd m
    d = diagRect 0 s r c
    r = rows m
    c = cols m

-- | Similar to 'thinSVD', returning only the nonzero singular values and the corresponding singular vectors.
compactSVD :: Field t  => Matrix t -> (Matrix t, Vector Double, Matrix t)
compactSVD m = (u', subVector 0 d s, v') where
    (u,s,v) = thinSVD m
    d = rankSVD (1*eps) m s `max` 1
    u' = takeColumns d u
    v' = takeColumns d v

-- | Singular values and all right singular vectors.
rightSV :: Field t => Matrix t -> (Vector Double, Matrix t)
rightSV m | vertical m = let (_,s,v) = thinSVD m in (s,v)
          | otherwise  = let (_,s,v) = svd m     in (s,v)

-- | Singular values and all left singular vectors.
leftSV :: Field t => Matrix t -> (Matrix t, Vector Double)
leftSV m  | vertical m = let (u,s,_) = svd m     in (u,s)
          | otherwise  = let (u,s,_) = thinSVD m in (u,s)

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

-- | Obtains the LU decomposition of a matrix in a compact data structure suitable for 'luSolve'.
luPacked :: Field t => Matrix t -> (Matrix t, [Int])
luPacked = {-# SCC "luPacked" #-} luPacked'

-- | Solution of a linear system (for several right hand sides) from the precomputed LU factorization obtained by 'luPacked'.
luSolve :: Field t => (Matrix t, [Int]) -> Matrix t -> Matrix t
luSolve = {-# SCC "luSolve" #-} luSolve'

-- | Solve a linear system (for square coefficient matrix and several right-hand sides) using the LU decomposition. For underconstrained or overconstrained systems use 'linearSolveLS' or 'linearSolveSVD'.
-- It is similar to 'luSolve' . 'luPacked', but @linearSolve@ raises an error if called on a singular system.
linearSolve :: Field t => Matrix t -> Matrix t -> Matrix t
linearSolve = {-# SCC "linearSolve" #-} linearSolve'

-- | Solve a symmetric or Hermitian positive definite linear system using a precomputed Cholesky decomposition obtained by 'chol'.
cholSolve :: Field t => Matrix t -> Matrix t -> Matrix t
cholSolve = {-# SCC "cholSolve" #-} cholSolve'

-- | Minimum norm solution of a general linear least squares problem Ax=B using the SVD. Admits rank-deficient systems but it is slower than 'linearSolveLS'. The effective rank of A is determined by treating as zero those singular valures which are less than 'eps' times the largest singular value.
linearSolveSVD :: Field t => Matrix t -> Matrix t -> Matrix t
linearSolveSVD = {-# SCC "linearSolveSVD" #-} linearSolveSVD'

-- | Least squared error solution of an overconstrained linear system, or the minimum norm solution of an underconstrained system. For rank-deficient systems use 'linearSolveSVD'.
linearSolveLS :: Field t => Matrix t -> Matrix t -> Matrix t
linearSolveLS = {-# SCC "linearSolveLS" #-} linearSolveLS'

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

-- | Eigenvalues and eigenvectors of a general square matrix.
--
-- If @(s,v) = eig m@ then @m \<> v == v \<> diag s@
eig :: Field t => Matrix t -> (Vector (Complex Double), Matrix (Complex Double))
eig = {-# SCC "eig" #-} eig'

-- | Eigenvalues of a general square matrix.
eigenvalues :: Field t => Matrix t -> Vector (Complex Double)
eigenvalues = {-# SCC "eigenvalues" #-} eigOnly

-- | Similar to 'eigSH' without checking that the input matrix is hermitian or symmetric. It works with the upper triangular part.
eigSH' :: Field t => Matrix t -> (Vector Double, Matrix t)
eigSH' = {-# SCC "eigSH'" #-} eigSH''

-- | Similar to 'eigenvaluesSH' without checking that the input matrix is hermitian or symmetric. It works with the upper triangular part.
eigenvaluesSH' :: Field t => Matrix t -> Vector Double
eigenvaluesSH' = {-# SCC "eigenvaluesSH'" #-} eigOnlySH

-- | Eigenvalues and Eigenvectors of a complex hermitian or real symmetric matrix.
--
-- If @(s,v) = eigSH m@ then @m == v \<> diag s \<> ctrans v@
eigSH :: Field t => Matrix t -> (Vector Double, Matrix t)
eigSH m | exactHermitian m = eigSH' m
        | otherwise = error "eigSH requires complex hermitian or real symmetric matrix"

-- | Eigenvalues of a complex hermitian or real symmetric matrix.
eigenvaluesSH :: Field t => Matrix t -> Vector Double
eigenvaluesSH m | exactHermitian m = eigenvaluesSH' m
                | otherwise = error "eigenvaluesSH requires complex hermitian or real symmetric matrix"

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

-- | QR factorization.
--
-- If @(q,r) = qr m@ then @m == q \<> r@, where q is unitary and r is upper triangular.
qr :: Field t => Matrix t -> (Matrix t, Matrix t)
qr = {-# SCC "qr" #-} qr'

-- | RQ factorization.
--
-- If @(r,q) = rq m@ then @m == r \<> q@, where q is unitary and r is upper triangular.
rq :: Field t => Matrix t -> (Matrix t, Matrix t)
rq m =  {-# SCC "rq" #-} (r,q) where
    (q',r') = qr $ trans $ rev1 m
    r = rev2 (trans r')
    q = rev2 (trans q')
    rev1 = flipud . fliprl
    rev2 = fliprl . flipud

-- | Hessenberg factorization.
--
-- If @(p,h) = hess m@ then @m == p \<> h \<> ctrans p@, where p is unitary
-- and h is in upper Hessenberg form (it has zero entries below the first subdiagonal).
hess        :: Field t => Matrix t -> (Matrix t, Matrix t)
hess = hess'

-- | Schur factorization.
--
-- If @(u,s) = schur m@ then @m == u \<> s \<> ctrans u@, where u is unitary
-- and s is a Shur matrix. A complex Schur matrix is upper triangular. A real Schur matrix is
-- upper triangular in 2x2 blocks.
--
-- \"Anything that the Jordan decomposition can do, the Schur decomposition
-- can do better!\" (Van Loan)
schur       :: Field t => Matrix t -> (Matrix t, Matrix t)
schur = schur'

-- | Similar to 'cholSH', but instead of an error (e.g., caused by a matrix not positive definite) it returns 'Nothing'.
mbCholSH :: Field t => Matrix t -> Maybe (Matrix t)
mbCholSH = {-# SCC "mbCholSH" #-} mbCholSH'

-- | Similar to 'chol', without checking that the input matrix is hermitian or symmetric. It works with the upper triangular part.
cholSH      :: Field t => Matrix t -> Matrix t
cholSH = {-# SCC "cholSH" #-} cholSH'

-- | Cholesky factorization of a positive definite hermitian or symmetric matrix.
--
-- If @c = chol m@ then @c@ is upper triangular and @m == ctrans c \<> c@.
chol :: Field t => Matrix t ->  Matrix t
chol m | exactHermitian m = cholSH m
       | otherwise = error "chol requires positive definite complex hermitian or real symmetric matrix"

-- | Joint computation of inverse and logarithm of determinant of a square matrix.
invlndet :: (Floating t, Field t)
         => Matrix t
         -> (Matrix t, (t, t)) -- ^ (inverse, (log abs det, sign or phase of det))
invlndet m | square m = (im,(ladm,sdm))
           | otherwise = error $ "invlndet of nonsquare "++ shSize m ++ " matrix"
  where
    lp@(lup,perm) = luPacked m
    s = signlp (rows m) perm
    dg = toList $ takeDiag $ lup
    ladm = sum $ map (log.abs) dg
    sdm = s* product (map signum dg)
    im = luSolve lp (ident (rows m))

-- | Determinant of a square matrix. To avoid possible overflow or underflow use 'invlndet'.
det :: Field t => Matrix t -> t
det m | square m = {-# SCC "det" #-} s * (product $ toList $ takeDiag $ lup)
      | otherwise = error $ "det of nonsquare "++ shSize m ++ " matrix"
    where (lup,perm) = luPacked m
          s = signlp (rows m) perm

-- | Explicit LU factorization of a general matrix.
--
-- If @(l,u,p,s) = lu m@ then @m == p \<> l \<> u@, where l is lower triangular,
-- u is upper triangular, p is a permutation matrix and s is the signature of the permutation.
lu :: Field t => Matrix t -> (Matrix t, Matrix t, Matrix t, t)
lu = luFact . luPacked

-- | Inverse of a square matrix. See also 'invlndet'.
inv :: Field t => Matrix t -> Matrix t
inv m | square m = m `linearSolve` ident (rows m)
      | otherwise = error $ "inv of nonsquare "++ shSize m ++ " matrix"

-- | Pseudoinverse of a general matrix.
pinv :: Field t => Matrix t -> Matrix t
pinv m = linearSolveSVD m (ident (rows m))

-- | Numeric rank of a matrix from the SVD decomposition.
rankSVD :: Element t
        => Double   -- ^ numeric zero (e.g. 1*'eps')
        -> Matrix t -- ^ input matrix m
        -> Vector Double -- ^ 'sv' of m
        -> Int      -- ^ rank of m
rankSVD teps m s = ranksv teps (max (rows m) (cols m)) (toList s)

-- | Numeric rank of a matrix from its singular values.
ranksv ::  Double   -- ^ numeric zero (e.g. 1*'eps')
        -> Int      -- ^ maximum dimension of the matrix
        -> [Double] -- ^ singular values
        -> Int      -- ^ rank of m
ranksv teps maxdim s = k where
    g = maximum s
    tol = fromIntegral maxdim * g * teps
    s' = filter (>tol) s
    k = if g > teps then length s' else 0

-- | The machine precision of a Double: @eps = 2.22044604925031e-16@ (the value used by GNU-Octave).
eps :: Double
eps =  2.22044604925031e-16

-- | 1 + 0.5*peps == 1,  1 + 0.6*peps /= 1
peps :: RealFloat x => x
peps = x where x = 2.0 ** fromIntegral (1 - floatDigits x)

-- | The imaginary unit: @i = 0.0 :+ 1.0@
i :: Complex Double
i = 0:+1

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

-- | The nullspace of a matrix from its SVD decomposition.
nullspaceSVD :: Field t
             => Either Double Int -- ^ Left \"numeric\" zero (eg. 1*'eps'),
                                  --   or Right \"theoretical\" matrix rank.
             -> Matrix t          -- ^ input matrix m
             -> (Vector Double, Matrix t) -- ^ 'rightSV' of m
             -> [Vector t]        -- ^ list of unitary vectors spanning the nullspace
nullspaceSVD hint a (s,v) = vs where
    tol = case hint of
        Left t -> t
        _      -> eps
    k = case hint of
        Right t -> t
        _       -> rankSVD tol a s
    vs = drop k $ toRows $ ctrans v

-- | The nullspace of a matrix. See also 'nullspaceSVD'.
nullspacePrec :: Field t
              => Double     -- ^ relative tolerance in 'eps' units (e.g., use 3 to get 3*'eps')
              -> Matrix t   -- ^ input matrix
              -> [Vector t] -- ^ list of unitary vectors spanning the nullspace
nullspacePrec t m = nullspaceSVD (Left (t*eps)) m (rightSV m)

-- | The nullspace of a matrix, assumed to be one-dimensional, with machine precision.
nullVector :: Field t => Matrix t -> Vector t
nullVector = last . nullspacePrec 1

orth :: Field t => Matrix t -> [Vector t]
-- ^ Return an orthonormal basis of the range space of a matrix
orth m = take r $ toColumns u
  where
    (u,s,_) = compactSVD m
    r = ranksv eps (max (rows m) (cols m)) (toList s)

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

{-  Pseudoinverse of a real matrix with the desired tolerance, expressed as a
multiplicative factor of the default tolerance used by GNU-Octave (see 'pinv').

@\> let m = 'fromLists' [[1,0,    0]
                    ,[0,1,    0]
                    ,[0,0,1e-10]]
\  --
\> 'pinv' m
1. 0.           0.
0. 1.           0.
0. 0. 10000000000.
\  --
\> pinvTol 1E8 m
1. 0. 0.
0. 1. 0.
0. 0. 1.@

-}
--pinvTol :: Double -> Matrix Double -> Matrix Double
pinvTol t m = v' `mXm` diag s' `mXm` trans u' where
    (u,s,v) = thinSVDRd m
    sl@(g:_) = toList s
    s' = fromList . map rec $ sl
    rec x = if x < g*tol then 1 else 1/x
    tol = (fromIntegral (max r c) * g * t * eps)
    r = rows m
    c = cols m
    d = dim s
    u' = takeColumns d u
    v' = takeColumns d v

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

-- many thanks, quickcheck!

haussholder :: (Field a) => a -> Vector a -> Matrix a
haussholder tau v = ident (dim v) `sub` (tau `scale` (w `mXm` ctrans w))
    where w = asColumn v

zh k v = fromList $ replicate (k-1) 0 ++ (1:drop k xs)
              where xs = toList v

zt 0 v = v
zt k v = join [subVector 0 (dim v - k) v, konst 0 k]

unpackQR :: (Field t) => (Matrix t, Vector t) -> (Matrix t, Matrix t)
unpackQR (pq, tau) =  {-# SCC "unpackQR" #-} (q,r)
    where cs = toColumns pq
          m = rows pq
          n = cols pq
          mn = min m n
          r = fromColumns $ zipWith zt ([m-1, m-2 .. 1] ++ repeat 0) cs
          vs = zipWith zh [1..mn] cs
          hs = zipWith haussholder (toList tau) vs
          q = foldl1' mXm hs

unpackHess :: (Field t) => (Matrix t -> (Matrix t,Vector t)) -> Matrix t -> (Matrix t, Matrix t)
unpackHess hf m
    | rows m == 1 = ((1><1)[1],m)
    | otherwise = (uH . hf) m

uH (pq, tau) = (p,h)
    where cs = toColumns pq
          m = rows pq
          n = cols pq
          mn = min m n
          h = fromColumns $ zipWith zt ([m-2, m-3 .. 1] ++ repeat 0) cs
          vs = zipWith zh [2..mn] cs
          hs = zipWith haussholder (toList tau) vs
          p = foldl1' mXm hs

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

-- | Reciprocal of the 2-norm condition number of a matrix, computed from the singular values.
rcond :: Field t => Matrix t -> Double
rcond m = last s / head s
    where s = toList (singularValues m)

-- | Number of linearly independent rows or columns.
rank :: Field t => Matrix t -> Int
rank m = rankSVD eps m (singularValues m)

{-
expm' m = case diagonalize (complex m) of
    Just (l,v) -> v `mXm` diag (exp l) `mXm` inv v
    Nothing -> error "Sorry, expm not yet implemented for non-diagonalizable matrices"
  where exp = vectorMapC Exp
-}

diagonalize m = if rank v == n
                    then Just (l,v)
                    else Nothing
    where n = rows m
          (l,v) = if exactHermitian m
                    then let (l',v') = eigSH m in (real l', v')
                    else eig m

-- | Generic matrix functions for diagonalizable matrices. For instance:
--
-- @logm = matFunc log@
--
matFunc :: (Complex Double -> Complex Double) -> Matrix (Complex Double) -> Matrix (Complex Double)
matFunc f m = case diagonalize m of
    Just (l,v) -> v `mXm` diag (mapVector f l) `mXm` inv v
    Nothing -> error "Sorry, matFunc requires a diagonalizable matrix"

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

golubeps :: Integer -> Integer -> Double
golubeps p q = a * fromIntegral b / fromIntegral c where
    a = 2^^(3-p-q)
    b = fact p * fact q
    c = fact (p+q) * fact (p+q+1)
    fact n = product [1..n]

epslist = [ (fromIntegral k, golubeps k k) | k <- [1..]]

geps delta = head [ k | (k,g) <- epslist, g<delta]

{- | Matrix exponential. It uses a direct translation of Algorithm 11.3.1 in Golub & Van Loan,
     based on a scaled Pade approximation.
-}
expm :: Field t => Matrix t -> Matrix t
expm = expGolub

expGolub :: ( Fractional t, Element t, Field t
            , Normed Matrix t
            , RealFrac (RealOf t)
            , Floating (RealOf t)
            ) => Matrix t -> Matrix t
expGolub m = iterate msq f !! j
    where j = max 0 $ floor $ logBase 2 $ pnorm Infinity m
          a = m */ fromIntegral ((2::Int)^j)
          q = geps eps -- 7 steps
          eye = ident (rows m)
          work (k,c,x,n,d) = (k',c',x',n',d')
              where k' = k+1
                    c' = c * fromIntegral (q-k+1) / fromIntegral ((2*q-k+1)*k)
                    x' = a <> x
                    n' = n |+| (c' .* x')
                    d' = d |+| (((-1)^k * c') .* x')
          (_,_,_,nf,df) = iterate work (1,1,eye,eye,eye) !! q
          f = linearSolve df nf
          msq x = x <> x

          (<>) = multiply
          v */ x = scale (recip x) v
          (.*) = scale
          (|+|) = add

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

{- | Matrix square root. Currently it uses a simple iterative algorithm described in Wikipedia.
It only works with invertible matrices that have a real solution. For diagonalizable matrices you can try @matFunc sqrt@.

@m = (2><2) [4,9
           ,0,4] :: Matrix Double@

@\>sqrtm m
(2><2)
 [ 2.0, 2.25
 , 0.0,  2.0 ]@
-}
sqrtm ::  Field t => Matrix t -> Matrix t
sqrtm = sqrtmInv

sqrtmInv x = fst $ fixedPoint $ iterate f (x, ident (rows x))
    where fixedPoint (a:b:rest) | pnorm PNorm1 (fst a |-| fst b) < peps   = a
                                | otherwise = fixedPoint (b:rest)
          fixedPoint _ = error "fixedpoint with impossible inputs"
          f (y,z) = (0.5 .* (y |+| inv z),
                     0.5 .* (inv y |+| z))
          (.*) = scale
          (|+|) = add
          (|-|) = sub

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

signlp r vals = foldl f 1 (zip [0..r-1] vals)
    where f s (a,b) | a /= b    = -s
                    | otherwise =  s

swap (arr,s) (a,b) | a /= b    = (arr // [(a, arr!b),(b,arr!a)],-s)
                   | otherwise = (arr,s)

fixPerm r vals = (fromColumns $ elems res, sign)
    where v = [0..r-1]
          s = toColumns (ident r)
          (res,sign) = foldl swap (listArray (0,r-1) s, 1) (zip v vals)

triang r c h v = (r><c) [el s t | s<-[0..r-1], t<-[0..c-1]]
    where el p q = if q-p>=h then v else 1 - v

luFact (l_u,perm) | r <= c    = (l ,u ,p, s)
                  | otherwise = (l',u',p, s)
  where
    r = rows l_u
    c = cols l_u
    tu = triang r c 0 1
    tl = triang r c 0 0
    l = takeColumns r (l_u |*| tl) |+| diagRect 0 (konst 1 r) r r
    u = l_u |*| tu
    (p,s) = fixPerm r perm
    l' = (l_u |*| tl) |+| diagRect 0 (konst 1 c) r c
    u' = takeRows c (l_u |*| tu)
    (|+|) = add
    (|*|) = mul

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

data NormType = Infinity | PNorm1 | PNorm2 | Frobenius

class (RealFloat (RealOf t)) => Normed c t where
    pnorm :: NormType -> c t -> RealOf t

instance Normed Vector Double where
    pnorm PNorm1    = norm1
    pnorm PNorm2    = norm2
    pnorm Infinity  = normInf
    pnorm Frobenius = norm2

instance Normed Vector (Complex Double) where
    pnorm PNorm1    = norm1
    pnorm PNorm2    = norm2
    pnorm Infinity  = normInf
    pnorm Frobenius = pnorm PNorm2

instance Normed Vector Float where
    pnorm PNorm1    = norm1
    pnorm PNorm2    = norm2
    pnorm Infinity  = normInf
    pnorm Frobenius = pnorm PNorm2

instance Normed Vector (Complex Float) where
    pnorm PNorm1    = norm1
    pnorm PNorm2    = norm2
    pnorm Infinity  = normInf
    pnorm Frobenius = pnorm PNorm2

instance Normed Matrix Double where
    pnorm PNorm1    = maximum . map (pnorm PNorm1) . toColumns
    pnorm PNorm2    = (@>0) . singularValues
    pnorm Infinity  = pnorm PNorm1 . trans
    pnorm Frobenius = pnorm PNorm2 . flatten

instance Normed Matrix (Complex Double) where
    pnorm PNorm1    = maximum . map (pnorm PNorm1) . toColumns
    pnorm PNorm2    = (@>0) . singularValues
    pnorm Infinity  = pnorm PNorm1 . trans
    pnorm Frobenius = pnorm PNorm2 . flatten

instance Normed Matrix Float where
    pnorm PNorm1    = maximum . map (pnorm PNorm1) . toColumns
    pnorm PNorm2    = realToFrac . (@>0) . singularValues . double
    pnorm Infinity  = pnorm PNorm1 . trans
    pnorm Frobenius = pnorm PNorm2 . flatten

instance Normed Matrix (Complex Float) where
    pnorm PNorm1    = maximum . map (pnorm PNorm1) . toColumns
    pnorm PNorm2    = realToFrac . (@>0) . singularValues . double
    pnorm Infinity  = pnorm PNorm1 . trans
    pnorm Frobenius = pnorm PNorm2 . flatten

-- | Approximate number of common digits in the maximum element.
relativeError :: (Normed c t, Container c t) => c t -> c t -> Int
relativeError x y = dig (norm (x `sub` y) / norm x)
    where norm = pnorm Infinity
          dig r = round $ -logBase 10 (realToFrac r :: Double)

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

-- | Generalized symmetric positive definite eigensystem Av = lBv,
-- for A and B symmetric, B positive definite (conditions not checked).
geigSH' :: Field t
        => Matrix t -- ^ A
        -> Matrix t -- ^ B
        -> (Vector Double, Matrix t)
geigSH' a b = (l,v')
  where
    u = cholSH b
    iu = inv u
    c = ctrans iu <> a <> iu
    (l,v) = eigSH' c
    v' = iu <> v
    (<>) = mXm


import Data.Complex
import Data.List (partition)
import Data.Map ((!))
import qualified Data.Map as M
import Data.Ratio

dft :: [Complex Double] -> [Complex Double]
dft x = matMult dftMat x
  where
    n = length x
    w = exp $ (-2) * pi * (0 :+ 1) / fromIntegral n
    dftMat = [[w ^ (j * k) | j <- [0 .. n - 1]] | k <- [0 .. n - 1]]
    matMult m x = map (sum . zipWith (*) x) m

fft :: [Complex Double] -> [Complex Double]
fft x = fft' x
  where
    n = length x
    w0 = exp ((-2) * pi * (0 :+ 1) / fromIntegral n)
    w = M.fromList [(k % n, w0 ^ k) | k <- [0 .. n - 1]]
    fft' [x] = [x]
    fft' x =
      let (evens, odds) = partition (even . fst) $ zip [0 ..] x
          e = fft' $ map snd evens
          o = fft' $ map snd odds
          x1 = zipWith3 (\e o k -> e + o * w ! (k %n)) e o [0 ..]
          x2 = zipWith3 (\e o k -> e - o * w ! (k %n)) e o [0 ..]
       in x1 ++ x2

main = do
  print $ dft [0, 1, 2, 3]
  print $ fft [0, 1, 2, 3]


{-# LANGUAGE TypeSynonymInstances #-}
{-# LANGUAGE FlexibleInstances #-}
{-# LANGUAGE GeneralizedNewtypeDeriving #-}

-- | This module defines the Syntax of the Quantum IO Monad, which is
-- an embedded language for writing quantum computations. It is an
-- alternative definition using the approach of defining F-Algebras.
module QIO.QioSynAlt where

import Data.Monoid as Monoid
import Data.Complex
import Control.Applicative (Applicative(..))
import Control.Monad       (liftM, ap)

-- | For Real numbers, we simply use the built in Double type
type RR = Double

-- | For Complex numbers, we use the built in Complex numbers, over our Real
-- number type (i.e. Double)
type CC = Complex RR

-- | The amplitude of a complex number is the magnitude squared.
amp :: CC -> RR
amp k = (magnitude k)*(magnitude k)

-- | The type of Qubits in QIO are simply integer references.
newtype Qbit = Qbit Int deriving (Num, Enum, Eq, Ord)

-- | We can display a qubit reference
instance Show Qbit where
    show (Qbit q) = "(Qbit:" ++ show q ++ ")"

-- | A rotation is in essence a two-by-two complex valued matrix
type Rotation = ((Bool,Bool) -> CC)

-- | We can display the matrix representation of a rotation
instance Show Rotation where
    show f = "(" ++ (show (f (False,False))) ++ "," ++ (show (f (False,True))) ++ "," ++ (show (f (True,False))) ++ "," ++ (show (f (True,True))) ++ ")"

-- | The identity rotation
rid :: Rotation
rid (x,y) = if x==y then 1 else 0

-- | The not rotation
rnot :: Rotation
rnot (x,y) = if x==y then 0 else 1

-- | The hadamard rotation
rhad :: Rotation
rhad (x,y) = if x && y then -h else h where h = (1/sqrt 2)

-- | The phase rotation
rphase :: RR -> Rotation
rphase _ (False,False)  = 1
rphase r (True,True)    = exp(0:+r)
rphase _ (_,_)          = 0

-- | Returns the inverse (or reverse) of the given rotation
rrev :: Rotation -> Rotation
rrev r (False,True)   = conjugate (r (True,False))
rrev r (True,False)   = conjugate (r (False,True))
rrev r xy             = conjugate (r xy)

-- | Rotations can be compared for equality.
-- They are equal if the define the same matrix.
instance Eq Rotation where
    f == g =    (f (False,False)  == g (False,False))
             && (f (False,True)   == g (False,True))
             && (f (True,False)   == g (True,False))
             && (f (True,True)    == g (True,True))
    f /= g =    (f (False,False)  /= g (False,False))
             || (f (False,True)   /= g (False,True))
             || (f (True,False)   /= g (True,False))
             || (f (True,True)    /= g (True,True))

-- | The non-recursive data type definition of a unitary operation
data UFunctor u = UReturn
                | Rot Qbit Rotation u
                | Swap Qbit Qbit u
                | Cond Qbit (Bool -> u) u
                | Ulet Bool (Qbit -> u) u

-- | In order to define an F-Algebra, 'UFunctor' must be a functor.
instance Functor UFunctor where
    fmap eval UReturn = UReturn
    fmap eval (Rot q r u) = Rot q r (eval u)
    fmap eval (Swap q1 q2 u) = Swap q1 q2 (eval u)
    fmap eval (Cond q f u) = Cond q (eval . f) (eval u)
    fmap eval (Ulet b f u) = Ulet b (eval . f) (eval u)

-- | The fix point type construtor.
newtype Fix f = Fx (f (Fix f))

-- | We can define the inverse of Fx
unFix :: Fix f -> f (Fix f)
unFix (Fx x) = x

-- | We fix the non-recursice data-type in order to get our type 'U'
-- of unitary operations.
type U = Fix UFunctor

-- | The type of an F-Algebra.
type Algebra f a = f a -> a

-- | The type of the initial algebra for UFunctor
type UInitialAlgebra = Algebra UFunctor U

-- | We can now define the initial algebra for U
uInitialAlgebra :: UInitialAlgebra
uInitialAlgebra = Fx

-- | We can use a catamorphism to abstract evaluation over a given
-- algebra
cata :: Functor f => Algebra f a -> Fix f -> a
cata algebra = algebra . fmap (cata algebra) . unFix

-- | The type "U" forms a Monoid.
instance Semigroup U where
    (Fx UReturn) <> u = u
    (Fx (Rot x a u)) <> u' = Fx $ Rot x a (u <> u')
    (Fx (Swap x y u)) <> u' = Fx $ Swap x y (u <> u')
    (Fx (Cond x br u')) <> u'' = Fx $ Cond x br (u' <> u'')
    (Fx (Ulet b f u)) <> u' = Fx $ Ulet b f (u <> u')

instance Monoid U where
    mempty = Fx UReturn

-- | Apply the given rotation to the given qubit
rot :: Qbit -> Rotation -> U
rot x r = Fx $ Rot x r mempty

-- | Swap the state of the two given qubits
swap :: Qbit -> Qbit -> U
swap x y = Fx $ Swap x y mempty

-- | Apply the conditional unitary, depending on the value of the given qubit
cond :: Qbit -> (Bool -> U) -> U
cond x br = Fx $ Cond x br mempty

-- | Introduce an Ancilla qubit in the given state, for use in the sub-unitary
ulet :: Bool -> (Qbit -> U) -> U
ulet b ux = Fx $ Ulet b ux mempty

-- | Apply a not rotation to the given qubit
unot :: Qbit -> U
unot x = rot x rnot

-- | Apply a hadamard rotation to the given qubit
uhad :: Qbit -> U
uhad x = rot x rhad

-- | Apply a phase rotation (of the given angle) to the given qubit
uphase :: Qbit -> RR -> U
uphase x r = rot x (rphase r)

-- | Returns the inverse (or reverse) of the given unitary operation,
-- using an F-Algebra
urev :: U -> U
urev = cata urev_algebra
 where
   urev_algebra :: UFunctor U -> U
   urev_algebra UReturn = Fx UReturn
   urev_algebra (Rot x r u) = u `mappend` rot x (rrev r)
   urev_algebra (Swap x y u) = u `mappend` swap x y
   urev_algebra (Cond x br u) = u `mappend` cond x br
   urev_algebra (Ulet b xu u) = u `mappend` ulet b xu

-- | We can display a representation of a unitary, using an F-Algebra
instance Show U where
  show = cata showU_algebra
    where
      showU_algebra :: UFunctor String -> String
      showU_algebra UReturn = ""
      showU_algebra (Rot q r u) =
       "Rotate " ++ show q ++ ":" ++ show r ++ "\n" ++ u
      showU_algebra (Swap q1 q2 u) =
       "Swap " ++ show q1 ++ " and " ++ show q2 ++ "\n" ++ u
      showU_algebra (Cond q br u) =
       "Cond " ++ show q ++ " \\b -> if b then (\n"
       ++ unlines (map (' ':) (lines $ br True))
       ++ ") else (\n"
       ++ unlines (map (' ':) (lines $ br False))
       ++ ")\n" ++ u
      showU_algebra (Ulet b xu u) =
       let fv = find_fv xu in
       "Ulet " ++ show fv ++ " = " ++ (if b then "1" else "0") ++ " in (\n"
       ++ unlines (map (' ':) (lines $ xu fv))
       ++ ")\n" ++ u
      -- this is currently a dummy function
      find_fv :: (Qbit -> String) -> Qbit
      find_fv _ = -1

-- | The non-recursive data type definition of a QIO computation
data QIOFunctor a q = QReturn a
                    | MkQbit Bool (Qbit -> q)
                    | ApplyU U q
                    | Meas Qbit (Bool -> q)

-- | In order to define an F-Algebra, 'UF' must be a functor.
instance Functor (QIOFunctor a) where
    fmap eval (QReturn a) = QReturn a
    fmap eval (MkQbit b f) = MkQbit b (eval . f)
    fmap eval (ApplyU u q) = ApplyU u (eval q)
    fmap eval (Meas q f) = Meas q (eval . f)

-- | We fix the non-recursice data-type in order to get our type 'U'
-- of unitary operations.
type QIOprim a = Fix (QIOFunctor a)

-- | The type of the initial algebra for UFunctor
type QIOInitialAlgebra a = Algebra (QIOFunctor a) (QIOprim a)

-- | We can now define the initial algebra for U
qioInitialAlgebra :: QIOInitialAlgebra a
qioInitialAlgebra = Fx

-- | The "QIO" type forms a Monad, by wrapping 'QIOprim'
data QIO a = Apply (Fix (QIOFunctor a))

-- | We can remove the wrapper.
primQIO :: QIO a -> QIOprim a
primQIO (Apply q) = q

instance Functor QIO where
    fmap = liftM

instance Applicative QIO  where
    pure  = Apply . Fx . QReturn
    (<*>) = ap

-- | The wrapper type 'ApplyFix' forms a Monad
instance Monad QIO where
    return = pure
    (Apply (Fx (QReturn a))) >>= f = f a
    (Apply (Fx (MkQbit b g))) >>= f = Apply . Fx $
      MkQbit b  (\q -> primQIO $ (Apply (g q)) >>= f)
    (Apply (Fx (ApplyU u q))) >>= f = Apply . Fx $
      ApplyU u $ primQIO (Apply q >>= f)
    (Apply (Fx (Meas x g))) >>= f = Apply . Fx $
      Meas x (\b -> primQIO $ (Apply (g b)) >>= f)

-- | Initialise a qubit in the given state (adding it to the overall quantum state)
mkQbit :: Bool -> QIO Qbit
mkQbit b = Apply . Fx $ MkQbit b (\q -> primQIO (return q))

-- | Apply the given unitary operation to the current quantum state
applyU :: U -> QIO ()
applyU u =  Apply . Fx $ ApplyU u $ primQIO (return ())

-- | Measure the given qubit, and return the measurement outcome (note that this
-- operation may affect the overall quantum state, as a measurement is destructive)
measQbit :: Qbit -> QIO Bool
measQbit x = Apply . Fx $ Meas x (\b -> primQIO (return b))

-- | We can show a QIO computation, using an F-Algebra
instance (Show a) => Show (QIO a) where
  show = (cata showQIO_algebra) . primQIO
   where
    showQIO_algebra :: (Show a) => Algebra (QIOFunctor a) String
    showQIO_algebra (QReturn a) =
       "Return: " ++ show a ++ "\n"
    showQIO_algebra (MkQbit b f) =
       "Init" ++ (if b then "1" else "0") ++ "\n"
       ++ f 0
    showQIO_algebra (ApplyU u qio) =
       "Apply Unitary: (\n"
       ++ unlines (map (' ':) (lines $ show u))
       ++ ")\n" ++ qio
    showQIO_algebra (Meas q f) =
       "Measure " ++ show q ++ " \\b -> if b then (\n"
       ++ unlines (map (' ':) (lines $ f True))
       ++ ") else (\n"
       ++ unlines (map (' ':) (lines $ f False))
       ++ ")\n"

-- | We can count the number of each primitive operation using an F-Algebra
count :: QIO a -> (Int,Int,Int)
count = (cata count_algebra) . primQIO
 where
  count_algebra :: Algebra (QIOFunctor a) (Int,Int,Int)
  count_algebra (QReturn _) = (0,0,0)
  count_algebra (MkQbit b f) = let (mk,ap,ms) = f 0 in
                               (mk+1,ap,ms)
  count_algebra (ApplyU _ (mk,ap,ms)) = (mk,ap+1,ms)
  count_algebra (Meas q f) = let (mk,ap,ms) = f False in
                             (mk,ap,ms+1)

toffoli :: Qbit -> Qbit -> Qbit -> U
toffoli q1 q2 q3 =
  cond q1 (\b1 -> if b1 then (
   cond q2 (\b2 -> if b2 then (unot q3)
    else mempty)) else mempty)

and :: Bool -> Bool -> QIO Bool
and a b = do
 q1 <- mkQbit a
 q2 <- mkQbit b
 q3 <- mkQbit False
 applyU (toffoli q1 q2 q3)
 measQbit q3


{-# LANGUAGE FlexibleContexts #-}
-- |
-- Module    : Statistics.Sample
-- Copyright : (c) 2008 Don Stewart, 2009 Bryan O'Sullivan
-- License   : BSD3
--
-- Maintainer  : bos@serpentine.com
-- Stability   : experimental
-- Portability : portable
--
-- Commonly used sample statistics, also known as descriptive
-- statistics.

module Statistics.Sample
    (
    -- * Statistics of location
      mean

    -- ** Two-pass functions (numerically robust)
    -- $robust
    , variance
    , varianceUnbiased
    , stdDev

    -- * References
    -- $references
    ) where

import Statistics.Sample.Internal (robustSumVar, sum)
import qualified Data.Vector as V
import qualified Data.Vector.Generic as G
import qualified Data.Vector.Unboxed as U

-- Operator ^ will be overriden
import Prelude hiding ((^), sum)

-- | /O(n)/ Arithmetic mean.  This uses Kahan-Babuška-Neumaier
-- summation, so is more accurate than 'welfordMean' unless the input
-- values are very large.
mean :: (G.Vector v Double) => v Double -> Double
mean xs = sum xs / fromIntegral (G.length xs)
{-# SPECIALIZE mean :: U.Vector Double -> Double #-}
{-# SPECIALIZE mean :: V.Vector Double -> Double #-}

-- $variance
--
-- The variance—and hence the standard deviation—of a
-- sample of fewer than two elements are both defined to be zero.

-- $robust
--
-- These functions use the compensated summation algorithm of Chan et
-- al. for numerical robustness, but require two passes over the
-- sample data as a result.
--
-- Because of the need for two passes, these functions are /not/
-- subject to stream fusion.

-- | Maximum likelihood estimate of a sample's variance.  Also known
-- as the population variance, where the denominator is /n/.
variance :: (G.Vector v Double) => v Double -> Double
variance samp
    | n > 1     = robustSumVar (mean samp) samp / fromIntegral n
    | otherwise = 0
    where
      n = G.length samp
{-# SPECIALIZE variance :: U.Vector Double -> Double #-}
{-# SPECIALIZE variance :: V.Vector Double -> Double #-}

-- | Unbiased estimate of a sample's variance.  Also known as the
-- sample variance, where the denominator is /n/-1.
varianceUnbiased :: (G.Vector v Double) => v Double -> Double
varianceUnbiased samp
    | n > 1     = robustSumVar (mean samp) samp / fromIntegral (n-1)
    | otherwise = 0
    where
      n = G.length samp
{-# SPECIALIZE varianceUnbiased :: U.Vector Double -> Double #-}
{-# SPECIALIZE varianceUnbiased :: V.Vector Double -> Double #-}

-- | Standard deviation.  This is simply the square root of the
-- unbiased estimate of the variance.
stdDev :: (G.Vector v Double) => v Double -> Double
stdDev = sqrt . varianceUnbiased
{-# SPECIALIZE stdDev :: U.Vector Double -> Double #-}
{-# SPECIALIZE stdDev :: V.Vector Double -> Double #-}

-- $cancellation
--
-- The functions prefixed with the name @fast@ below perform a single
-- pass over the sample data using Knuth's algorithm. They usually
-- work well, but see below for caveats. These functions are subject
-- to array fusion.
--
-- /Note/: in cases where most sample data is close to the sample's
-- mean, Knuth's algorithm gives inaccurate results due to
-- catastrophic cancellation.

-- $references
--
-- * Chan, T. F.; Golub, G.H.; LeVeque, R.J. (1979) Updating formulae
--   and a pairwise algorithm for computing sample
--   variances. Technical Report STAN-CS-79-773, Department of
--   Computer Science, Stanford
--   University. <ftp://reports.stanford.edu/pub/cstr/reports/cs/tr/79/773/CS-TR-79-773.pdf>
--
-- * Knuth, D.E. (1998) The art of computer programming, volume 2:
--   seminumerical algorithms, 3rd ed., p. 232.
--
-- * Welford, B.P. (1962) Note on a method for calculating corrected
--   sums of squares and products. /Technometrics/
--   4(3):419–420. <http://www.jstor.org/stable/1266577>
--
-- * West, D.H.D. (1979) Updating mean and variance estimates: an
--   improved method. /Communications of the ACM/
--   22(9):532–535. <http://doi.acm.org/10.1145/359146.359153>


module RandomExtra where

import Data.Complex (Complex(..))
import qualified Control.Monad.Random.Strict as R
import Numeric (log)

-- | Yield two randomly selected values which follow standard
-- normal distribution.
-- code stolen from https://hackage.haskell.org/package/moo-1.2/docs/src/Moo.GeneticAlgorithm.Random.html#getNormal2
getNormal2 :: (R.Random a, RealFloat a, R.MonadRandom m) => m (a, a)
getNormal2 = do
    -- Box-Muller method
    u <- R.getRandom
    v <- R.getRandom
    let (c :+ s) = exp (0 :+ (2 * pi * v))
    let r = sqrt $ (-2) * log u
    return (r * c, r * s)

-- | Yield one randomly selected value from standard normal distribution.
-- code stolen from https://hackage.haskell.org/package/moo-1.2/docs/src/Moo.GeneticAlgorithm.Random.html#getNormal2
-- code stolen from https://hackage.haskell.org/package/moo-1.2/docs/src/Moo.GeneticAlgorithm.Random.html#getNormal2
getNormal :: (R.Random a, RealFloat a, R.MonadRandom m) => m a
getNormal = fst <$> getNormal2

newtype Mu a = Mu a

newtype Sigma a = Sigma a

getScaledNormal
    :: (R.Random a, RealFloat a, R.MonadRandom m) => Mu a -> Sigma a -> m a
getScaledNormal (Mu mu) (Sigma sigma) = (\x -> x * sigma + mu) <$> getNormal

randomAngle :: (Floating a, R.Random a, R.MonadRandom m) => m a
randomAngle = R.getRandomR (0, 2 * pi)


{-# LANGUAGE CPP #-}
-----------------------------------------------------------------------------
-- |
-- Copyright   :  (c) Edward Kmett 2010-2015
-- License     :  BSD3
-- Maintainer  :  ekmett@gmail.com
-- Stability   :  experimental
-- Portability :  GHC only
--
-- Root finding using Halley's rational method (the second in
-- the class of Householder methods). Assumes the function is three
-- times continuously differentiable and converges cubically when
-- progress can be made.
--
-----------------------------------------------------------------------------

module Numeric.AD.Rank1.Halley
  (
  -- * Halley's Method (Tower AD)
    findZero
  , findZeroNoEq
  , inverse
  , inverseNoEq
  , fixedPoint
  , fixedPointNoEq
  , extremum
  , extremumNoEq
  ) where

import Prelude hiding (all)
import Numeric.AD.Internal.Forward (Forward)
import Numeric.AD.Internal.On
import Numeric.AD.Internal.Tower (Tower)
import Numeric.AD.Mode
import Numeric.AD.Rank1.Tower (diffs0)
import Numeric.AD.Rank1.Forward (diff)
import Numeric.AD.Internal.Combinators (takeWhileDifferent)

-- $setup
-- >>> import Data.Complex

-- | The 'findZero' function finds a zero of a scalar function using
-- Halley's method; its output is a stream of increasingly accurate
-- results.  (Modulo the usual caveats.) If the stream becomes constant
-- ("it converges"), no further elements are returned.
--
-- Examples:
--
-- >>> take 10 $ findZero (\x->x^2-4) 1
-- [1.0,1.8571428571428572,1.9997967892704736,1.9999999999994755,2.0]
--
-- >>> last $ take 10 $ findZero ((+1).(^2)) (1 :+ 1)
-- 0.0 :+ 1.0
findZero :: (Fractional a, Eq a) => (Tower a -> Tower a) -> a -> [a]
findZero f = takeWhileDifferent . findZeroNoEq f
{-# INLINE findZero #-}

-- | The 'findZeroNoEq' function behaves the same as 'findZero' except that it
-- doesn't truncate the list once the results become constant. This means it
-- can be used with types without an 'Eq' instance.
findZeroNoEq :: Fractional a => (Tower a -> Tower a) -> a -> [a]
findZeroNoEq f = iterate go where
  go x = xn where
    (y:y':y'':_) = diffs0 f x
    xn = x - 2*y*y'/(2*y'*y'-y*y'') -- 9.606671960457536 bits error
       -- = x - recip (y'/y - y''/ y') -- "improved error" = 6.640625e-2 bits
       -- = x - y' / (y'/y/y' - y''/2) -- "improved error" = 1.4
#ifdef HERBIE
{-# ANN findZeroNoEq "NoHerbie" #-}
#endif
{-# INLINE findZeroNoEq #-}

-- | The 'inverse' function inverts a scalar function using
-- Halley's method; its output is a stream of increasingly accurate
-- results.  (Modulo the usual caveats.) If the stream becomes constant
-- ("it converges"), no further elements are returned.
--
-- Note: the @take 10 $ inverse sqrt 1 (sqrt 10)@ example that works for Newton's method
-- fails with Halley's method because the preconditions do not hold!
inverse :: (Fractional a, Eq a) => (Tower a -> Tower a) -> a -> a -> [a]
inverse f x0 = takeWhileDifferent . inverseNoEq f x0
{-# INLINE inverse  #-}

-- | The 'inverseNoEq' function behaves the same as 'inverse' except that it
-- doesn't truncate the list once the results become constant. This means it
-- can be used with types without an 'Eq' instance.
inverseNoEq :: Fractional a => (Tower a -> Tower a) -> a -> a -> [a]
inverseNoEq f x0 y = findZeroNoEq (\x -> f x - auto y) x0
{-# INLINE inverseNoEq #-}

-- | The 'fixedPoint' function find a fixedpoint of a scalar
-- function using Halley's method; its output is a stream of
-- increasingly accurate results.  (Modulo the usual caveats.)
--
-- If the stream becomes constant ("it converges"), no further
-- elements are returned.
--
-- >>> last $ take 10 $ fixedPoint cos 1
-- 0.7390851332151607
fixedPoint :: (Fractional a, Eq a) => (Tower a -> Tower a) -> a -> [a]
fixedPoint f = takeWhileDifferent . fixedPointNoEq f
{-# INLINE fixedPoint #-}

-- | The 'fixedPointNoEq' function behaves the same as 'fixedPoint' except that
-- it doesn't truncate the list once the results become constant. This means it
-- can be used with types without an 'Eq' instance.
fixedPointNoEq :: Fractional a => (Tower a -> Tower a) -> a -> [a]
fixedPointNoEq f = findZeroNoEq (\x -> f x - x)
{-# INLINE fixedPointNoEq #-}

-- | The 'extremum' function finds an extremum of a scalar
-- function using Halley's method; produces a stream of increasingly
-- accurate results.  (Modulo the usual caveats.) If the stream becomes
-- constant ("it converges"), no further elements are returned.
--
-- >>> take 10 $ extremum cos 1
-- [1.0,0.29616942658570555,4.59979519460002e-3,1.6220740159042513e-8,0.0]
extremum :: (Fractional a, Eq a) => (On (Forward (Tower a)) -> On (Forward (Tower a))) -> a -> [a]
extremum f = takeWhileDifferent . extremumNoEq f
{-# INLINE extremum #-}

-- | The 'extremumNoEq' function behaves the same as 'extremum' except that it
-- doesn't truncate the list once the results become constant. This means it
-- can be used with types without an 'Eq' instance.
extremumNoEq :: Fractional a => (On (Forward (Tower a)) -> On (Forward (Tower a))) -> a -> [a]
extremumNoEq f = findZeroNoEq (diff (off . f . On))
{-# INLINE extremumNoEq #-}


{-# LANGUAGE OverloadedStrings #-}
{-# LANGUAGE DeriveDataTypeable, DeriveGeneric #-}
-- |
-- Module    : Statistics.Distribution.Uniform
-- Copyright : (c) 2011 Aleksey Khudyakov
-- License   : BSD3
--
-- Maintainer  : bos@serpentine.com
-- Stability   : experimental
-- Portability : portable
--
-- Variate distributed uniformly in the interval.
module Statistics.Distribution.Uniform
    (
      UniformDistribution
    -- * Constructors
    , uniformDistr
    , uniformDistrE
    -- ** Accessors
    , uniformA
    , uniformB
    ) where

import Control.Applicative
import Data.Data           (Data, Typeable)
import GHC.Generics        (Generic)
import qualified System.Random.MWC       as MWC

import qualified Statistics.Distribution as D
import Statistics.Internal

-- | Uniform distribution from A to B
data UniformDistribution = UniformDistribution {
      uniformA :: {-# UNPACK #-} !Double -- ^ Low boundary of distribution
    , uniformB :: {-# UNPACK #-} !Double -- ^ Upper boundary of distribution
    } deriving (Eq, Typeable, Data, Generic)

instance Show UniformDistribution where
  showsPrec i (UniformDistribution a b) = defaultShow2 "uniformDistr" a b i
instance Read UniformDistribution where
  readPrec = defaultReadPrecM2 "uniformDistr" uniformDistrE

-- | Create uniform distribution.
uniformDistr :: Double -> Double -> UniformDistribution
uniformDistr a b = maybe (error errMsg) id $ uniformDistrE a b

-- | Create uniform distribution.
uniformDistrE :: Double -> Double -> Maybe UniformDistribution
uniformDistrE a b
  | b < a     = Just $ UniformDistribution b a
  | a < b     = Just $ UniformDistribution a b
  | otherwise = Nothing
-- NOTE: failure is in default branch to guard against NaNs.

errMsg :: String
errMsg = "Statistics.Distribution.Uniform.uniform: wrong parameters"

instance D.Distribution UniformDistribution where
  cumulative (UniformDistribution a b) x
    | x < a     = 0
    | x > b     = 1
    | otherwise = (x - a) / (b - a)

instance D.ContDistr UniformDistribution where
  density (UniformDistribution a b) x
    | x < a     = 0
    | x > b     = 0
    | otherwise = 1 / (b - a)
  quantile (UniformDistribution a b) p
    | p >= 0 && p <= 1 = a + (b - a) * p
    | otherwise        =
      error $ "Statistics.Distribution.Uniform.quantile: p must be in [0,1] range. Got: "++show p
  complQuantile (UniformDistribution a b) p
    | p >= 0 && p <= 1 = b + (a - b) * p
    | otherwise        =
      error $ "Statistics.Distribution.Uniform.complQuantile: p must be in [0,1] range. Got: "++show p

instance D.Mean UniformDistribution where
  mean (UniformDistribution a b) = 0.5 * (a + b)

instance D.Variance UniformDistribution where
  -- NOTE: 1/sqrt 12 is not constant folded (#4101) so it's written as
  --       numerical constant. (Also FIXME!)
  stdDev   (UniformDistribution a b) = 0.2886751345948129 * (b - a)
  variance (UniformDistribution a b) = d * d / 12 where d = b - a

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

instance D.MaybeVariance UniformDistribution where
    maybeStdDev   = Just . D.stdDev

instance D.Entropy UniformDistribution where
  entropy (UniformDistribution a b) = log (b - a)

instance D.MaybeEntropy UniformDistribution where
  maybeEntropy = Just . D.entropy

instance D.ContGen UniformDistribution where
    genContVar (UniformDistribution a b) gen = MWC.uniformR (a,b) gen


-----------------------------------------------------------------------------
-- |
-- Module     : Numeric.LinearAlgebra.Packed.Statistics
-- Copyright  : Copyright (c) 2010, Patrick Perry <patperry@gmail.com>
-- License    : BSD3
-- Maintainer : Patrick Perry <patperry@gmail.com>
-- Stability  : experimental
--
-- Basic multivariate statistics.
--

module Numeric.LinearAlgebra.Packed.Statistics (
    defaultCovUplo,

    -- * Immutable interface
    cov,
    covWithMean,
    weightedCov,
    weightedCovWithMean,

    -- * Mutable interface
    covTo,
    covWithMeanTo,
    weightedCovTo,
    weightedCovWithMeanTo,

    ) where

import Control.Monad( when )
import Control.Monad.ST( ST )
import Data.List( foldl' )
import Text.Printf( printf )

import Numeric.LinearAlgebra.Types
import Numeric.LinearAlgebra.Packed.Base( Packed, STPacked )
import qualified Numeric.LinearAlgebra.Packed.Base as P

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

import qualified Numeric.LinearAlgebra.Matrix as M

-- | Returns the default storage scheme for covariance matrices.
defaultCovUplo :: Uplo
defaultCovUplo = Lower

-- | Returns the sample covariance matrix hermitian matrix (in packed form)
-- with storage scheme equal to 'defaultCovUplo'.  The first argument gives
-- the dimension of the vectors.
cov :: (BLAS2 e)
    => Int -> CovMethod -> [Vector e] -> Herm Packed e
cov p t xs = P.hermCreate $ do
    c <- Herm uplo `fmap` P.new_ p
    covTo c t xs
    return c
  where
    uplo = defaultCovUplo

-- | Given the pre-computed mean, returns the sample covariance matrix
-- (in packed form) with storage scheme equal to 'defaultCovUplo'.
covWithMean :: (BLAS2 e)
            => Vector e -> CovMethod -> [Vector e] -> Herm Packed e
covWithMean mu t xs = P.hermCreate $ do
    c <- Herm uplo `fmap` P.new_ p
    covWithMeanTo c mu t xs
    return c
  where
    p = V.dim mu
    uplo = defaultCovUplo

-- | Returns the weighed sample covariance matrix (in packed form) with
-- storage scheme equal to 'defaultCovUplo'. The first argument gives the
-- dimension of the vectors.
weightedCov :: (BLAS2 e)
            => Int -> CovMethod -> [(Double, Vector e)] -> Herm Packed e
weightedCov p t wxs = P.hermCreate $ do
    c <- Herm uplo `fmap` P.new_ p
    weightedCovTo c t wxs
    return c
  where
    uplo = defaultCovUplo

-- | Given the pre-computed mean, returns the weighed sample covariance matrix
-- (in packed form) with storage scheme equal to 'defaultCovUplo'.
weightedCovWithMean :: (BLAS2 e)
                    => Vector e -> CovMethod -> [(Double, Vector e)]
                    -> Herm Packed e
weightedCovWithMean mu t wxs = P.hermCreate $ do
    c <- Herm uplo `fmap` P.new_ p
    weightedCovWithMeanTo c mu t wxs
    return c
  where
    p = V.dim mu
    uplo = defaultCovUplo

-- | Computes and copies the sample covariance matrix (in packed form)
-- to the given destination.
covTo :: (RVector v, BLAS2 e)
      => Herm (STPacked s) e -> CovMethod -> [v e] -> ST s ()
covTo c@(Herm _ a) t xs = do
    p <- P.getDim a
    mu <- V.new p 1
    V.meanTo mu xs
    covWithMeanTo c mu t xs

-- | Given the pre-computed mean, computes and copies the sample covariance
-- matrix (in packed form) to the given destination.
covWithMeanTo :: (RVector v1, RVector v2, BLAS2 e)
              => Herm (STPacked s) e
              -> v1 e -> CovMethod -> [v2 e]
              -> ST s ()
covWithMeanTo c@(Herm _ a) mu t xs = do
    pa <- P.getDim a
    p <- V.getDim mu

    when (pa /= p) $ error $
        printf ("covWithMeanTo"
                ++ " (Herm _ <packed matrix with dim %d>)"
                ++ " <vector with dim %d>"
                ++ " _ _"
                ++ ": dimension mismatch")
               pa p

    xt <- M.new_ (p,n)
    M.withColsM xt $ \xs' ->
        sequence_ [ V.subTo x' mu x
                  | (x,x') <- zip xs xs'
                  ]
    P.withVectorM a V.clear
    M.withColsM xt $ \xs' ->
        sequence_ [ P.hermRank1UpdateM_ scale x' c | x' <- xs' ]
  where
    n = length xs
    df = fromIntegral $ case t of { MLCov -> n ; UnbiasedCov -> n - 1 }
    scale = 1/df

-- | Computes and copies the weighed sample covariance matrix (in packed
-- form) to the given destination.
weightedCovTo :: (RVector v, BLAS2 e)
              => Herm (STPacked s) e
              -> CovMethod -> [(Double, v e)]
              -> ST s ()
weightedCovTo c@(Herm _ a) t wxs = do
    p <- P.getDim a
    mu <- V.new p 1
    V.weightedMeanTo mu wxs
    weightedCovWithMeanTo c mu t wxs

-- | Given the pre-computed mean, computes and copies the weighed sample
-- covariance matrix (in packed form) to the given destination.
weightedCovWithMeanTo :: (RVector v1, RVector v2, BLAS2 e)
                      => Herm (STPacked s) e
                      -> v1 e -> CovMethod -> [(Double, v2 e)]
                      -> ST s ()
weightedCovWithMeanTo c@(Herm _ a) mu t wxs = do
    pa <- P.getDim a
    p <- V.getDim mu

    when (pa /= p) $ error $
        printf ("weightedCovWithMeanTo"
                ++ " (Herm _ <packed matrix with dim %d>)"
                ++ " <vector with dim %d>"
                ++ " _ _"
                ++ ": dimension mismatch")
               pa p

    xt <- M.new_ (p,n)
    M.withColsM xt $ \xs' ->
        sequence_ [  V.subTo x' mu x
                  >> V.scaleM_ (realToFrac $ sqrt (w / invscale)) x'
                  |  (w,x,x') <- zip3 ws xs xs'
                  ]
    P.withVectorM a V.clear
    M.withCols xt $ \xs' ->
        sequence_ [ P.hermRank1UpdateM_ 1 x' c | x' <- xs' ]
  where
    (ws0,xs) = unzip wxs
    w_sum = foldl' (+) 0 ws0
    ws = if w_sum == 0 then ws0 else map (/w_sum) ws0
    w2s_sum = foldl' (+) 0 $ map (^^(2::Int)) ws
    invscale = case t of
                   MLCov -> 1
                   UnbiasedCov -> (1 - w2s_sum)
    n = length ws0


module Fractal.Math.Mandelbrot(
  mandel,
  mandel'
)where
{-

Math packet kommt alles mathe bezogene rein

-}

import Data.Complex
import Fractal.Math.Kernel (evaluateFractal)

type Point = Complex Double
type Limit = Double
type Iteration = Complex Double
type Magnitude = Double
type MaxIter = Int

mandelEquation :: Point -> Iteration -> Iteration
mandelEquation c (0.0 :+ 0.0) = c
mandelEquation c (0.0 :+ 1) = 0
mandelEquation c (0.0 :+ (-1)) = 0
mandelEquation c z = z * z + c

mandel :: MaxIter -> Limit -> [Point] -> [Int]
mandel maxIter l field = map mandelEvaluator field
  where mandelEvaluator = evaluateFractal mandelEquation l maxIter

mandel' :: MaxIter -> Limit -> Point -> Int
mandel' maxIter l field = mandelEvaluator field
  where mandelEvaluator = evaluateFractal mandelEquation l maxIter


{-# LANGUAGE DataKinds             #-}
{-# LANGUAGE FlexibleContexts      #-}
{-# LANGUAGE FlexibleInstances     #-}
{-# LANGUAGE LambdaCase            #-}
{-# LANGUAGE MultiParamTypeClasses #-}
{-# LANGUAGE PolyKinds             #-}
{-# LANGUAGE RecordWildCards       #-}
{-# LANGUAGE ScopedTypeVariables   #-}
{-# LANGUAGE StandaloneDeriving    #-}
{-# LANGUAGE TypeApplications      #-}
{-# LANGUAGE TypeFamilies          #-}
{-# LANGUAGE TypeOperators         #-}
{-# LANGUAGE UndecidableInstances  #-}
module Numeric.Matrix.SVD
  ( MatrixSVD (..), SVD (..)
  , svd1, svd2, svd3, svd3q
  ) where

import Control.Monad
import Control.Monad.ST
import Data.Kind
import Numeric.Basics
import Numeric.DataFrame.Internal.PrimArray
import Numeric.DataFrame.ST
import Numeric.DataFrame.SubSpace
import Numeric.DataFrame.Type
import Numeric.Dimensions
import Numeric.Matrix.Bidiagonal
import Numeric.Matrix.Internal
import Numeric.Quaternion.Internal
import Numeric.Scalar.Internal
import Numeric.Subroutine.Sort
import Numeric.Tuple
import Numeric.Vector.Internal

-- | Result of SVD factorization
--   @ M = svdU %* asDiag svdS %* transpose svdV @.
--
--   Invariants:
--
--   * Singular values `svdS` are in non-increasing order and are non-negative.
--   * svdU and svdV are orthogonal matrices
--   * det svdU == 1
--
--   NB: <https://en.wikipedia.org/wiki/Singular_value_decomposition SVD on wiki>
data SVD (t :: Type) (n :: Nat) (m :: Nat)
  = SVD
  { svdU :: Matrix t n n
    -- ^ Left-singular basis matrix
  , svdS :: Vector t (Min n m)
    -- ^ Vector of singular values
  , svdV :: Matrix t m m
    -- ^ Right-singular basis matrix
  }

deriving instance ( Show t, PrimBytes t
                  , KnownDim n, KnownDim m, KnownDim (Min n m))
                  => Show (SVD t n m)
deriving instance ( Eq (Matrix t n n)
                  , Eq (Matrix t m m)
                  , Eq (Vector t (Min n m)))
                  => Eq (SVD t n m)

class RealFloatExtras t
    => MatrixSVD (t :: Type) (n :: Nat) (m :: Nat) where
    -- | Compute SVD factorization of a matrix
    svd :: IterativeMethod => Matrix t n m -> SVD t n m

-- | Obvious dummy implementation of SVD for 1x1 matrices
svd1 :: (PrimBytes t, Num t, Eq t) => Matrix t 1 1 -> SVD t 1 1
svd1 m = SVD
    { svdU = 1
    , svdS = broadcast $ abs x
    , svdV = broadcast $ if x == 0 then 1 else signum x
    }
  where
    x = ixOff 0 m

-- | SVD of a 2x2 matrix can be computed analytically
--
--   Related discussion:
--
--   https://scicomp.stackexchange.com/questions/8899/robust-algorithm-for-2-times-2-svd/
--
--   https://ieeexplore.ieee.org/document/486688
svd2 :: forall t . RealFloatExtras t => Matrix t 2 2 -> SVD t 2 2
svd2 (DF2 (DF2 m00 m01) (DF2 m10 m11)) =
    SVD
    { svdU = DF2 (DF2         uc  us)
                 (DF2 (negate us) uc)
    , svdS = DF2 sigma1 (abs sigma2)
    , svdV = DF2 (DF2 vc  (sg2 (negate vs)))
                 (DF2 vs  (sg2 vc         ))
    }
  where
    x1 = m00 - m11 -- 2F
    x2 = m00 + m11 -- 2E
    y1 = m01 + m10 -- 2G
    y2 = m01 - m10 -- 2H
    yy = y1*y2
    h1 = hypot x1 y1 -- h1 >= abs x1
    h2 = hypot x2 y2 -- h2 >= abs x2
    sigma1 = 0.5 * (h2 + h1)  -- sigma1 >= abs sigma2
    sigma2 = 0.5 * (h2 - h1)  -- can be negative, which is accounted by sg2
    sg2 = negateUnless (sigma2 >= 0)
    hx1 = h1 + x1
    hx2 = h2 + x2
    hxhx = hx1*hx2
    hxy  = hx1*y2
    yhx  = y1*hx2
    (uc', us', vc', vs') = case (x1 > 0 || y1 /= 0, x2 > 0 || y2 /= 0) of
      (True , True ) -> (hxhx + yy, hxy - yhx, hxhx - yy, hxy + yhx)
      (True , False) -> (y1,  hx1, -y1, hx1)
      (False, True ) -> (y2, -hx2, -y2, hx2)
      (False, False) -> (1, 0, -1, 0)
    ru = recip $ hypot uc' us'
    rv = recip $ hypot vc' vs'
    uc = uc' * ru
    us = us' * ru
    vc = vc' * rv
    vs = vs' * rv

-- | Get SVD decomposition of a 3x3 matrix using `svd3q` function.
--
--   This function reorders the singular components under the hood to make sure
--   @s1 >= s2 >= s3 >= 0@.
--   Thus, it has some overhead on top of `svd3q`.
svd3 :: forall t . (Quaternion t, RealFloatExtras t) => Matrix t 3 3 -> SVD t 3 3
svd3 m = SVD
        { svdU = toMatrix33 u
        , svdS = DF3 s1 s2 s3'
        , svdV = neg3If (s3 < 0) (toMatrix33 v)
        }
  where
    (u, (DF3 s1 s2 s3), v) = svd3q m
    s3' = abs s3
    neg3If :: Bool -> Matrix t 3 3 -> Matrix t 3 3
    neg3If False = id
    neg3If True  = ewmap @t @'[3] neg3
    neg3 :: Vector t 3 -> Vector t 3
    neg3 (DF3 a b c) = DF3 a b (negate c)

-- | Get SVD decomposition of a 3x3 matrix, with orthogonal matrices U and V
--   represented as quaternions.
--   Important: U and V are bound to be rotations at the expense of the last
--              singular value being possibly negative.
--
--   This is an adoptation of a specialized 3x3 SVD algorithm described in
--     "Computing the Singular Value Decomposition of 3x3 matrices
--      with minimal branching and elementary floating point operations",
--   by  A. McAdams, A. Selle, R. Tamstorf, J. Teran, E. Sifakis.
--
--   http://pages.cs.wisc.edu/~sifakis/papers/SVD_TR1690.pdf
svd3q :: forall t . (Quaternion t, RealFloatExtras t)
      => Matrix t 3 3 -> (Quater t, Vector t 3, Quater t)
svd3q m = (u, s, v)
  where
    v = jacobiEigenQ (transpose m %* m)
    (s, u) = uncurry fixSigns $ qrDecomposition3 (m %* toMatrix33 v)
    -- last bit: make sure    s1 >= s2 >= 0
    fixSigns :: Vector t 3 -> Quater t -> (Vector t 3, Quater t)
    fixSigns (DF3 s1 s2 s3) q@(Quater a b c d) = case (s1 >= 0, s2 >= 0) of
      (True , True ) -> (mk3 s1 s2 s3, q)
      (False, True ) -> (mk3 (negate s1) s2 (negate s3), Quater (-c)  d   a  (-b))
      (True , False) -> (mk3 s1 (negate s2) (negate s3), Quater   d   c (-b) (-a))
      (False, False) -> (mk3 (negate s1) (negate s2) s3, Quater   b (-a)  d  (-c))
    -- one more thing:
    --   the singular values are ordered, but may have small errors;
    --   as a result adjacent values may seem to be out of order by a very small number
    mk3 :: Scalar t -> Scalar t -> Scalar t -> Vector t 3
    mk3 s1 s2 s3' = case (s1 >= s2, s1 >= abs s3, s2 >= abs s3) of
        (True , True , True ) -> DF3 s1 s2     s3' -- s1 >= s2 >= s3
        (True , True , False) -> DF3 s1 s3 (cs s2) -- s1 >= s3 >  s2
        (True , False, _    ) -> DF3 s3 s1 (cs s2) -- s3 >  s1 >= s2
        (False, True , True ) -> DF3 s2 s1     s3' -- s2 >  s1 >= s3
        (False, _    , False) -> DF3 s3 s2 (cs s1) -- s3 >  s2 >  s1
        (False, False, True ) -> DF3 s2 s3 (cs s1) -- s2 >= s3 >  s1
      where
        s3 = abs s3'
        cs = negateUnless (s3' >= 0)

-- | Approximate values for cos (a/2) and sin (a/2) of a Givens rotation for
--    a 2x2 symmetric matrix. (Algorithm 2)
jacobiGivensQ :: forall t . RealFloatExtras t => t -> t -> t -> (t, t)
jacobiGivensQ aii aij ajj
    | g*sh*sh < ch*ch = (w * ch, w * sh)
    | otherwise       = (c', s')
  where
    ch = 2 * (aii-ajj)
    sh = aij
    w = recip $ hypot ch sh
    g  = 5.82842712474619 :: t  -- 3 + sqrt 8
    c' = 0.9238795325112867 :: t -- cos (pi/8)
    s' = 0.3826834323650898 :: t -- sin (pi/8)

-- | A quaternion for a QR Givens iteration
qrGivensQ :: forall t . RealFloatExtras t => t -> t -> (t, t)
qrGivensQ a1 a2
    | a1 < 0    = (sh * w, ch * w)
    | otherwise = (ch * w, sh * w)
  where
    rho2 = a1*a1 + a2*a2
    sh = if rho2 > M_EPS then a2 else 0
    ch = abs a1 + sqrt (max rho2 M_EPS)
    w = recip $ hypot ch sh -- TODO: consider something like a hypot

-- | One iteration of the Jacobi algorithm on a symmetric 3x3 matrix
--
--   The three words arguments are indices:
--     0 <= i /= j /= k <= 2
jacobiEigen3Iteration :: (Quaternion t, RealFloatExtras t)
                     => Int -> Int -> Int
                     -> STDataFrame s t '[3,3]
                     -> ST s (Quater t)
jacobiEigen3Iteration i j k sPtr = do
    sii <- readDataFrameOff sPtr ii
    sij <- readDataFrameOff sPtr ij
    sjj <- readDataFrameOff sPtr jj
    sik <- readDataFrameOff sPtr ik
    sjk <- readDataFrameOff sPtr jk
    -- Coefficients for a quaternion corresponding to a Givens rotation
    let (ch, sh) = jacobiGivensQ sii sij sjj
        a = ch*ch - sh*sh
        b = 2 * sh*ch
        aa = a * a
        ab = a * b
        bb = b * b
    -- update the matrix
    writeDataFrameOff sPtr ii $
      aa * sii + 2 * ab * sij + bb * sjj
    writeDataFrameOff sPtr ij $
      ab * (sjj - sii) + (aa - bb) * sij
    writeDataFrameOff sPtr jj $
      bb * sii - 2 * ab * sij + aa * sjj
    writeDataFrameOff sPtr ik $ a * sik + b * sjk
    writeDataFrameOff sPtr jk $ a * sjk - b * sik

    -- write the quaternion
    qPtr <- unsafeThawDataFrame 0
    writeDataFrameOff qPtr k (negate sh)
    writeDataFrameOff qPtr 3 ch
    fromVec4 <$> unsafeFreezeDataFrame qPtr
  where
    ii = i*3 + i
    ij = if i < j then i*3 + j else j*3 + i
    jj = j*3 + j
    ik = if i < k then i*3 + k else k*3 + i
    jk = if j < k then j*3 + k else k*3 + j

-- | Total number of the Givens rotations during the Jacobi eigendecomposition
--   part of the 3x3 SVD equals eigenItersX3*3.
--   Value `eigenItersX3 = 6` corresponds to 18 iterations and gives a good precision.
eigenItersX3 :: Int
eigenItersX3 = 12

-- | Run a few iterations of the Jacobi algorithm on a real-valued 3x3 symmetric matrix.
--   The eigenvectors basis of such matrix is orthogonal, and can be represented as
--   a quaternion.
jacobiEigenQ :: forall t
              . (Quaternion t, RealFloatExtras t)
             => Matrix t 3 3 -> Quater t
jacobiEigenQ m = runST $ do
    mPtr <- thawDataFrame m
    q  <- go eigenItersX3 mPtr 1
    s1 <- readDataFrameOff mPtr 0
    s2 <- readDataFrameOff mPtr 4
    s3 <- readDataFrameOff mPtr 8
    return $ sortQ s1 s2 s3 * q
  where
    go :: Int -> STDataFrame s t '[3,3] -> Quater t -> ST s (Quater t)
    go 0 _ q = pure q

    -- -- primitive cyclic iteration;
    -- --   fast, but the convergence is not perfect
    -- --
    -- -- set eigenItersX3 = 6 for good precision
    -- go n p q = do
    --   q1 <- jacobiEigen3Iteration 0 1 2 p
    --   q2 <- jacobiEigen3Iteration 1 2 0 p
    --   q3 <- jacobiEigen3Iteration 2 0 1 p
    --   go (n - 1) p (q3 * q2 * q1 * q)

    -- Pick the largest element on lower triangle;
    --   slow because of branching, but has a better convergence
    --
    --  set eigenItersX3 = 12 for good precision
    --    (slightly faster than the cyclic version with -O0)
    go n p q = do
      a10 <- abs <$> readDataFrameOff p 1
      a20 <- abs <$> readDataFrameOff p 2
      a21 <- abs <$> readDataFrameOff p 5
      q' <- jiter n p a10 a20 a21
      go (n - 1) p (q' * q)
    jiter :: Int -> STDataFrame s t '[3,3]
          -> Scalar t -> Scalar t -> Scalar t -> ST s (Quater t)
    jiter n p a10 a20 a21
      | gt2 a10 a20 a21
        = jacobiEigen3Iteration 0 1 2 p
      | gt2 a20 a10 a21
        = jacobiEigen3Iteration 2 0 1 p
      | gt2 a21 a10 a20
        = jacobiEigen3Iteration 1 2 0 p
      | otherwise
        = case mod n 3 of
            0 -> jacobiEigen3Iteration 0 1 2 p
            1 -> jacobiEigen3Iteration 2 0 1 p
            _ -> jacobiEigen3Iteration 1 2 0 p
    gt2 :: Scalar t -> Scalar t -> Scalar t -> Bool
    gt2 a b c = case compare a b of
                  GT -> a >= c
                  EQ -> a >  c
                  LT -> False

    -- Make such a quaternion that rotates the matrix so that:
    -- abs s1 >= abs s2 >= abs s3
    -- Note, the corresponding singular values may be negative, which must be
    -- taken into account later.
    sortQ :: Scalar t -> Scalar t -> Scalar t -> Quater t
    sortQ s1 s2 s3 = sortQ' (s1 >= s2) (s1 >= s3) (s2 >= s3)
    sortQ' :: Bool -> Bool -> Bool -> Quater t
    sortQ' True  True  True  = Quater 0 0 0 1                    -- s1 >= s2 >= s3
    sortQ' True  True  False = Quater M_SQRT1_2 0 0 (-M_SQRT1_2) -- s1 >= s3 >  s2
    sortQ' True  False _     = Quater 0.5 0.5 0.5 0.5            -- s3 >  s1 >= s2
    sortQ' False True  True  = Quater 0 0 M_SQRT1_2 (-M_SQRT1_2) -- s2 >  s1 >= s3
    sortQ' False _     False = Quater 0 M_SQRT1_2 0 (-M_SQRT1_2) -- s3 >  s2 >  s1
    sortQ' False False True  = Quater 0.5 0.5 0.5 (-0.5)         -- s2 >= s3 >  s1

-- | One Givens rotation for a QR algorithm on a 3x3 matrix
--
--   The three words arguments are indices:
--     0 <= i /= j /= k <= 2
--
--     if i < j then the eigen values are already sorted!
qrDecomp3Iteration :: (Quaternion t, RealFloatExtras t)
                   => Int -> Int -> Int
                   -> STDataFrame s t '[3,3]
                   -> ST s (Quater t)
qrDecomp3Iteration i j k sPtr = do
    sii <- readDataFrameOff sPtr ii
    sij <- readDataFrameOff sPtr ij
    sji <- readDataFrameOff sPtr ji
    sjj <- readDataFrameOff sPtr jj
    sik <- readDataFrameOff sPtr ik
    sjk <- readDataFrameOff sPtr jk
    -- Coefficients for a quaternion corresponding to a Givens rotation
    let (ch, sh) = qrGivensQ sii sji
        a = ch*ch - sh*sh
        b = 2 * sh*ch
    -- update the matrix
    writeDataFrameOff sPtr ii $ a * sii + b * sji
    writeDataFrameOff sPtr ij $ a * sij + b * sjj
    writeDataFrameOff sPtr ik $ a * sik + b * sjk
    writeDataFrameOff sPtr ji 0 --  a * sji - b * sii
    writeDataFrameOff sPtr jj $ a * sjj - b * sij
    writeDataFrameOff sPtr jk $ a * sjk - b * sik

    -- write the quaternion
    qPtr <- unsafeThawDataFrame 0
    writeDataFrameOff qPtr k (negateUnless leftTriple sh)
    writeDataFrameOff qPtr 3 ch
    fromVec4 <$> unsafeFreezeDataFrame qPtr
  where
    leftTriple = (j - i) /= 1 && (k - j) /= 1
    i3 = i*3
    j3 = j*3
    ii = i3 + i
    ij = i3 + j
    ik = i3 + k
    ji = j3 + i
    jj = j3 + j
    jk = j3 + k

-- | Run QR decomposition in context of 3x3 svd: AV = US = QR
--   The input here is matrix AV
--   The R upper-triangular matrix here is in fact a diagonal matrix Sigma;
--   The Q orthogonal matrix is matrix U in the svd decomposition,
--     represented here as a quaternion.
qrDecomposition3 :: (Quaternion t, RealFloatExtras t)
                 => Matrix t 3 3 -> (Vector t 3, Quater t)
qrDecomposition3 m = runST $ do
    mPtr <- thawDataFrame m
    q1 <- qrDecomp3Iteration 0 1 2 mPtr
    q2 <- qrDecomp3Iteration 0 2 1 mPtr
    q3 <- qrDecomp3Iteration 1 2 0 mPtr
    sig0 <- readDataFrameOff mPtr 0
    sig1 <- readDataFrameOff mPtr 4
    sig2 <- readDataFrameOff mPtr 8
    return (DF3 sig0 sig1 sig2, q3 * q2 * q1)

instance RealFloatExtras t => MatrixSVD t 1 1 where
    svd = svd1

instance RealFloatExtras t => MatrixSVD t 2 2 where
    svd = svd2

instance (RealFloatExtras t, Quaternion t) => MatrixSVD t 3 3 where
    svd = svd3

instance {-# INCOHERENT #-}
         ( RealFloatExtras t, KnownDim n, KnownDim m)
         => MatrixSVD t n m where
    svd a = runST $ do
      D <- pure dnm
      Dict <- pure $ minIsSmaller dn dm -- GHC is not convinced :(
      alphas <- unsafeThawDataFrame bdAlpha
      betas <- unsafeThawDataFrame bdBeta
      uPtr <- unsafeThawDataFrame bdU
      vPtr <- unsafeThawDataFrame bdV

      -- remove last beta if m > n
      bLast <- readDataFrameOff betas nm1
      when (abs bLast > M_EPS) $
        svdGolubKahanZeroCol alphas betas vPtr nm1

      -- main routine for a bidiagonal matrix
      let maxIter = 3*nm -- number of tries
      withinIters <- svdBidiagonalInplace alphas betas uPtr vPtr nm maxIter
      unless withinIters . tooManyIterations
        $ "SVD - Givens rotation sweeps for a bidiagonal matrix ("
           ++ show maxIter ++ " sweeps max)."

      -- sort singular values
      sUnsorted <- unsafeFreezeDataFrame alphas
      let sSorted :: Vector (Tuple '[t, Word]) (Min n m)
          sSorted = sortBy (\(S (x :! _)) (S (y :! _)) -> compare y x)
                  $ iwmap @_ @_ @'[] (\(Idx i :* U) (S x) -> S (abs x :! i :! U) ) sUnsorted
          svdS = ewmap @t @_ @'[] (\(S (x :! _)) -> S x) sSorted
          perm = ewmap @Word @_ @'[] (\(S (_ :! i :! U)) -> S i) sSorted
          pCount =
             if nm < 2
             then 0 :: Word
             else foldl (\s (i, j) -> if perm!i > perm!j then succ s else s)
                        0 [(i, j) | i <- [0..nm2w], j <- [i+1..nm2w+1]]
          pPositive = even pCount

      -- alphas and svdS are now out of sync, but that is not a problem

      -- make sure det U == 1
      when ((bdUDet < 0) == pPositive) $ do
        readDataFrameOff alphas 0 >>= writeDataFrameOff alphas 0 . negate
        forM_ [0..n - 1] $ \i ->
          readDataFrameOff uPtr (i*n) >>= writeDataFrameOff uPtr (i*n) . negate

      -- negate negative singular values
      forM_ [0..nm1] $ \i -> do
        s <- readDataFrameOff alphas i
        when (s < 0) $ do
          writeDataFrameOff alphas i $ negate s
          forM_ [0..m - 1] $ \j ->
            readDataFrameOff vPtr (j*m + i)
              >>= writeDataFrameOff vPtr (j*m + i) . negate

      -- apply permutations if necessary
      if pCount == 0
      then do
        svdU <- unsafeFreezeDataFrame uPtr
        svdV <- unsafeFreezeDataFrame vPtr
        return SVD {..}
      else do
        svdU' <- unsafeFreezeDataFrame uPtr
        svdV' <- unsafeFreezeDataFrame vPtr
        let svdU = iwgen @_ @_ @'[] $ \(i :* Idx j :* U) ->
              if j >= dimVal dnm
              then index (i :* Idx j :* U) svdU'
              else index (i :* Idx (unScalar $ perm!j) :* U) svdU'
            svdV = iwgen @_ @_ @'[] $ \(i :* Idx j :* U) ->
              if j >= dimVal dnm
              then index (i :* Idx j :* U) svdV'
              else index (i :* Idx (unScalar $ perm!j) :* U) svdV'
        return SVD {..}
      where
        n = fromIntegral $ dimVal dn :: Int
        m = fromIntegral $ dimVal dm :: Int
        dn = dim @n
        dm = dim @m
        dnm = minDim dn dm
        nm1 = nm - 1
        nm = fromIntegral (dimVal dnm) :: Int
        nm2w = fromIntegral (max (nm - 2) 0) :: Word
        -- compute the bidiagonal form b first, solve svd for b.
        BiDiag {..} = bidiagonalHouseholder a

{- Compute svd for a square bidiagonal matrix inplace
   \( B = U S V^\intercal \)

@
  B = | a1 b1 0     ... 0 |
      | 0  a2 b2 0  ... 0 |
      | 0  0 a3 b3  ... 0 |
      | ................. |
      | 0  0  ... an1 bn1 |
      | 0  0  ...  0  an  | bn? (in case if n > m)
@
 -}
svdBidiagonalInplace ::
       forall (s :: Type) (t :: Type) (n :: Nat) (m :: Nat) (nm :: Nat)
     . ( IterativeMethod, RealFloatExtras t
       , KnownDim n, KnownDim m, KnownDim nm, nm ~ Min n m)
    => STDataFrame s t '[nm] -- ^ the main diagonal of B and then the singular values.
    -> STDataFrame s t '[nm] -- ^ first upper diagonal of B
    -> STDataFrame s t '[n,n] -- ^ U
    -> STDataFrame s t '[m,m] -- ^ V
    -> Int -- ^ 0 < q <= nm -- size of a reduced matrix, such that leftover is diagonal
    -> Int -- iters
    -> ST s Bool -- whether the algorithm succeeds within the maxIter
svdBidiagonalInplace _ _ _ _ 0 _ = pure True
svdBidiagonalInplace _ _ _ _ 1 _ = pure True
svdBidiagonalInplace _ _ _ _ _ 0 = pure False
svdBidiagonalInplace aPtr bPtr uPtr vPtr q' iter = do
    Dict <- pure $ minIsSmaller (dim @n) (dim @m)
    (p, q) <- findCounters q'
    if (q /= 0)
    then do
      findZeroDiagonal p q >>= \case
        Just k
          | k == q-1  -> svdGolubKahanZeroCol aPtr bPtr vPtr (k-1)
          | otherwise -> svdGolubKahanZeroRow aPtr bPtr uPtr k
        Nothing -> svdGolubKahanStep aPtr bPtr uPtr vPtr p q
      svdBidiagonalInplace aPtr bPtr uPtr vPtr q (iter - 1)
    else return True
  where
    -- nm = fromIntegral $ dimVal' @nm :: Int

    -- Check if off-diagonal elements are close to zero and nullify them
    -- if they are small along the way.
    -- And find such indices p and q that satisfy condition in alg. 8.6.2
    -- on p. 492. of "Matrix Computations " (4-th edition).
    -- Except these are inverted:
    --   p -- is the starting index of B22 (last submatrix with non-zero superdiagonal)
    --   q -- is the starting index of B33 (diagonal submatrix)
    --
    -- that is, p and q determine the index and the size of next work piece.
    findCounters :: Int -> ST s (Int, Int)
    findCounters = goQ
      where
        checkEps :: Int -> ST s Bool
        checkEps k = do
          b <- abs <$> readDataFrameOff bPtr (k-1)
          if b == 0
          then return True
          else do
            a1 <- abs <$> readDataFrameOff aPtr (k-1)
            a2 <- abs <$> readDataFrameOff aPtr  k
            if b <= M_EPS * (max (a1 + a2) 1)
            then True <$ writeDataFrameOff bPtr (k-1) 0
            else return False
        goQ :: Int -> ST s (Int, Int)
        goQ 0 = pure (0, 0) -- guard against calling with q == 0
        goQ 1 = pure (0, 0) -- 1x1 matrix is always diagonal
        goQ k = checkEps (k-1) >>= \case
          True  -> goQ (k-1)
          False -> flip (,) k <$> goP (k-2)
        goP :: Int -> ST s Int
        goP 0 = pure 0
        goP k = checkEps k >>= \case
          True  -> return k
          False -> goP (k-1)

    -- For indices p and q (p < q), find the biggest index (< q) such that
    --  a[p] == 0
    findZeroDiagonal :: Int -> Int -> ST s (Maybe Int)
    findZeroDiagonal p q
      | k < p     = pure Nothing
      | otherwise = do
        ak <- readDataFrameOff aPtr k
        if ak == 0
        then pure $ Just k
        else if abs ak <= M_EPS
             then Just k <$ writeDataFrameOff aPtr k 0
             else findZeroDiagonal p k
      where
        k = q - 1

-- | Apply a series of column transformations to make b[k] (and whole column k+1) zero
--    (page 491, 1st paragraph) when a[k+1] == 0.
--   To make this element zero, I apply a series of Givens transforms on columns
--   (multiply on the right).
--
--   Prerequisites:
--     * a[k+1] == 0
--     * 0 <= k < min n (m-1)
--   Invariants:
--     * matrix \(B :: n \times m \) is bidiagonal, represented by two diagonals;
--     * matrix V is orthogonal
--     * \( B = A V^\intercal \), where \(A :: n \times m\) is an implicit original matrix
--   Results:
--     * Same bidiagonal matrix with b[k] == 0; i.e. (k+1)-th column is zero.
--     * matrix V is updated (multiplied on the right)
--
--   NB: All changes are made inplace.
--
svdGolubKahanZeroCol ::
       forall (s :: Type) (t :: Type) (n :: Nat) (m :: Nat)
     . (RealFloatExtras t, KnownDim n, KnownDim m, n <= m)
    => STDataFrame s t '[n] -- ^ the main diagonal of \(B\)
    -> STDataFrame s t '[n] -- ^ first upper diagonal of \(B\)
    -> STDataFrame s t '[m,m] -- ^ \(V\)
    -> Int -- ^ 0 <= k < min (n+1) m
    -> ST s ()
svdGolubKahanZeroCol aPtr bPtr vPtr k
  | k < 0 || k >= lim = error $ unwords
      [ "svdGolubKahanZeroCol: k =", show k
      , "is outside of a valid range 0 <= k <", show lim]
    -- this trick is to convince GHC that constraint (n <= m) is not redundant
  | Dict <- Dict @(n <= m) = do
    b <- readDataFrameOff bPtr k
    writeDataFrameOff bPtr k 0
    foldM_ goGivens b [k, k-1 .. 0]
  where
    n = fromIntegral $ dimVal' @n :: Int
    m = fromIntegral $ dimVal' @m :: Int
    lim = min n (m-1)
    goGivens :: Scalar t -> Int -> ST s (Scalar t)
    goGivens 0 _ = return 0 -- non-diagonal element is nullified prematurely
    goGivens b i = do
      ai <- readDataFrameOff aPtr i
      let rab = recip $ hypot b ai
          c = ai*rab
          s = b *rab
      updateGivensMat vPtr i (k+1) c s
      writeDataFrameOff aPtr i $ ai*c + b*s -- B[i,i]
      if i == 0
      then return 0
      else do
        bi1 <- readDataFrameOff bPtr (i - 1)  -- B[i,i-1]
        writeDataFrameOff bPtr (i - 1) $ bi1 * c
        return $ negate (bi1 * s)

-- | Apply a series of row transformations to make b[k] (and whole column k) zero
--    (page 490, last paragraph) when a[k] == 0.
--   To make this element zero, I apply a series of Givens transforms on rows
--   (multiply on the left).
--
--   Prerequisites:
--     * a[k] == 0
--     * 0 <= k < n - 1
--   Invariants:
--     * matrix \(B :: m \times n \) is bidiagonal, represented by two diagonals;
--     * matrix U is orthogonal
--     * \( B = U A \), where \(A :: m \times n \) is an implicit original matrix
--   Results:
--     * Same bidiagonal matrix with b[k] == 0; i.e. k-th column is zero.
--     * matrix U is updated (multiplied on the right)
--
--   NB: All changes are made inplace.
--
svdGolubKahanZeroRow ::
       forall (s :: Type) (t :: Type) (n :: Nat) (m :: Nat)
     . (RealFloatExtras t, KnownDim n, KnownDim m, n <= m)
    => STDataFrame s t '[n] -- ^ the main diagonal of B
    -> STDataFrame s t '[n] -- ^ first upper diagonal of B
    -> STDataFrame s t '[m,m] -- ^ U
    -> Int -- ^ 0 <= k < n - 1
    -> ST s ()
svdGolubKahanZeroRow aPtr bPtr uPtr k
  | k < 0 || k >= n1 = error $ unwords
      [ "svdGolubKahanZeroRow: k =", show k
      , "is outside of a valid range 0 <= k <", show n1]
    -- this trick is to convince GHC that constraint (n <= m) is not redundant
  | Dict <- Dict @(n <= m) = do
    b <- readDataFrameOff bPtr k
    writeDataFrameOff bPtr k 0
    foldM_ goGivens b [k+1..n1]
  where
    n = fromIntegral $ dimVal' @n :: Int
    n1 = n - 1
    goGivens :: Scalar t -> Int -> ST s (Scalar t)
    goGivens 0 _ = return 0 -- non-diagonal element is nullified prematurely
    goGivens b j = do
      aj <- readDataFrameOff aPtr j
      bj <- readDataFrameOff bPtr j
      let rab = recip $ hypot b aj
          c = aj*rab
          s =  b*rab
      updateGivensMat uPtr k j c (negate s)
      writeDataFrameOff aPtr j $ b*s + aj*c
      writeDataFrameOff bPtr j $ bj*c
      return $ negate (bj * s)

-- | A Golub-Kahan bidiagonal SVD step on an unreduced matrix
svdGolubKahanStep ::
       forall (s :: Type) (t :: Type) (n :: Nat) (m :: Nat) (nm :: Nat)
     . ( RealFloatExtras t
       , KnownDim n, KnownDim m, KnownDim nm, nm ~ Min n m)
    => STDataFrame s t '[nm] -- ^ the main diagonal of B and then the singular values.
    -> STDataFrame s t '[nm] -- ^ first upper diagonal of B
    -> STDataFrame s t '[n,n] -- ^ U
    -> STDataFrame s t '[m,m] -- ^ V
    -> Int -- ^ p : 0 <= p < q <= nm; p <= q - 2
    -> Int -- ^ q : 0 <= p < q <= nm; p <= q - 2
    -> ST s ()
svdGolubKahanStep aPtr bPtr uPtr vPtr p q
  | p > q - 2 || p < 0 || q > nm
    = error $ unwords
        [ "svdGolubKahanStep: p =", show p, "and q =", show q
        , "do not satisfy p <= q - 2 or 0 <= p < q <=", show nm]
  | Dict <- Dict @(nm ~ Min n m) = do
    (y,z) <- getWilkinsonShiftYZ
    goGivens2 y z p
  where
    nm = fromIntegral $ dimVal' @nm :: Int

    -- get initial values for one recursion sweep.
    -- Note, input must satisfy: q >= p+2
    getWilkinsonShiftYZ :: ST s (Scalar t, Scalar t)
    getWilkinsonShiftYZ  = do
      a1 <- readDataFrameOff aPtr p
      b1 <- readDataFrameOff bPtr p
      am <- readDataFrameOff aPtr (q-2)
      an <- readDataFrameOff aPtr (q-1)
      bm <- if q >= p + 3
            then readDataFrameOff bPtr (q-3)
            else pure 0
      bn <- readDataFrameOff bPtr (q-2)
      let t11 = a1*a1
          t12 = a1*b1
          tmm = am*am + bm*bm
          tnn = an*an + bn*bn
          tnm = am*bn
          d   = 0.5*(tmm - tnn)
          mu  = tnn + d - negateUnless (d >= 0) (hypot d tnm)
      return (t11 - mu, t12)

    -- yv = b[k-1]; zv = B[k-1,k+1] -- to be eliminated by 1st Givens r
    -- yu = a[k];   zu = B[k+1,k-1] -- to be eliminated by 2nd Givens r
    goGivens2 :: Scalar t -> Scalar t -> Int -> ST s ()
    goGivens2 yv zv k = do
          a1 <- readDataFrameOff aPtr k     -- B[k,k]
          a2 <- readDataFrameOff aPtr (k+1) -- B[k+1,k+1]
          b1 <- readDataFrameOff bPtr k     -- B[k,k+1]
          let a1' = a1*cv + b1*sv  -- B[k,k] == yu
              a2' = a2*cv          -- B[k+1,k+1]
              b0' = yv*cv + zv*sv  -- B[k-1,k]
              b1' = b1*cv - a1*sv  -- B[k,k+1]
              yu  = a1'            -- B[k,k]
              zu  = a2*sv          -- B[k+1,k]
              ryzu = recip $ hypot yu zu
              cu = yu * ryzu
              su = zu * ryzu
              a1'' = yu *cu + zu *su
              a2'' = a2'*cu - b1'*su
              b1'' = b1'*cu + a2'*su
          updateGivensMat vPtr k (k+1) cv sv
          updateGivensMat uPtr k (k+1) cu su

          when (k > p) $ writeDataFrameOff bPtr (k-1) b0'
          writeDataFrameOff bPtr k b1''
          writeDataFrameOff aPtr k a1''
          writeDataFrameOff aPtr (k+1) a2''
          when (k < q - 2) $ do
            b2 <- readDataFrameOff bPtr (k+1) -- B[k+1,k+2]
            let b2'' = b2*cu
                zvn  = b2*su
            writeDataFrameOff bPtr (k+1) b2''
            goGivens2 b1'' zvn (k+1)
        where
          ryzv = recip $ hypot yv zv
          cv = yv * ryzv
          sv = zv * ryzv

-- | Update a transformation matrix with a Givens transform (on the right)
updateGivensMat ::
       forall (s :: Type) (t :: Type) (n :: Nat)
     . (PrimBytes t, Num t, KnownDim n)
    => STDataFrame s t '[n,n]
    -> Int -> Int
    -> Scalar t -> Scalar t -> ST s ()
updateGivensMat p i j c s = forM_ [0..n-1] $ \k -> do
    let nk = n*k
        ioff = nk + i
        joff = nk + j
    uki <- readDataFrameOff p ioff
    ukj <- readDataFrameOff p joff
    writeDataFrameOff p ioff $ uki*c + ukj*s
    writeDataFrameOff p joff $ ukj*c - uki*s
  where
    n = fromIntegral $ dimVal' @n :: Int

minIsSmaller :: forall (n :: Nat) (m :: Nat)
              . Dim n -> Dim m -> Dict (Min n m <= n, Min n m <= m)
minIsSmaller dn dm
  | Just Dict <- lessOrEqDim dnm dn
  , Just Dict <- lessOrEqDim dnm dm
    = Dict
  | otherwise
    = error "minIsSmaller: impossible type-level comparison"
  where
    dnm = minDim dn dm


module Gates where

import           Protolude

import           Control.Exception (assert)
import           Data.Complex      (Complex (..), magnitude)
import           Data.List         ((!!))
import qualified Data.Matrix       as M
import qualified Data.Vector       as V
import           System.Random     (randomIO)

import           Instruction       (Operation, Qubit, Qubits)
import           State             (QuantumState)

-- Dimension in qubits by taking the log base 2
dimensionQubits :: Integral a => a -> Int
dimensionQubits = ceiling . logBase 2 . fromIntegral

i :: Operation
i = M.identity 2

x :: Operation
x = M.fromLists [[0, 1], [1, 0]]

h :: Operation
h = M.scaleMatrix (1 / sqrt 2) $ M.fromLists [[1, 1], [1, -1]]

cnot :: Operation
cnot = M.fromLists [[1, 0, 0, 0], [0, 1, 0, 0], [0, 0, 0, 1], [0, 0, 1, 0]]

swap :: Operation
swap = M.fromLists [[1, 0, 0, 0], [0, 0, 1, 0], [0, 1, 0, 0], [0, 0, 0, 1]]

-- Apply an operation to a quantum state by matrix multiplication
applyOperation :: Operation -> QuantumState -> QuantumState
applyOperation g s = M.getCol 1 $ g * M.colVector s

-- Compose operators by doing matrix multiplication
-- NOTE: (*) operator chosen over multStd for possible speed improvements for larger matrices
composeOperations :: Operation -> Operation -> Operation
composeOperations = (*)

-- Sample quantum state, randomly selecting a state, weighted by |phi_i|^2
sample :: MonadIO m => QuantumState -> m Int
sample s = do
  r <- liftIO $ randomIO
  return $ sampleIter r 0
  where
    sampleIter :: Double -> Int -> Int
    sampleIter r t =
      let newR = r - (magnitude (s V.! t)) ** 2
       in if newR < 0
            then t
            else sampleIter newR (t + 1)

-- Collapse the quantum state into the state at the given index
collapse :: QuantumState -> Int -> QuantumState
collapse s i =
  V.generate
    (V.length s)
    (\x ->
       if x == i
         then (1 :+ 0)
         else (0 :+ 0))

-- Observe the quantum state and collapse into that state
observe :: MonadIO m => QuantumState -> m QuantumState
observe s = do
  i <- sample s
  return $ collapse s i

-- Oh, to have dependent types
kroneckerMultiply :: (Num a) => M.Matrix a -> M.Matrix a -> M.Matrix a
kroneckerMultiply a b = M.flatten $ M.mapPos (\_ -> flip M.scaleMatrix b) a

-- TODO: re-eval constraints
kroneckerExpt :: (Num a) => M.Matrix a -> Int -> M.Matrix a
kroneckerExpt a n
  | n < 1 = M.fromLists [[1]]
  | n == 1 = a
  | otherwise = kroneckerMultiply (kroneckerExpt a (n - 1)) a

-- Lift an operation by padding with identity
liftOperation :: Operation -> Qubit -> Int -> Operation
liftOperation a index n =
  assert (index < n) kroneckerMultiply left (kroneckerMultiply a right)
  where
    left = kroneckerExpt i (n - index - (dimensionQubits $ M.nrows a))
    right = kroneckerExpt i index

-- Calculate needed transpositions from given permutation
permutationToTranspositions :: [Int] -> [(Int, Int)]
permutationToTranspositions p = permutationToTranspositionsIter 0
  where
    permutationToTranspositionsIter dest
      | dest == (length p) = []
      | otherwise =
        (case compare src dest of
           LT -> [(src, dest)]
           GT -> [(dest, src)]
           _  -> []) ++
        permutationToTranspositionsIter (dest + 1)
      where
        src = nextSrc (p !! dest) dest
    nextSrc s d
      | s < d = nextSrc (p !! s) d
      | otherwise = s

-- Convert transpositions to list of adjacent transpositions
transpositionToAdjacentTranspositions :: [(Int, Int)] -> [Int]
transpositionToAdjacentTranspositions ts = foldMap expandCons ts
  where
    expandCons (x, y)
      | y - x == 1 = [x]
      | otherwise =
        let trans = [x .. y - 1]
         in trans ++ (reverse (initSafe trans))

-- Apply 1 qubit gate
apply1QGate :: QuantumState -> Operation -> Qubit -> QuantumState
apply1QGate s a q =
  applyOperation (liftOperation a q (dimensionQubits (length s))) s

-- Apply multiple qubit gate
applynQGate :: QuantumState -> Operation -> Qubits -> QuantumState
applynQGate s a qubits = do
  applyOperation upq s
  where
    n :: Int
    n = dimensionQubits (length s)
    transpositionsToOperation :: [Int] -> Operation
    transpositionsToOperation ts =
      foldr composeOperations (M.identity (2 ^ n)) $
      map (\x -> liftOperation Gates.swap x n) ts
    u01 :: Operation
    u01 = liftOperation a 0 n
    fromSpace :: [Int]
    fromSpace =
      (reverse qubits) ++ (filter (\x -> not (elem x qubits)) [0 .. n - 1])
    trans :: [Int]
    trans =
      transpositionToAdjacentTranspositions . permutationToTranspositions $
      fromSpace
    toFrom :: Operation
    toFrom = transpositionsToOperation trans
    fromTo :: Operation
    fromTo = transpositionsToOperation (reverse trans)
    upq :: Operation
    upq = composeOperations toFrom (composeOperations u01 fromTo)

-- Apply gate, picking apply1QGate or applynQGate based on size
applyGate :: QuantumState -> Operation -> [Int] -> QuantumState
applyGate s a qubits =
  assert
    (length qubits == dimensionQubits (M.nrows a))
    (if length qubits == 1
       then apply1QGate s a (qubits !! 0)
       else applynQGate s a qubits)


{-
  An attempt to implement a basic case of Cooley-Tukey FFT algorithm.
 -}
{-# LANGUAGE ViewPatterns #-}
module Main
  ( main
  ) where

import Control.DeepSeq
import Control.Monad
import Control.Monad.State.Strict
import Data.Bits
import Data.Complex
import Data.Time.Clock
import Statistics.Sample
import System.Random
import Text.Printf

import qualified Data.Vector as V
import qualified Data.Vector.Mutable as VM

type Cpx = Complex Double

-- e^((c pi i) * (k / N))
expAuxCommon :: Int -> Int -> Int -> Cpx
expAuxCommon c k n = cis theta
  where
    theta = pi * fromIntegral (c * k) / fromIntegral n

-- Computes e^((-2 pi i) * (k / N))
expFft :: Int -> Int -> Cpx
expFft = expAuxCommon (-2)

-- Computes e^((2 pi i) * (k / N))
expIfft :: Int -> Int -> Cpx
expIfft = expAuxCommon 2

splitByParity :: V.Vector Cpx -> (V.Vector Cpx, V.Vector Cpx)
splitByParity vs = (evens, odds)
  where
    l = V.length vs
    oddCount = (l + 1) `quot` 2
    evenCount = l - oddCount
    evens =
      V.fromListN evenCount $ (vs V.!) <$> [0,2..]
    odds =
      V.fromListN oddCount $ (vs V.!) <$> [1,3..]

-- TODO: For now we assume that vector length is a power of two,
-- this can be easily extended to any length, by pretending out-of-range
-- values are 0.
gDitFft :: (Int -> Int -> Cpx) -> V.Vector Cpx -> V.Vector Cpx
gDitFft expF = fix $
  \impl vs ->
    let l = V.length vs
    in if l <= 1
      then vs
      else V.create $ do
        let (impl -> es, impl -> os) = splitByParity vs
            hf = l `quot` 2
        vec <- VM.unsafeNew l
        forM_ [0 .. hf - 1] $ \k -> do
          let a = es V.! k
              b = expF k l * os V.! k
          VM.unsafeWrite vec k (a + b)
          VM.unsafeWrite vec (hf + k) (a - b)
        pure vec

ditFft :: V.Vector Cpx -> V.Vector Cpx
ditFft = gDitFft expFft . rightPadZeros

iditFft :: V.Vector Cpx -> V.Vector Cpx
iditFft (rightPadZeros -> vs) = V.map (/ fromIntegral l) (iditFftAux vs)
  where
    l = V.length vs
    iditFftAux :: V.Vector Cpx -> V.Vector Cpx
    iditFftAux = gDitFft expIfft

{-
  right padding zeros in the end of a matrix to make the length of the vector
  a power of 2.

  special case: returns empty vector when the input is empty.
 -}
rightPadZeros :: V.Vector Cpx -> V.Vector Cpx
rightPadZeros vs
    | l <= 1 = vs
    | isPowerOf2 = vs
    | otherwise =
        let closestPowOf2Gt =
              {-
                Find the closest power of two that is greater or equal to input v.
                only works when v > 1 and v is not a power of two.
               -}
              unsafeShiftL 1 (finiteBitSize l - countLeadingZeros l)
        in vs <> V.fromListN (closestPowOf2Gt - l) (repeat 0)
  where
    l = V.length vs
    isPowerOf2 =
      countLeadingZeros l + countTrailingZeros l + 1 == finiteBitSize l

evaluateOnVector :: V.Vector Cpx -> IO ()
evaluateOnVector vs = do
  let result = iditFft (ditFft vs)
      diffs :: V.Vector Double
      diffs = V.zipWith (\x y -> magnitude (x - y)) vs result
  printf "StdDev: %.9f\n" (stdDev diffs)

type M = StateT StdGen IO

genRandomR :: Random a => (a, a) -> M a
genRandomR = state . randomR

genCpxVector :: M (V.Vector Cpx)
genCpxVector =
  genRandomR (1024, 8192) >>= genCpxVectorOfSize

genCpxVectorOfSize :: Int -> M (V.Vector Cpx)
genCpxVectorOfSize l = do
  let rDouble = genRandomR (-100,100)
  xs <- replicateM l $ (:+) <$> rDouble <*> rDouble
  pure $ V.fromListN l xs

evalRandomVector :: M ()
evalRandomVector = do
  vs <- genCpxVector
  liftIO $ do
    printf "Vector length: %d\n" (V.length vs)
    evaluateOnVector vs

directConvolve :: V.Vector Cpx -> V.Vector Cpx -> V.Vector Cpx
directConvolve xs ys = V.create $ do
  let lX = V.length xs
      lY = V.length ys
      lZ = lX + lY - 1
  vec <- VM.replicate lZ 0
  forM_ [0 .. lX - 1] $ \i -> do
    let x = V.unsafeIndex xs i
    forM_ [0 .. lY - 1] $ \j -> do
      let y = V.unsafeIndex ys j
          k = i + j
      v <- VM.unsafeRead vec k
      VM.unsafeWrite vec k $! v + x * y
  pure vec

fftConvolve :: V.Vector Cpx -> V.Vector Cpx -> V.Vector Cpx
fftConvolve xsPre ysPre = V.fromListN lZ (V.toList (iditFft rs))
  where
    lX = V.length xsPre
    lY = V.length ysPre
    lZ = lX + lY - 1
    -- TOOD: admittedly this is a bit wasteful...
    xs = V.fromListN lZ (V.toList xsPre <> repeat 0)
    ys = V.fromListN lZ (V.toList ysPre <> repeat 0)
    xs' = ditFft xs
    ys' = ditFft ys
    rs = V.zipWith (*) xs' ys'

evalRandomConvolve :: M ()
evalRandomConvolve = do
  xs <- genCpxVector
  ys <- genCpxVector
  let lX = V.length xs
      lY = V.length ys
      lZ = lX + lY - 1
      zsDirect = directConvolve xs ys
      zsFft = fftConvolve xs ys
      diffs :: V.Vector Double
      diffs = V.zipWith (\x y -> magnitude (x - y)) zsDirect zsFft
  liftIO $ do
    putStrLn $ "Vector lengths (xs,ys,xs * ys): " <> show (lX, lY, lZ)
    _ <- forceAndMeasure "Direct" zsDirect
    _ <- forceAndMeasure "FFT" zsFft
    printf "StdDev: %.9f\n" (stdDev diffs)

forceAndMeasure :: NFData a => String -> a -> IO a
forceAndMeasure tag r = do
  tBefore <- getCurrentTime
  tAfter <- r `deepseq` getCurrentTime
  printf "Time (%s): %.6f\n" tag ((realToFrac $ diffUTCTime tAfter tBefore) :: Double)
  pure r

main :: IO ()
main = do
  let cs = (\x -> x*2 :+ (x*2 + 1)) <$> [0..19]
      vs = V.fromList cs
  evaluateOnVector vs
  g <- newStdGen
  evalStateT (evalRandomVector >> evalRandomConvolve) g


{-|
Module      : kMeans
Description : k-Means
Copyright   : (c) Fabrício Olivetti, 2017
License     : GPL-3
Maintainer  : fabricio.olivetti@gmail.com

k-Means Clustering algorithm
-}

module KMeans where

import Data.List
import Data.Ord
--import qualified Numeric.LinearAlgebra as H

import Vector
import Dados

-- |'parseFile' parses a space separated file
-- to a list of lists of Double
parseFile :: String -> [[Double]]
parseFile file = map parseLine (lines file)
  where
    parseLine l = map toDouble (words l)
    toDouble  w = read w :: Double

firstClusters idxs points = [ points !! i | i <- idxs ]

kmeans it points clusters
  | it == 0               = clusters
  | clusters' == clusters = clusters
  | otherwise             = kmeans (it-1) points clusters'
  where
    clusters' = emStep points clusters

emStep points clusters = maximizationStep $ estimationStep points clusters

estimationStep :: [[Double]] -> [[Double]] -> [[[Double]]]
estimationStep points clusters = map unTuple
                               $ groupByValue
                               $ sortByValue
                               $ zip points
                               $ assign points clusters
  where
    unTuple = map fst

maximizationStep :: [[[Double]]] -> [[Double]]
maximizationStep candidates = map mean candidates
  where
    mean xs = (foldl' (.+.) (head xs) (tail xs)) ./ (fromIntegral $ length xs)

assign :: [[Double]] -> [[Double]] -> [Integer]
assign points clusters = map closestTo points
  where
    closestTo p = argmin $ map (euclid p) clusters
    euclid pi ci = sum $ map (^2) $ zipWith (-) pi ci
    argmin xs = fst $ head $ sortByValue $ zip [0..] xs