{-# LANGUAGE GeneralizedNewtypeDeriving #-}
{-# LANGUAGE TypeFamilies #-}
{-# LANGUAGE FlexibleContexts #-}
module ProductMixAuction.LinearAlgebra where

import Prelude
import Data.Vector (Vector)
import qualified Data.Vector as V
import GHC.Exts (IsList (..))

-- * Types

type Dimension = Int

class ComponentWise f where
  componentWise :: (a -> b -> c) -> f a -> f b -> f c

instance ComponentWise [] where
  componentWise = zipWith

instance ComponentWise Vector where
  componentWise = V.zipWith

-- * Linear algebra

-- | Check whether a vector is a zero vector (all components are exactly 0)
isZero :: (Foldable f, Eq a, Num a) => f a -> Bool
isZero v = length v == 0 || all (== 0) v

-- | The zero vector
zero :: (Num (Item (f a)), IsList (f a)) => Dimension -> f a
zero dim = fromList $ replicate dim 0

-- | A vector where all elements are 1
ones :: (Num (Item (f a)), IsList (f a)) => Dimension -> f a
ones dim = fromList $ replicate dim 1

-- | Addition of vectors, pointwise.
vplus :: (Num a, ComponentWise f) => f a -> f a -> f a
vplus = componentWise (+)

-- | Addition of vectors, pointwise (operator alias for 'vplus')
(<+>) :: (Num a, ComponentWise f) => f a -> f a -> f a
(<+>) = vplus

-- | Subtraction of vectors, pointwise.
vminus :: (Num a, ComponentWise f) => f a -> f a -> f a
vminus = componentWise (-)

-- | Subtraction of vectors, pointwise (operator alias for 'vminus')
(<->) :: (Num a, ComponentWise f) => f a -> f a -> f a
(<->) = vminus

-- | Multiplication of vectors, pointwise.
(<.>) :: (Num a, ComponentWise f) => f a -> f a -> f a
(<.>) = componentWise (*)

-- | Component-wise max
vmax :: (Ord a, ComponentWise f) => f a -> f a -> f a
vmax = componentWise max

-- | Component-wise min
vmin :: (Ord a, ComponentWise f) => f a -> f a -> f a
vmin = componentWise min

-- | Scale a vector (scalar multiplication)
vscale :: (Functor f, Num a) => a -> f a -> f a
vscale s = fmap (* s)

-- | Vector-less-than-or-equal: <= operator for vectors under Euclidean order
-- as described in the paper "Valid Combinations of Bids" (Baldwin, Goldberg &
-- Klemperer, validBids.pdf), section 1.1
--
-- Note that this is a partial ordering, so @a `vlte` b@ does not imply @not (a
-- `vlte` b)@, nor does @a `vlte` b && b `vlte` a@ imply @a == b@.
vlte :: (Ord a, Functor f, ComponentWise f, Foldable f) => f a -> f a -> Bool
vlte x y = and $ componentWise (<=) x y

-- | Dot product.
dot :: (Num a, Functor f, ComponentWise f, Foldable f)
    => f a -> f a -> a
dot = dotBy (*)

-- | Dot product with a custom multiplication operator. This is required
-- so that we can form the dot product of vectors from two different unit
-- spaces (e.g. price · quantity).
dotBy :: (Num c, Functor f, ComponentWise f, Foldable f)
      => (a -> b -> c)
      -> f a
      -> f b
      -> c
dotBy op x y = sum (componentWise op x y)

-- | Alias for 'csum': dot product with the unit vector, which amounts to
-- adding the components.
dot1 :: (Num a) => Vector a -> a
dot1 = V.sum