{-# OPTIONS_GHC -fno-warn-orphans #-}
--------------------------------------------------------------------------------
-- |
-- Module      :  Data.Vector.Circular.Util
-- Copyright   :  (C) Frank Staals, David Himmelstrup
-- License     :  see the LICENSE file
-- Maintainer  :  David Himmelstrup
--------------------------------------------------------------------------------
module Data.Vector.Circular.Util where

import           Algorithms.StringSearch.KMP (isSubStringOf)
import           Control.Lens
import           Data.Ext
import           Data.Maybe
import           Data.Semigroup.Foldable
import qualified Data.Vector as V
import           Data.Vector.Circular as CV
import qualified Data.Vector.NonEmpty as NV
import           Test.QuickCheck (Arbitrary (..), NonEmptyList (..))


-- FIXME: Upstream this to the non-empty vector library?
instance Foldable1 NV.NonEmptyVector

-- | Access the ith item in the CircularVector (w.r.t the rotation) as a lens
item   :: Int -> Lens' (CircularVector a) a
item :: Int -> Lens' (CircularVector a) a
item Int
i = (CircularVector a -> a)
-> (CircularVector a -> a -> CircularVector a)
-> Lens' (CircularVector a) a
forall s a b t. (s -> a) -> (s -> b -> t) -> Lens s t a b
lens (CircularVector a -> Int -> a
forall a. CircularVector a -> Int -> a
`CV.index` Int
i) (\CircularVector a
s a
x -> Vector a -> CircularVector a
forall a. Vector a -> CircularVector a
unsafeFromVector (CircularVector a -> Vector a
forall a. CircularVector a -> Vector a
toVector CircularVector a
s Vector a -> [(Int, a)] -> Vector a
forall a. Vector a -> [(Int, a)] -> Vector a
V.// [(Int
i,a
x)]))

-- | All elements, starting with the focus, going to the right
--
-- >>> rightElements $ unsafeFromList [3,4,5,1,2]
-- [3,4,5,1,2]
rightElements :: CircularVector a -> NV.NonEmptyVector a
rightElements :: CircularVector a -> NonEmptyVector a
rightElements = CircularVector a -> NonEmptyVector a
forall a. CircularVector a -> NonEmptyVector a
toNonEmptyVector

-- | All elements, starting with the focus, going to the left
--
-- >>> leftElements $ unsafeFromList [3,4,5,1,2]
-- [3,2,1,5,4]
leftElements :: CircularVector a -> NV.NonEmptyVector a
leftElements :: CircularVector a -> NonEmptyVector a
leftElements CircularVector a
v = Int -> (Int -> a) -> NonEmptyVector a
forall a. Int -> (Int -> a) -> NonEmptyVector a
NV.generate1 (CircularVector a -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length CircularVector a
v) (CircularVector a -> Int -> a
forall a. CircularVector a -> Int -> a
CV.index CircularVector a
v (Int -> a) -> (Int -> Int) -> Int -> a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Int -> Int
forall a. Num a => a -> a
negate)

-- | Finds an element in the CircularVector
--
-- >>> findRotateTo (== 3) $ unsafeFromList [1..5]
-- Just (CircularVector {vector = [1,2,3,4,5], rotation = 2})
-- >>> findRotateTo (== 7) $ unsafeFromList [1..5]
-- Nothing
findRotateTo   :: (a -> Bool) -> CircularVector a -> Maybe (CircularVector a)
findRotateTo :: (a -> Bool) -> CircularVector a -> Maybe (CircularVector a)
findRotateTo a -> Bool
p (CircularVector NonEmptyVector a
v Int
_rot) = NonEmptyVector a -> Int -> CircularVector a
forall a. NonEmptyVector a -> Int -> CircularVector a
CircularVector NonEmptyVector a
v (Int -> CircularVector a) -> Maybe Int -> Maybe (CircularVector a)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (a -> Bool) -> NonEmptyVector a -> Maybe Int
forall a. (a -> Bool) -> NonEmptyVector a -> Maybe Int
NV.findIndex a -> Bool
p NonEmptyVector a
v

