-----------------------------------------------------------------------------
--
-- Module      :  Math.Roots
-- Copyright   :  (c) 2014-16 Brian W Bush
-- License     :  MIT
--
-- Maintainer  :  Brian W Bush <consult@brianwbush.info>
-- Stability   :  Stable
-- Portability :  Portable
--
-- | Root finding.
--
-----------------------------------------------------------------------------


{-# LANGUAGE Safe #-}


module Math.Roots (
-- * Types.
  RootFinder
, Bracketer
-- * Functions.
, bracket
, bracketInward
, bracketOutward
) where


import Control.Monad (ap)


-- | A function for findng a root.
type RootFinder a =
     (a -> a)        -- ^ The function.
  -> (a, a)          -- ^ The interval bracketing the root.
  -> Either String a -- ^ The root, or an error message if it could not be found.


-- | A function for bracketing a root.
type Bracketer a =
     (a -> a)             -- ^ The function.
  -> (a, a)               -- ^ A guess at the interval bracketing a root.
  -> Either String (a, a) -- ^ The bracketing of the root, or an error message if it could not be found.


-- | Bracket a bracket around a root.
bracket :: (Fractional a, Num a, Ord a) =>
     Int         -- ^ The maximum number of function evaluations allowed.
  -> Bracketer a -- ^ The root bracketing function.
bracket n f xs =
  let
    inward = bracketInward n f xs
    outward = bracketOutward n f xs
  in
    either (const outward) Right inward


-- | Search for a bracket around a root outside an interval.
-- |
-- | Reference: William H. Press, Saul A. Teukolsky, William T. Vetterling, and Brian P. Flannery, Numerical Recipes in C: The Art of Scientific Computing>, Second Edition (New York: Cambridge Univ. Press, 1992).
bracketOutward :: (Fractional a, Num a, Ord a) =>
     Int         -- ^ The maximum number of function evaluations allowed.
  -> Bracketer a -- ^ The root bracketing function.
bracketOutward 0 _ _ = Left "Reached maximum number of iterations."
bracketOutward ntry f (x1, x2)
 | x1 == x2  = Left "Invalid initial range."
 | otherwise =
  do
    let
      factor = 1.6
      f1 = f x1
      f2 = f x2
      dx = x2 - x1
    if f1 * f2 < 0
      then return (x1, x2)
      else bracketOutward (ntry - 1) f $
        if abs f1 < abs f2
          then (x1 - factor * dx, x2)
          else (x1, x2 + factor * dx)


-- | Search for a bracket around a root inside an interval.
-- |
-- | Reference: William H. Press, Saul A. Teukolsky, William T. Vetterling, and Brian P. Flannery, Numerical Recipes in C: The Art of Scientific Computing>, Second Edition (New York: Cambridge Univ. Press, 1992).
bracketInward :: (Num a, Ord a) =>
     Int         -- ^ The maximum number of function evaluations allowed.
  -> Bracketer a -- ^ The root bracketing function.
bracketInward ntry f (x1, x2)
 | x1 == x2  = Left "Invalid initial range."
 | otherwise = bracketInward' $ map (ap (,) f . (x1 +) . ((x2 - x1) *) . fromIntegral) [0..ntry]
  where
    bracketInward' ((xa, fa) : zs@((xb, fb) : _)) =
      if fa * fb < 0
        then return (xa, xb)
        else bracketInward' zs
    bracketInward' _ = Left "Reached maximum number of iterations."