{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE NoImplicitPrelude #-}
{-# LANGUAGE ScopedTypeVariables #-}

-- | Gaussian quadrature over [-1, 1] using the Legendre polynomials.
--
module Integration.Gaussian (
  gaussian,
  gaussian',
  jacobi_matrix,
  nodes_and_weights )
where

import Algebra.Absolute ( abs )
import qualified Algebra.Absolute as Absolute ( C )
import qualified Algebra.Algebraic as Algebraic ( C )
import qualified Algebra.Field as Field ( C )
import qualified Algebra.Ring as Ring ( C )
import qualified Algebra.ToRational as ToRational ( C )
import Data.Vector.Fixed ( Arity, S )
import NumericPrelude hiding ( abs )

import Linear.Matrix (
  Col,
  Col10,
  Mat(..),
  construct,
  element_sum2,
  fromList,
  identity_matrix,
  map2,
  row,
  transpose,
  zipwith2 )
import Linear.QR ( eigenvectors_symmetric )


-- | Compute the Jacobi matrix for the Legendre polynomials over
--   [-1,1]. This matrix is derived from the three-term recurrence
--   relation which is expressable as a matrix equation. When the
--   polynomials are ortho/normal/, its Jacobi matrix should be
--   symmetric. But since the Legendre polynomials are not
--   ortho/normal/, we have to make a few modifications. The procedure
--   is described in Amparo, Gil, and Segura, \"Numerical Methods for
--   Special Functions\".
--
--   Examples:
--
--   >>> import Linear.Matrix ( Mat2, frobenius_norm )
--
--   >>> let c = 1/(sqrt 3)
--   >>> let expected = fromList [[0,c],[c, 0]] :: Mat2 Double
--   >>> let actual = jacobi_matrix :: Mat2 Double
--   >>> frobenius_norm (actual - expected) < 1e-14
--   True
--
jacobi_matrix :: forall m a. (Arity m, Algebraic.C a, Ring.C a)
              => Mat (S m) (S m) a
jacobi_matrix =
  construct lambda
  where
    coeff_a :: Int -> a
    coeff_a n = fromRational' $ (toRational (n + 1))/(toRational (2*n + 1))

    coeff_c :: Int -> a
    coeff_c n = fromRational' $ (toRational n)/(toRational (2*n + 1))

    max2 :: Ord b => b -> b -> b
    max2 x y = if x >= y then x else y

    lambda i j
      | j == i = fromInteger 0
      | abs (i-j) == 1 =
          let k = max2 i j
          in  sqrt $ (coeff_a (k-1)) * (coeff_c k)

      | otherwise = fromInteger 0



-- | Compute the nodes/weights for Gaussian quadrature over [-1,1]
--   using the Legendre polynomials. The test cases can all be found
--   on e.g. Wikipedia.
--
--   Examples:
--
--   >>> import Linear.Matrix ( Col2, frobenius_norm )
--
--   >>> let (ns,ws) = nodes_and_weights 1000 :: (Col2 Double, Col2 Double)
--   >>> let c = 1/(sqrt 3)
--   >>> let expected_ns = fromList [[-c],[c]] :: Col2 Double
--   >>> let expected_ws = fromList [[1],[1]] :: Col2 Double
--   >>> frobenius_norm (ns - expected_ns) < 1e-12
--   True
--   >>> frobenius_norm (ws - expected_ws) < 1e-12
--   True
--
nodes_and_weights :: forall m a.
                     (Arity m, Eq a, Algebraic.C a)
                  => Int
                  -> (Col (S m) a, Col (S m) a)
nodes_and_weights iterations =
  (nodes, weights)
  where
    -- A shift is needed to make sure the QR algorithm
    -- converges. Since the roots (and thus eigenvalues) of orthogonal
    -- polynomials are real and distinct, there does exist a shift
    -- that will make it work. We guess lambda=1 makes it work.

    shifted_jac = jacobi_matrix - identity_matrix :: Mat (S m) (S m) a
    (shifted_nodes, vecs) = eigenvectors_symmetric iterations shifted_jac

    ones :: Col (S m) a
    ones = construct (\_ _ -> fromInteger 1)

    nodes :: Col (S m) a
    nodes = shifted_nodes + ones -- unshift the nodes

    -- Get the first component of each column.
    first_components = row vecs 0

    -- Square it and multiply by 2; see the Golub-Welsch paper for
    -- this magic.
    weights_row = map2 (\x -> (fromInteger 2)*x^2) first_components

    weights = transpose $ weights_row


-- | Integrate a function over [-1,1] using Gaussian quadrature. This
--   is the \"simple\" version of the 'gaussian'' function; here the ten
--   nodes and weights are fixed.
--
--   Examples:
--
--   >>> let f x = x * (sin x)
--   >>> let actual = gaussian f :: Double
--   >>> let expected = 0.6023373578795136
--   >>> abs (actual - expected) < 1e-12
--   True
--
gaussian :: forall a.
            (Absolute.C a, Algebraic.C a, ToRational.C a, Field.C a)
         => (a -> a) -- ^ The function to integrate.
         -> a
gaussian f =
  gaussian' f nodes weights
  where
    coerce :: (ToRational.C b) => b -> a
    coerce = fromRational' . toRational

    nodes :: Col10 a
    nodes = fromList $ map (map coerce) node_list

    node_list :: [[Double]]
    node_list = [[-0.9739065285171726],
                 [-0.8650633666889886],
                 [-0.6794095682990235],
                 [-0.43339539412924943],
                 [-0.14887433898163027],
                 [0.14887433898163094],
                 [0.4333953941292481],
                 [0.6794095682990247],
                 [0.8650633666889846],
                 [0.9739065285171717]]

    weights :: Col10 a
    weights = fromList $ map (map coerce) weight_list

    weight_list :: [[Double]]
    weight_list =  [[0.06667134430868775],
                    [0.14945134915057925],
                    [0.21908636251598282],
                    [0.26926671930999607],
                    [0.295524224714752],
                    [0.295524224714754],
                    [0.26926671930999635],
                    [0.219086362515982],
                    [0.14945134915058048],
                    [0.06667134430868814]]



-- | Integrate a function over [-1,1] using Gaussian quadrature. The
--   nodes and the weights must be given (due to the design of the
--   library, the number of nodes/weights must be requested at the
--   type level so this function couldn't just compute e.g. @n@ of
--   them for you).
--
--   Examples:
--
--   >>> import Linear.Matrix ( Col5 )
--
--   >>> let (ns, ws) = nodes_and_weights 100 :: (Col5 Double, Col5 Double)
--   >>> let f x = x * (sin x)
--   >>> let actual = gaussian' f ns ws
--   >>> let expected = 0.6023373578795136
--   >>> abs (actual - expected) < 1e-8
--   True
--
gaussian' :: forall m a.
             (Arity m, ToRational.C a, Ring.C a)
          => (a -> a)    -- ^ The function @f@ to integrate.
          -> Col (S m) a -- ^ Column matrix of nodes
          -> Col (S m) a -- ^ Column matrix of weights
          -> a
gaussian' f nodes weights =
  -- The one norm is just the sum of the entries, which is what we
  -- want.
  element_sum2 weighted_values
  where
    function_values = map2 f nodes
    weighted_values = zipwith2 (*) weights function_values