{-# LANGUAGE BangPatterns         #-}
{-# LANGUAGE ScopedTypeVariables  #-}

module Math.HypergeoMatrix.HypergeoMatrix (hypergeomat) where
import           Control.Monad                (when)
import           Data.Array                   hiding (index)
import           Data.Array.IO                hiding (index)
import           Data.Sequence                (Seq, index, update, (!?), (|>))
import qualified Data.Sequence                as S
import           Math.HypergeoMatrix.Internal 

hypergeoI :: forall a. (Eq a, Fractional a, BaseFrac a)
  => Int -> BaseFracType a -> [a] -> [a] -> Int -> a -> a
hypergeoI :: forall a.
(Eq a, Fractional a, BaseFrac a) =>
Int -> BaseFracType a -> [a] -> [a] -> Int -> a -> a
hypergeoI Int
m BaseFracType a
alpha [a]
a [a]
b Int
n a
x =
  a
1 forall a. Num a => a -> a -> a
+ Fractional a => Int -> a -> Int -> [Int] -> a
summation' Int
0 a
1 Int
m []
  where
  summation' :: Fractional a => Int -> a -> Int -> [Int] -> a
  summation' :: Fractional a => Int -> a -> Int -> [Int] -> a
summation' Int
i a
z Int
j [Int]
kappa = Int -> a -> a -> a
go Int
1 a
z a
0
    where
    go :: Int -> a -> a -> a
    go :: Int -> a -> a -> a
go Int
kappai a
zz a
s
      | Int
i forall a. Eq a => a -> a -> Bool
== Int
0 Bool -> Bool -> Bool
&& Int
kappai forall a. Ord a => a -> a -> Bool
> Int
j Bool -> Bool -> Bool
|| Int
iforall a. Ord a => a -> a -> Bool
>Int
0 Bool -> Bool -> Bool
&& Int
kappai forall a. Ord a => a -> a -> Bool
> forall a. Ord a => a -> a -> a
min ([Int]
kappaforall a. [a] -> Int -> a
!!(Int
iforall a. Num a => a -> a -> a
-Int
1)) Int
j = a
s
      | Bool
otherwise = Int -> a -> a -> a
go (Int
kappai forall a. Num a => a -> a -> a
+ Int
1) a
z' a
s''
      where
      kappa' :: [Int]
kappa' = [Int]
kappa forall a. [a] -> [a] -> [a]
++ [Int
kappai]
      t :: a
t = forall a.
(Fractional a, Eq a, BaseFrac a) =>
BaseFracType a -> [a] -> [a] -> Seq Int -> a
_T BaseFracType a
alpha [a]
a [a]
b (forall a. [a] -> Seq a
S.fromList forall a b. (a -> b) -> a -> b
$ forall a. (a -> Bool) -> [a] -> [a]
filter (forall a. Ord a => a -> a -> Bool
> Int
0) [Int]
kappa') -- inutile de filtrer
      z' :: a
z' = a
zz forall a. Num a => a -> a -> a
* a
x forall a. Num a => a -> a -> a
*
        (forall a b. (Integral a, Num b) => a -> b
fromIntegral (Int
nforall a. Num a => a -> a -> a
-Int
i) forall a. Num a => a -> a -> a
+ forall a. BaseFrac a => BaseFracType a -> a
inject BaseFracType a
alpha forall a. Num a => a -> a -> a
* (forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
kappaiforall a. Num a => a -> a -> a
-a
1)) forall a. Num a => a -> a -> a
* a
t
      s' :: a
s' = if Int
j forall a. Ord a => a -> a -> Bool
> Int
kappai Bool -> Bool -> Bool
&& Int
i forall a. Ord a => a -> a -> Bool
<= Int
n
        then a
s forall a. Num a => a -> a -> a
+ Fractional a => Int -> a -> Int -> [Int] -> a
summation' (Int
iforall a. Num a => a -> a -> a
+Int
1) a
z' (Int
jforall a. Num a => a -> a -> a
-Int
kappai) [Int]
kappa'
        else a
s
      s'' :: a
s'' = a
s' forall a. Num a => a -> a -> a
+ a
z'


summation :: forall a. (Fractional a, Eq a, BaseFrac a)
  => [a] -> [a] -> [a] -> Seq (Maybe Int) -> Int -> BaseFracType a -> Int
     -> a -> Int -> Seq Int -> IOArray (Int, Int) a -> IO a
summation :: forall a.
(Fractional a, Eq a, BaseFrac a) =>
[a]
-> [a]
-> [a]
-> Seq (Maybe Int)
-> Int
-> BaseFracType a
-> Int
-> a
-> Int
-> Seq Int
-> IOArray (Int, Int) a
-> IO a
summation [a]
a [a]
b [a]
x Seq (Maybe Int)
dico Int
n BaseFracType a
alpha Int
i a
z Int
j Seq Int
kappa IOArray (Int, Int) a
jarray
  = if Int
i forall a. Eq a => a -> a -> Bool
== Int
n
    then
      forall (m :: * -> *) a. Monad m => a -> m a
return a
0
    else do
      let lkappa :: Int
lkappa = Seq Int
kappa forall a. Seq a -> Int -> a
`index` (forall a. Seq a -> Int
S.length Seq Int
kappa forall a. Num a => a -> a -> a
- Int
1)
      let go :: Int -> a -> a -> IO a
          go :: Int -> a -> a -> IO a
go Int
kappai !a
z' !a
s
            | Int
i forall a. Eq a => a -> a -> Bool
== Int
0 Bool -> Bool -> Bool
&& Int
kappai forall a. Ord a => a -> a -> Bool
> Int
j Bool -> Bool -> Bool
|| Int
i forall a. Ord a => a -> a -> Bool
> Int
0 Bool -> Bool -> Bool
&& Int
kappai forall a. Ord a => a -> a -> Bool
> forall a. Ord a => a -> a -> a
min Int
lkappa Int
j =
              forall (m :: * -> *) a. Monad m => a -> m a
return a
s
            | Bool
otherwise = do
              let kappa' :: Seq Int
kappa' = Seq Int
kappa forall a. Seq a -> a -> Seq a
|> Int
kappai
                  nkappa :: Int
nkappa = Seq (Maybe Int) -> Seq Int -> Int
_nkappa Seq (Maybe Int)
dico Seq Int
kappa'
                  z'' :: a
z'' = a
z' forall a. Num a => a -> a -> a
* forall a.
(Fractional a, Eq a, BaseFrac a) =>
BaseFracType a -> [a] -> [a] -> Seq Int -> a
_T BaseFracType a
alpha [a]
a [a]
b Seq Int
kappa'
                  lkappa' :: Int
lkappa' = forall a. Seq a -> Int
S.length Seq Int
kappa'
              forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when (Int
nkappa forall a. Ord a => a -> a -> Bool
> Int
1 Bool -> Bool -> Bool
&& (Int
lkappa' forall a. Eq a => a -> a -> Bool
== Int
1 Bool -> Bool -> Bool
|| Seq Int
kappa' forall a. Seq a -> Int -> Maybe a
!? Int
1 forall a. Eq a => a -> a -> Bool
== forall a. a -> Maybe a
Just Int
0)) forall a b. (a -> b) -> a -> b
$ do
                a
entry <- forall (a :: * -> * -> *) e (m :: * -> *) i.
(MArray a e m, Ix i) =>
a i e -> i -> m e
readArray IOArray (Int, Int) a
jarray (Int
nkappa forall a. Num a => a -> a -> a
- Int
1, Int
1)
                let kap0m1' :: a
kap0m1' = forall a b. (Integral a, Num b) => a -> b
fromIntegral (Seq Int
kappa' forall a. Seq a -> Int -> a
`index` Int
0 forall a. Num a => a -> a -> a
- Int
1)
                    newval :: a
newval = forall a. [a] -> a
head [a]
x forall a. Num a => a -> a -> a
* (a
1 forall a. Num a => a -> a -> a
+ forall a. BaseFrac a => BaseFracType a -> a
inject BaseFracType a
alpha forall a. Num a => a -> a -> a
* a
kap0m1') forall a. Num a => a -> a -> a
* a
entry
                forall (a :: * -> * -> *) e (m :: * -> *) i.
(MArray a e m, Ix i) =>
a i e -> i -> e -> m ()
writeArray IOArray (Int, Int) a
jarray (Int
nkappa, Int
1) a
newval
              let go' :: Int -> IO ()
                  go' :: Int -> IO ()
go' Int
t
                    | Int
t forall a. Eq a => a -> a -> Bool
== Int
n forall a. Num a => a -> a -> a
+ Int
1 = forall (m :: * -> *) a. Monad m => a -> m a
return ()
                    | Bool
otherwise = do
                      ()
_ <- forall a.
(Fractional a, BaseFrac a) =>
BaseFracType a
-> [a]
-> Seq (Maybe Int)
-> Int
-> a
-> Int
-> Int
-> Seq Int
-> IOArray (Int, Int) a
-> Seq Int
-> Int
-> IO ()
jack BaseFracType a
alpha [a]
x Seq (Maybe Int)
dico Int
0 a
1 Int
0 Int
t Seq Int
kappa' IOArray (Int, Int) a
jarray Seq Int
kappa' Int
nkappa
                      Int -> IO ()
go' (Int
t forall a. Num a => a -> a -> a
+ Int
1)
              ()
_ <- Int -> IO ()
go' Int
2
              a
entry' <- forall (a :: * -> * -> *) e (m :: * -> *) i.
(MArray a e m, Ix i) =>
a i e -> i -> m e
readArray IOArray (Int, Int) a
jarray (Int
nkappa, Int
n)
              let s' :: a
s' = a
s forall a. Num a => a -> a -> a
+ a
z'' forall a. Num a => a -> a -> a
* a
entry'
              if Int
j forall a. Ord a => a -> a -> Bool
> Int
kappai Bool -> Bool -> Bool
&& Int
i forall a. Ord a => a -> a -> Bool
<= Int
n
                then do
                  a
s'' <-
                    forall a.
(Fractional a, Eq a, BaseFrac a) =>
[a]
-> [a]
-> [a]
-> Seq (Maybe Int)
-> Int
-> BaseFracType a
-> Int
-> a
-> Int
-> Seq Int
-> IOArray (Int, Int) a
-> IO a
summation
                      [a]
a
                      [a]
b
                      [a]
x
                      Seq (Maybe Int)
dico
                      Int
n
                      BaseFracType a
alpha
                      (Int
i forall a. Num a => a -> a -> a
+ Int
1)
                      a
z''
                      (Int
j forall a. Num a => a -> a -> a
- Int
kappai)
                      Seq Int
kappa'
                      IOArray (Int, Int) a
jarray
                  Int -> a -> a -> IO a
go (Int
kappai forall a. Num a => a -> a -> a
+ Int
1) a
z'' (a
s' forall a. Num a => a -> a -> a
+ a
s'')
                else Int -> a -> a -> IO a
go (Int
kappai forall a. Num a => a -> a -> a
+ Int
1) a
z'' a
s'
      Int -> a -> a -> IO a
go Int
1 a
z a
0

jack :: (Fractional a, BaseFrac a)
  => BaseFracType a -> [a] -> Seq (Maybe Int) -> Int -> a -> Int -> Int
  -> Seq Int -> IOArray (Int, Int) a -> Seq Int -> Int -> IO ()
jack :: forall a.
(Fractional a, BaseFrac a) =>
BaseFracType a
-> [a]
-> Seq (Maybe Int)
-> Int
-> a
-> Int
-> Int
-> Seq Int
-> IOArray (Int, Int) a
-> Seq Int
-> Int
-> IO ()
jack BaseFracType a
alpha [a]
x Seq (Maybe Int)
dico Int
k a
beta Int
c Int
t Seq Int
mu IOArray (Int, Int) a
jarray Seq Int
kappa Int
nkappa = do
  let i0 :: Int
i0 = forall a. Ord a => a -> a -> a
max Int
k Int
1
      i1 :: Int
i1 = forall a. Seq a -> Int
S.length (Seq Int -> Seq Int
cleanPart Seq Int
mu) forall a. Num a => a -> a -> a
+ Int
1
      go :: Int -> IO ()
      go :: Int -> IO ()
go Int
i
        | Int
i forall a. Eq a => a -> a -> Bool
== Int
i1 = forall (m :: * -> *) a. Monad m => a -> m a
return ()
        | Bool
otherwise
         = do
          let u :: Int
u = Seq Int
mu forall a. Seq a -> Int -> a
`index` (Int
i forall a. Num a => a -> a -> a
- Int
1)
          forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when (forall a. Seq a -> Int
S.length Seq Int
mu forall a. Eq a => a -> a -> Bool
== Int
i Bool -> Bool -> Bool
|| Int
u forall a. Ord a => a -> a -> Bool
> Seq Int
mu forall a. Seq a -> Int -> a
`index` Int
i) forall a b. (a -> b) -> a -> b
$ do
            let gamma :: a
gamma = a
beta forall a. Num a => a -> a -> a
* forall a.
(Fractional a, BaseFrac a) =>
Seq Int -> Seq Int -> Int -> BaseFracType a -> a
_betaratio Seq Int
kappa Seq Int
mu Int
i BaseFracType a
alpha
                mu' :: Seq Int
mu' = Seq Int -> Seq Int
cleanPart forall a b. (a -> b) -> a -> b
$ forall a. Int -> a -> Seq a -> Seq a
update (Int
iforall a. Num a => a -> a -> a
-Int
1) (Int
u forall a. Num a => a -> a -> a
- Int
1) Seq Int
mu
                nmu :: Int
nmu = Seq (Maybe Int) -> Seq Int -> Int
_nkappa Seq (Maybe Int)
dico Seq Int
mu'
            if forall a. Seq a -> Int
S.length Seq Int
mu' forall a. Ord a => a -> a -> Bool
>= Int
i Bool -> Bool -> Bool
&& Int
u forall a. Ord a => a -> a -> Bool
> Int
1  -- "not (S.null mu')" useless because i>=1
              then
                forall a.
(Fractional a, BaseFrac a) =>
BaseFracType a
-> [a]
-> Seq (Maybe Int)
-> Int
-> a
-> Int
-> Int
-> Seq Int
-> IOArray (Int, Int) a
-> Seq Int
-> Int
-> IO ()
jack BaseFracType a
alpha [a]
x Seq (Maybe Int)
dico Int
i a
gamma (Int
c forall a. Num a => a -> a -> a
+ Int
1) Int
t Seq Int
mu' IOArray (Int, Int) a
jarray Seq Int
kappa Int
nkappa
              else
                forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when (Int
nkappa forall a. Ord a => a -> a -> Bool
> Int
1) forall a b. (a -> b) -> a -> b
$ do
                  a
entry' <- forall (a :: * -> * -> *) e (m :: * -> *) i.
(MArray a e m, Ix i) =>
a i e -> i -> m e
readArray IOArray (Int, Int) a
jarray (Int
nkappa, Int
t)
                  if Bool -> Bool
not (forall a. Seq a -> Bool
S.null Seq Int
mu') -- any (> 0) mu'
                    then do
                      a
entry <- forall (a :: * -> * -> *) e (m :: * -> *) i.
(MArray a e m, Ix i) =>
a i e -> i -> m e
readArray IOArray (Int, Int) a
jarray (Int
nmu, Int
t forall a. Num a => a -> a -> a
- Int
1)
                      forall (a :: * -> * -> *) e (m :: * -> *) i.
(MArray a e m, Ix i) =>
a i e -> i -> e -> m ()
writeArray
                        IOArray (Int, Int) a
jarray
                        (Int
nkappa, Int
t)
                        (a
entry' forall a. Num a => a -> a -> a
+ a
gamma forall a. Num a => a -> a -> a
* a
entry forall a. Num a => a -> a -> a
* [a]
x forall a. [a] -> Int -> a
!! (Int
t forall a. Num a => a -> a -> a
- Int
1) forall a b. (Num a, Integral b) => a -> b -> a
^ (Int
c forall a. Num a => a -> a -> a
+ Int
1))
                    else forall (a :: * -> * -> *) e (m :: * -> *) i.
(MArray a e m, Ix i) =>
a i e -> i -> e -> m ()
writeArray
                           IOArray (Int, Int) a
jarray
                           (Int
nkappa, Int
t)
                           (a
entry' forall a. Num a => a -> a -> a
+ a
gamma forall a. Num a => a -> a -> a
* [a]
x forall a. [a] -> Int -> a
!! (Int
t forall a. Num a => a -> a -> a
- Int
1) forall a b. (Num a, Integral b) => a -> b -> a
^ (Int
c forall a. Num a => a -> a -> a
+ Int
1))
          Int -> IO ()
go (Int
i forall a. Num a => a -> a -> a
+ Int
1)
  ()
_ <- Int -> IO ()
go Int
i0
  a
entry1 <- forall (a :: * -> * -> *) e (m :: * -> *) i.
(MArray a e m, Ix i) =>
a i e -> i -> m e
readArray IOArray (Int, Int) a
jarray (Int
nkappa, Int
t)
  if Int
k forall a. Eq a => a -> a -> Bool
== Int
0
    then
      forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when (Int
nkappa forall a. Ord a => a -> a -> Bool
> Int
1) forall a b. (a -> b) -> a -> b
$ do
        a
entry2 <- forall (a :: * -> * -> *) e (m :: * -> *) i.
(MArray a e m, Ix i) =>
a i e -> i -> m e
readArray IOArray (Int, Int) a
jarray (Int
nkappa, Int
t forall a. Num a => a -> a -> a
- Int
1)
        forall (a :: * -> * -> *) e (m :: * -> *) i.
(MArray a e m, Ix i) =>
a i e -> i -> e -> m ()
writeArray IOArray (Int, Int) a
jarray (Int
nkappa, Int
t) (a
entry1 forall a. Num a => a -> a -> a
+ a
entry2)
    else do
      a
entry2 <- forall (a :: * -> * -> *) e (m :: * -> *) i.
(MArray a e m, Ix i) =>
a i e -> i -> m e
readArray IOArray (Int, Int) a
jarray (Seq (Maybe Int) -> Seq Int -> Int
_nkappa Seq (Maybe Int)
dico Seq Int
mu, Int
t forall a. Num a => a -> a -> a
- Int
1)
      forall (a :: * -> * -> *) e (m :: * -> *) i.
(MArray a e m, Ix i) =>
a i e -> i -> e -> m ()
writeArray IOArray (Int, Int) a
jarray (Int
nkappa, Int
t) (a
entry1 forall a. Num a => a -> a -> a
+ a
beta forall a. Num a => a -> a -> a
* [a]
x forall a. [a] -> Int -> a
!! (Int
t forall a. Num a => a -> a -> a
- Int
1) forall a b. (Num a, Integral b) => a -> b -> a
^ Int
c forall a. Num a => a -> a -> a
* a
entry2)

-- | Hypergeometric function of a matrix argument.
-- Actually the matrix argument is given by the eigenvalues of the matrix.
-- For a type `a` of real numbers, `BaseFracType a = a`. If `a = Complex b` 
-- is a type of complex numbers, then `BaseFracType a = b`. Thus `alpha` 
-- parameter cannot be a complex number.
hypergeomat :: forall a. (Eq a, Fractional a, BaseFrac a)
  => Int -- ^ truncation weight
  -> BaseFracType a -- ^ alpha parameter (usually 2)
  -> [a] -- ^ upper parameters
  -> [a] -- ^ lower parameters
  -> [a] -- ^ variables (the eigenvalues)
  -> IO a
hypergeomat :: forall a.
(Eq a, Fractional a, BaseFrac a) =>
Int -> BaseFracType a -> [a] -> [a] -> [a] -> IO a
hypergeomat Int
m BaseFracType a
alpha [a]
a [a]
b [a]
x = do
  let n :: Int
n = forall (t :: * -> *) a. Foldable t => t a -> Int
length [a]
x
  if forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
all (forall a. Eq a => a -> a -> Bool
== forall a. [a] -> a
head [a]
x) [a]
x
    then
      forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ forall a.
(Eq a, Fractional a, BaseFrac a) =>
Int -> BaseFracType a -> [a] -> [a] -> Int -> a -> a
hypergeoI Int
m BaseFracType a
alpha [a]
a [a]
b Int
n (forall a. [a] -> a
head [a]
x)
    else do
      let pmn :: Int
pmn = Int -> Int -> Int
_P Int
m Int
n
          dico :: Seq (Maybe Int)
dico = Int -> Int -> Seq (Maybe Int)
_dico Int
pmn Int
m
          xrange :: [Int]
xrange = [Int
1 .. Int
n]
          line1 :: [((Int, Int), a)]
line1 = forall a b c. (a -> b -> c) -> [a] -> [b] -> [c]
zipWith (\Int
i a
u -> ((Int
1, Int
i), a
u)) [Int]
xrange (forall a. (a -> a -> a) -> [a] -> [a]
scanl1 forall a. Num a => a -> a -> a
(+) [a]
x)
          otherlines :: [((Int, Int), a)]
otherlines = forall (t :: * -> *) a b. Foldable t => (a -> [b]) -> t a -> [b]
concatMap (\Int
j -> [((Int
j, Int
i), a
0) | Int
i <- [Int]
xrange]) [Int
2 .. Int
pmn]
          arr0 :: Array (Int, Int) a
arr0 =
            forall i e. Ix i => (i, i) -> [(i, e)] -> Array i e
array ((Int
1, Int
1), (Int
pmn, Int
n)) ([((Int, Int), a)]
line1 forall a. [a] -> [a] -> [a]
++ [((Int, Int), a)]
otherlines)
      IOArray (Int, Int) a
jarray <- forall i (a :: * -> * -> *) e (b :: * -> * -> *) (m :: * -> *).
(Ix i, IArray a e, MArray b e m) =>
a i e -> m (b i e)
thaw Array (Int, Int) a
arr0
      a
s <- forall a.
(Fractional a, Eq a, BaseFrac a) =>
[a]
-> [a]
-> [a]
-> Seq (Maybe Int)
-> Int
-> BaseFracType a
-> Int
-> a
-> Int
-> Seq Int
-> IOArray (Int, Int) a
-> IO a
summation [a]
a [a]
b [a]
x Seq (Maybe Int)
dico Int
n BaseFracType a
alpha Int
0 a
1 Int
m forall a. Seq a
S.empty IOArray (Int, Int) a
jarray
      forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ a
s forall a. Num a => a -> a -> a
+ a
1