module Data.Paired.Vector
  ( module Data.Paired.Vector
  , module Data.Paired.Common
  ) where

import Data.Vector.Generic as VG

import Data.Paired.Common



-- | Upper triangular elements.

upperTriVG
  :: (Vector v a, Vector w (a,a))
  => OnDiag
  -> v a
  -> (Int, w (a,a))
upperTriVG :: OnDiag -> v a -> (Int, w (a, a))
upperTriVG OnDiag
d v a
as = (Int
z, Int
-> ((Int, Int) -> Maybe ((a, a), (Int, Int)))
-> (Int, Int)
-> w (a, a)
forall (v :: * -> *) a b.
Vector v a =>
Int -> (b -> Maybe (a, b)) -> b -> v a
unfoldrN Int
z (Int, Int) -> Maybe ((a, a), (Int, Int))
go (Int
0,if OnDiag
d OnDiag -> OnDiag -> Bool
forall a. Eq a => a -> a -> Bool
== OnDiag
OnDiag then Int
0 else Int
1))
  where la :: Int
la = v a -> Int
forall (v :: * -> *) a. Vector v a => v a -> Int
VG.length v a
as
        z :: Int
z  = Int
la Int -> Int -> Int
forall a. Num a => a -> a -> a
* (Int
la Int -> Int -> Int
forall a. Num a => a -> a -> a
+ if OnDiag
d OnDiag -> OnDiag -> Bool
forall a. Eq a => a -> a -> Bool
== OnDiag
OnDiag then Int
1 else Int
0) Int -> Int -> Int
forall a. Integral a => a -> a -> a
`div` Int
2
        go :: (Int, Int) -> Maybe ((a, a), (Int, Int))
go (Int
k,Int
l)
          | Int
k Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= Int
la   = Maybe ((a, a), (Int, Int))
forall a. Maybe a
Nothing
          | Int
l Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= Int
la   = (Int, Int) -> Maybe ((a, a), (Int, Int))
go (Int
kInt -> Int -> Int
forall a. Num a => a -> a -> a
+Int
1,Int
kInt -> Int -> Int
forall a. Num a => a -> a -> a
+Int
1 Int -> Int -> Int
forall a. Num a => a -> a -> a
+ if OnDiag
d OnDiag -> OnDiag -> Bool
forall a. Eq a => a -> a -> Bool
== OnDiag
OnDiag then Int
0 else Int
1)
          | Bool
otherwise = ((a, a), (Int, Int)) -> Maybe ((a, a), (Int, Int))
forall a. a -> Maybe a
Just ((v a
as v a -> Int -> a
forall (v :: * -> *) a. Vector v a => v a -> Int -> a
`VG.unsafeIndex` Int
k, v a
as v a -> Int -> a
forall (v :: * -> *) a. Vector v a => v a -> Int -> a
`VG.unsafeIndex` Int
l), (Int
k,Int
lInt -> Int -> Int
forall a. Num a => a -> a -> a
+Int
1))
{-# Inline upperTriVG #-}

-- | Outer pairing of all @as@ with all @bs@. This one is quasi-trivial,
-- but here for completeness.

rectangularVG
  :: (Vector va a, Vector vb b, Vector w (a,b))
  => va a
  -> vb b
  -> (Int, w (a,b))
rectangularVG :: va a -> vb b -> (Int, w (a, b))
rectangularVG va a
as vb b
bs = (Int
z, Int
-> ((Int, Int) -> Maybe ((a, b), (Int, Int)))
-> (Int, Int)
-> w (a, b)
forall (v :: * -> *) a b.
Vector v a =>
Int -> (b -> Maybe (a, b)) -> b -> v a
unfoldrN Int
z (Int, Int) -> Maybe ((a, b), (Int, Int))
go (Int
0,Int
0))
  where la :: Int
la = va a -> Int
forall (v :: * -> *) a. Vector v a => v a -> Int
VG.length va a
as
        lb :: Int
lb = vb b -> Int
forall (v :: * -> *) a. Vector v a => v a -> Int
VG.length vb b
bs
        z :: Int
z  = Int
la Int -> Int -> Int
forall a. Num a => a -> a -> a
* Int
lb
        go :: (Int, Int) -> Maybe ((a, b), (Int, Int))
go (Int
k,Int
l)
          | Int
k Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= Int
la   = Maybe ((a, b), (Int, Int))
forall a. Maybe a
Nothing
          | Int
l Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= Int
lb   = (Int, Int) -> Maybe ((a, b), (Int, Int))
go (Int
kInt -> Int -> Int
forall a. Num a => a -> a -> a
+Int
1,Int
0)
          | Bool
otherwise = ((a, b), (Int, Int)) -> Maybe ((a, b), (Int, Int))
forall a. a -> Maybe a
Just ((va a
as va a -> Int -> a
forall (v :: * -> *) a. Vector v a => v a -> Int -> a
`VG.unsafeIndex` Int
k, vb b
bs vb b -> Int -> b
forall (v :: * -> *) a. Vector v a => v a -> Int -> a
`VG.unsafeIndex` Int
l), (Int
k,Int
lInt -> Int -> Int
forall a. Num a => a -> a -> a
+Int
1))
{-# Inline rectangularVG #-}