module Numeric.LAPACK.Format (
   (##),
   Format(format),
   FormatArray(formatArray),
   ) where

import qualified Numeric.LAPACK.Matrix.Shape.Private as MatrixShape
import qualified Numeric.LAPACK.Matrix.Square as Square
import Numeric.LAPACK.Matrix.Shape.Private (Order(RowMajor, ColumnMajor))
import Numeric.LAPACK.Matrix.Private (General)

import qualified Numeric.Netlib.Class as Class

import qualified Data.Array.Comfort.Storable.Internal as Array
import qualified Data.Array.Comfort.Shape as Shape
import Data.Array.Comfort.Storable (Array)

import Text.Printf (PrintfArg, printf)

import qualified Data.List.HT as ListHT
import qualified Data.Complex as Complex
import Data.Monoid (Endo(Endo,appEndo))
import Data.List (mapAccumL, transpose)
import Data.Complex (Complex((:+)))


infix 0 ##

(##) :: (Format a) => a -> String -> IO ()
a ## fmt = putStr $ unlines $ format fmt a


class Format a where
   format :: String -> a -> [String]

instance Format Int where
   format _fmt a = [show a]

instance Format Float where
   format fmt a = [printf fmt a]

instance Format Double where
   format fmt a = [printf fmt a]

instance (PrintfArg a, Class.Real a) => Format (Complex a) where
   format fmt a = [printfComplex fmt a]

instance (Format a, Format b) => Format (a,b) where
   format fmt (a,b) = format fmt a ++ [""] ++ format fmt b

instance (Format a, Format b, Format c) => Format (a,b,c) where
   format fmt (a,b,c) =
      format fmt a ++ [""] ++ format fmt b ++ [""] ++ format fmt c

instance (FormatArray sh, Class.Floating a) => Format (Array sh a) where
   format = formatArray


class (Shape.C sh) => FormatArray sh where
   formatArray :: (Class.Floating a) => String -> Array sh a -> [String]

instance (Integral i) => FormatArray (Shape.ZeroBased i) where
   formatArray fmt m = [unwords $ map (printfFloating fmt) $ Array.toList m]

instance (Integral i) => FormatArray (Shape.OneBased i) where
   formatArray fmt m = [unwords $ map (printfFloating fmt) $ Array.toList m]

instance (Shape.C sh) => FormatArray (MatrixShape.Square sh) where
   formatArray fmt = formatGeneral fmt . Square.toGeneral

instance
   (Shape.C height, Shape.C width) =>
      FormatArray (MatrixShape.General height width) where
   formatArray = formatGeneral

formatGeneral ::
   (Shape.C height, Shape.C width, Class.Floating a) =>
   String -> General height width a -> [String]
formatGeneral fmt m =
   let MatrixShape.General order height width = Array.shape m
   in  formatAligned $ formatRows fmt order (height,width) $ Array.toList m

instance
   (Shape.C height, Shape.C width) =>
      FormatArray (MatrixShape.Householder height width) where
   formatArray = formatHouseholder

formatHouseholder ::
   (Shape.C height, Shape.C width, Class.Floating a) =>
   String -> Array (MatrixShape.Householder height width) a -> [String]
formatHouseholder fmt m =
   let MatrixShape.Householder order height width = Array.shape m
   in formatSeparateTriangle $
      formatRows fmt order (height,width) $ Array.toList m

instance (Shape.C size) => FormatArray (MatrixShape.Hermitian size) where
   formatArray = formatHermitian

formatHermitian ::
   (Shape.C size, Class.Floating a) =>
   String -> Array (MatrixShape.Hermitian size) a -> [String]
formatHermitian fmt m =
   let MatrixShape.Hermitian order size = Array.shape m
   in  formatSeparateTriangle $
       map (map (printfFloating fmt)) $
       complementTriangle order (Shape.size size) $ Array.toList m

complementTriangle :: (Class.Floating a) => Order -> Int -> [a] -> [[a]]
complementTriangle order n xs =
   let mergeTriangles lower upper =
         zipWith (++) (map (map conjugate . init) lower) upper
   in case order of
         RowMajor ->
            let tri = slice (take n $ iterate pred n) xs
                trans = reverse $ transpose $ map reverse tri
            in  mergeTriangles trans tri
         ColumnMajor ->
            let tri = slice (take n [1..]) xs
            in  mergeTriangles tri (transpose tri)

conjugate :: (Class.Floating a) => a -> a
conjugate =
   appEndo $
   Class.switchFloating
      (Endo id) (Endo id) (Endo Complex.conjugate) (Endo Complex.conjugate)

instance
   (MatrixShape.Uplo uplo, Shape.C size) =>
      FormatArray (MatrixShape.Triangular uplo size) where
   formatArray = formatTriangular

formatTriangular ::
   (MatrixShape.Uplo uplo, Shape.C size, Class.Floating a) =>
   String -> Array (MatrixShape.Triangular uplo size) a -> [String]
formatTriangular fmt m =
   let MatrixShape.Triangular uplo order size = Array.shape m
   in  formatAligned $
       MatrixShape.caseUplo uplo
         padLowerTriangle padUpperTriangle order (Shape.size size) $
       map (printfFloating fmt) $ Array.toList m

padUpperTriangle :: Order -> Int -> [String] -> [[String]]
padUpperTriangle order n xs =
   case order of
      RowMajor ->
         zipWith (++) (iterate ("":) []) (slice (take n $ iterate pred n) xs)
      ColumnMajor ->
         transpose $
         zipWith (++)
            (slice (take n [1..]) xs)
            (reverse $ take n $ iterate ("":) [])

padLowerTriangle :: Order -> Int -> [String] -> [[String]]
padLowerTriangle order n xs =
   case order of
      RowMajor ->
         map (take n) $ map (++ repeat "") $ slice (take n [1..]) xs
      ColumnMajor ->
         transpose $
         zipWith (++) (iterate ("":) []) (slice (take n $ iterate pred n) xs)

_padLowerTriangle :: Order -> Int -> [a] -> [[a]]
_padLowerTriangle order n xs =
   case order of
      RowMajor -> slice (take n [1..]) xs
      ColumnMajor ->
         foldr (\(y:ys) zs -> [y] : zipWith (:) ys zs) []
            (slice (take n $ iterate pred n) xs)

slice :: [Int] -> [a] -> [[a]]
slice ns xs =
   snd $ mapAccumL (\ys n -> let (vs,ws) = splitAt n ys in (ws,vs)) xs ns

formatSeparateTriangle :: [[String]] -> [String]
formatSeparateTriangle xss =
   let strWidths = columnWidths xss
   in  zipWith
         (\row xs ->
            concat $
            zipWith (\col cell -> (if row==col then '|' else ' '):cell) [0..] $
            zipWith (ListHT.padLeft ' ') strWidths xs)
         [(0::Int)..] xss

formatRows ::
   (Class.Floating a, Shape.C height, Shape.C width) =>
   String -> Order -> (height, width) -> [a] -> [[String]]
formatRows fmt order (height,width) =
   (case order of
      RowMajor -> ListHT.sliceVertical (Shape.size width)
      ColumnMajor -> ListHT.sliceHorizontal (Shape.size height)) .
   map (printfFloating fmt)

formatAligned :: [[String]] -> [String]
formatAligned xss =
   let strWidths = columnWidths xss
   in  map (unwords . zipWith (ListHT.padLeft ' ') strWidths) xss

columnWidths :: [[[a]]] -> [Int]
columnWidths xss =
   case map (map length) xss of
      [] -> []
      w:ws -> foldl (zipWith max) w ws


newtype Printf a = Printf {runPrintf :: String -> a -> String}

printfFloating :: (Class.Floating a) => String -> a -> String
printfFloating =
   runPrintf $
   Class.switchFloating
      (Printf printf)
      (Printf printf)
      (Printf printfComplex)
      (Printf printfComplex)

printfComplex :: (PrintfArg a, Class.Real a) => String -> Complex a -> String
printfComplex fmt (r:+i) =
   if i<0 || isNegativeZero i
     then printf (fmt ++ "-i" ++ fmt) r (-i)
     else printf (fmt ++ "+i" ++ fmt) r i