{-# OPTIONS_GHC -Wall #-}
{-# LANGUAGE ScopedTypeVariables #-}
-----------------------------------------------------------------------------
-- |
-- Module      :  ToySolver.Data.Polynomial.Interpolation.Hermite
-- Copyright   :  (c) Masahiro Sakai 2020
-- License     :  BSD-style
--
-- Maintainer  :  masahiro.sakai@gmail.com
-- Stability   :  provisional
-- Portability :  non-portable
--
-- References:
--
-- * Lagrange polynomial <https://en.wikipedia.org/wiki/Hermite_interpolation>
--
-----------------------------------------------------------------------------
module ToySolver.Data.Polynomial.Interpolation.Hermite
  ( interpolate
  ) where

import Data.List (unfoldr)

import ToySolver.Data.Polynomial (UPolynomial, X (..))
import qualified ToySolver.Data.Polynomial as P


-- | Compute a hermite Hermite interpolation from a list
-- \([(x_0, [y_0, y'_0, \ldots y^{(m)}_0]), (x_1, [y_1, y'_1, \ldots y^{(m)}_1]), \ldots]\).
interpolate :: (Eq k, Fractional k) => [(k,[k])] -> UPolynomial k
interpolate :: forall k. (Eq k, Fractional k) => [(k, [k])] -> UPolynomial k
interpolate [(k, [k])]
xyss = [UPolynomial k] -> UPolynomial k
forall a. Num a => [a] -> a
forall (t :: * -> *) a. (Foldable t, Num a) => t a -> a
sum ([UPolynomial k] -> UPolynomial k)
-> [UPolynomial k] -> UPolynomial k
forall a b. (a -> b) -> a -> b
$ (k -> UPolynomial k -> UPolynomial k)
-> [k] -> [UPolynomial k] -> [UPolynomial k]
forall a b c. (a -> b -> c) -> [a] -> [b] -> [c]
zipWith (\k
c UPolynomial k
p -> k -> UPolynomial k
forall k v. (Eq k, Num k) => k -> Polynomial k v
P.constant k
c UPolynomial k -> UPolynomial k -> UPolynomial k
forall a. Num a => a -> a -> a
* UPolynomial k
p) [k]
cs [UPolynomial k]
ps
  where
    x :: UPolynomial k
x = X -> UPolynomial k
forall a v. Var a v => v -> a
P.var X
X
    ps :: [UPolynomial k]
ps = (UPolynomial k -> UPolynomial k -> UPolynomial k)
-> UPolynomial k -> [UPolynomial k] -> [UPolynomial k]
forall b a. (b -> a -> b) -> b -> [a] -> [b]
scanl UPolynomial k -> UPolynomial k -> UPolynomial k
forall a. Num a => a -> a -> a
(*) (k -> UPolynomial k
forall k v. (Eq k, Num k) => k -> Polynomial k v
P.constant k
1) [(UPolynomial k
x UPolynomial k -> UPolynomial k -> UPolynomial k
forall a. Num a => a -> a -> a
- k -> UPolynomial k
forall k v. (Eq k, Num k) => k -> Polynomial k v
P.constant k
x') | (k
x', [k]
ys') <- [(k, [k])]
xyss, k
_ <- [k]
ys']
    cs :: [k]
cs = ([k] -> k) -> [[k]] -> [k]
forall a b. (a -> b) -> [a] -> [b]
map [k] -> k
forall a. HasCallStack => [a] -> a
head ([[k]] -> [k]) -> [[k]] -> [k]
forall a b. (a -> b) -> a -> b
$ [(k, [k])] -> [[k]]
forall a. Fractional a => [(a, [a])] -> [[a]]
dividedDifferenceTable [(k, [k])]
xyss


type D a = Either (a, Int, [a]) ((a, a), a)


dividedDifferenceTable :: forall a. Fractional a => [(a, [a])] -> [[a]]
dividedDifferenceTable :: forall a. Fractional a => [(a, [a])] -> [[a]]
dividedDifferenceTable [(a, [a])]
xyss = ([D a] -> Maybe ([a], [D a])) -> [D a] -> [[a]]
forall b a. (b -> Maybe (a, b)) -> b -> [a]
unfoldr [D a] -> Maybe ([a], [D a])
f [D a]
xyss'
  where
    xyss' :: [D a]
    xyss' :: [D a]
xyss' =
      [ (a, Int, [a]) -> D a
forall a b. a -> Either a b
Left (a
x, Int
len, (a -> Integer -> a) -> [a] -> [Integer] -> [a]
forall a b c. (a -> b -> c) -> [a] -> [b] -> [c]
zipWith (\a
num Integer
den -> a
num a -> a -> a
forall a. Fractional a => a -> a -> a
/ Integer -> a
forall a. Num a => Integer -> a
fromInteger Integer
den) [a]
ys ((Integer -> Integer -> Integer)
-> Integer -> [Integer] -> [Integer]
forall b a. (b -> a -> b) -> b -> [a] -> [b]
scanl Integer -> Integer -> Integer
forall a. Num a => a -> a -> a
(*) Integer
1 [Integer
1..]))
      | (a
x, [a]
ys) <- [(a, [a])]
xyss
      , let len :: Int
len = [a] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [a]
ys
      ]

    d :: D a -> [a]
    d :: D a -> [a]
d (Left (a
_x, Int
n, [a]
ys)) = Int -> a -> [a]
forall a. Int -> a -> [a]
replicate Int
n ([a] -> a
forall a. HasCallStack => [a] -> a
head [a]
ys)
    d (Right ((a, a)
_, a
y)) = [a
y]

    gx1, gx2, gy :: D a -> a
    gx1 :: D a -> a
gx1 (Left (a
x, Int
_n, [a]
_ys)) = a
x
    gx1 (Right ((a
x1, a
_x2), a
_y)) = a
x1
    gx2 :: D a -> a
gx2 (Left (a
x, Int
_n, [a]
_ys)) = a
x
    gx2 (Right ((a
_x1, a
x2), a
_y)) = a
x2
    gy :: D a -> a
gy (Left (a
_x, Int
_n, [a]
ys)) = [a] -> a
forall a. HasCallStack => [a] -> a
head [a]
ys
    gy (Right ((a, a)
_, a
y)) = a
y

    f :: [D a] -> Maybe ([a], [D a])
    f :: [D a] -> Maybe ([a], [D a])
f [] = Maybe ([a], [D a])
forall a. Maybe a
Nothing
    f [D a]
xs = ([a], [D a]) -> Maybe ([a], [D a])
forall a. a -> Maybe a
Just ((D a -> [a]) -> [D a] -> [a]
forall (t :: * -> *) a b. Foldable t => (a -> [b]) -> t a -> [b]
concatMap D a -> [a]
d [D a]
xs, [D a] -> [D a]
h [D a]
xs)
      where
        h :: [D a] -> [D a]
        h :: [D a] -> [D a]
h (D a
z1 : [D a]
zzs) =
          (case D a
z1 of
             Left (a
x, Int
n, a
_ : ys :: [a]
ys@(a
_ : [a]
_)) -> [(a, Int, [a]) -> D a
forall a b. a -> Either a b
Left (a
x, Int
nInt -> Int -> Int
forall a. Num a => a -> a -> a
-Int
1, [a]
ys)]
             D a
_ -> [])
          [D a] -> [D a] -> [D a]
forall a. [a] -> [a] -> [a]
++
          (case [D a]
zzs of
             D a
z2 : [D a]
_ -> [((a, a), a) -> D a
forall a b. b -> Either a b
Right ((D a -> a
gx1 D a
z1, D a -> a
gx2 D a
z2), (D a -> a
gy D a
z2 a -> a -> a
forall a. Num a => a -> a -> a
- D a -> a
gy D a
z1) a -> a -> a
forall a. Fractional a => a -> a -> a
/ (D a -> a
gx2 D a
z2 a -> a -> a
forall a. Num a => a -> a -> a
- D a -> a
gx1 D a
z1))]
             [] -> [])
          [D a] -> [D a] -> [D a]
forall a. [a] -> [a] -> [a]
++
          [D a] -> [D a]
h [D a]
zzs
        h [] = []