module Numeric.Ransac (ransac) where
import Control.Applicative
import Control.Monad (replicateM)
import Data.Vector.Generic ((!))
import qualified Data.Vector.Generic as V
import System.Random
randDistinct :: Int -> IO Int -> IO [Int]
randDistinct n gen = go 0 [] []
where go !i acc _ | i == n = return acc
go !i acc [] = replicateM (ni) gen >>= go i acc
go !i acc (r:rs) = if r `elem` acc
then go i acc rs
else go (i+1) (r:acc) rs
untilJust :: Monad m => m (Maybe b) -> m b
untilJust x = go
where go = x >>= maybe go return
ransac :: (V.Vector v a, V.Vector v d, Num d, Ord d) =>
Int -> Int -> Float ->
(v a -> Maybe c) -> (c -> a -> d) -> (d -> Bool) ->
v a -> IO (Maybe (c, v a))
ransac maxIter sampleSize agree fit residual goodFit pts = genModel >>= go 0
where go i r@(model, bestError, inliers)
| i == maxIter = if ratioInliers (V.length inliers) < agree
then return Nothing
else return (Just (model, inliers))
| otherwise = do r'@(_, err, inliers') <- genModel
if ratioInliers (V.length inliers') >= agree
&& err < bestError
then go (i+1) r'
else go (i+1) r
sample = V.fromList . map (pts !) <$>
randDistinct sampleSize ((`rem` n) . abs <$> randomIO)
genModel = do model <- untilJust (fit <$> sample)
let !errors = V.map (residual model) pts
!inliers = V.ifilter (const . goodFit . (errors !)) pts
case fit inliers of
Nothing -> genModel
Just model' -> let err = V.sum $ V.map (residual model') pts
in return (model', err, inliers)
n = V.length pts
ratioInliers n' = fromIntegral n' / fromIntegral n