{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE MultiWayIf #-}
module Data.Find.Interpolation
 ( search, searchM
 ) where

import qualified Data.Find.Linear as Linear
import           Data.Functor.Identity
import           Data.Ratio

search :: forall a
        . (Ord a)
       => (Integer -> a) -- ^ Read a value at
       -> (a -> Integer)
       -> (Integer, Integer) -- ^ (min, max)
       -> a -- ^ The value we're looking for
       -> Maybe Integer
search r j rng t = runIdentity $ searchM (pure . r) j rng t

searchM :: forall m a
         . (Monad m, Ord a)
        => (Integer -> m a) -- ^ Read a value at
        -> (a -> Integer)
        -> (Integer, Integer) -- ^ (min, max)
        -> a -- ^ The value we're looking for
        -> m (Maybe Integer)
searchM r j' (rmin, rmax) t =
    go rmin rmax
  where
    j :: a -> Rational
    j = fromIntegral . j'
    go :: Integer -> Integer -> m (Maybe Integer)
    go l h | (l+2) >= h = Linear.searchM r (l, h) t
    go l h = do
      vl <- r l
      if vl == t
      then pure $ Just l
      else do
        vh <- r h
        if | vh == t -> pure $ Just h
           | (vh == vl) -> pure Nothing
           | otherwise -> do
            let m = max l $ min h $ floor $ (l%1) + (((j t) - (j vl)) * ((h%1) - (l%1)) / ((j vh) - (j vl)))
            -- This reads the m value even if we can't find the result there.
            vm <- r m
            if | vm < t  -> go (m+1) h
               | t < vm  -> go l (m-1)
               | vm == t -> pure $ Just m
               | vm /= t -> pure Nothing