{-# LANGUAGE BangPatterns, DeriveDataTypeable, DeriveGeneric, CPP #-}
-- |
-- Module    : Numeric.RootFinding
-- Copyright : (c) 2011 Bryan O'Sullivan
-- License   : BSD3
--
-- Maintainer  : bos@serpentine.com
-- Stability   : experimental
-- Portability : portable
--
-- Haskell functions for finding the roots of real functions of real arguments.
module Numeric.RootFinding
    (
      Root(..)
    , fromRoot
    , ridders
    , newtonRaphson
    -- * References
    -- $references
    ) where

import Control.Applicative              (Alternative(..), Applicative(..))
import Control.Monad                    (MonadPlus(..), ap)
import Data.Data                        (Data, Typeable)
#if __GLASGOW_HASKELL__ > 704
import GHC.Generics                     (Generic)
#endif
import Numeric.MathFunctions.Comparison (within)


-- | The result of searching for a root of a mathematical function.
data Root a = NotBracketed
            -- ^ The function does not have opposite signs when
            -- evaluated at the lower and upper bounds of the search.
            | SearchFailed
            -- ^ The search failed to converge to within the given
            -- error tolerance after the given number of iterations.
            | Root a
            -- ^ A root was successfully found.
              deriving (Eq, Read, Show, Typeable, Data
#if __GLASGOW_HASKELL__ > 704
                       , Generic
#endif
                       )


instance Functor Root where
    fmap _ NotBracketed = NotBracketed
    fmap _ SearchFailed = SearchFailed
    fmap f (Root a)     = Root (f a)

instance Monad Root where
    NotBracketed >>= _ = NotBracketed
    SearchFailed >>= _ = SearchFailed
    Root a       >>= m = m a

    return = Root

instance MonadPlus Root where
    mzero = SearchFailed

    r@(Root _) `mplus` _ = r
    _          `mplus` p = p

instance Applicative Root where
    pure  = Root
    (<*>) = ap

instance Alternative Root where
    empty = SearchFailed

    r@(Root _) <|> _ = r
    _          <|> p = p

-- | Returns either the result of a search for a root, or the default
-- value if the search failed.
fromRoot :: a                   -- ^ Default value.
         -> Root a              -- ^ Result of search for a root.
         -> a
fromRoot _ (Root a) = a
fromRoot a _        = a


-- | Use the method of Ridders to compute a root of a function.
--
-- The function must have opposite signs when evaluated at the lower
-- and upper bounds of the search (i.e. the root must be bracketed).
ridders :: Double               -- ^ Absolute error tolerance.
        -> (Double,Double)      -- ^ Lower and upper bounds for the search.
        -> (Double -> Double)   -- ^ Function to find the roots of.
        -> Root Double
ridders tol (lo,hi) f
    | flo == 0    = Root lo
    | fhi == 0    = Root hi
    | flo*fhi > 0 = NotBracketed -- root is not bracketed
    | otherwise   = go lo flo hi fhi 0
  where
    go !a !fa !b !fb !i
        -- Root is bracketed within 1 ulp. No improvement could be made
        | within 1 a b       = Root a
        -- Root is found. Check that f(m) == 0 is nessesary to ensure
        -- that root is never passed to 'go'
        | fm == 0            = Root m
        | fn == 0            = Root n
        | d < tol            = Root n
        -- Too many iterations performed. Fail
        | i >= (100 :: Int)  = SearchFailed
        -- Ridder's approximation coincide with one of old
        -- bounds. Revert to bisection
        | n == a || n == b   = case () of
          _| fm*fa < 0 -> go a fa m fm (i+1)
           | otherwise -> go m fm b fb (i+1)
        -- Proceed as usual
        | fn*fm < 0          = go n fn m fm (i+1)
        | fn*fa < 0          = go a fa n fn (i+1)
        | otherwise          = go n fn b fb (i+1)
      where
        d    = abs (b - a)
        dm   = (b - a) * 0.5
        !m   = a + dm
        !fm  = f m
        !dn  = signum (fb - fa) * dm * fm / sqrt(fm*fm - fa*fb)
        !n   = m - signum dn * min (abs dn) (abs dm - 0.5 * tol)
        !fn  = f n
    !flo = f lo
    !fhi = f hi


-- | Solve equation using Newton-Raphson iterations.
--
-- This method require both initial guess and bounds for root. If
-- Newton step takes us out of bounds on root function reverts to
-- bisection.
newtonRaphson
  :: Double
     -- ^ Required precision
  -> (Double,Double,Double)
  -- ^ (lower bound, initial guess, upper bound). Iterations will no
  -- go outside of the interval
  -> (Double -> (Double,Double))
  -- ^ Function to finds roots. It returns pair of function value and
  -- its derivative
  -> Root Double
newtonRaphson !prec (!low,!guess,!hi) function
  = go low guess hi
  where
    go !xMin !x !xMax
      | f == 0              = Root x
      | abs (dx / x) < prec = Root x
      | otherwise           = go xMin' x' xMax'
      where
        (f,f') = function x
        -- Calculate Newton-Raphson step
        delta | f' == 0   = error "handle f'==0"
              | otherwise = f / f'
        -- Calculate new approximation and actual change of approximation
        (dx,x') | z <= xMin = let d = 0.5*(x - xMin) in (d, x - d)
                | z >= xMax = let d = 0.5*(x - xMax) in (d, x - d)
                | otherwise = (delta, z)
          where z = x - delta
        -- Update root bracket
        xMin' | dx < 0    = x
              | otherwise = xMin
        xMax' | dx > 0    = x
              | otherwise = xMax



-- $references
--
-- * Ridders, C.F.J. (1979) A new algorithm for computing a single
--   root of a real continuous function.
--   /IEEE Transactions on Circuits and Systems/ 26:979&#8211;980.