-- |
-- Module      : ConClusion.Array.Conversion
-- Description : Type castings and conversions of array types
-- Copyright   : Phillip Seeber, 2022
-- License     : AGPL-3
-- Maintainer  : phillip.seeber@googlemail.com
-- Stability   : experimental
-- Portability : POSIX, Windows
module ConClusion.Array.Conversion
  ( vecH2M,
    vecM2H,
    matH2M,
    matM2H,
  )
where

import Data.Massiv.Array as Massiv hiding (IndexException)
import Data.Massiv.Array.Manifest.Vector as Massiv
import Numeric.LinearAlgebra as LA hiding (magnitude, (<>))
import RIO
import qualified RIO.Vector.Storable as VectorS

-- | Converts a vector from the HMatrix package to the Massiv representation.
{-# SCC vecH2M #-}
vecH2M :: (Element e, Mutable r e, Load r Ix1 e) => VectorS.Vector e -> Massiv.Vector r e
vecH2M :: Vector e -> Vector r e
vecH2M Vector e
hVec = Comp -> Sz Int -> Vector e -> Vector r e
forall (v :: * -> *) a ix r.
(HasCallStack, Typeable v, Vector v a, Load (ARepr v) ix a,
 Load r ix a, Manifest r a) =>
Comp -> Sz ix -> v a -> Array r ix a
fromVector' Comp
Seq (Int -> Sz Int
forall ix. Index ix => ix -> Sz ix
Sz (Int -> Sz Int) -> Int -> Sz Int
forall a b. (a -> b) -> a -> b
$ Vector e -> Int
forall a. Storable a => Vector a -> Int
VectorS.length Vector e
hVec) Vector e
hVec

-- | Converts a vector from the Massiv representation to the HMatrix representation.
{-# SCC vecM2H #-}
vecM2H :: (Manifest r e, Load r Ix1 e, Element e) => Massiv.Vector r e -> LA.Vector e
vecM2H :: Vector r e -> Vector e
vecM2H = Vector r e -> Vector e
forall r ix e (v :: * -> *).
(Manifest r e, Load r ix e, Manifest (ARepr v) e, Vector v e,
 VRepr (ARepr v) ~ v) =>
Array r ix e -> v e
Massiv.toVector

-- | Converts a matrix from the HMatrix representation to the Massiv representation.
{-# SCC matH2M #-}
matH2M :: (Mutable r e, Load r Ix1 e, Element e) => LA.Matrix e -> Massiv.Matrix r e
matH2M :: Matrix e -> Matrix r e
matH2M Matrix e
hMat = Sz Ix2 -> Array r Int e -> Matrix r e
forall r ix ix' e.
(HasCallStack, Index ix', Index ix, Size r) =>
Sz ix' -> Array r ix e -> Array r ix' e
Massiv.resize' (Ix2 -> Sz Ix2
forall ix. Index ix => ix -> Sz ix
Sz (Ix2 -> Sz Ix2) -> Ix2 -> Sz Ix2
forall a b. (a -> b) -> a -> b
$ Int
nRows Int -> Int -> Ix2
:. Int
nCols) (Array r Int e -> Matrix r e)
-> (Matrix e -> Array r Int e) -> Matrix e -> Matrix r e
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Vector e -> Array r Int e
forall e r.
(Element e, Mutable r e, Load r Int e) =>
Vector e -> Vector r e
vecH2M (Vector e -> Array r Int e)
-> (Matrix e -> Vector e) -> Matrix e -> Array r Int e
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Matrix e -> Vector e
forall t. Element t => Matrix t -> Vector t
LA.flatten (Matrix e -> Matrix r e) -> Matrix e -> Matrix r e
forall a b. (a -> b) -> a -> b
$ Matrix e
hMat
  where
    nRows :: Int
nRows = Matrix e -> Int
forall t. Matrix t -> Int
LA.rows Matrix e
hMat
    nCols :: Int
nCols = Matrix e -> Int
forall t. Matrix t -> Int
LA.cols Matrix e
hMat

-- | Converts a matrix from Massiv to HMatrix representation.
{-# SCC matM2H #-}
matM2H :: (Element e, Manifest r e, Load r Ix1 e) => Massiv.Matrix r e -> LA.Matrix e
matM2H :: Matrix r e -> Matrix e
matM2H Matrix r e
mMat = Int -> Vector e -> Matrix e
forall t. Storable t => Int -> Vector t -> Matrix t
LA.reshape Int
nCols (Vector e -> Matrix e)
-> (Matrix r e -> Vector e) -> Matrix r e -> Matrix e
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Vector r e -> Vector e
forall r e.
(Manifest r e, Load r Int e, Element e) =>
Vector r e -> Vector e
vecM2H (Vector r e -> Vector e)
-> (Matrix r e -> Vector r e) -> Matrix r e -> Vector e
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Matrix r e -> Vector r e
forall r ix e. (Index ix, Size r) => Array r ix e -> Vector r e
Massiv.flatten (Matrix r e -> Matrix e) -> Matrix r e -> Matrix e
forall a b. (a -> b) -> a -> b
$ Matrix r e
mMat
  where
    Sz (Int
_nRows :. Int
nCols) = Matrix r e -> Sz Ix2
forall r ix e. Size r => Array r ix e -> Sz ix
Massiv.size Matrix r e
mMat