haskell vector spaces

Space.hs

{-# LANGUAGE TypeFamilies, FlexibleContexts, TypeOperators, UndecidableInstances, GeneralizedNewtypeDeriving #-}
module Numeric.Vector.Space where

import Data.Functor.Representable.Trie hiding (Entry)
import Data.Key (Adjustable(..), Key)
import qualified Data.Key as Key
import qualified Data.Heap as Heap
import Data.Heap (Heap, Entry(..))
import qualified Data.List.NonEmpty as NonEmpty
import Data.List.NonEmpty (NonEmpty(..))
import Data.Complex
import Data.List (foldl')
import qualified Data.Map as Map
import Control.Applicative
import Control.Monad
import Control.Monad.Free
import Control.Monad.Codensity
import Control.Monad.Trans

infixl 6 .*, ./
-- todo: rename or constrain on a?
class (Alternative v, MonadPlus v, Num (Scalar v)) => Vector v where
  type Scalar v :: *
  (.*) :: Scalar v -> v a -> v a
  linear :: [(Scalar v, a)] -> v a
  linear = foldr (\(p,a) b -> if p /= 0 then (p.* pure a) <|> b else b) empty

bias :: Vector v => Scalar v -> v Bool
bias p = linear [(p,True),(1-p,False)]

instance Vector v => Vector (Free v) where
  type Scalar (Free v) = Scalar v
  p .* Free as = Free (p .* as)
  p .* v = Free $ linear [(p,v)]
  linear = lift . linear 

instance Vector v => Vector (CodensityT v) where
  type Scalar (CodensityT v) = Scalar v
  p .* CodensityT m = CodensityT (\k -> p .* m k)
  linear = lift . linear

class Vector v => Eval v where
  eval :: Ord a => v a -> [(Scalar v,a)]
  flatten :: v a -> [(Scalar v, a)]

instance Eval v => Eval (Free v) where
  eval    = eval . retract
  flatten = flatten . retract

instance Eval v => Eval (CodensityT v) where
  eval = eval . lowerCodensityT
  flatten = flatten . lowerCodensityT

optimize :: (Eval v, Ord a) => v a -> v a
optimize = linear . eval

(./) :: (Vector v, Fractional (Scalar v)) => v a -> Scalar v -> v a 
v ./ s = recip s .* v

-- | conservative upper bound on the L1 norm. Accurate If all vectors are non-negative.
d :: Eval v => v a -> Scalar v
d = foldl' (\b pa -> b + abs (fst pa)) 0 . flatten

newtype Linear p a = Linear { runLinear :: [(p,a)] } 
  deriving (Read,Show)

instance Num p => Functor (Linear p) where
  fmap f (Linear as) = Linear [ (p, f a) | (p,a) <- as ]
  b <$ as = Linear [(d as, b)]

instance Num p => Applicative (Linear p) where
  pure a = Linear [(1, a)]
  Linear fs <*> Linear as = Linear [(p*q,f a) | (p,f) <- fs, (q,a) <- as]
  as <* bs = d bs .* as
  as *> bs = d as .* bs
  
instance Num p => Alternative (Linear p) where
  empty = Linear []
  Linear as <|> Linear bs = Linear (as <|> bs)

instance Num p => Monad (Linear p) where
  return a = Linear [(1, a)]
  (>>) = (*>)
  Linear as >>= f = Linear [ (p*q, b) | (p,a) <- as, (q,b) <- runLinear (f a) ]

instance Num p => MonadPlus (Linear p) where
  mzero = empty
  mplus = (<|>)

instance Num p => Vector (Linear p) where
  type Scalar (Linear p) = p
  p .* Linear as = Linear [ (p*q,a) | (q,a) <- as ]
  linear = Linear

instance Num p => Eval (Linear p) where
  eval = map swap . Map.toList . Map.fromListWith (+) . map swap . runLinear
    where 
      swap :: (a,b) -> (b,a)
      swap (a,b) = (b,a)
  flatten = runLinear

pl :: Linear Double a -> Linear Double a
pl = id

pf :: Free (Linear Double) a -> Free (Linear Double) a
pf = id

pc :: CodensityT (Free (Linear Double)) a -> CodensityT (Free (Linear Double)) a
pc = id

ql :: Linear (Complex Double) a -> Linear (Complex Double) a
ql = id

qf :: Free (Linear (Complex Double)) a -> Free (Linear (Complex Double)) a
qf = id

qc :: CodensityT (Free (Linear (Complex Double))) a -> CodensityT (Free (Linear (Complex Double))) a
qc = id

grassModel :: (Vector v, Fractional (Scalar v)) => v Bool
grassModel = do
  cloudy    <- bias 0.5
  rain      <- bias $ if cloudy then 0.8 else 0.2
  sprinkler <- bias $ if cloudy then 0.1 else 0.5
  _wetRoof  <- (&& rain) <$> bias 0.7
  wetGrass  <- (||) <$> ((&&) rain <$> bias 0.9)
                    <*> ((&&) sprinkler <$> bias 0.9)
  cup True wetGrass -- guard wetGrass
  return rain

class Ord a => Space a where
  cup :: Vector v => a -> a -> v ()
  cap :: Vector v => v (a,a)

instance Space () where
  cup () () = pure ()
  cap = pure ((),())

instance Space Bool where
  cup True True = pure ()
  cup False False = pure ()
  cup _ _ = empty
  cap = pure (True, True) <|> pure (False,False)

transpose :: (Vector v, Space a, Space b) => (a -> v b) -> b -> v a
transpose m i = do
  (k,l) <- cap
  j <- m k
  cup i j
  return l

{-
data Distribution p a = Distribution
  { pdf :: a -> p
  , cdf :: a -> p
  , qf  :: p -> a
  , lo  :: a
  , hi  :: a
  }

uniform :: (Real p, Fractional a) => a -> a -> Distribution p a
uniform a b = Distribution {..} where
  pdf x | a <= x && x <= b = recip (b - a)
        | otherwise = 0
  cdf x | x <= a = 0
        | x <= b = (x - a) / (x - b)
        | otherwise = 1
  qf p = a + realToFrac p (b - a)
  support = (a,b)

exponential :: (Real p, Fractional a) => p -> Distribution p p
exponential l = Distribution {..} where
  pdf x = l * exp (-l * x)
  cdf x = 1 - exp(-l*x)     -- expm1(-l*x)
  qf  p = - log (1 - p) / l -- -logp1(p)/l
  support = (0,1/0)

-- clamp a distribution between two values. yielding a new distribution
bounded :: Real p => a -> a -> Distribution p a -> Distribution p a
bounded a b d = Distribution {
    pdf = pdf'
  , cdf = cdf'
  , qf = qf'
  , lo = lo'
  , hi = hi'
  } where 
    lo'  = max (lo d) a
    hi'  = min (hi d) b
    base = cdf d lo'
    m    = cdf d hi' - base
    oom  = recip m
    qlo = qf d lo'
    qhi = qf 
    pdf' x | lo' <= x && x <= hi' = oopr * pdf d x
           | otherwise = 0
    cdf' x | x <= lo' = 0
           | x <= hi' = (cdf d x - b) * oom
           | otherwise = 1
    qf' p = qd d (m * p + base)
-}

newtype Qubit = Qubit Bool deriving (Eq,Ord,Space,Show,Read)

q0, q1 :: Qubit
q0 = Qubit False
q1 = Qubit True

-- invariant: must rotate by a unitary matrix
rot :: Vector v => Scalar v -> Scalar v -> Scalar v -> Scalar v -> Qubit -> v Qubit
rot f t _ _ (Qubit False) = linear [(f,q0), (t,q1)]
rot _ _ f t (Qubit True)  = linear [(f,q0), (t,q1)]

hadamard :: (Vector v, Floating (Scalar v)) => Qubit -> v Qubit
hadamard = rot h h 
               h (-h) 
 where h = recip(sqrt 2)

-- quantum negation = pauliX, but it can be executed classically.
qnot :: Qubit -> Qubit
qnot (Qubit a) = Qubit (not a)

pauliX :: Applicative v => Qubit -> v Qubit 
pauliX = pure . qnot -- pauliX = rot 0 1 1 0

pauliY, pauliZ, phase, pi_8 :: (Vector v, Scalar v ~ Complex r, RealFloat r) => Qubit -> v Qubit
pauliY = rot 0 (-i)
             i 0 where i = 0 :+ 1
pauliZ = rot 1 0 
             0 (-1)
phase = rot 1 0 
            0 i where i = 0 :+ 1
pi_8 = rot 1 0
           0 (cis $ pi/4)

qplus, qminus :: (Vector v, Scalar v ~ Complex r, RealFloat r) => v Qubit
qplus = hadamard q0
qminus = hadamard q1

qphase :: (Vector v, Scalar v ~ Complex r, RealFloat r) => r -> Qubit -> v Qubit
qphase  phi = rot 1 0 
                  0 c 
  where c = cis $ 2 * pi * phi

-- invariant: the qubit used for the condition cannot be used in the conditional matrix
conditional :: Applicative v => (a -> v a) -> Qubit -> a -> v a
conditional m (Qubit True) = m
conditional m (Qubit False) = pure

-- run a quantum computation, extracting classical bits
measure :: (Vector v, Eval u, Scalar u ~ Complex r, RealFloat r, Fractional (Scalar v), Ord a) => u a -> v a
measure qbits = linear [ (fromRational $ toRational $ magnitude q ^ 2, a) | (q,a) <- eval qbits ]

condense :: Int -> NonEmpty a -> NonEmpty a
condense n (a:|as0) | n >= 1 = go a as0 n
  where
    go a []     _ = a :| []
    go a (b:bs) 1 = a :| NonEmpty.toList (go b bs n)
    go _ (b:bs) k = go b bs $! (k - 1)

newtype Reverse a = Reverse a

instance Eq a => Eq (Reverse a) where
  Reverse a == Reverse b = a == b

instance Ord a => Ord (Reverse a) where
  compare (Reverse a) (Reverse b) = compare b a

buildHeap :: (Eval v, Ord p) => (Scalar v -> p) -> Scalar v -> v (Free v a) -> (Heap (Entry (Reverse p) (Scalar v, Free v a)), p)
buildHeap f q = convert . foldr act ([],0) . flatten where
      act (p,a) (xs,sz) = (x:xs, qp+sz) where 
          qp = q*p
          x = Entry (Reverse (f qp)) (qp,a)
      convert (xs,sz) = (Heap.fromList xs, f sz) 

-- used when we need no priority

data Trivial = Trivial deriving (Eq,Ord,Show)

instance Num Trivial where
  _ + _ = Trivial
  _ * _ = Trivial
  _ - _ = Trivial
  abs _ = Trivial
  signum _ = Trivial
  fromInteger _ = Trivial

-- always expands the tree in the direction that maximizes information.
-- returns a list of lower bounds on the proability with one sided error
proj_ :: (Eval v, Num p, Ord p) => (Scalar v -> p) -> Free v a -> NonEmpty (Scalar v, p)
proj_ f (Pure b)  = (1, f 0) :| []
proj_ f (Free bs) = condense 20 $ go 0 ub0 entries0
    where 
        (entries0, ub0) = buildHeap f 1 bs 
        go lo err heap = (lo,err) :| case Heap.viewMin heap of 
            Nothing -> []
            Just (Entry (Reverse q) (p, Pure c), heap') -> NonEmpty.toList $ go (lo + p) (err - q) heap'
            Just (Entry (Reverse q) (p, Free cs), heap') -> let (heap'', q) = buildHeap f p cs in NonEmpty.toList $ 
                go lo (err - f p + q) $ Heap.union heap' heap''

-- always expands the tree in the direction that maximizes information. Only yields valid Improving values when the 
-- probabilities always correctly lie within the interval 0..1
proj :: (Eval v, Num p, Ord p, Eq a, HasTrie a) => (Scalar v -> p) -> Free v a -> NonEmpty (a :->: Scalar v, p)
proj f (Pure b) = (adjust (const 1) b (pure 0), f 0) :| []
proj f (Free bs) = condense 20 $ go (pure 0) ub0 entries0
    where 
        (entries0, ub0) = buildHeap f 1 bs 
        go lo err heap = (lo,err) :| case Heap.viewMin heap of 
            Nothing -> []
            Just (Entry (Reverse q) (p, Pure c), heap') -> NonEmpty.toList $ go (adjust (+ p) c lo) (err - q) heap'
            Just (Entry (Reverse q) (p, Free cs), heap') -> let (heap'', q) = buildHeap f p cs in NonEmpty.toList $ 
                go lo (err - f p + q) $ Heap.union heap' heap''

bash-3.2$ 
bash-3.2$ pwd
/Users/ekmett/linear/Numeric/Vector
bash-3.2$ ls
Basis.hs	Dim.hs		Free.hs		Inference.hs	Linear.hs	Probability.hs	Space.hs
bash-3.2$ vi Space.hs 
^[bash-3.2$ cat Space.hs 
{-# LANGUAGE TypeFamilies, FlexibleContexts, TypeOperators, UndecidableInstances, GeneralizedNewtypeDeriving #-}
module Numeric.Vector.Space where

import Data.Functor.Representable.Trie hiding (Entry)
import Data.Key (Adjustable(..), Key)
import qualified Data.Key as Key
import qualified Data.Heap as Heap
import Data.Heap (Heap, Entry(..))
import qualified Data.List.NonEmpty as NonEmpty
import Data.List.NonEmpty (NonEmpty(..))
import Data.Complex
import Data.List (foldl')
import qualified Data.Map as Map
import Control.Applicative
import Control.Monad
import Control.Monad.Free
import Control.Monad.Codensity
import Control.Monad.Trans

infixl 6 .*, ./
-- todo: rename or constrain on a?
class (Alternative v, MonadPlus v, Num (Scalar v)) => Vector v where
  type Scalar v :: *
  (.*) :: Scalar v -> v a -> v a
  linear :: [(Scalar v, a)] -> v a
  linear = foldr (\(p,a) b -> if p /= 0 then (p.* pure a) <|> b else b) empty

bias :: Vector v => Scalar v -> v Bool
bias p = linear [(p,True),(1-p,False)]

instance Vector v => Vector (Free v) where
  type Scalar (Free v) = Scalar v
  p .* Free as = Free (p .* as)
  p .* v = Free $ linear [(p,v)]
  linear = lift . linear 

instance Vector v => Vector (CodensityT v) where
  type Scalar (CodensityT v) = Scalar v
  p .* CodensityT m = CodensityT (\k -> p .* m k)
  linear = lift . linear

class Vector v => Eval v where
  eval :: Ord a => v a -> [(Scalar v,a)]
  flatten :: v a -> [(Scalar v, a)]

instance Eval v => Eval (Free v) where
  eval    = eval . retract
  flatten = flatten . retract

instance Eval v => Eval (CodensityT v) where
  eval = eval . lowerCodensityT
  flatten = flatten . lowerCodensityT

optimize :: (Eval v, Ord a) => v a -> v a
optimize = linear . eval

(./) :: (Vector v, Fractional (Scalar v)) => v a -> Scalar v -> v a 
v ./ s = recip s .* v

-- | conservative upper bound on the L1 norm. Accurate If all vectors are non-negative.
d :: Eval v => v a -> Scalar v
d = foldl' (\b pa -> b + abs (fst pa)) 0 . flatten

newtype Linear p a = Linear { runLinear :: [(p,a)] } 
  deriving (Read,Show)

instance Num p => Functor (Linear p) where
  fmap f (Linear as) = Linear [ (p, f a) | (p,a) <- as ]
  b <$ as = Linear [(d as, b)]

instance Num p => Applicative (Linear p) where
  pure a = Linear [(1, a)]
  Linear fs <*> Linear as = Linear [(p*q,f a) | (p,f) <- fs, (q,a) <- as]
  as <* bs = d bs .* as
  as *> bs = d as .* bs
  
instance Num p => Alternative (Linear p) where
  empty = Linear []
  Linear as <|> Linear bs = Linear (as <|> bs)

instance Num p => Monad (Linear p) where
  return a = Linear [(1, a)]
  (>>) = (*>)
  Linear as >>= f = Linear [ (p*q, b) | (p,a) <- as, (q,b) <- runLinear (f a) ]

instance Num p => MonadPlus (Linear p) where
  mzero = empty
  mplus = (<|>)

instance Num p => Vector (Linear p) where
  type Scalar (Linear p) = p
  p .* Linear as = Linear [ (p*q,a) | (q,a) <- as ]
  linear = Linear

instance Num p => Eval (Linear p) where
  eval = map swap . Map.toList . Map.fromListWith (+) . map swap . runLinear
    where 
      swap :: (a,b) -> (b,a)
      swap (a,b) = (b,a)
  flatten = runLinear

pl :: Linear Double a -> Linear Double a
pl = id

pf :: Free (Linear Double) a -> Free (Linear Double) a
pf = id

pc :: CodensityT (Free (Linear Double)) a -> CodensityT (Free (Linear Double)) a
pc = id

ql :: Linear (Complex Double) a -> Linear (Complex Double) a
ql = id

qf :: Free (Linear (Complex Double)) a -> Free (Linear (Complex Double)) a
qf = id

qc :: CodensityT (Free (Linear (Complex Double))) a -> CodensityT (Free (Linear (Complex Double))) a
qc = id

grassModel :: (Vector v, Fractional (Scalar v)) => v Bool
grassModel = do
  cloudy    <- bias 0.5
  rain      <- bias $ if cloudy then 0.8 else 0.2
  sprinkler <- bias $ if cloudy then 0.1 else 0.5
  _wetRoof  <- (&& rain) <$> bias 0.7
  wetGrass  <- (||) <$> ((&&) rain <$> bias 0.9)
                    <*> ((&&) sprinkler <$> bias 0.9)
  cup True wetGrass -- guard wetGrass
  return rain

class Ord a => Space a where
  cup :: Vector v => a -> a -> v ()
  cap :: Vector v => v (a,a)

instance Space () where
  cup () () = pure ()
  cap = pure ((),())

instance Space Bool where
  cup True True = pure ()
  cup False False = pure ()
  cup _ _ = empty
  cap = pure (True, True) <|> pure (False,False)

transpose :: (Vector v, Space a, Space b) => (a -> v b) -> b -> v a
transpose m i = do
  (k,l) <- cap
  j <- m k
  cup i j
  return l

{-
data Distribution p a = Distribution
  { pdf :: a -> p
  , cdf :: a -> p
  , qf  :: p -> a
  , lo  :: a
  , hi  :: a
  }

uniform :: (Real p, Fractional a) => a -> a -> Distribution p a
uniform a b = Distribution {..} where
  pdf x | a <= x && x <= b = recip (b - a)
        | otherwise = 0
  cdf x | x <= a = 0
        | x <= b = (x - a) / (x - b)
        | otherwise = 1
  qf p = a + realToFrac p (b - a)
  support = (a,b)

exponential :: (Real p, Fractional a) => p -> Distribution p p
exponential l = Distribution {..} where
  pdf x = l * exp (-l * x)
  cdf x = 1 - exp(-l*x)     -- expm1(-l*x)
  qf  p = - log (1 - p) / l -- -logp1(p)/l
  support = (0,1/0)

-- clamp a distribution between two values. yielding a new distribution
bounded :: Real p => a -> a -> Distribution p a -> Distribution p a
bounded a b d = Distribution {
    pdf = pdf'
  , cdf = cdf'
  , qf = qf'
  , lo = lo'
  , hi = hi'
  } where 
    lo'  = max (lo d) a
    hi'  = min (hi d) b
    base = cdf d lo'
    m    = cdf d hi' - base
    oom  = recip m
    qlo = qf d lo'
    qhi = qf 
    pdf' x | lo' <= x && x <= hi' = oopr * pdf d x
           | otherwise = 0
    cdf' x | x <= lo' = 0
           | x <= hi' = (cdf d x - b) * oom
           | otherwise = 1
    qf' p = qd d (m * p + base)
-}

newtype Qubit = Qubit Bool deriving (Eq,Ord,Space,Show,Read)

q0, q1 :: Qubit
q0 = Qubit False
q1 = Qubit True

-- invariant: must rotate by a unitary matrix
rot :: Vector v => Scalar v -> Scalar v -> Scalar v -> Scalar v -> Qubit -> v Qubit
rot f t _ _ (Qubit False) = linear [(f,q0), (t,q1)]
rot _ _ f t (Qubit True)  = linear [(f,q0), (t,q1)]

hadamard :: (Vector v, Floating (Scalar v)) => Qubit -> v Qubit
hadamard = rot h h 
               h (-h) 
 where h = recip(sqrt 2)

-- quantum negation = pauliX, but it can be executed classically.
qnot :: Qubit -> Qubit
qnot (Qubit a) = Qubit (not a)

pauliX :: Applicative v => Qubit -> v Qubit 
pauliX = pure . qnot -- pauliX = rot 0 1 1 0

pauliY, pauliZ, phase, pi_8 :: (Vector v, Scalar v ~ Complex r, RealFloat r) => Qubit -> v Qubit
pauliY = rot 0 (-i)
             i 0 where i = 0 :+ 1
pauliZ = rot 1 0 
             0 (-1)
phase = rot 1 0 
            0 i where i = 0 :+ 1
pi_8 = rot 1 0
           0 (cis $ pi/4)

qplus, qminus :: (Vector v, Scalar v ~ Complex r, RealFloat r) => v Qubit
qplus = hadamard q0
qminus = hadamard q1

qphase :: (Vector v, Scalar v ~ Complex r, RealFloat r) => r -> Qubit -> v Qubit
qphase  phi = rot 1 0 
                  0 c 
  where c = cis $ 2 * pi * phi

-- invariant: the qubit used for the condition cannot be used in the conditional matrix
conditional :: Applicative v => (a -> v a) -> Qubit -> a -> v a
conditional m (Qubit True) = m
conditional m (Qubit False) = pure

-- run a quantum computation, extracting classical bits
measure :: (Vector v, Eval u, Scalar u ~ Complex r, RealFloat r, Fractional (Scalar v), Ord a) => u a -> v a
measure qbits = linear [ (fromRational $ toRational $ magnitude q ^ 2, a) | (q,a) <- eval qbits ]

condense :: Int -> NonEmpty a -> NonEmpty a
condense n (a:|as0) | n >= 1 = go a as0 n
  where
    go a []     _ = a :| []
    go a (b:bs) 1 = a :| NonEmpty.toList (go b bs n)
    go _ (b:bs) k = go b bs $! (k - 1)

newtype Reverse a = Reverse a

instance Eq a => Eq (Reverse a) where
  Reverse a == Reverse b = a == b

instance Ord a => Ord (Reverse a) where
  compare (Reverse a) (Reverse b) = compare b a

buildHeap :: (Eval v, Ord p) => (Scalar v -> p) -> Scalar v -> v (Free v a) -> (Heap (Entry (Reverse p) (Scalar v, Free v a)), p)
buildHeap f q = convert . foldr act ([],0) . flatten where
      act (p,a) (xs,sz) = (x:xs, qp+sz) where 
          qp = q*p
          x = Entry (Reverse (f qp)) (qp,a)
      convert (xs,sz) = (Heap.fromList xs, f sz) 

-- used when we need no priority

data Trivial = Trivial deriving (Eq,Ord,Show)

instance Num Trivial where
  _ + _ = Trivial
  _ * _ = Trivial
  _ - _ = Trivial
  abs _ = Trivial
  signum _ = Trivial
  fromInteger _ = Trivial

-- always expands the tree in the direction that maximizes information.
-- returns a list of lower bounds on the proability with one sided error
proj_ :: (Eval v, Num p, Ord p) => (Scalar v -> p) -> Free v a -> NonEmpty (Scalar v, p)
proj_ f (Pure b)  = (1, f 0) :| []
proj_ f (Free bs) = condense 20 $ go 0 ub0 entries0
    where 
        (entries0, ub0) = buildHeap f 1 bs 
        go lo err heap = (lo,err) :| case Heap.viewMin heap of 
            Nothing -> []
            Just (Entry (Reverse q) (p, Pure c), heap') -> NonEmpty.toList $ go (lo + p) (err - q) heap'
            Just (Entry (Reverse q) (p, Free cs), heap') -> let (heap'', q) = buildHeap f p cs in NonEmpty.toList $ 
                go lo (err - f p + q) $ Heap.union heap' heap''

-- always expands the tree in the direction that maximizes information. Only yields valid Improving values when the 
-- probabilities always correctly lie within the interval 0..1
proj :: (Eval v, Num p, Ord p, Eq a, HasTrie a) => (Scalar v -> p) -> Free v a -> NonEmpty (a :->: Scalar v, p)
proj f (Pure b) = (adjust (const 1) b (pure 0), f 0) :| []
proj f (Free bs) = condense 20 $ go (pure 0) ub0 entries0
    where 
        (entries0, ub0) = buildHeap f 1 bs 
        go lo err heap = (lo,err) :| case Heap.viewMin heap of 
            Nothing -> []
            Just (Entry (Reverse q) (p, Pure c), heap') -> NonEmpty.toList $ go (adjust (+ p) c lo) (err - q) heap'
            Just (Entry (Reverse q) (p, Free cs), heap') -> let (heap'', q) = buildHeap f p cs in NonEmpty.toList $ 
                go lo (err - f p + q) $ Heap.union heap' heap''