{-# LANGUAGE ConstraintKinds  #-}
{-# LANGUAGE FlexibleContexts #-}

module Data.Array.Accelerate.KullbackLiebler ( kullbackLiebler
                                             , entropy
                                             , dropZeroes
                                             , scale
                                             , hellinger
                                             , fDivergence
                                             , alphaDivergence
                                             ) where

import qualified Data.Array.Accelerate as A

-- | \( D_f(p \| q) = \displaystyle\int p(x) f\left(\frac{p(x)}{q(x)}\right) dx \)
--
-- @since 0.1.2.0
fDivergence :: (A.Floating e) => (A.Exp e -> A.Exp e) -- ^ \(f\)
            -> A.Acc (A.Vector e)
            -> A.Acc (A.Vector e)
            -> A.Acc (A.Scalar e)
fDivergence f ps qs = A.sum (A.zipWith (\p q -> p * f (p / q)) ps qs)

-- | \( D^{(\alpha)}(p\| q) = \frac{4}{1 - \alpha^2}\left(1 - \displaystyle\int p(x)^{\frac{1-\alpha}{2}} q(x)^{\frac{1+\alpha}{2}} dx\right)\)
-- for \( \alpha \neq \pm 1\)
--
-- @since 0.1.2.0
alphaDivergence :: A.Floating e => A.Exp e -> A.Acc (A.Vector e) -> A.Acc (A.Vector e) -> A.Acc (A.Scalar e)
alphaDivergence α ps qs = A.map (\x -> (4 / (1 - α ** 2)) * (1 - x)) integrand
    where integrand = A.sum (A.zipWith (\p q -> p ** ((1 - α)/2) * q ** ((1 + α)/2)) ps qs)

-- | Hellinger distance
--
-- @since 0.1.2.0
hellinger :: (A.Floating e) => A.Acc (A.Vector e) -> A.Acc (A.Vector e) -> A.Acc (A.Scalar e)
hellinger ps qs = A.map (A.sqrt . (2*)) $ A.sum (A.zipWith (\p q -> (A.sqrt p - A.sqrt q) ** 2) ps qs)

-- | Assumes input is nonzero
kullbackLiebler :: (A.Floating e) => A.Acc (A.Vector e) -> A.Acc (A.Vector e) -> A.Acc (A.Scalar e)
kullbackLiebler ps qs = A.sum (A.zipWith (\p q -> p * log (p / q)) ps qs)

-- | Assumes input is nonzero
entropy :: (A.Floating e) => A.Acc (A.Vector e) -> A.Acc (A.Scalar e)
entropy = A.sum . A.map (\p -> p * log p)

dropZeroes :: (A.Eq e, Num (A.Exp e)) => A.Acc (A.Vector e) -> A.Acc (A.Vector e)
dropZeroes = A.afst . A.filter (A./= 0)

-- | Doesn't check for negative values
--
-- @since 0.1.1.0
scale :: A.Floating e => A.Acc (A.Vector e) -> A.Acc (A.Vector e)
scale xs =
    let tot = A.the $ A.sum xs
    in A.map (/tot) xs