module Main where

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

import qualified Data.Packed.ST as PackST
import qualified Data.Packed.Matrix as Matrix
import qualified Data.Packed.Vector as Vector
import Data.Packed.Matrix (Matrix)
import Data.Packed.Vector (Vector)

import qualified Numeric.LinearAlgebra.Banded as Banded
import qualified Numeric.Container as Container
import Numeric.Container ((<\>))

import qualified Graphics.Gnuplot.Advanced as GP
import qualified Graphics.Gnuplot.Plot.TwoDimensional as Plot2D
import qualified Graphics.Gnuplot.Graph.TwoDimensional as Graph2D

import System.Random (randomRs, mkStdGen)
import Control.Monad.HT (void)
import Control.Monad (when, zipWithM_, forM_)

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


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.fromList txs0
   in  Matrix.fromColumns $
       map (flip Container.cmap txs . Piecewise.interpolateConstantExt typ) $
       Type.basisFunctions typ xs

basisMatrixSparse ::
   Type.T Double Double ny -> [Double] -> [Double] -> Matrix Double
basisMatrixSparse typ xs txs = PackST.runSTMatrix $ do
   mat <- PackST.newMatrix 0 (length txs) (length $ Type.basisFunctions typ xs)
   zipWithM_
      (\k -> mapM_ (uncurry (PackST.writeMatrix mat k))) [0..] $
      map (Type.sampleBasisFunctions typ xs) txs
   return mat

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 <\> Vector.fromList tys

matrixDiff ::
   Type.T Double Double ny ->
   [Double] -> [(Double, Double)] -> Double
matrixDiff typ xs target =
   let txs = map fst target
   in  Container.maxElement $ Container.cmap abs $
       Container.sub
          (basisMatrixFull typ xs txs)
          (basisMatrixSparse typ xs txs)


mulSparseMatrixVector ::
   Int -> [[(Int, Double)]] -> [Double] -> Vector Double
mulSparseMatrixVector size samples tys = PackST.runSTVector $ do
   vec <- PackST.newVector 0 size
   forM_ (zip samples tys) $ \(row,ty) ->
      forM_ row $ \(k,y) ->
         PackST.modifyVector vec k (+y*ty)
   return vec

bandedGramian ::
   Int -> Int -> [[(Int, Double)]] -> Banded.SymmetricMatrix Double
bandedGramian size width samples =
      Banded.SymmetricMatrix $ PackST.runSTMatrix $ do
   mat <- PackST.newMatrix 0 size width
   forM_ samples $ \row ->
      forM_ row $ \(k,yk) ->
      forM_ row $ \(j,yj) ->
         when (k<=j) $ PackST.modifyMatrix mat k (j-k) (+yk*yj)
   return mat

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
       matrix =
          Banded.choleskyDecompose $
          bandedGramian size (Type.basisOverlap typ) samples
   in  Type.coefficientsToInterpolator typ xs $ Vector.toList $
       Banded.choleskySolve matrix $ 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 ::
   Type.T Double Double ny -> [Double] -> Plot2D.T Double Double
plotBasisFunctions nodeType xs =
   let abscissa = Plot2D.linearScale 1000 (minimum xs, maximum xs)
   in  Plot2D.functions Graph2D.lines abscissa $
       map (Piecewise.interpolateConstantExt nodeType) $
       Type.basisFunctions nodeType xs


main :: IO ()
main = do
   let xs = [0, 1, 3, 4, 6, 7]
       exs = (-1) : xs ++ [8]
   void $ GP.plotDefault $ plotBasisFunctions Type.linear xs
   void $ GP.plotDefault $ plotBasisFunctions Type.hermite1 xs
   void $ GP.plotDefault $ plotBasisFunctions Type.cubicLinear exs
   void $ GP.plotDefault $ plotBasisFunctions Type.cubicParabola 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 $
      Plot2D.list Graph2D.points noisy
      <>
      (Plot2D.functions Graph2D.lines (Plot2D.linearScale 1000 (-2,10)) $
       Piecewise.interpolateConstantExt Type.linear linearNodes :
       Piecewise.interpolateConstantExt Type.hermite1 hermite1Nodes :
       Piecewise.interpolateConstantExt Type.cubicLinear cubicLinearNodes :
       Piecewise.interpolateConstantExt Type.cubicParabola cubicParabolaNodes :
       [])

   putStrLn "differences between matrices should be almost zero:"
   putStrLn $ "linear: " ++ show (matrixDiff Type.linear xs noisy)
   putStrLn $ "hermite1: " ++ show (matrixDiff Type.hermite1 xs noisy)
   putStrLn $ "cubicLinear: " ++ show (matrixDiff Type.cubicLinear exs noisy)
   putStrLn $ "cubicParabola: " ++ show (matrixDiff Type.cubicParabola exs noisy)

   putStrLn "differences between samples should be almost zero:"
   putStrLn $ "linear: " ++ show (bandedDiff absDiffSingle Type.linear xs noisy)
   putStrLn $ "hermite1: " ++ show (bandedDiff absDiffPair Type.hermite1 xs noisy)
   putStrLn $ "cubicLinear: " ++ show (bandedDiff absDiffPair Type.cubicLinear exs noisy)
   putStrLn $ "cubicParabola: " ++ show (bandedDiff absDiffPair Type.cubicParabola exs noisy)