content: -- |
  -- Module:      Math.NumberTheory.Utils.Hyperbola
  -- Copyright:   (c) 2018 Andrew Lelechenko
  -- Licence:     MIT
  -- Maintainer:  Andrew Lelechenko <andrew.lelechenko@gmail.com>
  --
  -- Highest points under hyperbola.
  --
  
  module Math.NumberTheory.Utils.Hyperbola
    ( pointsUnderHyperbola
    ) where
  
  import Data.Bits
  
  import Math.NumberTheory.Roots
  
  -- | Straightforward computation of
  -- [ n `quot` x | x <- [hi, hi - 1 .. lo] ].
  -- Unfortunately, such list generator performs poor,
  -- so we fall back to manual recursion.
  pointsUnderHyperbola0 :: Int -> Int -> Int -> [Int]
  pointsUnderHyperbola0 n lo hi
    | n < 0     = error "pointsUnderHyperbola0: first argument must be non-negative"
    | lo <= 0   = error "pointsUnderHyperbola0: second argument must be positive"
    | otherwise = go hi
      where
        go x
          | x < lo    = []
          | otherwise = n `quot` x : go (x - 1)
  
  data Bresenham = Bresenham
    {  bresX       :: !Int
    ,  bresBeta    :: !Int
    , _bresGamma   :: !Int
    , _bresDelta1  :: !Int
    , _bresEpsilon :: !Int
    }
  
  initBresenham :: Int -> Int -> Bresenham
  initBresenham n x = Bresenham x beta gamma delta1 epsilon
    where
      beta    = n `quot` x
      epsilon = n `rem` x
      delta1  = n `quot` (x - 1) - beta
      gamma   = beta - (x - 1) * delta1
  
  -- | bresenham(x+1) -> bresenham(x) for x >= (2n)^1/3
  stepBack :: Bresenham -> Bresenham
  stepBack (Bresenham x' beta' gamma' delta1' epsilon')
    | eps >= x `shiftL` 1 {- delta2 = 2 -}
    = let delta1 = delta1' + 2 in Bresenham x (beta' + delta1) (gamma' + delta1 `shiftL` 1 - x `shiftL` 1) delta1 (eps - x `shiftL` 1)
    | eps >= x {- delta1 = 1 -}
    = let delta1 = delta1' + 1 in Bresenham x (beta' + delta1) (gamma' + delta1 `shiftL` 1 - x) delta1 (eps - x)
    | eps >= 0 {- delta2 =  0 -}
    = Bresenham x (beta' + delta1') (gamma' + delta1' `shiftL` 1) delta1' eps
    | otherwise {- delta2 = -1 -}
    = let delta1 = delta1' - 1 in Bresenham x (beta' + delta1) (gamma' + delta1 `shiftL` 1 + x) delta1 (eps + x)
    where
      x       = x' - 1
      eps     = epsilon' + gamma'
  {-# INLINE stepBack #-}
  
  -- | Division-free computation of
  -- [ n `quot` x | x <- [hi, hi - 1 .. lo] ].
  -- In other words, we compute y-coordinates of highest integral points
  -- under hyperbola @x * y = n@ between @x = lo@ and @x = hi@ in reverse order.
  --
  -- The implementation follows section 5 of <https://arxiv.org/pdf/1206.3369.pdf A successive approximation algorithm for computing the divisor summatory function>
  -- by R. Sladkey.
  -- It is 2x faster than a trivial implementation for 'Int'.
  pointsUnderHyperbola :: Int -> Int -> Int -> [Int]
  pointsUnderHyperbola n lo hi
    | n < 0        = error "pointsUnderHyperbola: first argument must be non-negative"
    | lo <= 0      = error "pointsUnderHyperbola: second argument must be positive"
    | hi <  lo     = []
    | hi == lo     = [n `quot` lo]
    | otherwise    = go (initBresenham n hi)
    where
      mid = (integerCubeRoot (2 * n) + 1) `max` lo
      go h
        | bresX h < mid = pointsUnderHyperbola0 n lo ((mid - 1) `min` hi)
        | otherwise = bresBeta h : go (stepBack h)
  