{-# LANGUAGE TypeFamilies #-}
module Numeric.LAPACK.Matrix.Shape.Private where

import qualified Data.Array.Comfort.Shape as Shape

import Data.Tuple.HT (swap)


data Order = RowMajor | ColumnMajor
   deriving (Eq, Show)

flipOrder :: Order -> Order
flipOrder RowMajor = ColumnMajor
flipOrder ColumnMajor = RowMajor

charFromOrder :: Order -> Char
charFromOrder RowMajor = 'T'
charFromOrder ColumnMajor = 'N'


data General height width =
   General {
      generalOrder :: Order,
      generalHeight :: height,
      generalWidth :: width
   } deriving (Eq, Show)

instance (Shape.C height, Shape.C width) => Shape.C (General height width) where
   type Index (General height width) = (Shape.Index height, Shape.Index width)
   indices (General _ height width) = Shape.indices (height,width)

   offset (General RowMajor height width) =
      Shape.offset (height,width)
   offset (General ColumnMajor height width) =
      Shape.offset (width,height) . swap
   uncheckedOffset (General RowMajor height width) =
      Shape.uncheckedOffset (height,width)
   uncheckedOffset (General ColumnMajor height width) =
      Shape.uncheckedOffset (width,height) . swap

   sizeOffset (General RowMajor height width) =
      Shape.sizeOffset (height,width)
   sizeOffset (General ColumnMajor height width) =
      Shape.sizeOffset (width,height) . swap
   uncheckedSizeOffset (General RowMajor height width) =
      Shape.uncheckedSizeOffset (height,width)
   uncheckedSizeOffset (General ColumnMajor height width) =
      Shape.uncheckedSizeOffset (width,height) . swap

   inBounds (General _ height width) = Shape.inBounds (height,width)
   size (General _ height width) = Shape.size (height,width)
   uncheckedSize (General _ height width) = Shape.uncheckedSize (height,width)


transpose :: General height width -> General width height
transpose (General order height width) = General (flipOrder order) width height

dimensions ::
   (Shape.C height, Shape.C width) => General height width -> (Int, Int)
dimensions (General order height width) =
   case order of
      RowMajor -> (Shape.size width, Shape.size height)
      ColumnMajor -> (Shape.size height, Shape.size width)


data Householder height width =
   Householder {
      householderOrder :: Order,
      householderHeight :: height,
      householderWidth :: width
   } deriving (Eq, Show)

data Reflector = Reflector deriving (Eq)
data Triangular = Triangular deriving (Eq)

householderPart ::
   (Shape.C height, Shape.C width) =>
   Householder height width ->
   (Shape.Index height, Shape.Index width) -> Either Reflector Triangular
householderPart (Householder _ height width) (r,c) =
   if Shape.offset height r > Shape.offset width c
     then Left Reflector
     else Right Triangular

instance
   (Shape.C height, Shape.C width) =>
      Shape.C (Householder height width) where

   type Index (Householder height width) =
            (Either Reflector Triangular,
             (Shape.Index height, Shape.Index width))

   indices sh@(Householder _ height width) =
      map (\ix -> (householderPart sh ix, ix)) $
      Shape.indices (height,width)

   offset sh@(Householder order height width) (part,ix) =
      if part == householderPart sh ix
        then
            case order of
               RowMajor -> Shape.offset (height,width) ix
               ColumnMajor -> Shape.offset (width,height) (swap ix)
        else error "Shape.Householder.offset: wrong matrix part"
   uncheckedOffset (Householder RowMajor height width) =
      Shape.uncheckedOffset (height,width) . snd
   uncheckedOffset (Householder ColumnMajor height width) =
      Shape.uncheckedOffset (width,height) . swap . snd

   sizeOffset sh@(Householder order height width) (part,ix) =
      if part == householderPart sh ix
        then
            case order of
               RowMajor -> Shape.sizeOffset (height,width) ix
               ColumnMajor -> Shape.sizeOffset (width,height) (swap ix)
        else error "Shape.Householder.sizeOffset: wrong matrix part"
   uncheckedSizeOffset (Householder RowMajor height width) =
      Shape.uncheckedSizeOffset (height,width) . snd
   uncheckedSizeOffset (Householder ColumnMajor height width) =
      Shape.uncheckedSizeOffset (width,height) . swap . snd

   size (Householder _ height width) = Shape.size (height,width)
   uncheckedSize (Householder _ height width) =
      Shape.uncheckedSize (height,width)
   inBounds sh@(Householder _ height width) (part,ix) =
      Shape.inBounds (height,width) ix
      &&
      part == householderPart sh ix