{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE TypeFamilies #-}
module Numeric.ADDual.VectorOps where

import Data.Kind (Type)
import qualified Data.Vector as V
import qualified Data.Vector.Strict as VSr
import qualified Data.Vector.Storable as VS
import qualified Data.Vector.Unboxed as VU
import Foreign.Storable (Storable)


class VectorOps v where
  type VectorOpsScalar v :: Type
  vfromListN :: Int -> [VectorOpsScalar v] -> v
  vfromList :: [VectorOpsScalar v] -> v
  vtoList :: v -> [VectorOpsScalar v]
  vlength :: v -> Int
  vreplicate :: Int -> VectorOpsScalar v -> v
  vselect :: VS.Vector Bool -> v -> v -> v  -- ^ True selects the first argument, False the second

class (VectorOps v, Num (VectorOpsScalar v)) => VectorOpsNum v where
  vadd :: v -> v -> v
  vsub :: v -> v -> v
  vmul :: v -> v -> v
  vsum :: v -> VectorOpsScalar v

class (VectorOpsNum v, Floating (VectorOpsScalar v)) => VectorOpsFloating v where
  vexp :: v -> v

class (VectorOps v, Ord (VectorOpsScalar v)) => VectorOpsOrd v where
  vcmpLE :: v -> v -> VS.Vector Bool
  vmaximum :: v -> VectorOpsScalar v

  vcmpLT, vcmpGT, vcmpGE :: v -> v -> VS.Vector Bool
  vcmpLT a b = VS.map not (vcmpLE b a)
  vcmpGT a b = VS.map not (vcmpLE a b)
  vcmpGE a b = vcmpLE b a

instance VectorOps (V.Vector a) where
  type VectorOpsScalar (V.Vector a) = a
  vfromListN = V.fromListN
  vfromList = V.fromList
  vtoList = V.toList
  vlength = V.length
  vreplicate = V.replicate
  vselect bs a b = V.fromListN (VS.length bs) [if bs VS.! i then a V.! i else b V.! i
                                              | i <- [0 .. VS.length bs - 1]]

instance Num a => VectorOpsNum (V.Vector a) where
  vadd = V.zipWith (+)
  vsub = V.zipWith (-)
  vmul = V.zipWith (*)
  vsum = V.sum

instance Floating a => VectorOpsFloating (V.Vector a) where
  vexp = V.map exp

instance Ord a => VectorOpsOrd (V.Vector a) where
  vcmpLE a b = VS.generate (V.length a) (\i -> a V.! i <= b V.! i)
  vmaximum = V.maximum

instance VectorOps (VSr.Vector a) where
  type VectorOpsScalar (VSr.Vector a) = a
  vfromListN = VSr.fromListN
  vfromList = VSr.fromList
  vtoList = VSr.toList
  vlength = VSr.length
  vreplicate = VSr.replicate
  vselect bs a b = VSr.fromListN (VS.length bs) [if bs VS.! i then a VSr.! i else b VSr.! i
                                                | i <- [0 .. VS.length bs - 1]]

instance Num a => VectorOpsNum (VSr.Vector a) where
  vadd = VSr.zipWith (+)
  vsub = VSr.zipWith (-)
  vmul = VSr.zipWith (*)
  vsum = VSr.sum

instance Floating a => VectorOpsFloating (VSr.Vector a) where
  vexp = VSr.map exp

instance Ord a => VectorOpsOrd (VSr.Vector a) where
  vcmpLE a b = VS.generate (VSr.length a) (\i -> a VSr.! i <= b VSr.! i)
  vmaximum = VSr.maximum

instance Storable a => VectorOps (VS.Vector a) where
  type VectorOpsScalar (VS.Vector a) = a
  vfromListN = VS.fromListN
  vfromList = VS.fromList
  vtoList = VS.toList
  vlength = VS.length
  vreplicate = VS.replicate
  vselect bs a b = VS.fromListN (VS.length bs) [if bs VS.! i then a VS.! i else b VS.! i
                                               | i <- [0 .. VS.length bs - 1]]

instance (Storable a, Num a) => VectorOpsNum (VS.Vector a) where
  vadd = VS.zipWith (+)
  vsub = VS.zipWith (-)
  vmul = VS.zipWith (*)
  vsum = VS.sum

instance (Storable a, Floating a) => VectorOpsFloating (VS.Vector a) where
  vexp = VS.map exp

instance (Storable a, Ord a) => VectorOpsOrd (VS.Vector a) where
  vcmpLE a b = VS.generate (VS.length a) (\i -> a VS.! i <= b VS.! i)
  vmaximum = VS.maximum

instance VU.Unbox a => VectorOps (VU.Vector a) where
  type VectorOpsScalar (VU.Vector a) = a
  vfromListN = VU.fromListN
  vfromList = VU.fromList
  vtoList = VU.toList
  vlength = VU.length
  vreplicate = VU.replicate
  vselect bs a b = VU.fromListN (VS.length bs) [if bs VS.! i then a VU.! i else b VU.! i
                                               | i <- [0 .. VS.length bs - 1]]

instance (VU.Unbox a, Num a) => VectorOpsNum (VU.Vector a) where
  vadd = VU.zipWith (+)
  vsub = VU.zipWith (-)
  vmul = VU.zipWith (*)
  vsum = VU.sum

instance (VU.Unbox a, Floating a) => VectorOpsFloating (VU.Vector a) where
  vexp = VU.map exp

instance (VU.Unbox a, Ord a) => VectorOpsOrd (VU.Vector a) where
  vcmpLE a b = VS.generate (VU.length a) (\i -> a VU.! i <= b VU.! i)
  vmaximum = VU.maximum