-- | Test if the circular list is a cyclic shift of the second
-- list.
--
-- Running time: \(O(n+m)\), where \(n\) and \(m\) are the sizes of
-- the lists.
isShiftOf         :: Eq a => CircularVector a -> CircularVector a -> Bool
CircularVector a
xs isShiftOf :: CircularVector a -> CircularVector a -> Bool
`isShiftOf` CircularVector a
ys = let twice :: CircularVector a -> NonEmptyVector a
twice CircularVector a
zs    = let zs' :: NonEmptyVector a
zs' = CircularVector a -> NonEmptyVector a
forall a. CircularVector a -> NonEmptyVector a
leftElements CircularVector a
zs in NonEmptyVector a
zs' NonEmptyVector a -> NonEmptyVector a -> NonEmptyVector a
forall a. Semigroup a => a -> a -> a
<> NonEmptyVector a
zs'
                        once :: CircularVector a -> NonEmptyVector a
once        = CircularVector a -> NonEmptyVector a
forall a. CircularVector a -> NonEmptyVector a
leftElements
                        check :: CircularVector a -> CircularVector a -> Bool
check CircularVector a
as CircularVector a
bs = Maybe Int -> Bool
forall a. Maybe a -> Bool
isJust (Maybe Int -> Bool) -> Maybe Int -> Bool
forall a b. (a -> b) -> a -> b
$ CircularVector a -> NonEmptyVector a
forall a. CircularVector a -> NonEmptyVector a
once CircularVector a
as NonEmptyVector a -> NonEmptyVector a -> Maybe Int
forall a (p :: * -> *) (t :: * -> *).
(Eq a, Foldable p, Foldable t) =>
p a -> t a -> Maybe Int
`isSubStringOf` CircularVector a -> NonEmptyVector a
forall a. CircularVector a -> NonEmptyVector a
twice CircularVector a
bs
                    in CircularVector a -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length CircularVector a
xs Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== CircularVector a -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length CircularVector a
ys Bool -> Bool -> Bool
&& CircularVector a -> CircularVector a -> Bool
forall a. Eq a => CircularVector a -> CircularVector a -> Bool
check CircularVector a
xs CircularVector a
ys

instance Arbitrary a => Arbitrary (CircularVector a) where
  arbitrary :: Gen (CircularVector a)
arbitrary = [a] -> CircularVector a
forall a. [a] -> CircularVector a
unsafeFromList ([a] -> CircularVector a) -> Gen [a] -> Gen (CircularVector a)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (NonEmptyList a -> [a]
forall a. NonEmptyList a -> [a]
getNonEmpty (NonEmptyList a -> [a]) -> Gen (NonEmptyList a) -> Gen [a]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Gen (NonEmptyList a)
forall a. Arbitrary a => Gen a
arbitrary)

-- | label the circular vector with indices, starting from zero at the
-- current focus, going right.
--
-- Running time: \(O(n)\)
withIndicesRight                      :: CircularVector a -> CircularVector (Int :+ a)
withIndicesRight :: CircularVector a -> CircularVector (Int :+ a)
withIndicesRight (CircularVector NonEmptyVector a
v Int
s) = NonEmptyVector (Int :+ a) -> Int -> CircularVector (Int :+ a)
forall a. NonEmptyVector a -> Int -> CircularVector a
CircularVector NonEmptyVector (Int :+ a)
v' Int
s
  where
    n :: Int
n  = NonEmptyVector a -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length NonEmptyVector a
v
    v' :: NonEmptyVector (Int :+ a)
v' = (Int -> a -> Int :+ a)
-> NonEmptyVector a -> NonEmptyVector (Int :+ a)
forall a b. (Int -> a -> b) -> NonEmptyVector a -> NonEmptyVector b
NV.imap (\Int
i a
x -> ((Int
iInt -> Int -> Int
forall a. Num a => a -> a -> a
-Int
s) Int -> Int -> Int
forall a. Integral a => a -> a -> a
`mod` Int
n) Int -> a -> Int :+ a
forall core extra. core -> extra -> core :+ extra
:+ a
x) NonEmptyVector a
v