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

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

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

import Linear.Matrix (
  Col,
  Mat(..),
  construct,
  identity_matrix,
  matmap,
  row',
  transpose )
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, fromList )
--
--   >>> 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, fromList )
--
--   >>> 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
    jac = jacobi_matrix

    -- 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 = jac - identity_matrix
    (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 = matmap (\x -> (fromInteger 2)*x^2) first_components

    weights = transpose $ weights_row