{-# LANGUAGE TypeFamilies #-}
{-# LANGUAGE TypeOperators #-}
{-# LANGUAGE GADTs #-}
{-# LANGUAGE FlexibleInstances #-}
module Numeric.LAPACK.Matrix.Plain.Format where

import qualified Numeric.LAPACK.Matrix.Layout.Private as Layout
import qualified Numeric.LAPACK.Matrix.Extent.Private as Extent
import qualified Numeric.LAPACK.Vector as Vector
import qualified Numeric.LAPACK.Output as Output
import qualified Numeric.LAPACK.Shape as ExtShape
import Numeric.LAPACK.Output (Output, formatRow, (<+>))
import Numeric.LAPACK.Matrix.Layout.Private
         (Order(RowMajor, ColumnMajor), UnaryProxy)
import Numeric.LAPACK.Matrix.Private (Full)
import Numeric.LAPACK.Matrix.Extent.Private (Extent)
import Numeric.LAPACK.Scalar (conjugate)
import Numeric.LAPACK.Private (caseRealComplexFunc)

import qualified Numeric.Netlib.Class as Class

import qualified Type.Data.Num.Unary.Literal as TypeNum
import qualified Type.Data.Num.Unary as Unary
import Type.Data.Num (integralFromProxy)

import qualified Data.Array.Comfort.Storable.Unchecked as Array
import qualified Data.Array.Comfort.Boxed as BoxedArray
import qualified Data.Array.Comfort.Shape as Shape
import Data.Array.Comfort.Storable (Array)
import Data.Array.Comfort.Shape ((::+))

import Text.Printf (PrintfArg, printf)

import qualified Data.List.Match as Match
import qualified Data.List.HT as ListHT
import Data.Foldable (Foldable, fold)
import Data.List (mapAccumL, transpose)
import Data.Complex (Complex((:+)))
import Data.Ix (Ix)

import Control.Applicative ((<$>))


deflt :: String
deflt = "%.4g"


data Config =
   Config {
      configFormat :: String,
      configEmpty :: String
   }

defltConfig :: Config
defltConfig =
   Config {
      configFormat = deflt,
      configEmpty = ""
   }


class (Shape.C sh) => FormatArray sh where
   {-
   We use constraint @(Class.Floating a)@ and not @(Format a)@
   because it allows us to align the components of complex numbers.
   -}
   formatArray :: (Class.Floating a, Output out) => Config -> Array sh a -> out

instance (Integral i) => FormatArray (Shape.ZeroBased i) where
   formatArray = formatVector

instance (Integral i) => FormatArray (Shape.OneBased i) where
   formatArray = formatVector

instance (Ix i) => FormatArray (Shape.Range i) where
   formatArray = formatVector

instance (Integral i) => FormatArray (Shape.Shifted i) where
   formatArray = formatVector

instance (Enum enum, Bounded enum) => FormatArray (Shape.Enumeration enum) where
   formatArray = formatVector

instance (FormatArray sh) => FormatArray (Shape.Deferred sh) where
   formatArray cfg =
      formatArray cfg . Array.mapShape (\(Shape.Deferred sh) -> sh)

instance (FormatArray sh0, FormatArray sh1) => FormatArray (sh0::+sh1) where
   formatArray cfg v =
      formatArray cfg (Vector.takeLeft v)
      <+>
      formatArray cfg (Vector.takeRight v)

instance (Shape.C sh) => FormatArray (ExtShape.IntIndexed sh) where
   formatArray = formatVector

formatVector ::
   (Shape.C sh, Class.Floating a, Output out) => Config -> Array sh a -> out
formatVector cfg =
   formatRow . map (Output.text . fold . printfFloating cfg) . Array.toList

instance
   (Extent.Measure meas, Extent.C vert, Extent.C horiz,
    Shape.C height, Shape.C width) =>
      FormatArray (Layout.Full meas vert horiz height width) where
   formatArray cfg = formatAligned (printfFloating cfg) . layoutFull

layoutFull ::
   (Extent.Measure meas, Extent.C vert, Extent.C horiz,
    Shape.C height, Shape.C width, Class.Floating a) =>
   Full meas vert horiz height width a -> [[a]]
layoutFull m =
   let Layout.Full order extent = Array.shape m
   in  splitRows order (Extent.dimensions extent) $ Array.toList m

instance
   (Eq lower, Extent.Measure meas, Extent.C vert, Extent.C horiz,
    Shape.C height, Shape.C width) =>
      FormatArray (Layout.Split lower meas vert horiz height width) where
   formatArray cfg =
      formatSeparateTriangle (printfFloating cfg) . layoutSplit

layoutSplit ::
   (Eq lower, Extent.Measure meas, Extent.C vert, Extent.C horiz,
    Shape.C height, Shape.C width, Class.Floating a) =>
   Array (Layout.Split lower meas vert horiz height width) a -> [[a]]
layoutSplit m =
   let Layout.Split _ order extent = Array.shape m
   in splitRows order (Extent.dimensions extent) $ Array.toList m

instance (Shape.C size) => FormatArray (Layout.Hermitian size) where
   formatArray = formatMirrored conjugate

instance (Shape.C size) => FormatArray (Layout.Symmetric size) where
   formatArray = formatMirrored id

formatMirrored ::
   (Shape.C size, Class.Floating a, Output out) =>
   (a -> a) ->
   Config ->
   Array (Layout.Mosaic Layout.Packed mirror Shape.Upper size) a ->
   out
formatMirrored adapt cfg m =
   formatSeparateTriangle (printfFloating cfg) $ layoutMirrored adapt m

layoutMirrored ::
   (Shape.C size, Class.Floating a) =>
   (a -> a) ->
   Array (Layout.Mosaic Layout.Packed mirror Shape.Upper size) a ->
   [[a]]
layoutMirrored adapt m =
   let Layout.Mosaic Layout.Packed _mirror Layout.Upper order size
         = Array.shape m
   in  complementTriangle adapt order (Shape.size size) $ Array.toList m

complementTriangle ::
   (Class.Floating a) => (a -> a) -> Order -> Int -> [a] -> [[a]]
complementTriangle adapt order n xs =
   let mergeTriangles lower upper =
         zipWith (++) (map (map adapt . 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)


formatDiagonal ::
   (Shape.C size, Class.Floating a, Output out) =>
   Config -> Order -> size -> [a] -> out
formatDiagonal cfg order size xs =
   let n0 = Unary.unary TypeNum.u0
   in formatAligned (printfFloatingMaybe cfg) $
      padBanded (n0,n0) order (size,size) xs


instance
   (Layout.UpLo uplo, Shape.C size) =>
      FormatArray (Layout.Triangular uplo size) where
   formatArray cfg = formatAligned (printfFloatingMaybe cfg) . layoutTriangular

layoutTriangular ::
   (Layout.UpLo uplo, Shape.C size, Class.Floating a) =>
   Array (Layout.Triangular uplo size) a -> [[Maybe a]]
layoutTriangular m =
   let Layout.Mosaic Layout.Packed Layout.NoMirror uplo order size
         = Array.shape m
   in padTriangle uplo order (Shape.size size) $ Array.toList m

padTriangle ::
   Layout.UpLoSingleton uplo -> Order -> Int -> [a] -> [[Maybe a]]
padTriangle uplo =
   case uplo of
      Layout.Lower -> padLowerTriangle
      Layout.Upper -> padUpperTriangle

padUpperTriangle :: Order -> Int -> [a] -> [[Maybe a]]
padUpperTriangle order n xs =
   let mxs = map Just xs
       nothings = iterate (Nothing:) []
   in case order of
         RowMajor ->
            zipWith (++) nothings (slice (take n $ iterate pred n) mxs)
         ColumnMajor ->
            transpose $
            zipWith (++)
               (slice (take n [1..]) mxs)
               (reverse $ take n nothings)

padLowerTriangle :: Order -> Int -> [a] -> [[Maybe a]]
padLowerTriangle order n xs =
   map (map Just) $
   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


instance
   (Unary.Natural sub, Unary.Natural super,
    Extent.Measure meas, Extent.C vert, Extent.C horiz,
    Shape.C height, Shape.C width) =>
      FormatArray (Layout.Banded sub super meas vert horiz height width) where
   formatArray cfg = formatAligned (printfFloatingMaybe cfg) . layoutBanded

layoutBanded ::
   (Unary.Natural sub, Unary.Natural super,
    Extent.Measure meas, Extent.C vert, Extent.C horiz,
    Shape.C height, Shape.C width, Class.Floating a) =>
   Array (Layout.Banded sub super meas vert horiz height width) a -> [[Maybe a]]
layoutBanded m =
   let Layout.Banded offDiag order extent = Array.shape m
   in  padBanded offDiag order (Extent.dimensions extent) $ Array.toList m

padBanded ::
   (Shape.C height, Shape.C width, Unary.Natural sub, Unary.Natural super) =>
   (UnaryProxy sub, UnaryProxy super) -> Order ->
   (height, width) -> [a] -> [[Maybe a]]
padBanded (sub,super) order (height,width) xs =
   let slices =
         ListHT.sliceVertical (Layout.bandedBreadth (sub,super)) xs
       m = Shape.size height
       n = Shape.size width
   in case order of
         RowMajor ->
            map (take n) $
            zipWith (shiftRow Nothing)
               (iterate (1+) (- integralFromProxy sub))
               (map (map Just) slices)
         ColumnMajor ->
            let ku = integralFromProxy super
            in take m $ drop ku $
               foldr
                  (\col mat ->
                     zipWith (:) (map Just col ++ repeat Nothing) ([]:mat))
                  (replicate (ku + m - n) [])
                  slices


instance
   (Unary.Natural offDiag, Shape.C size) =>
      FormatArray (Layout.BandedHermitian offDiag size) where
   formatArray cfg =
      formatSeparateTriangle (printfFloatingMaybe cfg) . layoutBandedHermitian

layoutBandedHermitian ::
   (Unary.Natural offDiag, Shape.C sh, Class.Floating a) =>
   Array (Layout.BandedHermitian offDiag sh) a -> [[Maybe a]]
layoutBandedHermitian m =
   let Layout.BandedHermitian offDiag order size = Array.shape m
   in padBandedHermitian offDiag order size $ Array.toList m

padBandedHermitian ::
   (Unary.Natural offDiag, Shape.C size, Class.Floating a) =>
   UnaryProxy offDiag -> Order -> size -> [a] -> [[Maybe a]]
padBandedHermitian offDiag order _size xs =
   let k = integralFromProxy offDiag
       slices = ListHT.sliceVertical (k + 1) xs
   in case order of
         RowMajor ->
            foldr
               (\row square ->
                  Match.take ([]:square) (map Just row)
                  :
                  zipWith (:)
                     (tail $ map (Just . conjugate) row ++ repeat Nothing)
                     square)
               [] slices
         ColumnMajor ->
            zipWith (shiftRow Nothing) (iterate (1+) (-k)) $ map (map Just) $
            zipWith (++)
               (map (map conjugate . init) slices)
               (drop k $
                foldr
                  (\column band ->
                     zipWith (++) (map (:[]) column ++ repeat []) ([]:band))
                  (replicate k [])
                  slices)

shiftRow :: a -> Int -> [a] -> [a]
shiftRow pad k = if k<=0 then drop (-k) else (replicate k pad ++)

splitRows ::
   (Shape.C height, Shape.C width) =>
   Order -> (height, width) -> [a] -> [[a]]
splitRows order (height,width) =
   case order of
      RowMajor -> ListHT.sliceVertical (Shape.size width)
      ColumnMajor -> ListHT.sliceHorizontal (Shape.size height)


formatAligned ::
   (Functor f, Foldable f, Output out) =>
   (a -> f String) -> [[a]] -> out
formatAligned printFmt =
   Output.formatAligned . map (map (fmap Output.text . printFmt))

formatSeparateTriangle ::
   (Functor f, Foldable f, Output out) =>
   (a -> f String) -> [[a]] -> out
formatSeparateTriangle printFmt =
   Output.formatSeparateTriangle . map (map (fmap Output.text . printFmt))


arrayFromList2 ::
   (Extent.Measure meas, Extent.C vert, Extent.C horiz) =>
   (Shape.C height, Shape.C width) =>
   Extent meas vert horiz height width -> [[a]] ->
   BoxedArray.Array (height, width) (Output.Separator, Maybe (Output.Style, a))
arrayFromList2 extent =
   incompleteArrayFromList2 extent . map (map Just)

incompleteArrayFromList2 ::
   (Extent.Measure meas, Extent.C vert, Extent.C horiz) =>
   (Shape.C height, Shape.C width) =>
   Extent meas vert horiz height width -> [[Maybe a]] ->
   BoxedArray.Array (height, width) (Output.Separator, Maybe (Output.Style, a))
incompleteArrayFromList2 extent =
   let (height, width) = Extent.dimensions extent in
   let n = Shape.size width in
   BoxedArray.fromList (height, width) . concat .
   map (map ((,) Output.Space) .
        ListHT.padRight Nothing n . map (fmap ((,) Output.Stored)))

splitArrayFromList2 ::
   (Extent.Measure meas, Extent.C vert, Extent.C horiz) =>
   (Shape.C height, Shape.C width) =>
   Extent meas vert horiz height width -> [[a]] ->
   BoxedArray.Array (height, width) (Output.Separator, Maybe (Output.Style, a))
splitArrayFromList2 extent =
   incompleteSplitArrayFromList2 extent . map (map Just)

incompleteSplitArrayFromList2 ::
   (Extent.Measure meas, Extent.C vert, Extent.C horiz) =>
   (Shape.C height, Shape.C width) =>
   Extent meas vert horiz height width -> [[Maybe a]] ->
   BoxedArray.Array (height, width) (Output.Separator, Maybe (Output.Style, a))
incompleteSplitArrayFromList2 extent =
   let (height, width) = Extent.dimensions extent in
   let n = Shape.size width in
   BoxedArray.fromList (height, width) . concat .
   Output.decorateTriangle
      (\seps styles row ->
         zip seps $ ListHT.padRight Nothing n $
         zipWith (\s r -> (,) s <$> r) styles row)

-- ToDo: could be in BoxedArray
toRows ::
   (Shape.C height, Shape.C width) =>
   BoxedArray.Array (height,width) a -> [[a]]
toRows arr =
   case BoxedArray.toList arr of
      [] -> replicate (Shape.size $ fst $ BoxedArray.shape arr) []
      xs -> ListHT.sliceVertical (Shape.size $ snd $ BoxedArray.shape arr) xs

toColumns ::
   (Shape.C height, Shape.C width) =>
   BoxedArray.Array (height,width) a -> [[a]]
toColumns arr =
   ListHT.sliceHorizontal (Shape.size $ snd $ BoxedArray.shape arr) $
   BoxedArray.toList arr



data TupleShape a = TupleShape

instance (Class.Floating a) => Shape.C (TupleShape a) where
   size sh = caseRealComplexFunc sh 1 2

type Tuple a = BoxedArray.Array (TupleShape a)

fillTuple :: (Class.Floating a) => b -> Tuple a b
fillTuple = BoxedArray.replicate TupleShape


newtype ToTuple a = ToTuple {getToTuple :: a -> Tuple a String}

printfFloatingPlain :: (Class.Floating a) => String -> a -> Tuple a String
printfFloatingPlain fmt =
   getToTuple $
   Class.switchFloating
      (ToTuple $ fillTuple . printf fmt)
      (ToTuple $ fillTuple . printf fmt)
      (ToTuple $ printfComplex fmt)
      (ToTuple $ printfComplex fmt)

printfFloating :: (Class.Floating a) => Config -> a -> Tuple a String
printfFloating cfg = printfFloatingPlain (configFormat cfg)

printfFloatingMaybe ::
   (Class.Floating a) => Config -> Maybe a -> Tuple a String
printfFloatingMaybe cfg =
   maybe (fillTuple $ configEmpty cfg) (printfFloating cfg)

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

complexTuple :: (Class.Real a) => b -> b -> Tuple (Complex a) b
complexTuple b0 b1 = BoxedArray.fromList TupleShape [b0,b1]