{-# LANGUAGE TypeFamilies, TypeOperators #-}
module NoSlow.Backend.Uvector (
  UA, UArr,

  null, length, replicate, cons, head, tail, map, imap,
  zip, zip3, zipWith, unzip,
  sum, and, prescanl',
  enumFromTo_Int,
  filter, index, append, take, drop, slice, backpermute, update, update_,
  minIndex, maxIndex, unstablePartition,
  prescanr',
  pair, from2, fst, snd, triple, from3
) where

import NoSlow.Util.Computation
import NoSlow.Util.Base ( Unsupported(..) )

import Data.Array.Vector
import qualified Data.Array.Vector.UArr as UArr ( indexU )
import qualified Prelude
import Prelude ( Ord(..), Num, Int, Bool, (.), ($), fmap, seq, flip, not )

instance DeepSeq (UArr a)

instance (TestData a, UA a) => TestData (UArr a) where
  testData = testList

instance UA a => ListLike UArr a where
  fromList = toU

null :: UA a => UArr a -> Bool
{-# INLINE null #-}
null = nullU

length :: UA a => UArr a -> Int
{-# INLINE length #-}
length = lengthU

replicate :: UA a => Int -> a -> UArr a
{-# INLINE replicate #-}
replicate = replicateU

cons :: UA a => a -> UArr a -> UArr a
{-# INLINE cons #-}
cons = consU

head :: UA a => UArr a -> a
{-# INLINE head #-}
head = headU

tail :: UA a => UArr a -> UArr a
{-# INLINE tail #-}
tail = tailU

map :: (UA a, UA b) => (a -> b) -> UArr a -> UArr b
{-# INLINE map #-}
map = mapU

imap :: (UA a, UA b) => (Int -> a -> b) -> UArr a -> UArr b
{-# INLINE imap #-}
imap f xs = mapU (uncurryS f) (indexedU xs)

zip :: (UA a, UA b) => UArr a -> UArr b -> UArr (a :*: b)
{-# INLINE zip #-}
zip = zipU

zip3 :: (UA a, UA b, UA c) => UArr a -> UArr b -> UArr c -> UArr (a :*: b :*: c)
{-# INLINE zip3 #-}
zip3 = zip3U

zipWith :: (UA a, UA b, UA c) => (a -> b -> c) -> UArr a -> UArr b -> UArr c
{-# INLINE zipWith #-}
zipWith = zipWithU

unzip :: (UA a, UA b) => UArr (a :*: b) -> UArr a :*: UArr b
{-# INLINE unzip #-}
unzip = unzipU

sum :: (Num a, UA a) => UArr a -> a
{-# INLINE sum #-}
sum = sumU

and :: UArr Bool -> Bool
{-# INLINE and #-}
and = andU

prescanl' :: (UA a, UA b) => (a -> b -> a) -> a -> UArr b -> UArr a
{-# INLINE prescanl' #-}
prescanl' = scanlU

prescanr' = Unsupported

enumFromTo_Int :: Int -> Int -> UArr Int
{-# INLINE enumFromTo_Int #-}
enumFromTo_Int = enumFromToU

filter :: UA a => (a -> Bool) -> UArr a -> UArr a
{-# INLINE filter #-}
filter = filterU

index :: UA a => UArr a -> Int -> a
{-# INLINE index #-}
index = indexU

append :: UA a => UArr a -> UArr a -> UArr a
{-# INLINE append #-}
append = appendU

take :: UA a => Int -> UArr a -> UArr a
{-# INLINE take #-}
take = takeU

drop :: UA a => Int -> UArr a -> UArr a
{-# INLINE drop #-}
drop = dropU

slice :: UA a => UArr a -> Int -> Int -> UArr a
{-# INLINE slice #-}
slice = sliceU

-- Sigh! Where did my backpermute go?
backpermute :: UA a => UArr a -> UArr Int -> UArr a
{-# INLINE backpermute #-}
backpermute xs is = xs `seq` mapU (UArr.indexU xs) is

update = Unsupported
update_ = Unsupported

minIndex :: (Ord e, UA e) => UArr e -> Int
{-# INLINE minIndex #-}
minIndex = fstS . minimumByU (\p q -> compare (sndS p) (sndS q)) . indexedU

maxIndex :: (Ord e, UA e) => UArr e -> Int
{-# INLINE maxIndex #-}
maxIndex = fstS . maximumByU (\p q -> compare (sndS p) (sndS q)) . indexedU

unstablePartition :: UA a => (a -> Bool) -> UArr a -> (UArr a, UArr a)
{-# INLINE unstablePartition #-}
unstablePartition f xs = (filterU f xs, filterU (not . f) xs)

pair :: a -> b -> a :*: b
{-# INLINE pair #-}
pair = (:*:)

from2 :: a :*: b -> (a,b)
{-# INLINE from2 #-}
from2 (a :*: b) = (a,b)

fst :: a :*: b -> a
{-# INLINE fst #-}
fst = fstS

snd :: a :*: b -> b
{-# INLINE snd #-}
snd = sndS

triple :: a -> b -> c -> a :*: b :*: c
{-# INLINE triple #-}
triple a b c = a :*: b :*: c

from3 :: a :*: b :*: c -> (a,b,c)
{-# INLINE from3 #-}
from3 (a :*: b :*: c) = (a,b,c)