{-# LANGUAGE Rank2Types #-}
module Main where

import qualified Plot2DExtra

import qualified Numeric.Interpolation.NodeList as Nodes
import qualified Numeric.Interpolation.Piecewise as Piecewise
import qualified Numeric.Interpolation.Type as Type

import qualified Numeric.LAPACK.Singular as Singular
import qualified Numeric.LAPACK.ShapeStatic as ShapeStatic
import qualified Numeric.LAPACK.Matrix.Shape as MatrixShape
import qualified Numeric.LAPACK.Matrix.BandedHermitianPositiveDefinite
                                                       as BandedSPD
import qualified Numeric.LAPACK.Matrix.BandedHermitian as BandedHermitian
import qualified Numeric.LAPACK.Matrix.Array as ArrMatrix
import qualified Numeric.LAPACK.Matrix as Matrix
import qualified Numeric.LAPACK.Vector as Vector

import qualified Type.Data.Num.Unary as Unary
import Type.Base.Proxy (Proxy)

import qualified Data.Array.Comfort.Storable as Array
import qualified Data.Array.Comfort.Shape as Shape

import qualified Graphics.Gnuplot.Advanced as GP
import qualified Graphics.Gnuplot.Frame as Frame
import qualified Graphics.Gnuplot.Frame.OptionSet as Opts
import qualified Graphics.Gnuplot.Plot.TwoDimensional as Plot2D
import qualified Graphics.Gnuplot.Graph.TwoDimensional as Graph2D
import qualified Graphics.Gnuplot.LineSpecification as LineSpec

import System.Random (randomRs, mkStdGen)
import Control.Monad.HT (void)
import Control.Applicative ((<$>))

import qualified Data.Foldable as Fold
import Data.Tuple.HT (mapFst)
import Data.Monoid ((<>))


type BandedHermitianMatrix k =
       BandedHermitian.BandedHermitian k Matrix.ShapeInt
type Matrix = Matrix.General Matrix.ShapeInt Matrix.ShapeInt
type Vector = Vector.Vector Matrix.ShapeInt
type ShortVector k = Vector.Vector (ShapeStatic.ZeroBased k)


noisy :: [(Double, Double)]
noisy =
   take 100 $
   zipWith
      (\x d -> (x, sin x + d))
      (randomRs (0,7) (mkStdGen 23))
      (randomRs (-0.2,0.2) (mkStdGen 42))

basisMatrixFull ::
   Type.T Double Double ny -> [Double] -> [Double] -> Matrix Double
basisMatrixFull typ xs txs0 =
   let txs = Vector.autoFromList txs0
   in  Matrix.fromColumns (Array.shape txs) $
       map (flip Array.map txs . Piecewise.interpolateConstantExt typ) $
       Type.basisFunctions typ xs

zipRowsWith :: (a -> b -> c) -> [a] -> [[b]] -> [c]
zipRowsWith f as bs = concat $ zipWith (map . f) as bs

-- ToDo: generalize shapes
basisMatrixSparse ::
   Type.T Double Double ny -> [Double] -> [Double] -> Matrix Double
basisMatrixSparse typ xs txs =
   Matrix.fromRowMajor $
   Array.fromAssociations 0
      (Shape.ZeroBased $ length txs,
       Shape.ZeroBased $ length $ Type.basisFunctions typ xs) $
   zipRowsWith (\k (j,x) -> ((k,j),x)) [0..] $
   map (Type.sampleBasisFunctions typ xs) txs

fit ::
   Type.T Double Double ny ->
   [Double] -> [(Double, Double)] -> Nodes.T Double ny
fit typ xs target =
   let (txs, tys) = unzip target
       matrix = basisMatrixSparse typ xs txs
   in  Type.coefficientsToInterpolator typ xs $
       Vector.toList $
       Matrix.unliftColumn MatrixShape.ColumnMajor
         (snd . Singular.leastSquaresMinimumNormRCond 1e-5 matrix) $
       Vector.autoFromList tys

matrixDiff ::
   Type.T Double Double ny -> [Double] -> [(Double, Double)] -> Double
matrixDiff typ xs target =
   let txs = map fst target
   in  Vector.normInf $ ArrMatrix.toVector $
       Matrix.sub
          (basisMatrixFull typ xs txs)
          (basisMatrixSparse typ xs txs)


mulSparseMatrixVector ::
   Int -> [[(Int, Double)]] -> [Double] -> Vector Double
mulSparseMatrixVector size samples tys =
   Array.accumulate (+)
      (Vector.zero (Matrix.shapeInt size))
      (zipRowsWith (\ty (k,y) -> (k, y*ty)) tys samples)

{- ToDo:
use index type (ZeroInt, Enumeration Order)
with Order = Absolute | Derivative
Problem: not all interpolation types use derivatives
-}
shortVector ::
   (Unary.Natural k) =>
   Proxy k -> [(Int, Double)] -> (Int, ShortVector k Double)
shortVector width xs =
   let i0 = minimum $ map fst xs
   in (i0,
--       Array.fromAssociations 0 Shape.static $
       Array.reshape Shape.static $
       Array.fromAssociations 0 (Matrix.shapeInt $ Unary.integralFromProxy width) $
       map (mapFst (subtract i0)) xs)

bandedGramian ::
   (Unary.Natural k) =>
   Int -> Proxy (Unary.Succ k) ->
   [[(Int, Double)]] -> BandedHermitianMatrix k Double
bandedGramian size width samples =
   BandedHermitian.sumRank1 MatrixShape.ColumnMajor (Matrix.shapeInt size) $
   map ((,) 1 . shortVector width) samples

reifyPositive ::
   Integer -> (forall s. Unary.Natural s => Proxy (Unary.Succ s) -> w) -> w
reifyPositive n f = Unary.reifyNatural (n-1) (f . Unary.succ)

fitBanded ::
   Type.T Double Double ny ->
   [Double] -> [(Double, Double)] -> Nodes.T Double ny
fitBanded typ xs target =
   let size = length $ Type.basisFunctions typ xs
       (txs, tys) = unzip target
       samples = map (Type.sampleBasisFunctions typ xs) txs
   in reifyPositive (toInteger $ Type.basisOverlap typ) 
         (\width ->
            Type.coefficientsToInterpolator typ xs $ Vector.toList $
            Matrix.unliftColumn MatrixShape.ColumnMajor
               (BandedSPD.solve (bandedGramian size width samples)) $
            mulSparseMatrixVector size samples tys)

bandedDiff ::
   (ny -> ny -> Double) ->
   Type.T Double Double ny ->
   [Double] -> [(Double, Double)] -> Double
bandedDiff absDiff typ xs target =
   maximum $
   zipWith absDiff
      (Fold.toList $ fit typ xs target)
      (Fold.toList $ fitBanded typ xs target)

absDiffSingle :: Double -> Double -> Double
absDiffSingle x y = abs (x-y)

absDiffPair :: (Double,Double) -> (Double,Double) -> Double
absDiffPair (x,dx) (y,dy) = max (abs (x-y)) (abs (dx-dy))


plotBasisFunctions ::
   (String, Type.T Double Double ny) -> [Double] ->
   Frame.T (Graph2D.T Double Double)
plotBasisFunctions (title, nodeType) xs =
   let abscissa = Plot2D.linearScale 1000 (minimum xs, maximum xs)
   in  Frame.cons (Opts.title title $ Opts.key False Opts.deflt) $
       Plot2D.functions Graph2D.lines abscissa $
       map (Piecewise.interpolateConstantExt nodeType) $
       Type.basisFunctions nodeType xs


typeLinear :: (String, Type.T Double Double Double)
typeLinear        = ("linear",        Type.linear)
typeHermite1, typeCubicLinear, typeCubicParabola ::
   (String, Type.T Double Double (Double, Double))
typeHermite1      = ("hermite1",      Type.hermite1)
typeCubicLinear   = ("cubicLinear",   Type.cubicLinear)
typeCubicParabola = ("cubicParabola", Type.cubicParabola)


main :: IO ()
main = do
   let xs = [0, 1, 3, 4, 6, 7]
       exs = (-1) : xs ++ [8]
   void $ GP.plotDefault $ plotBasisFunctions typeLinear xs
   void $ GP.plotDefault $ plotBasisFunctions typeHermite1 xs
   void $ GP.plotDefault $ plotBasisFunctions typeCubicLinear exs
   void $ GP.plotDefault $ plotBasisFunctions typeCubicParabola exs
   let linearNodes = fit Type.linear xs noisy
       hermite1Nodes = fit Type.hermite1 xs noisy
       cubicLinearNodes = fit Type.cubicLinear exs noisy
       cubicParabolaNodes = fit Type.cubicParabola exs noisy
   void $ GP.plotDefault $
      (Graph2D.lineSpec (LineSpec.title "target" LineSpec.deflt)
         <$> Plot2D.list Graph2D.points noisy)
      <>
      (let interpolate (name,typ) nodes =
             (name, Piecewise.interpolateConstantExt typ nodes)
       in Plot2DExtra.functions (Plot2D.linearScale 1000 (-2,10)) $
            interpolate typeLinear linearNodes :
            interpolate typeHermite1 hermite1Nodes :
            interpolate typeCubicLinear cubicLinearNodes :
            interpolate typeCubicParabola cubicParabolaNodes :
            [])

   putStrLn "differences between matrices should be almost zero:"
   let printMatrixDiff (name,typ) ps =
         putStrLn $ name ++ ": " ++ show (matrixDiff typ ps noisy)
   printMatrixDiff typeLinear xs
   printMatrixDiff typeHermite1 xs
   printMatrixDiff typeCubicLinear exs
   printMatrixDiff typeCubicParabola exs

   putStrLn "differences between samples should be almost zero:"
   let printBandedDiff diff (name,typ) ps =
         putStrLn $ name ++ ": " ++ show (bandedDiff diff typ ps noisy)
   printBandedDiff absDiffSingle typeLinear xs
   printBandedDiff absDiffPair typeHermite1 xs
   printBandedDiff absDiffPair typeCubicLinear exs
   printBandedDiff absDiffPair typeCubicParabola exs