{-# LANGUAGE TypeOperators #-}
module Data.Packed.Matrix.Lens where


import Data.Yall.Lens
import Data.Packed.Matrix
import Data.Packed.Vector
import Data.Packed.ST

import Data.Packed.Matrix.Lens.Util


rowL :: Element t => Int -> (Matrix t :-> Vector t)
rowL i = lens get' set'
  where
  get' matrix = buildVector (cols matrix) (\j -> matrix @@> (i,j))
  set' matrix row' = withSTMatrix matrix $ \matrix_st -> 
      mapVectorWithIndexM_ (\j t -> writeMatrix matrix_st i j t) row'

colL :: Element t => Int -> (Matrix t :-> Vector t)
colL j = lens get' set'
  where
  get' matrix = buildVector (rows matrix) (\i -> matrix @@> (i,j))
  set' matrix col' = withSTMatrix matrix $ \matrix_st -> 
      mapVectorWithIndexM_ (\i t -> writeMatrix matrix_st i j t) col'

diagL :: Element t => (Matrix t :-> Vector t)
diagL = lens get' set'
  where
  get' matrix = takeDiag matrix
  set' matrix diag' = withSTMatrix matrix $ \matrix_st -> 
    mapVectorWithIndexM_ (\i t -> writeMatrix matrix_st i i t) diag'

entryL :: Element t => Int -> Int -> (Matrix t :-> t)
entryL i j = lens get' set'
  where
  get' = (@@> (i,j))
  set' matrix entry' = withSTMatrix matrix $ \matrix_st ->
    writeMatrix matrix_st i j entry'
