{-# LANGUAGE TypeOperators #-}
module Numeric.BLAS.Subobject.Shape where

import qualified Numeric.BLAS.Subobject.View as View
import qualified Numeric.BLAS.Subobject.Layout as Layout
import qualified Data.Array.Comfort.Shape.SubSize as SubSize
import qualified Data.Array.Comfort.Shape as Shape
import Data.Array.Comfort.Shape ((::+)((::+)))


{-
FixMe:
This is misleading.
@size@ measures the whole underlying Array,
whereas the pair @start, shape@ indicate the subarray.
-}
data T lay sh =
   Cons {
      start, size :: Int,
      layout :: lay,
      shape :: sh
   }
   deriving (Show)

instance (Shape.C sh) => Shape.C (T lay sh) where
   size = size



fromVector :: (View.Slice lay, Shape.C sh) => sh -> T lay sh
fromVector sh =
   Cons {
      start = 0,
      size = Shape.size sh,
      layout = View.defltLayout,
      shape = sh
   }

fromVector_ :: (Shape.C sh) => lay -> sh -> T lay sh
fromVector_ layout_ sh =
   Cons {
      start = 0,
      size = Shape.size sh,
      layout = layout_,
      shape = sh
   }


focus :: (View.T lay sh0 -> View.T lay sh1) -> T lay sh0 -> T lay sh1
focus f (Cons s0 n lay0 sh0) =
   case f (View.Cons s0 lay0 sh0) of
      View.Cons s1 lay1 sh1 -> Cons s1 n lay1 sh1

focusMany :: (Functor f) =>
   (View.T lay sh0 -> f (View.T lay sh1)) -> T lay sh0 -> f (T lay sh1)
focusMany f (Cons s0 n lay0 sh0) =
   fmap (\(View.Cons s1 lay1 sh1) -> Cons s1 n lay1 sh1) $
   f $ View.Cons s0 lay0 sh0


type Subvector = T Layout.Subvector

subvectorFromVector :: (Shape.C sh) => sh -> Subvector sh
subvectorFromVector = fromVector

type Slice = T Layout.Slice

sliceInc :: Slice sh -> Int
sliceInc = Layout.sliceInc . layout

sliceFromVector :: (Shape.C sh) => sh -> Slice sh
sliceFromVector = fromVector

sliceFromSubvector :: Subvector sh -> Slice sh
sliceFromSubvector sv =
   Cons {
      start = start sv,
      size = size sv,
      layout = Layout.Slice 1,
      shape = shape sv
   }


{- |
@sh@ can be @(height, width)@ or @Shape.Square sh@.
-}
type Submatrix = T Layout.Submatrix

submatrixLeadingDim :: Submatrix sh -> Int
submatrixLeadingDim = Layout.submatrixLeadingDim . layout

submatrixFromMatrix ::
   (Shape.C height, Shape.C width) =>
   (height,width) -> Submatrix (height,width)
submatrixFromMatrix sh@(height,width) =
   let n = Shape.size height in
   let m = Shape.size width in
   Cons {
      start = 0,
      size = n*m,
      layout = Layout.Submatrix m,
      shape = sh
   }

submatrixFromAppend ::
   (Shape.C prefix, Shape.C suffix, Shape.C leftPad, Shape.C rightPad) =>
   (Shape.C height, Shape.C width) =>
   (prefix ::+ (height, leftPad::+width::+rightPad) ::+ suffix) ->
   Submatrix (height,width)
submatrixFromAppend sh@(_ ::+ (height, _::+width::+_) ::+ _) =
   let (size_,
         SubSize.Atom prefixSize
         ::+
         (SubSize.Atom _, SubSize.Sub widthSize widthShape)
         ::+
         SubSize.Atom _)
            = SubSize.evaluate sh in
   let (SubSize.Atom leftSize ::+ SubSize.Atom _ ::+ SubSize.Atom _)
            = widthShape in
   Cons {
      start = prefixSize+leftSize,
      size = size_,
      layout = Layout.Submatrix widthSize,
      shape = (height, width)
   }

submatrixFromAppend_ ::
   (Shape.C prefix, Shape.C suffix, Shape.C leftPad, Shape.C rightPad) =>
   (Shape.C height, Shape.C width) =>
   (prefix ::+ (height, leftPad::+width::+rightPad) ::+ suffix) ->
   Submatrix (height,width)
submatrixFromAppend_
      sh@(prefix ::+ (height, leftPad::+width::+rightPad) ::+ _suffix) =
--   let n = Shape.size height in
--   let m = Shape.size width in
   Cons {
      start = Shape.size $ prefix ::+ leftPad,
      size = Shape.size sh,
      layout = Layout.Submatrix $ Shape.size $ leftPad::+width::+rightPad,
      shape = (height, width)
   }

submatrixFromSubvector ::
   (Shape.C width) =>
   Subvector (height,width) -> Submatrix (height,width)
submatrixFromSubvector sv =
   Cons {
      start = start sv,
      size = size sv,
      layout = Layout.Submatrix $ Shape.size $ snd $ shape sv,
      shape = shape sv
   }

{- ToDo:
Shape.Index (Shape.SubVector sh) = Either Int (Shape.Index sh)
Left - access all elements, this is what Shape.indices returns
Right - access only elements in the subobject
That is, there are two admissible indices for elements in the subobject, both Left and Right
-}