{-# LANGUAGE RebindableSyntax #-}

-- | The Roots.Simple module contains root-finding algorithms. That
--   is, procedures to (numerically) find solutions to the equation,
--
--   > f(x) = 0
--
--   where f is assumed to be continuous on the interval of interest.
--

module Roots.Simple
where

import Data.List (find)

import Normed

import qualified Roots.Fast as F

import NumericPrelude hiding (abs)
import Algebra.Absolute (abs)
import qualified Algebra.Additive as Additive
import qualified Algebra.Algebraic as Algebraic
import qualified Algebra.Field as Field
import qualified Algebra.RealField as RealField
import qualified Algebra.RealRing as RealRing

-- | Does the (continuous) function @f@ have a root on the interval
--   [a,b]? If f(a) <] 0 and f(b) ]> 0, we know that there's a root in
--   [a,b] by the intermediate value theorem. Likewise when f(a) >= 0
--   and f(b) <= 0.
--
--   Examples:
--
--   >>> let f x = x**3
--   >>> has_root f (-1) 1 Nothing
--   True
--
--   This fails if we don't specify an @epsilon@, because cos(-2) ==
--   cos(2) doesn't imply that there's a root on [-2,2].
--
--   >>> has_root cos (-2) 2 Nothing
--   False
--   >>> has_root cos (-2) 2 (Just 0.001)
--   True
--
has_root :: (RealField.C a, RealRing.C b)
         => (a -> b) -- ^ The function @f@
         -> a       -- ^ The \"left\" endpoint, @a@
         -> a       -- ^ The \"right\" endpoint, @b@
         -> Maybe a -- ^ The size of the smallest subinterval
                   --   we'll examine, @epsilon@
         -> Bool
has_root f a b epsilon =
  F.has_root f a b epsilon Nothing Nothing


-- | We are given a function @f@ and an interval [a,b]. The bisection
--   method finds a root by splitting [a,b] in half repeatedly.
--
--   If one is found within some prescribed tolerance @epsilon@, it is
--   returned. Otherwise, the interval [a,b] is split into two
--   subintervals [a,c] and [c,b] of equal length which are then both
--   checked via the same process.
--
--   Returns 'Just' the value x for which f(x) == 0 if one is found,
--   or Nothing if one of the preconditions is violated.
--
--   Examples:
--
--   >>> let actual = 1.5707963267948966
--   >>> let Just root = bisect cos 1 2 0.001
--   >>> root
--   1.5712890625
--   >>> abs (root - actual) < 0.001
--   True
--
--   >>> bisect sin (-1) 1 0.001
--   Just 0.0
--
bisect :: (RealField.C a, RealRing.C b)
       => (a -> b) -- ^ The function @f@ whose root we seek
       -> a       -- ^ The \"left\" endpoint of the interval, @a@
       -> a       -- ^ The \"right\" endpoint of the interval, @b@
       -> a       -- ^ The tolerance, @epsilon@
       -> Maybe a
bisect f a b epsilon =
  F.bisect f a b epsilon Nothing Nothing


-- | We are given a function @f@ and an interval [a,b]. The trisection
--   method finds a root by splitting [a,b] into three
--   subintervals repeatedly.
--
--   If one is found within some prescribed tolerance @epsilon@, it is
--   returned. Otherwise, the interval [a,b] is split into two
--   subintervals [a,c] and [c,b] of equal length which are then both
--   checked via the same process.
--
--   Returns 'Just' the value x for which f(x) == 0 if one is found,
--   or Nothing if one of the preconditions is violated.
--
--   Examples:
--
--   >>> let actual = 1.5707963267948966
--   >>> let Just root = trisect cos 1 2 0.001
--   >>> root
--   1.5713305898491083
--   >>> abs (root - actual) < 0.001
--   True
--
--   >>> trisect sin (-1) 1 0.001
--   Just 0.0
--
trisect :: (RealField.C a, RealRing.C b)
       => (a -> b) -- ^ The function @f@ whose root we seek
       -> a       -- ^ The \"left\" endpoint of the interval, @a@
       -> a       -- ^ The \"right\" endpoint of the interval, @b@
       -> a       -- ^ The tolerance, @epsilon@
       -> Maybe a
trisect f a b epsilon =
  F.trisect f a b epsilon Nothing Nothing


-- | Find a fixed point of the function @f@ with the search starting
--   at x0. We delegate to the version that returns the number of
--   iterations and simply discard the number of iterations.
--
fixed_point :: (Normed a, Additive.C a, Algebraic.C b, RealField.C b)
            => (a -> a) -- ^ The function @f@ to iterate.
            -> b       -- ^ The tolerance, @epsilon@.
            -> a       -- ^ The initial value @x0@.
            -> a       -- ^ The fixed point.
fixed_point f epsilon x0 =
  snd $ F.fixed_point_with_iterations f epsilon x0


-- | Return the number of iterations required to find a fixed point of
--   the function @f@ with the search starting at x0 and tolerance
--   @epsilon@. We delegate to the version that returns the number of
--   iterations and simply discard the fixed point.
fixed_point_iteration_count :: (Normed a,
                                Additive.C a,
                                RealField.C b,
                                Algebraic.C b)
                            => (a -> a) -- ^ The function @f@ to iterate.
                            -> b       -- ^ The tolerance, @epsilon@.
                            -> a       -- ^ The initial value @x0@.
                            -> Int       -- ^ The fixed point.
fixed_point_iteration_count f epsilon x0 =
  fst $ F.fixed_point_with_iterations f epsilon x0


