{-# LANGUAGE AllowAmbiguousTypes #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE TypeApplications #-}

{- | Ball arithmetic

Also known as mid-rad interval arithmetic, is an extension of floating-point arithmetic
in which an error bound is attached to each variable.
-}

module Support.Ball
  ( Ball(..)
  , exactly
  , balls
  , balls44
  , widen
  , narrow
  , isExactly
  , sumAnyOrder
  , libmUlps
  , Outside(..)
  , outside
  , enclosed
  , Disjoint(..)
  , disjoint
  , consistent
  ) where

import Data.Foldable (toList)
import Linear (M44)
import Linear.Epsilon (Epsilon(..))

data Ball p = Ball
  { center :: {-# UNPACK #-} !Double
  , radius :: {-# UNPACK #-} !Double
  }
  deriving (Eq, Show)

exactly :: Real p => p -> Ball p
exactly x = Ball (realToFrac x) 0

balls :: (Functor f, Real p) => f p -> f (Ball p)
balls = fmap exactly

balls44 :: Real p => M44 p -> M44 (Ball p)
balls44 = fmap balls

widen :: Ball p -> Ball Double
widen (Ball m r) = Ball m r

narrow :: Ball Double -> Ball Float
narrow (Ball m r)
  | r == 0, representable @Float m = Ball m 0
  | otherwise = rounded m r

isExactly :: Double -> Ball p -> Bool
isExactly x (Ball m r) = r == 0 && m == x

-- | Assume `sin`, `cos` and `tan` are accurate to 1 ulp.
libmUlps :: Double
libmUlps = 1

unitRoundoff :: forall p. RealFloat p => Double
unitRoundoff = 2 ^^ negate (floatDigits (0 :: p))

smallestSubnormal :: forall p. RealFloat p => Double
smallestSubnormal = encodeFloat 1 (fst (floatRange (0 :: p)) - floatDigits (0 :: p))

representable :: forall p. RealFloat p => Double -> Bool
representable x = realToFrac (realToFrac x :: p) == x

gamma :: Int -> Double -> Double
gamma n u = nu / (1 - nu)
  where
    nu = fromIntegral n * u

radiusSlack :: Double -> Double
radiusSlack r = r * (1 + 16 * unitRoundoff @Double)

rounded :: forall p. RealFloat p => Double -> Double -> Ball p
rounded m propagated =
  Ball m $ radiusSlack (propagated + midpointRounding + simulatedRounding)
  where
    midpointRounding = unitRoundoff @Double * abs m + smallestSubnormal @Double
    simulatedRounding =
      unitRoundoff @p * (abs m + propagated + midpointRounding) + smallestSubnormal @p

libmRounded :: forall p. RealFloat p => Double -> Double -> Ball p
libmRounded m propagated =
  Ball m $ radiusSlack (propagated + midpointError + simulatedError)
  where
    midpointError = libmError @Double (abs m)
    simulatedError = libmError @p (abs m + propagated + midpointError)

libmError :: forall p. RealFloat p => Double -> Double
libmError magnitude = libmUlps * (2 * unitRoundoff @p * magnitude + smallestSubnormal @p)

twoSum :: Double -> Double -> (Double, Double)
twoSum a b = (s, (a - (s - b')) + (b - b'))
  where
    s = a + b
    b' = s - a

exactProduct :: Double -> Double -> Double -> Bool
exactProduct a b prod =
  snd (twoProduct a b) == 0 && toRational a * toRational b == toRational prod

twoProduct :: Double -> Double -> (Double, Double)
twoProduct a b = (p, ((ah * bh - p) + ah * bl + al * bh) + al * bl)
  where
    p = a * b
    (ah, al) = split a
    (bh, bl) = split b
    split x = (hi, x - hi)
      where
        c = 134217729 * x
        hi = c - (c - x)

plus :: forall p. RealFloat p => Ball p -> Ball p -> Ball p
plus a@(Ball ma ra) b@(Ball mb rb)
  | isExactly 0 a = b
  | isExactly 0 b = a
  | ra == 0, rb == 0, roundoff == 0, representable @p s = Ball s 0
  | otherwise = rounded s (ra + rb)
  where
    (s, roundoff) = twoSum ma mb

times :: forall p. RealFloat p => Ball p -> Ball p -> Ball p
times a@(Ball ma ra) b@(Ball mb rb)
  | isExactly 0 a || isExactly 0 b = Ball 0 0
  | isExactly 1 a = b
  | isExactly 1 b = a
  | isExactly (-1) a = negate b
  | isExactly (-1) b = negate a
  | ra == 0, rb == 0, exactProduct ma mb prod, representable @p prod = Ball prod 0
  | otherwise = rounded prod (abs ma * rb + abs mb * ra + ra * rb)
  where
    prod = ma * mb

sumAnyOrder :: forall p. RealFloat p => [Ball p] -> Ball p
sumAnyOrder summands =
  case filter (not . isExactly 0) summands of
    [] -> Ball 0 0
    [single] -> single
    nonZero -> anyOrder nonZero
  where
    anyOrder xs = Ball m $ radiusSlack (propagated + midpointRounding + simulatedRounding)
      where
        m = sum $ map center xs
        propagated = (1 + gamma additions (unitRoundoff @Double)) * sum (map radius xs)
        magnitudes = sum $ map (abs . center) xs
        additions = length xs - 1
        midpointRounding =
          gamma additions (unitRoundoff @Double) * magnitudes +
          fromIntegral additions * smallestSubnormal @Double
        simulatedRounding =
          gamma additions (unitRoundoff @p) * (magnitudes + propagated + midpointRounding) +
          fromIntegral additions * smallestSubnormal @p

literal :: forall p. RealFloat p => Rational -> Ball p
literal q
  | exact = Ball m 0
  | otherwise = roundedLiteral m (realToFrac (fromRational q :: p))
  where
    m = fromRational q
    exact = toRational m == q && toRational (fromRational q :: p) == q

roundedLiteral :: Double -> Double -> Ball p
roundedLiteral m simulated =
  Ball m $ radiusSlack (max (abs (simulated - m)) midpointRounding)
  where
    midpointRounding = unitRoundoff @Double * abs m + smallestSubnormal @Double

instance RealFloat p => Num (Ball p) where
  (+) = plus
  a - b = plus a (negate b)
  (*) = times
  negate (Ball m r) = Ball (negate m) r
  abs (Ball m r) = Ball (abs m) r
  signum = error "Ball.signum: no enclosure for a discontinuous function"
  fromInteger = literal . fromInteger

instance RealFloat p => Fractional (Ball p) where
  a@(Ball ma ra) / b@(Ball mb rb)
    | isExactly 1 b = a
    | abs mb <= 2 * rb = error "Ball./: divisor is too close to zero to bound the quotient"
    | ra == 0, rb == 0, exactQuotient = Ball quotient 0
    | otherwise = rounded quotient (((abs quotient + smallestSubnormal @Double) * rb + ra) / (abs mb - rb))
    where
      quotient = ma / mb
      exactQuotient =
        toRational quotient == toRational ma / toRational mb && representable @p quotient
  fromRational = literal

instance RealFloat p => Floating (Ball p) where
  pi = roundedLiteral pi (realToFrac (pi :: p))

  sqrt (Ball m r)
    | r == 0, exactProduct root root m, representable @p root = Ball root 0
    | m > 2 * r = rounded root (r / root)
    | m + r < 0 = error "Ball.sqrt: negative argument"
    | otherwise = rounded (upper / 2) (upper / 2)
    where
      root = sqrt m
      upper = sqrt (m + r)

  sin x@(Ball m r)
    | isExactly 0 x = Ball 0 0
    | otherwise = libmRounded (sin m) r
  cos x@(Ball m r)
    | isExactly 0 x = Ball 1 0
    | otherwise = libmRounded (cos m) r
  tan x
    | isExactly 0 x = Ball 0 0
    | otherwise = Ball m $ radiusSlack (r + libmError @p (abs m + r))
    where
      Ball m r = sin (widen x) / cos (widen x)

  exp = unsupported "exp"
  log = unsupported "log"
  asin = unsupported "asin"
  acos = unsupported "acos"
  atan = unsupported "atan"
  sinh = unsupported "sinh"
  cosh = unsupported "cosh"
  asinh = unsupported "asinh"
  acosh = unsupported "acosh"
  atanh = unsupported "atanh"

unsupported :: String -> a
unsupported name = error $ "Ball." <> name <> ": no enclosure implemented"

instance (RealFloat p, Epsilon p) => Epsilon (Ball p) where
  nearZero (Ball m r)
    | nearest == farthest = nearest
    | otherwise = error "Ball.nearZero: the ball straddles the threshold"
    where
      nearest = nearZero (realToFrac (max 0 (abs m - r)) :: p)
      farthest = nearZero (realToFrac (abs m + r) :: p)

data Outside p = Outside
  { actual :: p
  , enclosure :: Ball p
  }
  deriving (Eq, Show)

outside :: RealFloat p => Ball p -> p -> Maybe (Outside p)
outside ball@(Ball m r) x
  | isNaN x || isInfinite x = Just $ Outside x ball
  | abs (realToFrac x - m) <= r = Nothing
  | otherwise = Just $ Outside x ball

enclosed :: (Foldable f, Foldable g, RealFloat p) => f (Ball p) -> g p -> [Maybe (Outside p)]
enclosed bs xs = zipWith outside (toList bs) (toList xs)

data Disjoint p = Disjoint (Ball p) (Ball p)
  deriving (Eq, Show)

disjoint :: Ball p -> Ball p -> Maybe (Disjoint p)
disjoint a b
  | abs (center a - center b) <= radius a + radius b = Nothing
  | otherwise = Just $ Disjoint a b

consistent :: (Foldable f, Foldable g) => f (Ball p) -> g (Ball p) -> [Maybe (Disjoint p)]
consistent as bs = zipWith disjoint (toList as) (toList bs)