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

import qualified Data.Array.Comfort.Shape as Shape

import Control.Applicative (Const(Const, getConst))

import Data.Functor.Identity (Identity(Identity), runIdentity)
import Data.List (tails)
import Data.Tuple.HT (swap)


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

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

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


type family HeightOf shape
type family WidthOf shape


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

type instance HeightOf (General height width) = height
type instance WidthOf (General height width) = width

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 Square size =
   Square {
      squareOrder :: Order,
      squareSize :: size
   } deriving (Eq, Show)

type instance HeightOf (Square size) = size
type instance WidthOf (Square size) = size

generalFromSquare :: Square size -> General size size
generalFromSquare (Square order sh) = General order sh sh

transposeSquare :: Square sh -> Square sh
transposeSquare (Square order size) = Square (flipOrder order) size

instance (Shape.C size) => Shape.C (Square size) where
   type Index (Square size) = (Shape.Index size, Shape.Index size)
   indices (Square _ size) = Shape.indices (size,size)

   offset (Square RowMajor size) =
      Shape.offset (size,size)
   offset (Square ColumnMajor size) =
      Shape.offset (size,size) . swap
   uncheckedOffset (Square RowMajor size) =
      Shape.uncheckedOffset (size,size)
   uncheckedOffset (Square ColumnMajor size) =
      Shape.uncheckedOffset (size,size) . swap

   sizeOffset (Square RowMajor size) =
      Shape.sizeOffset (size,size)
   sizeOffset (Square ColumnMajor size) =
      Shape.sizeOffset (size,size) . swap
   uncheckedSizeOffset (Square RowMajor size) =
      Shape.uncheckedSizeOffset (size,size)
   uncheckedSizeOffset (Square ColumnMajor size) =
      Shape.uncheckedSizeOffset (size,size) . swap

   inBounds (Square _ size) = Shape.inBounds (size,size)
   size (Square _ size) = Shape.size (size,size)
   uncheckedSize (Square _ size) = Shape.uncheckedSize (size,size)


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

type instance HeightOf (Householder height width) = height
type instance WidthOf (Householder height width) = width

