{-# LANGUAGE CPP #-}
{-# LANGUAGE FlexibleContexts #-}

{-# OPTIONS_GHC -fno-warn-missing-signatures #-}

-----------------------------------------------------------------------------
{- |
Module      :  Numeric.LinearAlgebra
Copyright   :  (c) Alberto Ruiz 2006-15
License     :  BSD3
Maintainer  :  Alberto Ruiz
Stability   :  provisional

-}
-----------------------------------------------------------------------------
module Numeric.LinearAlgebra (

    -- * Basic types and data manipulation
    -- | This package works with 2D ('Matrix') and 1D ('Vector')
    -- arrays of real ('R') or complex ('C') double precision numbers.
    -- Single precision and machine integers are also supported for
    -- basic arithmetic and data manipulation.
    module Numeric.LinearAlgebra.Data,

    -- * Numeric classes
    -- |
    -- The standard numeric classes are defined elementwise (commonly referred to
    -- as the Hadamard product or the Schur product):
    --
    -- >>>  vector [1,2,3] * vector [3,0,-2]
    -- [3.0,0.0,-6.0]
    -- it :: Vector R
    --
    -- >>> matrix 3 [1..9] * ident 3
    -- (3><3)
    --  [ 1.0, 0.0, 0.0
    --  , 0.0, 5.0, 0.0
    --  , 0.0, 0.0, 9.0 ]

    -- * Autoconformable dimensions
    -- |
    -- In most operations, single-element vectors and matrices
    -- (created from numeric literals or using 'scalar'), and matrices
    -- with just one row or column, automatically
    -- expand to match the dimensions of the other operand:
    --
    -- >>> 5 + 2*ident 3 :: Matrix Double
    -- (3><3)
    --  [ 7.0, 5.0, 5.0
    --  , 5.0, 7.0, 5.0
    --  , 5.0, 5.0, 7.0 ]
    --
    -- >>> (4><3) [1..] + row [10,20,30]
    -- (4><3)
    --  [ 11.0, 22.0, 33.0
    --  , 14.0, 25.0, 36.0
    --  , 17.0, 28.0, 39.0
    --  , 20.0, 31.0, 42.0 ]
    --

    -- * Products
    -- ** Dot
    dot, (<.>),
    -- ** Matrix-vector
    (#>), (<#), (!#>),
    -- ** Matrix-matrix
    (<>),
    -- | The matrix product is also implemented in the "Data.Monoid" instance, where
    -- single-element matrices (created from numeric literals or using 'scalar')
    -- are used for scaling.
    --
    -- >>> import Data.Monoid as M
    -- >>>  let m = matrix 3 [1..6]
    -- >>> m M.<> 2 M.<> diagl[0.5,1,0]
    -- (2><3)
    --  [ 1.0,  4.0, 0.0
    --  , 4.0, 10.0, 0.0 ]
    --
    -- 'mconcat' uses 'optimiseMult' to get the optimal association order.

    -- ** Other
    outer, kronecker, cross,
    scale, add,
    sumElements, prodElements,

    -- * Linear systems
    -- ** General
    (<\>),
    linearSolveLS,
    linearSolveSVD,
    -- ** Determined
    linearSolve,
    luSolve, luPacked,
    luSolve', luPacked',
    -- ** Symmetric indefinite
    ldlSolve, ldlPacked,
    -- ** Positive definite
    cholSolve,
    -- ** Triangular
    UpLo(..),
    triSolve,
    -- ** Tridiagonal
    triDiagSolve,
    -- ** Sparse
    cgSolve,
    cgSolve',

    -- * Inverse and pseudoinverse
    inv, pinv, pinvTol,

    -- * Determinant and rank
    rcond, rank,
    det, invlndet,

    -- * Norms
    Normed(..),
    norm_Frob, norm_nuclear,

    -- * Nullspace and range
    orth,
    nullspace, null1, null1sym,

    -- * Singular value decomposition
    svd,
    thinSVD,
    compactSVD,
    compactSVDTol,
    singularValues,
    leftSV, rightSV,

    -- * Eigendecomposition
    eig, geig, eigSH,
    eigenvalues, geigenvalues, eigenvaluesSH,
    geigSH,

    -- * QR
    qr, thinQR, rq, thinRQ, qrRaw, qrgr,

    -- * Cholesky
    chol, mbChol,

    -- * LU
    lu, luFact,

    -- * Hessenberg
    hess,

    -- * Schur
    schur,

    -- * Matrix functions
    expm,
    sqrtm,
    matFunc,

    -- * Correlation and convolution
    corr, conv, corrMin, corr2, conv2,

    -- * Random arrays

    Seed, RandDist(..), randomVector, rand, randn, gaussianSample, uniformSample,

    -- * Misc
    meanCov, rowOuters, pairwiseD2, normalize, peps, relativeError, magnit,
    haussholder, optimiseMult, udot, nullspaceSVD, orthSVD, ranksv,
    iC, sym, mTm, trustSym, unSym,
    -- * Auxiliary classes
    Element, Container, Product, Numeric, LSDiv, Herm,
    Complexable, RealElement,
    RealOf, ComplexOf, SingleOf, DoubleOf,
    IndexOf,
    Field, Linear(), Additive(),
    Transposable,
    LU(..),
    LDL(..),
    QR(..),
    CGState(..),
    Testable(..)
) where

import Numeric.LinearAlgebra.Data

import Numeric.Matrix()
import Numeric.Vector()
import Internal.Matrix
import Internal.Container hiding ((<>))
import Internal.Numeric hiding (mul)
import Internal.Algorithms hiding (linearSolve,Normed,orth,luPacked',linearSolve',luSolve',ldlPacked')
import qualified Internal.Algorithms as A
import Internal.Util
import Internal.Random
import Internal.Sparse((!#>))
import Internal.CG
import Internal.Conversion
#if MIN_VERSION_base(4,11,0)
import Prelude hiding ((<>))
#endif

{- | dense matrix product

>>> let a = (3><5) [1..]
>>> a
(3><5)
 [  1.0,  2.0,  3.0,  4.0,  5.0
 ,  6.0,  7.0,  8.0,  9.0, 10.0
 , 11.0, 12.0, 13.0, 14.0, 15.0 ]

>>> let b = (5><2) [1,3, 0,2, -1,5, 7,7, 6,0]
>>> b
(5><2)
 [  1.0, 3.0
 ,  0.0, 2.0
 , -1.0, 5.0
 ,  7.0, 7.0
 ,  6.0, 0.0 ]

>>> a <> b
(3><2)
 [  56.0,  50.0
 , 121.0, 135.0
 , 186.0, 220.0 ]

-}
(<>) :: Numeric t => Matrix t -> Matrix t -> Matrix t
(<>) = mXm
infixr 8 <>

{- | Solve a linear system (for square coefficient matrix and several right-hand sides) using the LU decomposition, returning Nothing for a singular system. For underconstrained or overconstrained systems use 'linearSolveLS' or 'linearSolveSVD'.

@
a = (2><2)
 [ 1.0, 2.0
 , 3.0, 5.0 ]
@

@
b = (2><3)
 [  6.0, 1.0, 10.0
 , 15.0, 3.0, 26.0 ]
@

>>> linearSolve a b
Just (2><3)
 [ -1.4802973661668753e-15,     0.9999999999999997, 1.999999999999997
 ,       3.000000000000001, 1.6653345369377348e-16, 4.000000000000002 ]

>>> let Just x = it
>>> disp 5 x
2x3
-0.00000  1.00000  2.00000
 3.00000  0.00000  4.00000

>>> a <> x
(2><3)
 [  6.0, 1.0, 10.0
 , 15.0, 3.0, 26.0 ]

-}
linearSolve m b = A.mbLinearSolve m b

-- | return an orthonormal basis of the null space of a matrix. See also 'nullspaceSVD'.
nullspace m = nullspaceSVD (Left (1*eps)) m (rightSV m)

-- | return an orthonormal basis of the range space of a matrix. See also 'orthSVD'.
orth m = orthSVD (Left (1*eps)) m (leftSV m)


{-# LANGUAGE FlexibleInstances #-}
-----------------------------------------------------------------------------
-- |
-- Module     : Foreign.VMath.VFloating
-- Copyright  : Copyright (c) 2010, Patrick Perry <patperry@gmail.com>
-- License    : BSD3
-- Maintainer : Patrick Perry <patperry@gmail.com>
-- Stability  : experimental
--
-- Vector Floating operations.
--

module Foreign.VMath.VFloating (
    VFloating(..)
    ) where

import Foreign( Ptr, Storable, peek, poke, advancePtr )
import Data.Complex( Complex(..) )

import Foreign.VMath.VFractional
import Foreign.VMath.Double
import Foreign.VMath.Zomplex

-- | Types with vectorized 'Floating' operations.
class (VFractional a, Floating a) => VFloating a where
    vExp :: Int -> Ptr a -> Ptr a -> IO ()
    vSqrt :: Int -> Ptr a -> Ptr a -> IO ()
    vLog :: Int -> Ptr a -> Ptr a -> IO ()
    vPow :: Int -> Ptr a -> Ptr a -> Ptr a -> IO ()
    vSin :: Int -> Ptr a -> Ptr a -> IO ()
    vCos :: Int -> Ptr a -> Ptr a -> IO ()
    vTan :: Int -> Ptr a -> Ptr a -> IO ()
    vASin :: Int -> Ptr a -> Ptr a -> IO ()
    vACos :: Int -> Ptr a -> Ptr a -> IO ()
    vATan :: Int -> Ptr a -> Ptr a -> IO ()
    vSinh :: Int -> Ptr a -> Ptr a -> IO ()
    vCosh :: Int -> Ptr a -> Ptr a -> IO ()
    vTanh :: Int -> Ptr a -> Ptr a -> IO ()
    vASinh :: Int -> Ptr a -> Ptr a -> IO ()
    vACosh :: Int -> Ptr a -> Ptr a -> IO ()
    vATanh :: Int -> Ptr a -> Ptr a -> IO ()

    vExp = vop exp
    vSqrt = vop sqrt
    vLog = vop log
    vPow = vop2 (**)
    vSin = vop sin
    vCos = vop cos
    vTan = vop tan
    vASin = vop asin
    vACos = vop acos
    vATan = vop atan
    vSinh = vop sinh
    vCosh = vop cosh
    vTanh = vop tanh
    vASinh = vop asinh
    vACosh = vop acosh
    vATanh = vop atanh

vop :: (Storable a, Storable b)
    => (a -> b) -> Int -> Ptr a -> Ptr b -> IO ()
vop f n src dst | n <= 0    = return ()
                | otherwise = do
    a <- peek src
    poke dst $ f a
    vop f (n-1) (src `advancePtr` 1) (dst `advancePtr` 1)
{-# INLINE vop #-}

vop2 :: (Storable a1, Storable a2, Storable b)
     => (a1 -> a2 -> b) -> Int -> Ptr a1 -> Ptr a2 -> Ptr b -> IO ()
vop2 f n src1 src2 dst | n <= 0    = return ()
                       | otherwise = do
    a1 <- peek src1
    a2 <- peek src2
    poke dst $ f a1 a2
    vop2 f (n-1) (src1 `advancePtr` 1) (src2 `advancePtr` 1)
         (dst `advancePtr` 1)
{-# INLINE vop2 #-}

instance VFloating Double where
    vExp = vdExp
    {-# INLINE vExp #-}
    vSqrt = vdSqrt
    {-# INLINE vSqrt #-}
    vLog = vdLog
    {-# INLINE vLog #-}
    vPow = vdPow
    {-# INLINE vPow #-}
    vSin = vdSin
    {-# INLINE vSin #-}
    vCos = vdCos
    {-# INLINE vCos #-}
    vTan = vdTan
    {-# INLINE vTan #-}
    vASin = vdASin
    {-# INLINE vASin #-}
    vACos = vdACos
    {-# INLINE vACos #-}
    vATan = vdATan
    {-# INLINE vATan #-}
    vSinh = vdSinh
    {-# INLINE vSinh #-}
    vCosh = vdCosh
    {-# INLINE vCosh #-}
    vTanh = vdTanh
    {-# INLINE vTanh #-}
    vASinh = vdASinh
    {-# INLINE vASinh #-}
    vACosh = vdACosh
    {-# INLINE vACosh #-}
    vATanh = vdATanh
    {-# INLINE vATanh #-}

instance VFloating (Complex Double) where
    vSqrt = vzSqrt
    {-# INLINE vSqrt #-}
    vLog = vzLog
    {-# INLINE vLog #-}

    {- These functions have branch cuts in the wrong places
    vExp = vzExp
    {-# INLINE vExp #-}
    vPow = vzPow
    {-# INLINE vPow #-}
    vSin = vzSin
    {-# INLINE vSin #-}
    vCos = vzCos
    {-# INLINE vCos #-}
    vTan = vzTan
    {-# INLINE vTan #-}
    vASin = vzASin
    {-# INLINE vASin #-}
    vACos = vzACos
    {-# INLINE vACos #-}
    vATan = vzATan
    {-# INLINE vATan #-}
    vSinh = vzSinh
    {-# INLINE vSinh #-}
    vCosh = vzCosh
    {-# INLINE vCosh #-}
    vTanh = vzTanh
    {-# INLINE vTanh #-}
    vASinh = vzASinh
    {-# INLINE vASinh #-}
    vACosh = vzACosh
    {-# INLINE vACosh #-}
    vATanh = vzATanh
    {-# INLINE vATanh #-}
    -}


module Parser (
  readExpr
              )
  where

import Text.ParserCombinators.Parsec hiding (spaces)
import System.Environment
import Control.Monad
import Numeric (readOct, readHex, readFloat)
import Data.Char (digitToInt)
import Data.Ratio ((%), Rational)
import Data.Complex (Complex (..))

data LispVal = Atom String
             | List [LispVal]
             | DottedList [LispVal] LispVal
             | Integer Integer
             | String String
             | Bool Bool
             | Ratio Rational
             | Complex (Complex Float)
             | Float Float
             | Character Char
  deriving (Eq, Show)

symbol :: Parser Char
symbol = oneOf "!$%&|*+-/:<=>?@^_~"

spaces :: Parser ()
spaces = skipMany1 space

myEscapes :: Parser Char
myEscapes = do
  char '\\'
  char_escaped <- oneOf ['"', 'n', 't', 'r', '\\']
  return $ case char_escaped of
    '\\' -> '\\'
    '"'  -> '"'
    'n'  -> '\n'
    'r'  -> '\r'
    't'  -> '\t'

-- test escaping quotes with: "\"\\\"\"" hahaha
parseString :: Parser LispVal
parseString = do
  char '"'
  x <- many (myEscapes <|> noneOf ['\\', '"'])
  char '"'
  return $ String x

parseAtom :: Parser LispVal
parseAtom = do
  first <- letter <|> symbol
  rest <- many (letter <|> digit <|> symbol)
  let atom = first:rest
  return $ case atom of
    "#t" -> Bool True
    "#f" -> Bool False
    _    -> Atom atom

parseDecimal :: Parser LispVal
parseDecimal = liftM (Integer . read) $ many1 digit

readBinRec :: [Int] -> Integer
readBinRec (lsb : rest) = (toInteger lsb) + 2 * readBinRec rest
readBinRec [] = 0

readBinReal :: String -> Integer
readBinReal = readBinRec . reverse . map digitToInt

parseBinary :: Parser LispVal
parseBinary = liftM (Integer. readBinReal) $ many1 (oneOf "01")

readOctReal :: String -> Integer
readOctReal s = fst ((readOct s) !! 0)
parseOctal :: Parser LispVal
parseOctal = liftM (Integer . readOctReal) $ many1 (oneOf "01234567")

readHexReal :: String -> Integer
readHexReal s = fst ((readHex s) !! 0)
parseHex :: Parser LispVal
parseHex = liftM (Integer . readHexReal) $ many1 (oneOf "0123456789ABCDEFabcdef")

parsePrefixedNumber :: Parser LispVal
parsePrefixedNumber = do
  char '#'
  radix <- oneOf ['b', 'o', 'd', 'x']
  case radix of
    'b' -> parseBinary
    'o' -> parseOctal
    'd' -> parseDecimal
    'x' -> parseHex

parseFloat :: Parser LispVal
parseFloat = do
  whole_num <- many1 digit
  char '.'
  fractional_bit <- many1 digit
  return . Float . fst $ (readFloat (whole_num ++ "." ++ fractional_bit) !! 0 )

parseRational :: Parser LispVal
parseRational = do
  top <- many1 digit
  char '/'
  bottom <- many1 digit
  return $ Ratio (read top % read bottom)

floatCast :: LispVal -> Float
floatCast (Float f) = f
floatCast (Integer i) = fromIntegral i

parseComplex :: Parser LispVal
parseComplex = do
  first_bit <- try parseFloat <|> parseDecimal
  char '+'
  second_bit <- try parseFloat <|> parseDecimal
  char 'i'
  (return . Complex) ((floatCast first_bit) :+ (floatCast second_bit))

-- #b (binary), #o (octal), #d (decimal), and #x (hexadecimal)
parseNumber :: Parser LispVal
parseNumber =
  try parseRational
  <|> try parseComplex
  <|> try parseFloat
  <|> try parsePrefixedNumber
  <|> parseDecimal

{-
parseNumber' :: Parser LispVal
parseNumber' = do
  digits <- many1 digit
  return ((Number . read) digits)

parseNumber'' :: Parser LispVal
parseNumber'' = many1 digit >>= (return . Number . read)
-}

{-Characters are objects that represent printed characters such as letters and digits. Characters are written using the notation #\<character> or #\<character name>. For example:

#\a	; lower case letter
#\A	; upper case letter
#\(	; left parenthesis
#\	; the space character
#\space	; the preferred way to write a space
#\newline	; the newline character
Case is significant in #\<character>, but not in #\<character name>. If <character> in #\<character> is alphabetic, then the character following <character> must be a delimiter character such as a space or parenthesis. This rule resolves the ambiguous case where, for example, the sequence of characters ``#\space'' could be taken to be either a representation of the space character or a representation of the character ``#\s'' followed by a representation of the symbol ``pace.''

Characters written in the #\ notation are self-evaluating. That is, they do not have to be quoted in programs. Some of the procedures that operate on characters ignore the difference between upper case and lower case. The procedures that ignore case have ``-ci'' (for ``case insensitive'') embedded in their names.
-}
parseChar :: Parser LispVal
parseChar = do
  char '#' >> char '\\'
  rest <- (string "space") <|> (string "newline") <|> (anyChar >>= (\x -> return  [x]))
  return (Character $ case rest of
    "space" -> ' '
    "newline" -> '\n'
    c -> c !! 0)

parseList :: Parser LispVal
parseList = liftM List $ sepBy parseExpr spaces

parseDottedList :: Parser LispVal
parseDottedList = do
    head <- endBy parseExpr spaces
    tail <- char '.' >> spaces >> parseExpr
    return $ DottedList head tail

parseQuoted :: Parser LispVal
parseQuoted = do
    char '\''
    x <- parseExpr
    return $ List [Atom "quote", x]

{-parseQuasiQuoted :: Parser LispVal
parseQuasiQuoted = do
    char '`'
    x <- parseExpr
    return $ List [Atom "quote", x]-}
parseExpr :: Parser LispVal
parseExpr = try parseAtom
         <|> try parseChar
         <|> try parseNumber
         <|> try parseString
         <|> try parseQuoted
         <|> do char '('
                x <- try parseList <|> parseDottedList
                char ')'
                return x

readExpr :: String -> String
readExpr input = case parse parseExpr "lisp" input of
    Left err -> "No match: " ++ show err ++ ", input: " ++ input
    Right val -> "Found value: " ++ show val


{-# LANGUAGE FlexibleContexts #-}
module DFTPinwheel where

import           Control.Monad                  as M
import qualified Data.Array.Accelerate          as A
import           Data.Array.Accelerate.LLVM.PTX
import           Data.Array.Repa                as R
import           Data.Complex
import           Data.List                      as L
import           Data.Vector.Storable           as VS
import           Data.Vector.Unboxed            as VU
import           DFT.Plan
import           Filter.Utils
import           Foreign.CUDA.Driver            as CUDA
import           FourierMethod.BlockCudaMatrix
import           FourierMethod.FourierSeries2D
import           Image.IO
import           Pinwheel.Base
import           Pinwheel.FourierSeries2D
import           System.Directory
import           System.Environment
import           System.FilePath
import           System.Random
import           Utils.SimpsonRule
import Utils.Array
import Numeric.GSL.Special.Bessel
import           Utils.Distribution

main = do
  args@(deviceIDsStr:numPointsStr:numR2FreqStr:deltaStr:deltaReconStr:periodR2Str:angularFreqStr:radialFreqStr:sigmaStr:periodEnvelopeStr:stdStr:numBatchStr:_) <-
    getArgs
  let deviceIDs = read deviceIDsStr :: [Int]
      numPoints = read numPointsStr :: Int
      numR2Freqs = read numR2FreqStr :: Int
      delta = read deltaStr :: Double
      deltaRecon = read deltaReconStr :: Double
      periodR2 = read periodR2Str :: Double
      angularFreq = read angularFreqStr :: Int
      radialFreq = read radialFreqStr :: Int
      sigma = read sigmaStr :: Double
      periodEnvelope = read periodEnvelopeStr :: Double
      std = read stdStr :: Double
      numBatch = read numBatchStr :: Int
      folderPath = "output/test/DFTPinwheel"
  removePathForcibly folderPath
  createDirectoryIfMissing True folderPath
  initialise []
  devs <- M.mapM device deviceIDs
  ctxs <- M.mapM (\dev -> CUDA.create dev []) devs
  ptxs <- M.mapM createTargetFromContext ctxs
  let center = div numPoints 2
      pinwheel =
        computeUnboxedS . R.fromFunction (Z :. numPoints :. numPoints) $ \(Z :. x :. y) ->
          fourierMellinInvPeriod
            sigma
            (periodEnvelope * sqrt 2)
            angularFreq
            radialFreq
            ( deltaRecon * fromIntegral (x - center)
            , deltaRecon * fromIntegral (y - center))
      pinwheelMat =
        CuMat (numPoints ^ 2) 1 . CuVecHost . VU.convert . toUnboxed $ pinwheel
      centerFreq = div numR2Freqs 2
      -- pinwheelFreq =
      --   computeUnboxedS . R.fromFunction (Z :. numR2Freqs :. numR2Freqs) $ \(Z :. x :. y) ->
      --     fourierMellin
      --       sigma
      --       angularFreq
      --       radialFreq
      --       ( delta * fromIntegral (x - centerFreq)
      --       , delta * fromIntegral (y - centerFreq))
      -- pinwheelFreqMat =
      --   CuMat 1 (numR2Freqs ^ 2) . CuVecHost . VU.convert . toUnboxed $ -- . computeS . weightedArray
      --   pinwheelFreq
  -- invHarmonics <-
  --   createInverseHarmonicMatriesGPU ptxs 1 numPoints numR2Freqs periodR2 delta
  -- pinwheelCoef <-
  --   computeFourierCoefficientsR2Stream
  --     deviceIDs
  --     ptxs
  --     numR2Freqs
  --     numPoints
  --     periodR2
  --     delta
  --     1
  --     numBatch
  --     [pinwheelMat]
  -- pinwheelFreqSeries <-
  --   computeFourierSeriesR2
  --     deviceIDs
  --     numR2Freqs
  --     numPoints
  --     periodR2
  --     invHarmonics
  --     [pinwheelFreqMat]
  plotImageRepaComplex (folderPath </> "Pinwheel.png") .
    ImageRepa 8 .
    computeS . extend (Z :. (1 :: Int) :. All :. All) . R.traverse pinwheel id $ \f idx@(Z :. i :. j) ->
    let x = fromIntegral $ i - center
        y = fromIntegral $ j - center
        r = sqrt $ x ^ 2 + y ^ 2
       -- if r <= 6
       --   then 0
       --   else
     in f idx
  -- plotImageRepaComplex (folderPath </> "PinwheelTest.png") .
  --   ImageRepa 8 . computeS . extend (Z :. (1 :: Int) :. All :. All) $
  --   pinwheelTest
  -- plotImageRepaComplex (folderPath </> "PinwheelCoefficients.png") .
  --   ImageRepa 8 . computeS . R.traverse pinwheelCoef id $ \f idx@(Z :. _ :. i :. j) ->
  --   let x = fromIntegral $ i - div numR2Freqs 2
  --       y = fromIntegral $ j - div numR2Freqs 2
  --       r = sqrt $ x ^ 2 + y ^ 2
  --   in if r <= 0
  --        then 0
  --        else f idx
  -- plotImageRepaComplex (folderPath </> "PinwheelFrequency.png") .
  --   ImageRepa 8 . computeS . extend (Z :. (1 :: Int) :. All :. All) $
  --   pinwheelFreq
  -- plotImageRepaComplex (folderPath </> "PinwheelFrequencySeries.png") .
  --   ImageRepa 8 $
  --   pinwheelFreqSeries
  initVec <-
    (VS.fromList . L.map (\x -> x :+ 0)) <$>
    M.replicateM (numPoints ^ 2) randomIO
  lock <- getFFTWLock
  plan <-
    fst <$> idft1dGPlan lock emptyPlan [numPoints, numPoints] [0, 1] initVec
  -- dftPinwheel <-
  --   dftExecute plan (DFTPlanID IDFT1DG [numPoints, numPoints] [0, 1]) .
  --   VU.convert . toUnboxed . computeUnboxedS . makeFilter2D $
  --   pinwheel
  -- plotImageRepaComplex (folderPath </> "DFTPinwheel.png") .
  --   ImageRepa 8 .
  --   computeUnboxedS .
  --   makeFilter2D .
  --   fromUnboxed (Z :. (1 :: Int) :. numPoints :. numPoints) . VS.convert $
  --   dftPinwheel
  let centerR2Freq = div numR2Freqs 2
      gaussian2D =
        computeUnboxedS . fromFunction (Z :. numR2Freqs :. numR2Freqs) $ \(Z :. i' :. j') ->
          let i = i' - centerR2Freq
              j = j' - centerR2Freq
              rho = (sqrt . fromIntegral $ i ^ 2 + j ^ 2) / periodR2 * 2*pi
           in (-- rho ^ (abs angularFreq) *
               exp
                 (pi * fromIntegral (i ^ 2 + j ^ 2) /
                  ((-1) * periodR2 ^ 2 * std ^ 2)) /
               (2 * pi * std ^ 2)) :+
              0
      -- pinwheelFreqAnaticalNonGaussian =
      --   centerHollowArray numR2Freqs . computeUnboxedS $
      --   analyticalFourierCoefficients2
      --     numR2Freqs
      --     delta
      --     angularFreq
      --     radialFreq
      --     (-sigma) --(sigma - 1)
      --     periodR2
      --     (periodEnvelope * sqrt 2)
      pinwheelFreqAnatical =
        centerHollowArray numR2Freqs .
        computeUnboxedS
        -- . R.zipWith (*) gaussian2D
        $
        analyticalFourierCoefficients2
          numR2Freqs
          delta
          angularFreq
          radialFreq
          (-sigma) --(sigma - 1)
          periodR2
          (periodEnvelope * sqrt 2)
      -- pinwheelFreqAnaticalMat =
      --   A.use .
      --   A.fromList (A.Z A.:. (numR2Freqs ^ 2) A.:. (1 :: Int)) . R.toList $
      --   pinwheelFreqAnatical
      -- pinwheelFreqAnaticalInner =
      --   centerHollowArray numR2Freqs . computeUnboxedS $
      --   analyticalFourierCoefficients2'
      --     numR2Freqs
      --     (1 / fromIntegral (div numR2Freqs 2))
      --     angularFreq
      --     radialFreq
      --     (-sigma)
      --     periodR2
      --     (periodEnvelope * sqrt 2)
      -- pinwheelFreqAnaticalInnerMat =
      --   A.use .
      --   A.fromList (A.Z A.:. (numR2Freqs ^ 2) A.:. (1 :: Int)) . R.toList $
      --   pinwheelFreqAnaticalInner
      -- pinwheelFreqSeriesAnatical =
      --   computeUnboxedS $
      --   analyticalFourierSeries1
      --     numPoints
      --     1 -- delta
      --     angularFreq
      --     radialFreq
      --     (sigma - 1)
      --     periodR2
  plotImageRepaComplex (folderPath </> "PinwheelCoefficientsAnalytical.png") .
    ImageRepa 8 . computeS . extend (Z :. (1 :: Int) :. All :. All) $
    pinwheelFreqAnatical
  -- plotImageRepaComplex
  --   (folderPath </> "PinwheelCoefficientsAnalyticalNonGaussian.png") .
  --   ImageRepa 8 . computeS . extend (Z :. (1 :: Int) :. All :. All) $
  --   pinwheelFreqAnaticalNonGaussian
  -- plotImageRepaComplex
  --   (folderPath </> "PinwheelCoefficientsAnalyticalInner.png") .
  --   ImageRepa 8 . computeS . extend (Z :. (1 :: Int) :. All :. All) $
  --   pinwheelFreqAnaticalInner
  -- plotImageRepaComplex (folderPath </> "PinwheelFrequencySeriesAnalytical.png") .
  --   ImageRepa 8 . computeS . extend (Z :. (1 :: Int) :. All :. All) $
  --   pinwheelFreqSeriesAnatical
  -- pinwheelFreqSeriesStream <-
  --   computeFourierSeriesR2Stream
  --     deviceIDs
  --     ptxs
  --     numR2Freqs
  --     numPoints
  --     periodR2
  --     delta
  --     numBatch
  --     [transposeCuMat pinwheelFreqMat]
  -- plotImageRepaComplex (folderPath </> "PinwheelFrequencySeriesStream.png") .
  --   ImageRepa 8 . computeS . R.traverse pinwheelFreqSeriesStream id $ \f idx@(Z :. _ :. i :. j) ->
  --   let x = fromIntegral $ i - div numPoints 2
  --       y = fromIntegral $ j - div numPoints 2
  --       r = sqrt $ x ^ 2 + y ^ 2
  --   in if r <= 0
  --        then 0
  --        else f idx
  -- pinwheelCoefficientsAnalyticalSeriesStream <-
  --   computeFourierSeriesR2StreamAcc
  --     ptxs
  --     numR2Freqs
  --     numPoints
  --     1
  --     periodR2
  --     deltaRecon
  --     numBatch
  --     pinwheelFreqAnaticalMat
  --     -- [ transposeCuMat . createCuMat $
  --     --   [VU.convert . toUnboxed $ pinwheelFreqAnatical]
  --     -- ]
  -- plotImageRepaComplex
  --   (folderPath </> "PinwheelCoefficientsAnalyticalSeries.png") .
  --   ImageRepa 8 .
  --   computeS .
  --   extend (Z :. (1 :: Int) :. All :. All) .
  --   R.traverse (sumS . rotate3D $ pinwheelCoefficientsAnalyticalSeriesStream) id $ \f idx@(Z :. i :. j) ->
  --   let x = fromIntegral $ i - center
  --       y = fromIntegral $ j - center
  --       r = sqrt $ x ^ 2 + y ^ 2
  --      -- if r <= 6
  --      --   then 0
  --      --   else
  --    in f idx
  pinwheelCoefficientsAnalyticalIDFT <-
    dftExecute plan (DFTPlanID IDFT1DG [numPoints, numPoints] [0, 1]) .
    VU.convert . toUnboxed . computeS . makeFilter2D $
    pinwheelFreqAnatical
  plotImageRepaComplex (folderPath </> "PinwheelCozefficientsIDFT.png") .
    ImageRepa 8 .
    computeS .
    makeFilter2DInverse .
    fromUnboxed (Z :. (1 :: Int) :. numPoints :. numPoints) . VS.convert $
    pinwheelCoefficientsAnalyticalIDFT
  -- pinwheelCoefficientsAnalyticalInnerSeriesStream <-
  --   computeFourierSeriesR2StreamAcc'
  --     ptxs
  --     numR2Freqs
  --     numPoints
  --     1
  --     periodR2
  --     1 -- (1 / fromIntegral (div numR2Freqs 2))
  --     numBatch
  --     pinwheelFreqAnaticalMat
  -- plotImageRepaComplex
  --   (folderPath </> "PinwheelCoefficientsAnalyticalInnverSeries.png") .
  --   ImageRepa 8 .
  --   computeS . extend (Z :. (1 :: Int) :. All :. All) . sumS . rotate3D $
  --   pinwheelCoefficientsAnalyticalInnerSeriesStream
  -- plotImageRepaComplex
  --   (folderPath </> "PinwheelCoefficientsAnalyticalSumSeries.png") .
  --   ImageRepa 8 .
  --   computeS . extend (Z :. (1 :: Int) :. All :. All) . sumS . rotate3D $
  --   R.zipWith
  --     (+)
  --     pinwheelCoefficientsAnalyticalSeriesStream
  --     pinwheelCoefficientsAnalyticalInnerSeriesStream
  --   computeS . R.traverse pinwheelCoefficientsAnalyticalSeriesStream id $ \f idx@(Z :. _ :. i :. j) ->
  --   let x = fromIntegral $ i - div numPoints 2
  --       y = fromIntegral $ j - div numPoints 2
  --       r = sqrt $ x ^ 2 + y ^ 2
  --   in if r <= 0
  --        then 0
  --        else f idx
  -- -- pinwheelCoefficientsSeriesStream <-
  --   computeFourierSeriesR2Stream
  --     deviceIDs
  --     ptxs
  --     numR2Freqs
  --     numPoints
  --     periodR2
  --     delta
  --     numBatch
  --     [transposeCuMat . createCuMat $ [VU.convert . toUnboxed $ pinwheelCoef]]
  -- plotImageRepaComplex (folderPath </> "PinwheelCoefficientsSeriesStream.png") .
  --   ImageRepa 8 . computeS . R.traverse pinwheelCoefficientsSeriesStream id $ \f idx@(Z :. _ :. i :. j) ->
  --   let x = fromIntegral $ i - div numPoints 2
  --       y = fromIntegral $ j - div numPoints 2
  --       r = sqrt $ x ^ 2 + y ^ 2
  --   in if r <= 0
  --        then 0
  --        else f idx
  -- let besselFilter =
  --       computeUnboxedS . fromFunction (Z :. numR2Freqs :. numR2Freqs) $ \(Z :. i :. j) ->
  --         let xFreq = fromIntegral $ i - centerFreq
  --             yFreq = fromIntegral $ j - centerFreq
  --             rho = sqrt $ xFreq ^ 2 + yFreq ^ 2
  --         in if rho == 0
  --              then 1
  --              else (bessel_J1 (2 * pi * rho / periodR2)) / rho :+ 0
  -- plotImageRepaComplex (folderPath </> "BesselFilter.png") .
  --   ImageRepa 8 . computeS . extend (Z :. (1 :: Int) :. All :. All) $
  --   besselFilter
  -- initVec1 <-
  --   (VS.fromList . L.map (\x -> x :+ 0)) <$>
  --   M.replicateM (numR2Freqs ^ 2) randomIO
  -- lock <- getFFTWLock
  -- plan <-
  --   fst <$>
  --   (dft1dGPlan lock emptyPlan [numR2Freqs, numR2Freqs] [0, 1] initVec1 >>= \(plan, vec) ->
  --      idft1dGPlan lock plan [numR2Freqs, numR2Freqs] [0, 1] vec)
  -- let besselVec =
  --       VU.convert . toUnboxed . computeS . makeFilter2D $ besselFilter
  --     pinwheelFreqAnaticalVec = VU.convert . toUnboxed $ pinwheelFreqAnatical
  -- besselVecF <-
  --   dftExecute plan (DFTPlanID DFT1DG [numR2Freqs, numR2Freqs] [0, 1]) besselVec
  -- pinwheelFreqAnaticalVecF <-
  --   dftExecute
  --     plan
  --     (DFTPlanID DFT1DG [numR2Freqs, numR2Freqs] [0, 1])
  --     pinwheelFreqAnaticalVec
  -- pinwheelFreqAnatical1 <-
  --   fmap (fromUnboxed (Z :. numR2Freqs :. numR2Freqs) . VS.convert) .
  --   dftExecute plan (DFTPlanID IDFT1DG [numR2Freqs, numR2Freqs] [0, 1]) $
  --   VS.zipWith (*) besselVecF pinwheelFreqAnaticalVecF
  -- plotImageRepaComplex
  --   (folderPath </> "BesselFilterConveledPinwheelFreqAnatical.png") .
  --   ImageRepa 8 . computeS . extend (Z :. (1 :: Int) :. All :. All) $
  --   pinwheelFreqAnatical1
  -- let pinwheelFreqAnatical2 =
  --       R.zipWith (-) pinwheelFreqAnatical pinwheelFreqAnatical1
  --     pinwheelFreqAnaticalMat2 =
  --       A.use .
  --       A.fromList (A.Z A.:. (numR2Freqs ^ 2) A.:. (1 :: Int)) . R.toList $
  --       pinwheelFreqAnatical2
  -- plotImageRepaComplex (folderPath </> "pinwheelFreqAnatical2.png") .
  --   ImageRepa 8 . computeS . extend (Z :. (1 :: Int) :. All :. All) $
  --   pinwheelFreqAnatical2
  -- pinwheelCoefficientsAnalyticalSeriesStream2 <-
  --   computeFourierSeriesR2StreamAcc
  --     ptxs
  --     numR2Freqs
  --     numPoints
  --     1
  --     periodR2
  --     delta
  --     numBatch
  --     pinwheelFreqAnaticalMat2
  -- plotImageRepaComplex
  --   (folderPath </> "PinwheelCoefficientsAnalyticalSeries2.png") .
  --   ImageRepa 8 .
  --   computeS . extend (Z :. (1 :: Int) :. All :. All) . sumS . rotate3D $
  --   pinwheelCoefficientsAnalyticalSeriesStream2
  -- dftPinwheel <-
  --   fmap (fromUnboxed (Z :. numR2Freqs :. numR2Freqs) . VS.convert) .
  --   dftExecute plan (DFTPlanID DFT1DG [numR2Freqs, numR2Freqs] [0, 1]) .
  --   VU.convert . toUnboxed . computeS . makeFilter2D $
  --   pinwheel
  -- let zeroArr =
  --       fromFunction (Z :. numR2Freqs :. numR2Freqs) $ \(Z :. i :. j) ->
  --         if (sqrt . fromIntegral $ (i - centerFreq) ^ 2 + (j - centerFreq) ^ 2) <
  --            1
  --           then 0
  --           else 1
  -- plotImageRepaComplex (folderPath </> "DFTPinwheel.png") .
  --   ImageRepa 8 .
  --   computeS .
  --   extend (Z :. (1 :: Int) :. All :. All) .
  --   R.zipWith (*) zeroArr . makeFilter2D $
  --   dftPinwheel
  -- let normFunc arr =
  --       let m = VU.maximum . toUnboxed . computeS . R.map magnitude $ arr
  --       in R.map (/ (m :+ 0)) arr
  -- plotImageRepaComplex (folderPath </> "Diff.png") .
  --   ImageRepa 8 . computeS . extend (Z :. (1 :: Int) :. All :. All) $
  --   R.zipWith
  --     (\a b ->
  --        if magnitude b == 0
  --          then 0
  --          else a / b)
  --     (normFunc . R.zipWith (*) zeroArr . makeFilter2D $ dftPinwheel)
  --     (normFunc . R.zipWith (*) zeroArr $ pinwheelFreqAnatical)


{-# LANGUAGE Trustworthy #-}
-- |
-- Module    : Statistics.Distribution.Hypergeometric.GenVar
-- Copyright : (c) 2015 Sam Rijs
--             (c) 2005 Robert Kern
--             (c) 1998 Ivan Frohne
-- License   : MIT
--
-- Maintainer  : srijs@airpost.net
-- Stability   : experimental
--

module Statistics.Distribution.Hypergeometric.GenVar (genVar) where

import Control.Monad.Primitive (PrimMonad, PrimState)

import Statistics.Distribution (ContGen(..), DiscreteGen(..))
import Statistics.Distribution.Hypergeometric (HypergeometricDistribution, hdM, hdL, hdK)

import Numeric.SpecFunctions (logGamma)

import System.Random.MWC (Gen, uniform)

d1 = 1.7155277699214135 -- 2*sqrt(2/e)
d2 = 0.8989161620588988 -- 3 - 2*sqrt(3/e)

instance ContGen HypergeometricDistribution
  where genContVar d g = return . fromIntegral =<< genVar (hdM d, hdL d, hdK d) g

instance DiscreteGen HypergeometricDistribution
  where genDiscreteVar d g = genVar (hdM d, hdL d, hdK d) g

-- | Random variates from the hypergeometric distribution /(m, l, k)/.
-- Returns the number of white balls drawn when /k/ balls
-- are drawn at random from an urn containing /m/ white
-- and /l-m/ black balls.
genVar :: (PrimMonad m, Integral a) => (a, a, a) -> Gen (PrimState m) -> m a
genVar (good, popsize, sample) gen = return . fix2 . fix1 =<< loop
  where fix1 z = if good > bad then m - z else z
        fix2 z = if m < sample then good - z else z
        mingoodbad = min good bad
        bad = popsize - good
        maxgoodbad = max good bad
        m = min sample (popsize - sample)
        gamma z = sum $ map (logGamma . fromIntegral)
          [ z + 1,     mingoodbad - z + 1
          , m - z + 1, maxgoodbad - m + z + 1 ]
        d4 = fromIntegral mingoodbad / fromIntegral popsize
        d5 = 1 - d4
        d6 = fromIntegral m * d4 + 0.5
        d7 = sqrt $ fromIntegral (popsize - m) * fromIntegral sample * d4 * d5 / fromIntegral (popsize - 1) + 0.5
        d8 = d1 * d7 + d2
        d9 = floor $ fromIntegral (m + 1) * fromIntegral (mingoodbad + 1) / fromIntegral (popsize + 2)
        d10 = gamma d9
        d11 = fromIntegral $ min (min m mingoodbad + 1) (floor (d6 + 16 * d7))
        -- 16 for 16-decimal-digit precision in d1 and d2
        loop = do
          x <- uniform gen
          y <- uniform gen
          case () of
            _  | w < 0 || w >= d11    -> loop -- fast rejection
               | x * (4 - x) - 3 <= t -> return z -- fast acceptance
               | x * (x - t) >= 1     -> loop -- fast rejection
               | 2 * (log x) <= t     -> return z -- acceptance
               | otherwise            -> loop -- rejection
               where w = d6 + d8 * (y - 0.5) / x
                     z = floor w
                     t = d10 - gamma z


{-# LANGUAGE DeriveDataTypeable, DeriveGeneric, FlexibleContexts,
    FlexibleInstances, UndecidableInstances #-}
-- |
-- Module    : Statistics.Distribution.Transform
-- Copyright : (c) 2013 John McDonnell;
-- License   : BSD3
--
-- Maintainer  : bos@serpentine.com
-- Stability   : experimental
-- Portability : portable
--
-- Transformations over distributions
module Statistics.Distribution.Transform
    (
      LinearTransform (..)
    , linTransFixedPoint
    , scaleAround
    ) where

import Data.Aeson (FromJSON, ToJSON)
import Control.Applicative ((<*>))
import Data.Binary (Binary)
import Data.Binary (put, get)
import Data.Data (Data, Typeable)
import Data.Functor ((<$>))
import GHC.Generics (Generic)
import qualified Statistics.Distribution as D

-- | Linear transformation applied to distribution.
--
-- > LinearTransform μ σ _
-- > x' = μ + σ·x
data LinearTransform d = LinearTransform
  { linTransLocation :: {-# UNPACK #-} !Double
    -- ^ Location parameter.
  , linTransScale    :: {-# UNPACK #-} !Double
    -- ^ Scale parameter.
  , linTransDistr    :: d
    -- ^ Distribution being transformed.
  } deriving (Eq, Show, Read, Typeable, Data, Generic)

instance (FromJSON d) => FromJSON (LinearTransform d)
instance (ToJSON d) => ToJSON (LinearTransform d)

instance (Binary d) => Binary (LinearTransform d) where
    get = LinearTransform <$> get <*> get <*> get
    put (LinearTransform x y z) = put x >> put y >> put z

-- | Apply linear transformation to distribution.
scaleAround :: Double           -- ^ Fixed point
            -> Double           -- ^ Scale parameter
            -> d                -- ^ Distribution
            -> LinearTransform d
scaleAround x0 sc = LinearTransform (x0 * (1 - sc)) sc

-- | Get fixed point of linear transformation
linTransFixedPoint :: LinearTransform d -> Double
linTransFixedPoint (LinearTransform loc sc _) = loc / (1 - sc)

instance Functor LinearTransform where
  fmap f (LinearTransform loc sc dist) = LinearTransform loc sc (f dist)

instance D.Distribution d => D.Distribution (LinearTransform d) where
  cumulative (LinearTransform loc sc dist) x = D.cumulative dist $ (x-loc) / sc

instance D.ContDistr d => D.ContDistr (LinearTransform d) where
  density    (LinearTransform loc sc dist) x = D.density    dist ((x-loc) / sc) / sc
  logDensity (LinearTransform loc sc dist) x = D.logDensity dist ((x-loc) / sc) - log sc
  quantile      (LinearTransform loc sc dist) p = loc + sc * D.quantile      dist p
  complQuantile (LinearTransform loc sc dist) p = loc + sc * D.complQuantile dist p

instance D.MaybeMean d => D.MaybeMean (LinearTransform d) where
  maybeMean (LinearTransform loc _ dist) = (+loc) <$> D.maybeMean dist

instance (D.Mean d) => D.Mean (LinearTransform d) where
  mean (LinearTransform loc _ dist) = loc + D.mean dist

instance D.MaybeVariance  d => D.MaybeVariance (LinearTransform d) where
  maybeVariance (LinearTransform _ sc dist) = (*(sc*sc)) <$> D.maybeVariance dist
  maybeStdDev   (LinearTransform _ sc dist) = (*sc)      <$> D.maybeStdDev dist

instance (D.Variance d) => D.Variance (LinearTransform d) where
  variance (LinearTransform _ sc dist) = sc * sc * D.variance dist
  stdDev   (LinearTransform _ sc dist) = sc * D.stdDev dist

instance (D.MaybeEntropy d) => D.MaybeEntropy (LinearTransform d) where
  maybeEntropy (LinearTransform _ _ dist) = D.maybeEntropy dist

instance (D.Entropy d) => D.Entropy (LinearTransform d) where
  entropy (LinearTransform _ _ dist) = D.entropy dist

instance D.ContGen d => D.ContGen (LinearTransform d) where
  genContVar (LinearTransform loc sc d) g = do
    x <- D.genContVar d g
    return $! loc + sc * x


{-# LANGUAGE FlexibleContexts #-}

module WeightDecay
  ( input
  , testdata
  , traindata
  , testY
  , testX
  , trainLinRegUnkownTarget
  , trainAndTestWithData
  , trainAndTestWithRegularization
  , makeTransformedMatrix
  , listToTuple'
  ) where

import           Data.Maybe
import           LinearRegression
import           Network.HTTP
import           NonlinearTransform
import           Numeric.LinearAlgebra.Data
import           Numeric.LinearAlgebra.HMatrix as HMatrix

--------------------------------------------------------------
-- Solutions to homework 6 of "Learning from data" question 5-6
--------------------------------------------------------------
input http = do
  response <- simpleHTTP $ getRequest http
  let body = fmap rspBody response
  let rows = fmap (map (map readDouble . words) . lines) body
  return rows

readDouble :: String -> Double
readDouble = read

listToTuple' :: [a] -> Maybe (a, a)
listToTuple' (x:y:xs) = Just (x, y)
listToTuple' _        = Nothing

-- | downloads the test data
testdata = input "http://work.caltech.edu/data/out.dta"

-- | downloads the training data
traindata = input "http://work.caltech.edu/data/in.dta"

testY = fmap (fmap makeVectorY) testdata
  where
    makeVectorY lists = vector $ map last lists

testX = makeDataMatrix testdata

trainY = fmap (fmap makeVectorY) traindata
  where
    makeVectorY lists = vector $ map last lists

trainX = makeDataMatrix traindata

makeLabelVector inputdata = fmap (fmap makeVectorY) inputdata
  where
    makeVectorY lists = vector $ map last lists

makeDataMatrix inputdata = fmap (fmap makeMatrixX) inputdata
  where
    makeMatrixX lists = matrix 3 $ concatMap xvectors lists
    xvectors list = 1 : init list

makeTransformedMatrix2 inputdata = fmap (fmap makeMatrixX) inputdata
  where
    makeMatrixX lists =
      createTransformedX2 $ map (fromMaybe (0, 0) . listToTuple') lists

makeTransformedMatrix k inputdata = fmap (fmap makeMatrixX) inputdata
  where
    makeMatrixX lists =
      createTransformedXk k $ map (fromMaybe (0, 0) . listToTuple') lists

trainXTransformed2 = makeTransformedMatrix2 traindata

testXTransformed2 = makeTransformedMatrix2 testdata

-- | Trains linear Regression from the data on the course website
trainLinRegUnkownTarget = do
  matrixXTrain <- trainXTransformed2
  vectorYTrain <- trainY
  let weights = linearRegressionWeight <$> matrixXTrain <*> vectorYTrain
  let inSampleError =
        linRegClassificationError <$> matrixXTrain <*> vectorYTrain <*> weights
  return (weights, inSampleError)

-- | Tests linear Regression from the data on the course website
testLinRegUnkownTarget weights = do
  matrixXTest <- testXTransformed2
  vectorYTest <- testY
  let outOfSampleError =
        linRegClassificationError <$> matrixXTest <*> vectorYTest <*> weights
  return outOfSampleError

-- | solution for question 2: train and test with given data, with nonlinear transformation but without regularization and print the result
trainAndTestWithData :: IO ()
trainAndTestWithData = do
  (weights, inSampleError) <- trainLinRegUnkownTarget
  outOfSampleError <- testLinRegUnkownTarget weights
  putStr "In Sample Error: "
  print inSampleError
  putStr "Out of Sample Error: "
  print outOfSampleError

-- | returns the weights determined by linear regression on the training data with regularization
linearRegressionWeightReg :: R -> Matrix R -> Vector R -> Vector R
linearRegressionWeightReg regParam matrixOfInputData vectorOfLabels =
  inv
    (tr matrixOfInputData HMatrix.<> matrixOfInputData +
     scalar regParam * (ident $fst $size $tr matrixOfInputData)) HMatrix.#>
  (tr matrixOfInputData #> vectorOfLabels)

-- | Trains linear Regression from the data on the course website with regularization
trainLinRegWithRegularization regParam = do
  matrixXTrain <- trainXTransformed2
  vectorYTrain <- trainY
  let weights =
        linearRegressionWeightReg regParam <$> matrixXTrain <*> vectorYTrain
  let inSampleError =
        linRegClassificationError <$> matrixXTrain <*> vectorYTrain <*> weights
  return (weights, inSampleError)

-- | performs training and testing with the regularization parameter provided as an input and prints the in-sample and out-of-sample error
trainAndTestWithRegularization :: R -> IO ()
trainAndTestWithRegularization regParam = do
  (weights, inSampleError) <- trainLinRegWithRegularization regParam
  outOfSampleError <- testLinRegUnkownTarget weights
  putStr "In Sample Error: "
  print inSampleError
  putStr "Out of Sample Error: "
  print outOfSampleError


{-# OPTIONS_HADDOCK prune #-}
module ForSyDe.Atom.Skel.FastVector.DSP where

import Data.Complex
import qualified Data.Number.FixedFunctions as F
import ForSyDe.Atom.MoC as MoC
import ForSyDe.Atom.Skel.FastVector.Lib as V hiding (duals, unduals)
import ForSyDe.Atom.Skel.FastVector.Matrix as M
import ForSyDe.Atom ((><))

-- | See 'ForSyDe.Atom.Skel.Vector.DSP.taylor'.
taylor :: (Eq a, Floating a)
       => Int      -- ^ number of points in the output window.
       -> Int      -- ^ Number of nearly constant level sidelobes adjacent to the mainlobe
       -> a        -- ^ Desired peak sidelobe level in decibels (db) relative to the mainlobe
       -> Vector a -- ^ The window, with the center value normalized to one (the value
                   -- one appears only if the number of samples is odd).
taylor n nbar level = V.farm11 (*scale) w
  where
    -- explicit conversions to floating
    bN   = fromIntegral n
    nBar = fromIntegral nbar
    ma   = map fromIntegral [1 .. nbar-1]
    -- calculate intermediate values
    b  = 10**((-level) / 20)
    a  = log(b + sqrt(b**2 - 1)) / pi
    s2 = nBar ** 2 / (a**2 + (nBar - 0.5)**2)
    -- functions for calculating coefficients
    fmcalc m  = let numer = (-1)**(m+1) * product[1.0 - m**2/s2/(a**2 + (j - 0.5)**2) | j <- ma]
                    denom = 2 * product[1 - m**2/j**2 | j <- ma, j /= m]
                in numer / denom
    ccalc m x = cos(2 * pi * x * (m - bN/2 + 1/2) / bN)
    wcalc m   = 2 * dotvv (vector $ map fmcalc ma) (vector $ map (ccalc m) ma) + 1
    -- calculate window coefficients
    w  = vector $ map (wcalc . fromIntegral) [0..n-1]
    -- normalize (Note that this is not described in the original text [1])
    scale = 1 / wcalc (bN - 1) / 2

-- | See 'ForSyDe.Atom.Skel.Vector.DSP.taylor''.
taylor' n = taylor 4 (-30)

-- | See 'ForSyDe.Atom.Skel.Vector.DSP.dotvm'.
dotvm :: Num a => Vector a -> Vector (Vector a) -> Vector a
dotvm = dotvm' (+) (*)

-- | See 'ForSyDe.Atom.Skel.Vector.DSP.dotvv'.
dotvv :: Num a => Vector a -> Vector a -> a
dotvv a b
  | V.length a == V.length b = V.reduce (+) (V.farm21 (*) a b)
  | otherwise                = error "Vector sizes must match"

-- | See 'ForSyDe.Atom.Skel.Vector.DSP.dotvv''.
dotvv' :: Num a
     => ((a -> a -> a) -> f a -> f a -> f a)
     -- ^ higher-order function that can wrap the (*) operation.
     -> Vector (f a) -> Vector (f a) -> f a
dotvv' wrap a b
  | V.length a == V.length b = V.reduce (wrap (+)) $ V.farm21 (wrap (*)) a b
  | otherwise                = error "Vector sizes must match"

-- | See 'ForSyDe.Atom.Skel.Vector.DSP.dotmv''.
dotmv' :: (a -> a -> a) -- ^ kernel function for a row/column reduction, e.g. @(+)@ for dot product
       -> (b -> a -> a) -- ^ binary operation for pair-wise elements, e.g. @(*)@ for dot product
       -> Vector (Vector b) -- ^ /size/ = @(xa,ya)@
       -> Vector a      -- ^ /length/ = @xa@
       -> Vector a      -- ^ /length/ = @xa@
dotmv' f g mA y = V.farm11 (\x -> V.reduce f $ V.farm21 g x y) mA

-- | See 'ForSyDe.Atom.Skel.Vector.DSP.dotvm''.
dotvm' :: (b -> b -> b)     -- ^ kernel function for a row/column reduction, e.g. @(+)@ for dot product
       -> (a -> b -> b)     -- ^ binary operation for pair-wise elements, e.g. @(*)@ for dot product
       -> Vector a          -- ^ /length/ = @xa@
       -> Vector (Vector b) -- ^ /size/ = @(xa,ya)@
       -> Vector b          -- ^ /length/ = @ya@
dotvm' f g vs = V.reduce (V.farm21 f) . V.farm21 (\x -> V.farm11 (g x)) vs

-- | See 'ForSyDe.Atom.Skel.Vector.DSP.hanning'.
hanning :: (Floating n)
        => Int -- ^ The length of the window
        -> Vector n
hanning size = V.farm11 func $ V.vector [0..size-1]
  where
    func idx = let i = fromIntegral idx
                   n = fromIntegral size
               in 0.5 * (1 - cos((2 * pi * i) / (n - 1)))

-- | See 'ForSyDe.Atom.Skel.Vector.DSP.hamming'.
hamming :: (Floating n)
        => Int -- ^ The length of the window
        -> Vector n
hamming size = V.farm11 func $ V.vector [0..size-1]
  where
    func idx = let i = fromIntegral idx
                   n = fromIntegral size
               in 0.54 - 0.46 * cos((2 * pi * i) / (n - 1))

-- | See 'ForSyDe.Atom.Skel.Vector.DSP.blackman'.
blackman :: (Floating n)
        => Int -- ^ The length of the window
        -> Vector n
blackman size = V.farm11 func $ V.vector [0..size-1]
  where
    func idx = let i = fromIntegral idx
                   n = fromIntegral size
               in 0.42 - 0.5 * cos((2 * pi * i) / (n - 1)) + 0.08 * cos((4 * pi * i) / (n - 1))

-- | See 'ForSyDe.Atom.Skel.Vector.DSP.fir'.
fir :: Num a
    => Vector a  -- ^ vector of coefficients
    -> Vector a  -- ^ input vector of numbers; /size/ = @n@
    -> Vector a  -- ^ output vector of numbers; /size/ = @n@
fir coefs = V.reverse . V.farm11 applyFilter . tails . V.reverse
  where
    applyFilter = V.reduce (+) . V.farm21 (*) coefs

-- | See 'ForSyDe.Atom.Skel.Vector.DSP.fir''.
fir' :: (a -> a -> a)  -- ^ process/operation replacing '+'
     -> (c -> a -> a)  -- ^ process/operation replacing '*'
     -> (a -> a)       -- ^ delay process
     -> Vector c       -- ^ vector of coefficients
     -> a              -- ^ input signal/structure
     -> a              -- ^ output signal/structure
fir' plus times delay coefs =
  -- V.reduce plus . V.farm21 (\c -> times c) coefs . V.recur (V.fanoutn n delay <: id)
  V.reduce plus . V.farm21 times coefs . V.recuri (V.fanoutn n delay)
  where n = V.length coefs - 1

-- | See 'ForSyDe.Atom.Skel.Vector.DSP.twiddles'.
twiddles :: Floating a => Int -> V.Vector (Complex a)
twiddles bN = (bitrev . V.take (bN `div` 2)) (V.farm11 bW $ vector [0..])
  where bW x = (cis . negate) (-2 * pi * fromIntegral x / fromIntegral bN)

-- | See 'ForSyDe.Atom.Skel.Vector.DSP.fft'.
fft :: RealFloat a
     => Int -> V.Vector (Complex a) -> V.Vector (Complex a)
fft k vs | n == 2^k = V.reverse $ bitrev $ (stage `V.pipe1` V.iterate k (*2) 2) vs
         | otherwise = error $ "n=" ++ show n ++ "; k=" ++ show k
  where
    stage   w = V.concat . V.farm21 segment (twiddles n) . V.group w
    segment t = (><) unduals . (><) (V.farm22 (butterfly t)) . duals
    n         = V.length vs        -- length of input
    -------------------------------------------------
    butterfly w x0 x1 = let t = w * x1 in (x0 + t, x0 - t) -- kernel function

-- | See 'ForSyDe.Atom.Skel.Vector.DSP.fft''.
fft' :: Floating a
     => (Complex a -> a -> a -> (a, a))
     -> Int -> V.Vector a -> V.Vector a
fft' butterfly k vs | n == 2^k = bitrev $ (stage `V.pipe1` (V.iterate k (*2) 2)) vs
  where
    stage   w = V.concat . V.farm21 segment (twiddles n) . V.group w
    segment t = (><) unduals . (><) (V.farm22 (butterfly t)) . duals
    n         = V.length vs        -- length of input

-- | See 'ForSyDe.Atom.Skel.Vector.DSP.duals'.
duals    :: Vector a -> (Vector a, Vector a)
duals v  = (V.take k v, V.drop k v)
  where
    k = V.length v `div` 2

-- | See 'ForSyDe.Atom.Skel.Vector.DSP.unduals'.
unduals  :: Vector a -> Vector a -> Vector a
unduals x y =  x <++> y

-- | See 'ForSyDe.Atom.Skel.Vector.DSP.bitrev'.
bitrev = V.unsafeLift bitrevF
  where
    bitrevF [x] = [x]
    bitrevF xs  = bitrevF (V.evensF xs) ++ bitrevF (V.oddsF xs)


module Carlson
  (carlsonRF, carlsonRF',
  carlsonRD, carlsonRD',
  carlsonRJ, carlsonRJ',
  carlsonRC, carlsonRC',
  carlsonRG, carlsonRG')
  where
import           Data.Complex
import           Internal

rf_ :: Cplx -> Cplx -> Cplx -> Double -> ((Double,Double,Double), Cplx)
rf_ x y z err =
  let a = (x+y+z)/3 in
  let delta = map (\u -> magnitude(1-u/a)) [x,y,z] in
  if maximum delta < err
    then ((delta !! 0, delta !! 1, delta !! 2), a)
    else
      let (sqrtx, sqrty, sqrtz) = (sqrt x, sqrt y, sqrt z) in
      let lambda = sqrtx*sqrty + sqrty*sqrtz + sqrtz*sqrtx in
      rf_ ((x+lambda)/4) ((y+lambda)/4) ((z+lambda)/4) err

carlsonRF' :: Double -> Cplx -> Cplx -> Cplx -> Cplx
carlsonRF' err x y z =
  if zeros > 1
    then error "At most one of x, y, z can be 0"
    else
      let ((dx,dy,dz), a) = rf_ x y z err in
      let (e2,e3) = (dx*dy + dy*dz + dz*dx, dx*dy*dz) in
      toCplx(1 - e2/10 + e3/14 + e2*e2/24 - 3*e2*e3/44 - 5*e2*e2*e2/208 +
          3*e3*e3/104 + e2*e2*e3/16) / sqrt a
  where
    zeros = sum (map (\u -> fromEnum (u == 0)) [x,y,z])

carlsonRF :: Cplx -> Cplx -> Cplx -> Cplx
carlsonRF = carlsonRF' 1e-15

rd_ :: Cplx -> Cplx -> Cplx -> Cplx -> Cplx -> Double ->
       ((Double,Double,Double), Cplx, Cplx, Cplx)
rd_ x y z s fac err =
  let a = (x+y+z+z+z)/5 in
  let delta = map (\u -> magnitude(1-u/a)) [x,y,z] in
  if maximum delta < err
    then ((delta !! 0, delta !! 1, delta !! 2), a, s, fac)
    else
      let (sqrtx, sqrty, sqrtz) = (sqrt x, sqrt y, sqrt z) in
      let lambda = sqrtx*sqrty + sqrty*sqrtz + sqrtz*sqrtx in
      let s' = s + fac / (sqrt z * (z + lambda)) in
      rd_ ((x+lambda)/4) ((y+lambda)/4) ((z+lambda)/4) s' (fac/4) err

carlsonRD' :: Double -> Cplx -> Cplx -> Cplx -> Cplx
carlsonRD' err x y z =
  if zeros > 1
    then error "At most one of x, y, z can be 0"
    else
      let ((dx,dy,dz), a, s, fac) = rd_ x y z 0 1 err in
      let
        (e2,e3,e4,e5) = (dx*dy + dy*dz + 3*dz*dz + 2*dz*dx + dx*dz + 2*dy*dz,
          dz*dz*dz + dx*dz*dz + 3*dx*dy*dz + 2*dy*dz*dz + dy*dz*dz + 2*dx*dz*dz,
          dy*dz*dz*dz + dx*dz*dz*dz + dx*dy*dz*dz + 2*dx*dy*dz*dz,
          dx*dy*dz*dz*dz) in
      3*s + fac * toCplx(1 - 3*e2/14 + e3/6 + 9*e2*e2/88 - 3*e4/22 - 9*e2*e3/52 +
        3*e5/26 - e2*e2*e2/16 + 3*e3*e3/40 + 3*e2*e4/20 + 45*e2*e2*e3/272 -
        9*(e3*e4 + e2*e5)/68) / a / sqrt a
  where
    zeros = sum (map (\u -> fromEnum (u == 0)) [x,y,z])

carlsonRD :: Cplx -> Cplx -> Cplx -> Cplx
carlsonRD = carlsonRD' 1e-15

rj_ :: Cplx -> Cplx -> Cplx -> Cplx -> Cplx -> Double -> Cplx -> Int ->
       Double -> [Cplx] -> [Cplx] -> Double -> (Cplx, Int, [Cplx], [Cplx])
rj_ x y z p a maxmagns delta f fac d e err =
  let q = (4/err)**(1/6) * maxmagns / fromIntegral f in
  if magnitude a > q
    then (a, f, d, e)
    else
      let dnew = (sqrt p + sqrt x)*(sqrt p + sqrt y)*(sqrt p + sqrt z)
          d' = (fromIntegral f * dnew) : d
          e' = (toCplx fac * delta / dnew / dnew) : e
          lambda = sqrt x * sqrt y + sqrt y * sqrt z + sqrt z * sqrt x
          x' = (x + lambda) / 4
          y' = (y + lambda) / 4
          z' = (z + lambda) / 4
          p' = (p + lambda) / 4
          a' = (a + lambda) / 4
      in
      rj_ x' y' z' p' a' maxmagns delta (4*f) (fac/64) d' e' err

carlsonRJ' :: Double -> Cplx -> Cplx -> Cplx -> Cplx -> Cplx
carlsonRJ' err x y z p =
  if zeros > 1
    then error "At most one of x, y, z, p can be 0"
    else
      let a0 = (x + y + z + p + p) / 5
          maxmagns = maximum $ map (\u -> magnitude(a0-u)) [x, y, z, p]
          delta = (p-x)*(p-y)*(p-z)
      in
      let (a, f, d, e) = rj_ x y z p a0 maxmagns delta 1 1 [] [] err
          f' = fromIntegral f
      in
      let x' = (a0 - x) / f' / a
          y' = (a0 - y) / f' / a
          z' = (a0 - z) / f' / a
          p' = -(x'+y'+z') / 2
          e2 = x'*y' + x'*z' + y'*z' - 3*p'*p'
          e3 = x'*y'*z' + 2*e2*p' + 4*p'*p'*p'
          e4 = p'*(2*x'*y'*z' + e2*p' + 3*p'*p'*p')
          e5 = x'*y'*z'*p'*p'
          h = zipWith (\u v -> atanx_over_x(sqrt u) / v) e d
      in
      (1 - 3*e2/14 + e3/6 + 9*e2*e2/88 - 3*e4/22 - 9*e2*e3/52 + 3*e5/26) /
        f' / a / sqrt a + 6 * sum h
  where
    zeros = sum (map (\u -> fromEnum (u == 0)) [x,y,z,p])
    atanx_over_x w = if w == 0 then 1 else atanC w / w

carlsonRJ :: Cplx -> Cplx -> Cplx -> Cplx -> Cplx
carlsonRJ = carlsonRJ' 1e-15

rc_ :: Cplx -> Cplx -> Cplx -> Double -> Int -> Double -> (Cplx, Int)
rc_ x y a magn f err =
  let q = (1/3/err)**(1/8) * magn / fromIntegral f in
  if magnitude a > q
    then (a, f)
    else
      let lambda = 2 * sqrt x * sqrt y + y
          a' = (a + lambda) / 4
          x' = (x + lambda) / 4
          y' = (y + lambda) / 4
      in
      rc_ x' y' a' magn (4*f) err

carlsonRC' :: Double -> Cplx -> Cplx -> Cplx
carlsonRC' err x y =
  if y == 0
    then error "y cannot be 0"
    else
      let a0 = (x + y + y) / 3
          magn = magnitude(a0-x)
      in
      let (a, f) = rc_ x y a0 magn 1 err
          f' = fromIntegral f
      in
      let s = (y - a0) / f' / a in
      (1 + 3*s*s/10 + s*s*s/7 + 3*s*s*s*s/8 + 9*s*s*s*s*s/22 +
        159*s*s*s*s*s*s/208 + 9*s*s*s*s*s*s*s/8) / sqrt a

carlsonRC :: Cplx -> Cplx -> Cplx
carlsonRC = carlsonRC' 1e-15

carlsonRG' :: Double -> Cplx -> Cplx -> Cplx -> Cplx
carlsonRG' err x y z =
  if zeros > 1
    then sqrt(x+y+z) / 2
    else
      if z == 0
        then carlsonRG' err z x y
        else
          (z * carlsonRF' err x y z -
            (x-z) * (y-z) * carlsonRD' err x y z / 3 +
            sqrt x * sqrt y / sqrt z) / 2
  where
    zeros = sum (map (\u -> fromEnum (u == 0)) [x,y,z])

carlsonRG :: Cplx -> Cplx -> Cplx -> Cplx
carlsonRG = carlsonRG' 1e-15


{-# LANGUAGE OverloadedStrings, DuplicateRecordFields, RecordWildCards, NamedFieldPuns #-}
{-|
 - Module      : Numeric.GLP
 - Description : Specifiy a General Linear Problem and transform to a Standardized LP
 - Copyright   : (c) Torsten Kemps-Benedix, 2019
 - This module 'Numeric.GLP' specifiies a [Generalized Linear Problem](https://en.wikipedia.org/wiki/Linear_programming) (or LP in short) and provides means to tranform this into a Standardized LP (see module 'Numeric.LP') that may be solved with e.g. the simplex solver in module 'Numeric.LP.Simplex'.
-}
module Numeric.GLP
-- (GLP, Numeric.GLP.a, c, m1, m2, m3, dir, mkGLP, standardizeGLP)
where

import qualified Data.List                   as L
import           Control.Monad
import qualified Data.Set                    as S
import           Data.Set                    (member, isSubsetOf)
import           Numeric.LinearAlgebra       (I, Element, Matrix, Vector, ident, diagl,
                                              cols, rows, find, toColumns,
                                              fromColumns, cmap, fromList, toList,
                                              (!), (><),
                                              )
import qualified Numeric.LinearAlgebra       as LA
import           Numeric.LinearAlgebra.Devel (modifyMatrix, runSTMatrix,
                                              thawMatrix, mapVectorWithIndex, fi)
import           Numeric.LinearAlgebra       (vjoin, (|||), (===))

-- | This record defines a General Linear Problem containing ≤
-- and ≥ inequalities and == equalities. The sum of the numbers of
-- inequalities ('m1' times ≤ and 'm2' times ≥) and equalities ('m3'
-- times ==) must match one less than the number of
-- rows of the matrix 'a'. The first (i.e. the zero-th) row of the matrix
-- contains the coefficients of the objective function.
--
-- We assume that \(a_{00}\) is zero. The matrix \(a\) shall have \(m+1\) rows
-- and \(n+1\) columns.
--
-- The formal specification of the problem we intend to solve is as follows:
--
-- Maximize resp. minimize (according to 'dir') the following linear objective
-- function:
--
-- \[ z_0 = \sum_{i=1}^{n} a_{0,i} x_i \]
--
-- Subject to the following linear constraints:
--
-- \[ z_j = a_{j,0} + \sum_{i=1}^{n} a_{j,i} x_i,\, j=1,\ldots,m\]
-- \[ z_j \leq 0,\, j=1,\ldots,m_1\]
-- \[ z_j \geq 0,\, j=m_1+1,\ldots,m_1+m_2\]
-- \[ z_j = 0,\, j=m_1+m_2+1,\ldots,m=m1+m_2+m_3\]
--
-- The original variables \(x_1, \ldots, x_n\) may be constraint to be non-negative,
-- non-positive or they may be unconstrained according to the n-vector
-- \(c\in\{-1,0,+1\}^n\).
--
-- \[ c_j = -1 \Rightarrow x_j \leq 0\]
-- \[ c_j = 1 \Rightarrow x_j \geq 0\]
-- \[ c_j = 0 \Rightarrow x_j \text{ unrestricted} \]
data GLP = GLP {
    a   :: Matrix Double, -- ^ Tableau of coefficients, see above for a formal
                          -- description of their meaning
    c   :: Vector I,      -- ^ Type of constraint for \(x_j\), \(j=1,\ldots,n\),
                          -- see 'GLP' for further details. The element \(c_i\)
                          -- corresponds to column \(i+1\) of the matrix \(a\).
    n :: Int,             -- ^ Number of original variables \(x_1, \ldots, x_n\)
    m1  :: Int,           -- ^ Number of ≤ inequalities
    m2  :: Int,           -- ^ Number of ≥ inequalities
    m3  :: Int,           -- ^ Number of = equalities
    dir :: Dir,           -- ^ Whether the objective function is to be maximized
                          -- or minimzed
    lhs, rhs :: Vector I  -- ^ Indices of left and right hand variables respectively.
                          -- The original variables \(x_1,\ldots,x_n\) have indices
                          -- \(1,\ldots,n\), the slack variables get indices
                          -- \(n+1, \ldots, n+m_1+m_2+1\), and the auxiliary variables
                          -- finally get indices \(n+m_1+m_2+2, \ldots, n+m_1+m_2+m+2\).
    } deriving (Show)

-- | This smart constructor checks all the pre-conditions for 'GLP' stated in the
-- description of 'GLP' and initializes the variables 'lhs' and 'rhs'. Initially only
-- the \(m\) auxiliary variables \(z_1,\ldots,z_m\) are on the left hand side and all the
-- original variables \(x_1,\ldots,x_n\) are on the right hand side. Initially there are
-- no slack variables \(y_1,\ldots,y_{m_1+m_2+1}\). These will be introdced by the function
-- 'introduceSlackVariables' below.
mkGLP :: Matrix Double -> Vector I -> (Int, Int, Int) -> Dir -> Either String GLP
mkGLP a c (m1,m2,m3) dir =
    let m = rows a - 1
        n = cols a - 1
    in if (m /= m1+m2+m3)
       then Left $ "rows a - 1 == "++show m++" /= m1+m2+m3 == "++show (m1+m2+m3)
       else if (m <= 0)
            then Left "The number of constraints is m <= 0"
            else if (n <= 0)
                 then Left "The number of variables n is <= 0"
                 else if not (S.fromList (LA.toList c) `isSubsetOf` (S.fromList [-1,0,1]))
                      then Left "Vector c has elements other than -1, 0, +1"
                      else if LA.size c /= n
                           then Left $ "Vector c has not the proper length "++show n
                           else Right (GLP a c n m1 m2 m3 dir
                                           (LA.fromList [fi (n+m1+m2+2)..fi (n+m1+m2+m+2)])
                                           (LA.fromList [1..fi n]))

data Dir = Maximize | Minimize deriving Show

-- | Transform the Generalized Linear Problem to a Standardized Linear Problem
-- by means of the following steps:
--
-- 1. If the objective function is to be minimized multiply it by -1.
-- 2. Turn around the ≥-inequalities by negation of the corresponding rows. After this
-- step there are 'm1'+'m2' inequalities of the <= type.
-- 3. For those variables \(x_i\) that have a constraint \(x_i \leq 0\) substitute
-- \(x_i \rightarrow -x_i\). This means multiplying the column \(i\) of the matrix
-- \(a\) by \(-1\).
-- 4. For those variables \(x_i\) that are unconstrained introduce new non-negative
-- variables \(x_i = x_i⁺ - x_i⁻\). We give these the indices \(x_i⁺ = x_i\) and
-- \(x_⁻ = x_{i+1}\). The indices of all other variables \(x_k\), \(k>i\) have to be
-- increased by \(1\). The matrix \(a\) therefore grows by one column.
-- 5. For the 'm1'+'m2' ≤-inequalities introduce as many slack variables on the right
-- hand side and make these equalities instead. After this step there are only
-- equalities left as constraints. We do not represent these slack variables in the
-- matrix 'a' and therefore nothing really changes in this step.
-- 6. Introduce artificial variables \(z_j\) ( \(j=1,\ldots,m\) ) for each equation.
-- This changes the optimization problem and the solver has to take this into account.
--
-- The function returns a modified and standardized GLP and a matrix \(b\) that maps
-- solution vectors (including the objective functions value) of the modified GLP back
-- to solution vectors of the original GLP.
standardizeGLP :: GLP -> (GLP, Matrix Double)
standardizeGLP glp =
    let (glp1, b1) = resolveMin2Max glp
        (glp2, b2) = resolveGeqConstraints glp1
        (glp3, b3) = resolveNonPositiveVars glp2
        (glp4, b4) = resolveUnconstrainedVars glp3
        glp5 = introduceSlackVariables glp4
    in (glp5, b1 <> b2 <> b3 <> b4)

resolveMin2Max :: GLP -> (GLP, Matrix Double)
resolveMin2Max glp@GLP{..} = case dir of
        Maximize -> (glp, ident (1+m1+m2+m3))
        Minimize -> (glp{a = mapRow a 0 negate, dir = Maximize},
                     diagl (-1:(replicate (m1+m2+m3) 1)))

resolveGeqConstraints :: GLP -> (GLP, Matrix Double)
resolveGeqConstraints glp@GLP{..} =
    let a' = foldl (\a j -> mapRow a j negate) a [m1+1..m1+m2] -- negate rows from m1+1 to m1+m2
    in (glp{a = a', m1 = m1+m2, m2 = 0},
        ident (1+m1+m2+m3))

resolveNonPositiveVars :: GLP -> (GLP, Matrix Double)
resolveNonPositiveVars glp@GLP{..} =
    let varsToBeNegated = find (== -1) c
        varsToBeNegatedS = S.fromList varsToBeNegated
        (a', c') = foldl (\(a,c) i ->
                            (mapCol a (i+1) negate,
                             mapVectorWithIndex (\k x -> if k==i
                                                         then 1
                                                         else c!k)
                                                c))
                         (a, c)
                         (find (== -1) c)
    in (glp{a = a', c = c'},
        diagl $ 1:[if i `member` varsToBeNegatedS
                   then -1
                   else 1 | i <- [0..m1+m2+m3-1]])

resolveUnconstrainedVars :: GLP -> (GLP, Matrix Double)
resolveUnconstrainedVars glp@GLP{..} =
    let m = 1+m1+m2+m3
        zCol = vec1 m 0 1
        (a', c', b) =
            let listOfColsAndCs =
                   fmap (\(c1,ci,i1) ->
                            if ci==0
                            then ([c1, cmap negate c1],
                                  [1, 1],
                                  [vec1 m (i1+1) 1,
                                   vec1 m (i1+1) (-1)]) -- duplicate and adjust
                            else ([c1], [ci],
                                  [vec1 m (i1+1) 1])) -- return unchanged
                        (L.zip3 (toColumns a) (toList c) [0..(cols a)-1])
                columns = L.concat $ fmap (\(x,_,_) -> x) listOfColsAndCs
                cs = L.concat $ fmap (\(_,x,_) -> x) listOfColsAndCs
                bs = L.concat $ fmap (\(_,_,x) -> x) listOfColsAndCs
            in (fromColumns columns, fromList cs, fromColumns $ zCol:bs)
    in (glp{a=a', c=c'}, b)

introduceSlackVariables :: GLP -> GLP
introduceSlackVariables glp@GLP{..} =
    let m3' = m1+m2+m3
        n = cols a - 1
        a' = a ||| ((1><(m1+m2)) (replicate (m1+m2) 0) ===
                    diagl ((replicate m1 1)++(replicate m2 (-1))) ===
                    (m3><(m1+m2)) (replicate (m3*(m1+m2)) 0))
        rhs' = vjoin [rhs, LA.fromList [fi (n+1)..fi (n+1+m1+m2)]]
    in glp{a=a', m1=0, m2=0, m3=m3', rhs=rhs'}

vec1 :: Int -> Int -> Double -> Vector Double
vec1 n i x = fromList $ L.replicate i 0++x:L.replicate (n-i-1) 0

mapRow :: (Element t) => Matrix t -> Int -> (t -> t) -> Matrix t
mapRow a k f = runSTMatrix $ do
                    a' <- thawMatrix a
                    forM_ [0..cols a-1] (\i -> modifyMatrix a' k i f)
                    return a'

mapCol :: (Element t) => Matrix t -> Int -> (t -> t) -> Matrix t
mapCol a k f = runSTMatrix $ do
                    a' <- thawMatrix a
                    forM_ [0..rows a-1] (\i -> modifyMatrix a' i k f)
                    return a'

prettyPrintAsFormulas :: GLP -> String
prettyPrintAsFormulas glp@GLP{..} =
    let rs = LA.toRows a
        m = cols a - 1
        n = rows a - 1
        objFun = (show dir)++
                    " z_0 = "++linExp 0
        zRows = fmap (\k -> "z_"++show k++" = "++linExp k) [1..n]
        linExp k = L.intercalate " + " (
                    L.filter (/= "") $
                    fmap (\(a,i) ->
                            if a /= 0
                            then show a++(if i>0 then " "++varName i else "")
                            else "") (L.zip (toList $ rs!!k) [0..m]))
        varName i = if i <= n then "x_"++show i  else "y_"++show (i-n)
        zConstr  = L.intercalate ", " $ filter (/= "") $
                   [if m1 == 1
                    then "z_1 ≤ 0"
                    else if m1 == 2
                         then "z_1, z_2 ≤ 0"
                         else if m1 > 2
                              then "z_1, ..., z_"++show m1++" ≤ 0"
                              else "",
                    if m2 == 1
                    then "z_"++show (m1+1)++" ≥ 0"
                    else if m2 > 1
                         then "z_"++show (m1+1)++", ..., z_"++show (m1+m2)++" ≥ 0"
                         else "",
                    if m3 == 1
                    then "z_"++show (m1+m2+1)++" = 0"
                    else if m3 > 1
                         then "z_"++show (m1+m2+1)++", ..., z_"++show (m1+m2+m3)++" = 0"
                         else ""]
        xConstr ck k = "x_"++show k++if ck == -1
                                     then " ≤ 0"
                                     else if ck == 1
                                          then " ≥ 0"
                                          else " unconstrained"
        xConstrs = fmap (uncurry xConstr) (L.zip (toList c) [1..m])
    in L.intercalate "\n" (
            objFun:
            "Subject to":
            (fmap (\line -> "\t"++line) (zRows++[""]++
                                         [zConstr, ""]++
                                         xConstrs)))
            ++"\n"

(Right ex1') = ex1
ex1 = let a = (5><5)
                 [   0,    1,  1,  3, -0.5
                 , 740,   -1,  0, -2,  0
                 ,   0,    0, -2,  0,  7
                 ,   0.5,  0, -1,  1, -2
                 ,   9,   -1, -1, -1, -1]
          c = LA.fromList [-1, 1, 0, -1]
      in mkGLP a c (2, 1, 1) Minimize

allEx1 = do
    let g1 = ex1'
    putStr $ prettyPrintAsFormulas g1
    putStrLn $ "lhs = "++show (lhs g1)
    putStrLn $ "rhs = "++show (rhs g1)
    putStrLn "\nResolve min/max"
    let (g2, b1) = resolveMin2Max g1
    putStr $ prettyPrintAsFormulas g2
    putStrLn $ "lhs = "++show (lhs g2)
    putStrLn $ "rhs = "++show (rhs g2)
    print b1
    putStrLn "\nResove geq constraints"
    let (g3, b2) = resolveGeqConstraints g2
    putStr $ prettyPrintAsFormulas g3
    putStrLn $ "lhs = "++show (lhs g3)
    putStrLn $ "rhs = "++show (rhs g3)
    print b2
    putStrLn "\nResove non-positive variables"
    let (g4, b3) = resolveNonPositiveVars g3
    putStr $ prettyPrintAsFormulas g4
    putStrLn $ "lhs = "++show (lhs g4)
    putStrLn $ "rhs = "++show (rhs g4)
    print b3
    putStrLn "\nResove unconstrained variables"
    let (g5, b4) = resolveUnconstrainedVars g4
    putStr $ prettyPrintAsFormulas g5
    putStrLn $ "lhs = "++show (lhs g5)
    putStrLn $ "rhs = "++show (rhs g5)
    print b4
    putStrLn "\nIntroduce slack variables"
    let g6 = introduceSlackVariables g5
    putStr $ prettyPrintAsFormulas g6
    putStrLn $ "lhs = "++show (lhs g6)
    putStrLn $ "rhs = "++show (rhs g6)
    putStrLn "\n"
    return (g6, b1 <> b2 <> b3 <> b4)


{-# LANGUAGE PatternGuards #-}
-- |
-- Module    : Statistics.Matrix
-- Copyright : 2011 Aleksey Khudyakov, 2014 Bryan O'Sullivan
-- License   : BSD3
--
-- Basic matrix operations.
--
-- There isn't a widely used matrix package for Haskell yet, so
-- we implement the necessary minimum here.

module Statistics.Matrix
    ( -- * Data types
      Matrix(..)
    , Vector
      -- * Conversion from/to lists/vectors
    , fromVector
    , fromList
    , fromRowLists
    , fromRows
    , fromColumns
    , toVector
    , toList
    , toRows
    , toColumns
    , toRowLists
      -- * Other
    , generate
    , generateSym
    , ident
    , diag
    , dimension
    , center
    , multiply
    , multiplyV
    , transpose
    , power
    , norm
    , column
    , row
    , map
    , for
    , unsafeIndex
    , hasNaN
    , bounds
    , unsafeBounds
    ) where

import Prelude hiding (exponent, map)
import Control.Applicative ((<$>))
import Control.Monad.ST
import qualified Data.Vector.Unboxed as U
import           Data.Vector.Unboxed   ((!))
import qualified Data.Vector.Unboxed.Mutable as UM

import Numeric.Sum         (sumVector,kbn)
import Statistics.Matrix.Function
import Statistics.Matrix.Types
import Statistics.Matrix.Mutable  (unsafeNew,unsafeWrite,unsafeFreeze)

----------------------------------------------------------------
-- Conversion to/from vectors/lists
----------------------------------------------------------------

-- | Convert from a row-major list.
fromList :: Int                 -- ^ Number of rows.
         -> Int                 -- ^ Number of columns.
         -> [Double]            -- ^ Flat list of values, in row-major order.
         -> Matrix
fromList r c = fromVector r c . U.fromList

-- | create a matrix from a list of lists, as rows
fromRowLists :: [[Double]] -> Matrix
fromRowLists = fromRows . fmap U.fromList

-- | Convert from a row-major vector.
fromVector :: Int               -- ^ Number of rows.
           -> Int               -- ^ Number of columns.
           -> U.Vector Double   -- ^ Flat list of values, in row-major order.
           -> Matrix
fromVector r c v
  | r*c /= len = error "input size mismatch"
  | otherwise  = Matrix r c v
  where len    = U.length v

-- | create a matrix from a list of vectors, as rows
fromRows :: [Vector] -> Matrix
fromRows xs
  | [] <- xs        = error "Statistics.Matrix.fromRows: empty list of rows!"
  | any (/=nCol) ns = error "Statistics.Matrix.fromRows: row sizes do not match"
  | nCol == 0       = error "Statistics.Matrix.fromRows: zero columns in matrix"
  | otherwise       = fromVector nRow nCol (U.concat xs)
  where
    nCol:ns = U.length <$> xs
    nRow    = length xs

-- | create a matrix from a list of vectors, as columns
fromColumns :: [Vector] -> Matrix
fromColumns = transpose . fromRows

-- | Convert to a row-major flat vector.
toVector :: Matrix -> U.Vector Double
toVector (Matrix _ _ v) = v

-- | Convert to a row-major flat list.
toList :: Matrix -> [Double]
toList = U.toList . toVector

-- | Convert to a list of lists, as rows
toRowLists :: Matrix -> [[Double]]
toRowLists (Matrix _ nCol v)
  = chunks $ U.toList v
  where
    chunks [] = []
    chunks xs = case splitAt nCol xs of
      (rowE,rest) -> rowE : chunks rest

-- | Convert to a list of vectors, as rows
toRows :: Matrix -> [Vector]
toRows (Matrix _ nCol v) = chunks v
  where
    chunks xs
      | U.null xs = []
      | otherwise = case U.splitAt nCol xs of
          (rowE,rest) -> rowE : chunks rest

-- | Convert to a list of vectors, as columns
toColumns :: Matrix -> [Vector]
toColumns = toRows . transpose

----------------------------------------------------------------
-- Other
----------------------------------------------------------------

-- | Generate matrix using function
generate :: Int                 -- ^ Number of rows
         -> Int                 -- ^ Number of columns
         -> (Int -> Int -> Double)
            -- ^ Function which takes /row/ and /column/ as argument.
         -> Matrix
generate nRow nCol f
  = Matrix nRow nCol $ U.generate (nRow*nCol) $ \i ->
      let (r,c) = i `quotRem` nCol in f r c

-- | Generate symmetric square matrix using function
generateSym
  :: Int                 -- ^ Number of rows and columns
  -> (Int -> Int -> Double)
     -- ^ Function which takes /row/ and /column/ as argument. It must
     --   be symmetric in arguments: @f i j == f j i@
  -> Matrix
generateSym n f = runST $ do
  m <- unsafeNew n n
  for 0 n $ \r -> do
    unsafeWrite m r r (f r r)
    for (r+1) n $ \c -> do
      let x = f r c
      unsafeWrite m r c x
      unsafeWrite m c r x
  unsafeFreeze m

-- | Create the square identity matrix with given dimensions.
ident :: Int -> Matrix
ident n = diag $ U.replicate n 1.0

-- | Create a square matrix with given diagonal, other entries default to 0
diag :: Vector -> Matrix
diag v
  = Matrix n n $ U.create $ do
      arr <- UM.replicate (n*n) 0
      for 0 n $ \i ->
        UM.unsafeWrite arr (i*n + i) (v ! i)
      return arr
  where
    n = U.length v

-- | Return the dimensions of this matrix, as a (row,column) pair.
dimension :: Matrix -> (Int, Int)
dimension (Matrix r c _) = (r, c)

-- | Matrix-matrix multiplication. Matrices must be of compatible
-- sizes (/note: not checked/).
multiply :: Matrix -> Matrix -> Matrix
multiply m1@(Matrix r1 _ _) m2@(Matrix _ c2 _) =
  Matrix r1 c2 $ U.generate (r1*c2) go
  where
    go t = sumVector kbn $ U.zipWith (*) (row m1 i) (column m2 j)
      where (i,j) = t `quotRem` c2

-- | Matrix-vector multiplication.
multiplyV :: Matrix -> Vector -> Vector
multiplyV m v
  | cols m == c = U.generate (rows m) (sumVector kbn . U.zipWith (*) v . row m)
  | otherwise   = error $ "matrix/vector unconformable " ++ show (cols m,c)
  where c = U.length v

-- | Raise matrix to /n/th power. Power must be positive
-- (/note: not checked).
power :: Matrix -> Int -> Matrix
power mat 1 = mat
power mat n = res
  where
    mat2 = power mat (n `quot` 2)
    pow  = multiply mat2 mat2
    res | odd n     = multiply pow mat
        | otherwise = pow

-- | Element in the center of matrix (not corrected for exponent).
center :: Matrix -> Double
center mat@(Matrix r c _) =
    unsafeBounds U.unsafeIndex mat (r `quot` 2) (c `quot` 2)

-- | Calculate the Euclidean norm of a vector.
norm :: Vector -> Double
norm = sqrt . sumVector kbn . U.map square

-- | Return the given column.
column :: Matrix -> Int -> Vector
column (Matrix r c v) j= U.generate r (\i -> v `U.unsafeIndex` (j + i * c))
{-# INLINE column #-}

-- | Return the given row.
row :: Matrix -> Int -> Vector
row (Matrix _ c v) i = U.slice (c*i) c v

unsafeIndex :: Matrix
            -> Int              -- ^ Row.
            -> Int              -- ^ Column.
            -> Double
unsafeIndex = unsafeBounds U.unsafeIndex

-- | Apply function to every element of matrix
map :: (Double -> Double) -> Matrix -> Matrix
map f (Matrix r c v) = Matrix r c (U.map f v)

-- | Indicate whether any element of the matrix is @NaN@.
hasNaN :: Matrix -> Bool
hasNaN = U.any isNaN . toVector

-- | Given row and column numbers, calculate the offset into the flat
-- row-major vector.
bounds :: (Vector -> Int -> r) -> Matrix -> Int -> Int -> r
bounds k (Matrix rs cs v) r c
  | r < 0 || r >= rs = error "row out of bounds"
  | c < 0 || c >= cs = error "column out of bounds"
  | otherwise        = k v $! r * cs + c
{-# INLINE bounds #-}

-- | Given row and column numbers, calculate the offset into the flat
-- row-major vector, without checking.
unsafeBounds :: (Vector -> Int -> r) -> Matrix -> Int -> Int -> r
unsafeBounds k (Matrix _ cs v) r c = k v $! r * cs + c
{-# INLINE unsafeBounds #-}

transpose :: Matrix -> Matrix
transpose m@(Matrix r0 c0 _) = Matrix c0 r0 . U.generate (r0*c0) $ \i ->
  let (r,c) = i `quotRem` r0
  in unsafeIndex m c r


{-# LANGUAGE BangPatterns #-}
module ListStats (mean, var, stDiv, nCr,  varM, qf,unStdNorm,unStdNormSplits, -- cov, corr,
                  stDivM, stNorm, erf, cn, erfInv, binomial, cdf, cdF, meanBy, stDivBy,
                  normal, gammaln,  adjustedMatrix, adjustedMatrixN, adjustedMatrixPN, normalP,
                  module ListAlgebra,
                 ) where

import Numeric
import Numeric.LinearAlgebra
import qualified Data.Vector as V (map, fromList)
import Data.List
import Control.DeepSeq (NFData)
import LINreg2 (foldl'Rnf )
import Data.Ord (comparing)
import ListAlgebra

----------------------------------------------------------------------------------------
--- mean, variance and correlations
----------------------------------------------------------------------------------------

insertZeros :: Num a =>  [a] -> [Int] -> [a]
insertZeros xs   []   = xs
insertZeros xs (a:as) =
        let   (front,back) = splitAt (a-1) xs  in
            (1 :front) ++ (0 : insertZeros back  as)

-- map a function then calculate the standard deviation
stDivBy :: (a -> Double) -> [a] -> Double
stDivBy f xs = (sumBy (sqr . (\a -> a - ns) . f) xs) / k
    where
       (sm,k) = foldl' (\(n,m) a -> (f a + n, m + 1)) (0,0) xs
       ns     = sm / k
       sqr x  = x * x

-- map a function then calculate the mean
meanBy :: (a -> Double) -> [a] -> Double
meanBy g xs = sm / k
   where (sm,k) = foldl' (\(n,m) a -> (g a + n, m + 1)) (0,0) xs

----
mean  :: (NFData a, Fractional a) => [a] -> a
mean     xs                    = fst suml / (snd suml - 1)
                    where suml = foldl'Rnf  kF (0, 1) xs
                          kF (!n , !r) b =  (n + b, 1 + r)
---
var  :: (NFData a, Floating a) => [a] -> a
var     xs                    =  (fst suml / snd suml)
                   where suml = foldl'Rnf  kF (0, 0) xs
                         kF (!n , !r) b = (n + (b - k)^2 , 1 + r)
                         k    = mean xs
-- variance of a list of values with the mean given
varM  :: (NFData a , Floating a) => a -> [a] -> a
varM     k  xs              =  fst suml / snd suml
                 where suml = foldl'Rnf  kF (0, 0) xs
                       kF (!n , !r) b = (n + (b - k)^2 , 1 + r)
--
--- the standard deviation without the Mean given
stDiv   :: (NFData a , Floating a) =>  [a] -> a
stDiv                       = sqrt . var

---- standard deviation with the mean given
stDivM        k             = sqrt . (varM k)

--- standardise by subtracting the mean and dividing by the standard deviation
--stNorm :: [Double] -> [Double]
stNorm       xs             = map (\ a -> (a - m) / std) xs
                 where m    = mean xs
                       std  = stDivM m xs
---
unStdNorm  xs ys  = map (\a -> std * a + m)  ys
                where  m    = mean xs
                       std  = stDivM m xs

unStdNormSplits  xs zs = unStdNorm (insertZeros xs zs)

---   selecting the 'best' sublist of lenght n of a list of values
bestN :: (Ord a , Ord b)  => (a -> b) -> Int -> [a] -> [a]
bestN  comp   lim  xs
    | length xs <= lim   =  xs
    | length back <= lim =  myBest lim comp xs
    | otherwise          =  min (sortBy (comparing comp) front) (bestN comp lim back)
    where
        (front, back) = splitAt lim xs
        myBest n  f   = myTake n .  sortBy (comparing f)
        mn            = lim `div` 2

--- computed the weighted average of the residuals of values, which represent
--- points on the y-axis plotted against x-values [1 .. n], where n is the
-- length of the list. The weight biased towards is determind as the
wgtAvgRd ::  (NFData a, Enum a, Fractional a) => Int -> [[a]]  -> a
wgtAvgRd    n   xs     = weightsSum / len
    where
        (weightsSum , len) = foldl'Rnf (getWeightedAvg (2 * n)) (0,0) xs
        getWeightedAvg k (wa,ln) ys  =   (wa + wVal , ln  + 1)
            where
                f     = fromIntegral
                len   = f $ length ys
                ks    = [1 .. len]
                sumSqrs  as bs = foldl'Rnf (\a (b,c) -> a + b * c) 0 (zip as bs)
                wVal  =  (sumSqrs ks ys) / (sumSqrs ks ks)   -- (f k / len)  *

--------------------------------------------------------------------------------------------
--  Defining the X matrix for the MML function
--------------------------------------------------------------------------------------------
-- We combine adjustMatrix, fillZeros and allPieces to
-- give the final adjusted matrix. This matrix is standardised, then transposed,
-- but the transposeition can be removed if we require that to occur later on
adjustedMatrix  xs   =  [( tS . adjustMatrix . fillZeroes_aux n 0 1) yss | yss <- allPieces xs]
                         where n  = length xs
                               tS =  fromLists . map (1:) . transpose . map stNorm

adjustedMatrixN m xs =  [tS $ adjustMatrix $ fillZeroes_aux n 0 1 yss | yss <- nPieces m xs]
                         where n  = length xs
                               tS = map (1:). transpose . map stNorm
adjustedMatrixPN k m xs = [(tS $ adjustMatrix $ fillZeroes_aux n 0 1 yss, cuts) |
                                (yss,cuts) <- nPiecesP k m Nothing xs, not $ null cuts]
                         where n  = length xs
                               tS = map (1:). transpose . map stNorm
--
-- test function
joins :: Eq a => [a] -> [[a]]
joins  xs | null xs = [] -- && (length (concat a) > 5)
          | otherwise = joins_aux xs ++ joins (tail xs)
    where
        joins_aux (x : y : ys) = ([x,y] :  (map (x :) $ joins_aux (y : ys)))
        joins_aux  _           = []
--------------------------------------------------------------------------------------------
-- calculating the binomial coefficient, and normal prob density
-------------------------------------------------------------------------------------------
nCr ::  (Integral a, Ord a) => a -> a -> a
nCr   n    r   | n < 0 || r < 0  =  error "can find combinations for negative numbers"
nCr   n   0                      =  1
nCr   0   _                      =  0
nCr   !n   !k                    =  nCr (n-1) (k-1) * n `div` k

--- binomial distribution
binomial :: (Integral a, Ord b, Fractional b, Ord a) => a -> a -> b -> b
binomial   n   r   p  | p > 1 || p < 0   = error "Invalid probability value" ---n nCr r *
binomial   n   k   p  | k >= 0 && k <= n = let  nck = fromIntegral (nCr n k) in nck * (p^k) * (1 - p)^(n - k)
                      | otherwise        = 0

-- value of x from a normal distribution with mean 'mn' and standard deviation "std"
normal x mn std = (1 / (std * sqrt (2 * pi))) * exp(- 0.5 * ((x - mn)^2) / (std ^2))
--normal x mn std = cdf (x + eps) mn std  - cdf (x - eps) mn std
    where
        eps = 0.0101095
        --z   = (x - mn) / std

-- calculating the normal with a given precision
normalP eps x mn std = cdf (x + eps) mn std  - cdf (x - eps) mn std
------------------------------------------------------------------------------------------------
--- approximating the error functions and its inverse
----------------------------------------------------------------------------------------------
erf	x	| x  <  0    =  -1 *  erf (- 1 * x)
		| otherwise  = 1 - t * exp k
    where
        t  = 1 / (1 + 0.5 * abs x)
	t3 = t * (- 0.82215223 +  t *  0.17087277)
        t2 = t * ( 0.27886807 +  t * (-1.13520398 +  t *  1.48851587 + t3))
        t4 = t * ( 0.37409196 +  t * ( 0.09678418 +  t *  (- 0.18628806) + t2))
	k  = - x * x -  1.26551223 + t *  1.00002368 +  t4
--

--cn  :: Int -> Double
cn	 0	=	1
cn	 n	=	foldr (\k  b -> b + (cn k  * cn (n - 1 - k)) / ((k + 1) * (2 * k + 1))) 0 [0.. n - 1]

-- inverse of the error function
--erfInv :: (Fractional a, Floating a, Enum a) => a -> a
erfInv	y | y <  -1 || y > 1  = error "erfInv argument out of range"
          | abs y == 1        = pol (-y* log (0.0))
          | y < - 0.7         = pol (- 1 * folD z (head c) (tail c) / folD z 1 d)
          | y <   0.7         = pol ( y *  folD (y^2) (head a) (tail a) / folD (y^2) 1 b)
          | otherwise         = pol (folD z1 (head c) (tail c)  / folD z1 1 d)
                  where a     = [0.886226899, -1.645349621,  0.914624893, -0.140543331]
                        b     = [-2.118377725,  1.442710462, -0.329097515,  0.012229801]
                        c     = [-1.970840454, -1.624906493,  3.429567803,  1.641345311]
                        d     = [3.543889200,  1.637067800]
                        z     = sqrt(- log ((1.0+y)/2.0))
                        z1    = sqrt(- log ((1.0+y)/2.0))
                        pol x = x - (erf(x) - y) / (2.0 /(sqrt pi)  * exp(-x*x))
                        folD  zz a ls = foldl'Rnf  (\a b -> a + zz * (fst b + snd b * zz)) a (zipp ls)

-- commulative density function
cdf   x mean std | std < 0      = error"Negative variance to CDF"
cdf   x mean std | otherwise    =  0.5 * (1 + erf ((x - mean) / (std * sqrt 2)))

-- the CDF for the standard normal distribution
cdF  n  =  cdf  n 0 1

-- the q function, which is 1 - CDF (x) (here given for the standard normal dist)
qf  x  =  1 - cdF x
--------------------------------------------------------------------------------------
--- The gamm function takel  from wikipedia
--------------------------------------------------------------------------------------
gammaln :: Double -> Double
gammaln xx = -tmp' + log(2.5066282746310005 * ser' / xx)
			where
			tmp' = (xx+5.5) - (xx+0.5)*log(xx+5.5)
			ser' = ser + sum (zipWith (/) cof [xx+1..])
			cof  = [76.18009172947146,-86.50532032941677,24.01409824083091,-1.231739572450155,0.001208650973866179,-0.000005395239384953]
			ser  = 1.000000000190015


-- |
-- Module    : Statistics.Matrix
-- Copyright : 2011 Aleksey Khudyakov, 2014 Bryan O'Sullivan
-- License   : BSD3
--
-- Basic matrix operations.
--
-- There isn't a widely used matrix package for Haskell yet, so
-- we implement the necessary minimum here.

module Statistics.Matrix
    (
      Matrix(..)
    , Vector
    , fromList
    , fromVector
    , toVector
    , toList
    , dimension
    , center
    , multiply
    , multiplyV
    , transpose
    , power
    , norm
    , column
    , row
    , map
    , for
    , unsafeIndex
    , hasNaN
    , bounds
    , unsafeBounds
    ) where

import Prelude hiding (exponent, map, sum)
import Statistics.Function (for, square)
import Statistics.Matrix.Types
import Statistics.Sample.Internal (sum)
import qualified Data.Vector.Unboxed as U

-- | Convert from a row-major list.
fromList :: Int                 -- ^ Number of rows.
         -> Int                 -- ^ Number of columns.
         -> [Double]            -- ^ Flat list of values, in row-major order.
         -> Matrix
fromList r c = fromVector r c . U.fromList

-- | Convert from a row-major vector.
fromVector :: Int               -- ^ Number of rows.
           -> Int               -- ^ Number of columns.
           -> U.Vector Double   -- ^ Flat list of values, in row-major order.
           -> Matrix
fromVector r c v
  | r*c /= len = error "input size mismatch"
  | otherwise  = Matrix r c 0 v
  where len    = U.length v

-- | Convert to a row-major flat vector.
toVector :: Matrix -> U.Vector Double
toVector (Matrix _ _ _ v) = v

-- | Convert to a row-major flat list.
toList :: Matrix -> [Double]
toList = U.toList . toVector

-- | Return the dimensions of this matrix, as a (row,column) pair.
dimension :: Matrix -> (Int, Int)
dimension (Matrix r c _ _) = (r, c)

-- | Avoid overflow in the matrix.
avoidOverflow :: Matrix -> Matrix
avoidOverflow m@(Matrix r c e v)
  | center m > 1e140 = Matrix r c (e + 140) (U.map (* 1e-140) v)
  | otherwise        = m

-- | Matrix-matrix multiplication. Matrices must be of compatible
-- sizes (/note: not checked/).
multiply :: Matrix -> Matrix -> Matrix
multiply m1@(Matrix r1 _ e1 _) m2@(Matrix _ c2 e2 _) =
  Matrix r1 c2 (e1 + e2) $ U.generate (r1*c2) go
  where
    go t = sum $ U.zipWith (*) (row m1 i) (column m2 j)
      where (i,j) = t `quotRem` c2

-- | Matrix-vector multiplication.
multiplyV :: Matrix -> Vector -> Vector
multiplyV m v
  | cols m == c = U.generate (rows m) (sum . U.zipWith (*) v . row m)
  | otherwise   = error $ "matrix/vector unconformable " ++ show (cols m,c)
  where c = U.length v

-- | Raise matrix to /n/th power. Power must be positive
-- (/note: not checked).
power :: Matrix -> Int -> Matrix
power mat 1 = mat
power mat n = avoidOverflow res
  where
    mat2 = power mat (n `quot` 2)
    pow  = multiply mat2 mat2
    res | odd n     = multiply pow mat
        | otherwise = pow

-- | Element in the center of matrix (not corrected for exponent).
center :: Matrix -> Double
center mat@(Matrix r c _ _) =
    unsafeBounds U.unsafeIndex mat (r `quot` 2) (c `quot` 2)

-- | Calculate the Euclidean norm of a vector.
norm :: Vector -> Double
norm = sqrt . sum . U.map square

-- | Return the given column.
column :: Matrix -> Int -> Vector
column (Matrix r c _ v) i = U.backpermute v $ U.enumFromStepN i c r
{-# INLINE column #-}

-- | Return the given row.
row :: Matrix -> Int -> Vector
row (Matrix _ c _ v) i = U.slice (c*i) c v

unsafeIndex :: Matrix
            -> Int              -- ^ Row.
            -> Int              -- ^ Column.
            -> Double
unsafeIndex = unsafeBounds U.unsafeIndex

map :: (Double -> Double) -> Matrix -> Matrix
map f (Matrix r c e v) = Matrix r c e (U.map f v)

-- | Indicate whether any element of the matrix is @NaN@.
hasNaN :: Matrix -> Bool
hasNaN = U.any isNaN . toVector

-- | Given row and column numbers, calculate the offset into the flat
-- row-major vector.
bounds :: (Vector -> Int -> r) -> Matrix -> Int -> Int -> r
bounds k (Matrix rs cs _ v) r c
  | r < 0 || r >= rs = error "row out of bounds"
  | c < 0 || c >= cs = error "column out of bounds"
  | otherwise        = k v $! r * cs + c
{-# INLINE bounds #-}

-- | Given row and column numbers, calculate the offset into the flat
-- row-major vector, without checking.
unsafeBounds :: (Vector -> Int -> r) -> Matrix -> Int -> Int -> r
unsafeBounds k (Matrix _ cs _ v) r c = k v $! r * cs + c
{-# INLINE unsafeBounds #-}

transpose :: Matrix -> Matrix
transpose m@(Matrix r0 c0 e _) = Matrix c0 r0 e . U.generate (r0*c0) $ \i ->
  let (r,c) = i `quotRem` r0
  in unsafeIndex m c r