-- | Returns a list of ratios,
--
--     ||x^{*} - x_{n+1}|| / ||x^{*} - x_{n}||^{p}
--
--   of fixed point iterations for the function @f@ with initial guess
--   @x0@ and @p@ some positive power.
--
--   This is used to determine the rate of convergence.
--
fixed_point_error_ratios :: (Normed a,
                             Additive.C a,
                             RealField.C b,
                             Algebraic.C b)
                   => (a -> a) -- ^ The function @f@ to iterate.
                   -> a       -- ^ The initial value @x0@.
                   -> a       -- ^ The true solution, @x_star@.
                   -> Integer -- ^ The power @p@.
                   -> [b]     -- ^ The resulting sequence of x_{n}.
fixed_point_error_ratios f x0 x_star p =
  zipWith (/) en_plus_one en_exp
  where
    xn = F.fixed_point_iterations f x0
    en = map (\x -> norm (x_star - x)) xn
    en_plus_one = tail en
    en_exp = map (^p) en



-- | The sequence x_{n} of values obtained by applying Newton's method
--   on the function @f@ and initial guess @x0@.
--
--   Examples:
--
--   Atkinson, p. 60.
--   >>> let f x = x^6 - x - 1
--   >>> let f' x = 6*x^5 - 1
--   >>> tail $ take 4 $ newton_iterations f f' 2
--   [1.6806282722513088,1.4307389882390624,1.2549709561094362]
--
newton_iterations :: (Field.C a)
                    => (a -> a) -- ^ The function @f@ whose root we seek
                    -> (a -> a) -- ^ The derivative of @f@
                    -> a       -- ^ Initial guess, x-naught
                    -> [a]
newton_iterations f f' =
  iterate next
  where
  next xn =
    xn - ( (f xn) / (f' xn) )


-- | Use Newton's method to find a root of @f@ near the initial guess
--   @x0@. If your guess is bad, this will recurse forever!
--
--   Examples:
--
--   Atkinson, p. 60.
--
--   >>> let f x = x^6 - x - 1
--   >>> let f' x = 6*x^5 - 1
--   >>> let Just root = newtons_method f f' (1/1000000) 2
--   >>> root
--   1.1347241385002211
--   >>> abs (f root) < 1/100000
--   True
--
--   >>> import Data.Number.BigFloat
--   >>> let eps = 1/(10^20) :: BigFloat Prec50
--   >>> let Just root = newtons_method f f' eps 2
--   >>> root
--   1.13472413840151949260544605450647284028100785303643e0
--   >>> abs (f root) < eps
--   True
--
newtons_method :: (RealField.C a)
                 => (a -> a) -- ^ The function @f@ whose root we seek
                 -> (a -> a) -- ^ The derivative of @f@
                 -> a       -- ^ The tolerance epsilon
                 -> a       -- ^ Initial guess, x-naught
                 -> Maybe a
newtons_method f f' epsilon x0 =
  find (\x -> abs (f x) < epsilon) x_n
  where
    x_n = newton_iterations f f' x0


-- | Takes a function @f@ of two arguments and repeatedly applies @f@
--   to the previous two values. Returns a list containing all
--   generated values, f(x0, x1), f(x1, x2), f(x2, x3)...
--
--   Examples:
--
--   >>> let fibs = iterate2 (+) 0 1
--   >>> take 15 fibs
--   [0,1,1,2,3,5,8,13,21,34,55,89,144,233,377]
--
iterate2 :: (a -> a -> a) -- ^ The function @f@
         -> a           -- ^ The initial value @x0@
         -> a           -- ^ The second value, @x1@
         -> [a]         -- ^ The result list, [x0, x1, ...]
iterate2 f x0 x1 =
  x0 : x1 : (go x0 x1)
  where
    go prev2 prev1 =
      let next = f prev2 prev1 in
        next : go prev1 next


-- | The sequence x_{n} of values obtained by applying the secant
--   method on the function @f@ and initial guesses @x0@, @x1@.
--
--   The recursion more or less implements a two-parameter 'iterate',
--   although one list is passed to the next iteration (as opposed to
--   one function argument, with iterate). At each step, we peel the
--   first two elements off the list and then compute/append elements
--   three, four... onto the end of the list.
--
--   Examples:
--
--   Atkinson, p. 67.
--   >>> let f x = x^6 - x - 1
--   >>> take 4 $ secant_iterations f 2 1
--   [2.0,1.0,1.0161290322580645,1.190577768676638]
--
secant_iterations :: (Field.C a)
                    => (a -> a) -- ^ The function @f@ whose root we seek
                    -> a       -- ^ Initial guess, x-naught
                    -> a       -- ^ Second initial guess, x-one
                    -> [a]
secant_iterations f =
  iterate2 g
  where
  g prev2 prev1 =
    let x_change = prev1 - prev2
        y_change = (f prev1) - (f prev2)
    in
      (prev1 - (f prev1 * (x_change / y_change)))


-- | Use the secant method to find a root of @f@ near the initial guesses
--   @x0@ and @x1@. If your guesses are bad, this will recurse forever!
--
--   Examples:
--
--   Atkinson, p. 67.
--   >>> let f x = x^6 - x - 1
--   >>> let Just root = secant_method f (1/10^9) 2 1
--   >>> root
--   1.1347241384015196
--   >>> abs (f root) < (1/10^9)
--   True
--
secant_method :: (RealField.C a)
                 => (a -> a) -- ^ The function @f@ whose root we seek
                 -> a       -- ^ The tolerance epsilon
                 -> a       -- ^ Initial guess, x-naught
                 -> a       -- ^ Second initial guess, x-one
                 -> Maybe a
secant_method f epsilon x0 x1
  = find (\x -> abs (f x) < epsilon) x_n
  where
    x_n = secant_iterations f x0 x1