data Reflector = Reflector deriving (Eq)
data Triangle = Triangle deriving (Eq)

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

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

   type Index (Householder height width) =
            (Either Reflector Triangle,
             (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


{- |
Store the upper triangular half of a real symmetric or complex Hermitian matrix.
-}
data Hermitian size =
   Hermitian {
      hermitianOrder :: Order,
      hermitianSize :: size
   } deriving (Eq, Show)

type instance HeightOf (Hermitian size) = size
type instance WidthOf (Hermitian size) = size

uploFromOrder :: Order -> Char
uploFromOrder RowMajor = 'L'
uploFromOrder ColumnMajor = 'U'

instance (Shape.C size) => Shape.C (Hermitian size) where
   type Index (Hermitian size) = (Shape.Index size, Shape.Index size)

   indices (Hermitian _ size) =
      let ixs = Shape.indices size
      in  concat $ zipWith (\r cs -> map ((,) r) cs) ixs $ tails ixs

   uncheckedOffset sh ix =
      snd $ Shape.uncheckedSizeOffset sh ix

   sizeOffset sh ix =
      if Shape.inBounds sh ix
        then Shape.uncheckedSizeOffset sh ix
        else error "Shape.Hermitian.sizeOffset: wrong matrix part"

   uncheckedSizeOffset (Hermitian RowMajor size) (rs,cs) =
      let (s,r) = Shape.uncheckedSizeOffset size rs
          c = Shape.uncheckedOffset size cs
      in  (s, triangleSize s - triangleSize (s-r) + c-r)
   uncheckedSizeOffset (Hermitian ColumnMajor size) (rs,cs) =
      let (s,r) = Shape.uncheckedSizeOffset size rs
          c = Shape.uncheckedOffset size cs
      in  (s, triangleSize c + r)

   size (Hermitian _ size) = triangleSize $ Shape.size size
   uncheckedSize (Hermitian _ size) = triangleSize $ Shape.uncheckedSize size
   inBounds (Hermitian _ size) ix@(r,c) =
      Shape.inBounds (size,size) ix
      &&
      Shape.offset size r <= Shape.offset size c


data Triangular uplo size =
   Triangular {
      triangularUplo :: uplo,
      triangularOrder :: Order,
      triangularSize :: size
   } deriving (Eq, Show)

type instance HeightOf (Triangular uplo size) = size
type instance WidthOf (Triangular uplo size) = size

data Lower = Lower deriving (Eq, Show)
data Upper = Upper deriving (Eq, Show)

type LowerTriangular = Triangular Lower
type UpperTriangular = Triangular Upper

triangularTransposeUp :: LowerTriangular sh -> UpperTriangular sh
triangularTransposeUp (Triangular Lower order size) =
   Triangular Upper (flipOrder order) size

triangularTransposeDown :: UpperTriangular sh -> LowerTriangular sh
triangularTransposeDown (Triangular Upper order size) =
   Triangular Lower (flipOrder order) size


class Uplo uplo where switchUplo :: f Lower -> f Upper -> f uplo
instance Uplo Lower where switchUplo f _ = f
instance Uplo Upper where switchUplo _ f = f

autoUplo :: Uplo uplo => uplo
autoUplo = runIdentity $ switchUplo (Identity Lower) (Identity Upper)

uploOrder :: Uplo uplo => uplo -> Order -> Order
uploOrder uplo order = caseUplo uplo (flipOrder order) order

getUploConst :: uplo -> Const a uplo -> a
getUploConst _ = getConst

caseUplo :: Uplo uplo => uplo -> a -> a -> a
caseUplo uplo lo up =
   getUploConst uplo $ switchUplo (Const lo) (Const up)

instance (Uplo uplo, Shape.C size) => Shape.C (Triangular uplo size) where
   type Index (Triangular uplo size) = (Shape.Index size, Shape.Index size)

   indices (Triangular uplo _ size) =
      let ixs = Shape.indices size
          rcs = concat $ zipWith (\r cs -> map ((,) r) cs) ixs $ tails ixs
      in  caseUplo uplo (map swap rcs) rcs

   uncheckedOffset sh ix =
      snd $ Shape.uncheckedSizeOffset sh ix

   sizeOffset sh ix =
      if Shape.inBounds sh ix
        then Shape.uncheckedSizeOffset sh ix
        else error "Shape.Triangular.sizeOffset: wrong matrix part"

   uncheckedSizeOffset (Triangular uplo RowMajor size) (rs,cs) =
      let (s,r) = Shape.uncheckedSizeOffset size rs
          c = Shape.uncheckedOffset size cs
      in  (s,
           caseUplo uplo
               (triangleSize r + c)
               (triangleSize s - triangleSize (s-r) + c-r))
   uncheckedSizeOffset (Triangular uplo ColumnMajor size) (rs,cs) =
      let (s,r) = Shape.uncheckedSizeOffset size rs
          c = Shape.uncheckedOffset size cs
      in  (s,
           caseUplo uplo
               (triangleSize s - triangleSize (s-c) + r-c)
               (triangleSize c + r))

   size (Triangular _ _ size) = triangleSize $ Shape.size size
   uncheckedSize (Triangular _ _ size) = triangleSize $ Shape.uncheckedSize size
   inBounds (Triangular uplo _ size) ix@(r,c) =
      Shape.inBounds (size,size) ix
      &&
      caseUplo uplo
         (Shape.offset size r >= Shape.offset size c)
         (Shape.offset size r <= Shape.offset size c)

triangleSize :: Int -> Int
triangleSize n = div (n*(n+1)) 2

triangleRoot :: Floating a => a -> a
triangleRoot size = (sqrt (8*size+1)-1)/2

triangleExtent :: String -> Int -> Int
triangleExtent name size =
   let n = round (triangleRoot (fromIntegral size :: Double))
   in if size == triangleSize n
        then n
        else error (name ++ ": no triangular number of elements")