{-# LANGUAGE ScopedTypeVariables #-}

module Piecewise (
  Piecewise(..),
  evaluate,
  evaluate',
  from_intervals,
  list_plot' )
where

import qualified Algebra.Additive as Additive ( C )
import qualified Algebra.Field as Field ( C )
import Control.Arrow ( first )
import Misc ( partition )
import NumericPrelude
import Prelude ()

-- | A predicate is basically a function that returns True or
--   False. In this case, the predicate is used to determine which
--   piece of a 'Piecewise' function we want to use.
type Predicate a = (a -> Bool)


-- | One piece of a piecewise function. The predicate determines
--   whether or not the point lies within this piece, and the second
--   component (a function) is what we'll evaluate if it does.
type Piece a = (Predicate a, (a -> a))

-- | A representation of piecewise functions of one variable. A first
--   approach might use a list of intervals paired with the function's
--   value on those intervals. Rather than limit ourselves to a list
--   of intervals, we allow the use of any predicate. The simple
--   intervals case can be implemented with the predicate, \"is x
--   between x1 and x2?\".
data Piecewise a =
  Piecewise { pieces :: [Piece a] }


-- | Evaluate a piecewise function at a point. If the point is within
--   the domain of the function, the function is evaluated at the
--   point and the result is returned wrapped in a 'Just'. If the
--   point is outside of the domain, 'Nothing' is returned.
evaluate :: Piecewise a -> a -> Maybe a
evaluate (Piecewise ps) x =
  foldl f Nothing ps
  where
    f (Just y) _ = Just y
    f Nothing (p,g) = if p x then Just (g x) else Nothing


-- | Evaluate a piecewise function at a point. If the point is within
--   the domain of the function, the function is evaluated at the
--   point and the result is returned. If the point is outside of the
--   domain, zero is returned.
evaluate' :: (Additive.C a) => Piecewise a -> a -> a
evaluate' pw x =
  case evaluate pw x of
    Nothing -> zero
    Just result -> result


-- | Construct a piecewise function from a list of intervals paired
--   with functions. This is a convenience function that automatically
--   converts intervals to \"contained in\" predicates.
--
--   Examples:
--
--   >>> let p1 = ((-1,0), \x -> -x) :: ((Double,Double), Double->Double)
--   >>> let p2 = ((0,1), \x -> x) :: ((Double,Double), Double->Double)
--   >>> let ivs = [p1, p2]
--   >>> let pw = from_intervals ivs
--   >>> evaluate' pw (-0.5)
--   0.5
--   >>> evaluate' pw 0.5
--   0.5
--
from_intervals :: forall a. (Ord a) => [((a,a), (a -> a))] -> Piecewise a
from_intervals =
  Piecewise . map (first f)
  where
    f :: (a,a) -> Predicate a
    f (x1,x2) x = x1 <= x && x <= x2


-- | Plot a piecewise function over an interval @(x0,x1)@ using
--   'evaluate''. Return the result a list of @(x,y)@ tuples.
--
--   Examples:
--
--   >>> let p1 = ((-1,0), \x -> -x) :: ((Double,Double), Double->Double)
--   >>> let p2 = ((0,1), \x -> x) :: ((Double,Double), Double->Double)
--   >>> let ivs = [p1, p2]
--   >>> let pw = from_intervals ivs
--   >>> list_plot' pw (-2) 2 4
--   [(-2.0,0.0),(-1.0,1.0),(0.0,-0.0),(1.0,1.0),(2.0,0.0)]
--
list_plot' :: (Additive.C a, Field.C a)
           => Piecewise a
           -> a
           -> a
           -> Int
           -> [(a,a)]
list_plot' pw x0 x1 subintervals =
  [ (x, evaluate' pw x) | x <- nodes ]
  where
    nodes = x0 : (map snd $ partition subintervals x0 x1)