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

import qualified Numeric.BLAS.Subobject.View as View
import qualified Numeric.BLAS.Subobject.Layout.Class as LayoutClass
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 ((::+)((::+)))


{- |
Attention:
@size@ measures the whole underlying Array,
whereas the pair @start, shape@ determines position and size of the subarray.
@layout@ describes the gaps between the used elements.
-}
data T lay sh =
   Cons {
      size :: Int,
      start :: Int,
      layout :: lay,
      shape :: sh
   }
   deriving (Show)

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



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

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


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

focusMany :: (Functor f) =>
   (View.T lay0 sh0 -> f (View.T lay1 sh1)) -> T lay0 sh0 -> f (T lay1 sh1)
focusMany f (Cons n s0 lay0 sh0) =
   fmap (\(View.Cons s1 lay1 sh1) -> Cons n s1 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 {
      size = size sv,
      start = start 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

submatrixLeadingDim ::
   (LayoutClass.Submatrix lay, Shape.C width) =>
   T lay (height,width) -> Int
submatrixLeadingDim sm =
   LayoutClass.leadingDim (layout sm) (snd $ shape sm)

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 {
      size = n*m,
      start = 0,
      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 {
      size = size_,
      start = prefixSize+leftSize,
      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 {
      size = Shape.size sh,
      start = Shape.size $ prefix ::+ leftPad,
      layout = Layout.Submatrix $ Shape.size $ leftPad::+width::+rightPad,
      shape = (height, width)
   }

submatrixFromSubvector ::
   (Shape.C width) =>
   Subvector (height,width) -> Submatrix (height,width)
submatrixFromSubvector sv =
   Cons {
      size = size sv,
      start = start 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
-}