import Data.Ord
import Statistics.Quantile
import qualified Data.Vector as V
import qualified Control.Seq as Seq
import qualified Control.Parallel.Strategies as Strat

collatz :: Integer -> Integer
collatz n | even n    = n `div` 2
collatz n | otherwise = 3*n + 1 

collatzSeries :: Integer -> [Integer]
collatzSeries = takeWhile (/= 1) . iterate collatz

streaks :: [Integer] -> [[Integer]]
streaks [] = []
streaks (x:xs) = let (this,rest) = oneStreak (x:xs)
                 in this:streaks rest

oneStreak :: [Integer] -> ([Integer], [Integer])
oneStreak [x] = ([x],[])
oneStreak l@(x:y:_) = splitWhile2 (\a b -> a `compare` b == x `compare` y) l

splitWhile2 :: (Integer -> Integer -> Bool) -> [Integer] -> ([Integer], [Integer])
splitWhile2 p [x] = ([x],[])
splitWhile2 p (x:y:xs) | p x y = let (s,r) = splitWhile2 p (y:xs) in (x:s,r)
                       | otherwise = ([x],y:xs) 

rawData :: Integer -> V.Vector Double
rawData =
    V.fromList .
    -- No modification: 429MB
    -- withStrategy (rdeepseq) . -- 140MB
    -- valueStrictList . 67MB
    -- Seq.withStrategy (Seq.seqList Seq.rseq) . -- 140MB
    -- Strat.withStrategy (Strat.evalList Strat.rseq) . -- Stack overflow
    -- Strat.withStrategy (Strat.seqList Strat.rseq) . -- Stack overflow
    map fromIntegral .
    map length .
    concatMap (streaks . collatzSeries) .
    (\n -> [1..n])

quantile :: V.Vector Double -> Double
quantile = weightedAvg 9 10 

valueStrictList :: [a] -> [a]
valueStrictList [] = []
valueStrictList (x:xs) = x `seq` (x:valueStrictList xs)

main = do
    putStrLn "Collatz quantile calculation"
    let vec = rawData 100000
    putStrLn $ "Number of streaks: " ++ show (V.length vec)
    putStrLn $ "Quantile: " ++ show (quantile vec)