{-# LANGUAGE BangPatterns #-}
module Main where

import qualified Data.Vector.Unboxed.Mutable as MU
import Data.Vector.Unboxed.Mutable (IOVector, unsafeRead, unsafeWrite, new)
import qualified Data.Vector.Algorithms.Intro as I

import Control.Monad (when)
import System.Environment (getArgs)

countNaNs :: IOVector Double -> IO Int
countNaNs a = go 0 0
  where
    len = MU.length a
    go !ct i
        | i < len = do
            x <- unsafeRead a i
            go (if isNaN x then ct+1 else ct) (i+1)
        | otherwise = return ct

sample :: Int -> IO (IOVector Double)
sample k = do
    a <- new k
    let foo :: Double -> Double
        foo x = 1.0 + sin x / x
        fill i x
            | i < k = do
                unsafeWrite a i (foo x)
                fill (i+1) (x+1.0)
            | otherwise = return a
    fill 0 (fromIntegral k * 10)

main :: IO ()
main = do
    args <- getArgs
    let k = case args of
              (arg:_) -> read arg
              _       -> 10000
    a <- sample k
    b <- countNaNs a
    when (b /= 0) (putStrLn $ "Before sorting: " ++ show b ++ " NaNs.")
    I.sort a
    c <- countNaNs a
    when (c /= 0) (putStrLn $ "After sorting: " ++ show c ++ " NaNs.")