module PAVA where

data Block a = Block ([a] -> a) [a]

instance (Fractional a, Eq a) => Eq (Block a) where
  x == y = mu x == mu y

instance (Fractional a, Ord a) => Ord (Block a) where
  x <= y = mu x <= mu y

mu :: Fractional a => Block a -> a
mu (Block f xs) = f xs

mean :: Fractional a => [a] -> a
mean = (\(a, b) -> b / a) . foldl (\(a, b) c -> (a + 1, b + c)) (0, 0)

pava' :: (Show a, Eq a, Ord a, Fractional a) => ([a] -> a) -> [a] -> [a]
pava' f = concatMap unblock . untilFixed fixViol . map (Block f . pure)
  where
    untilFixed f x =
      let y = f x
       in if y == x then y else untilFixed f y

    fixViol (a : b : xs)
      | a > b = merge a b : xs
      | otherwise = a : fixViol (b : xs)
    fixViol x = x

    merge (Block _ a) (Block _ b) = Block f (a <> b)

pava = pava' mean

unblock :: Fractional a => Block a -> [a]
unblock (Block f xs) = replicate (length xs) (f xs)