{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE GADTs #-}
{-# LANGUAGE InstanceSigs #-}
{-# LANGUAGE OverloadedStrings #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE TypeApplications #-}

module DataFrame.Operations.Permutation where

import qualified Data.List as L
import qualified Data.Text as T
import qualified Data.Vector as V
import qualified Data.Vector.Algorithms.Merge as VA
import qualified Data.Vector.Generic as VG
import qualified Data.Vector.Unboxed as VU
import qualified Data.Vector.Unboxed.Mutable as VUM

import Control.Exception (throw)
import Control.Monad.ST (runST)
import Data.Type.Equality (testEquality, (:~:) (Refl))
import Data.Vector.Internal.Check (HasCallStack)
import DataFrame.Errors (DataFrameException (..))
import DataFrame.Internal.Column (Column (..), Columnable, atIndicesStable)
import DataFrame.Internal.DataFrame (
    DataFrame (..),
    columnNames,
    unsafeGetColumn,
 )
import DataFrame.Internal.Expression (Expr (Col), getColumns)
import DataFrame.Operations.Core (dimensions)
import DataFrame.Operations.Transformations (derive)
import System.Random (Random (randomR), RandomGen)
import Type.Reflection (typeRep)

-- | Sort order taken as a parameter by the 'sortBy' function.
data SortOrder where
    Asc :: (Columnable a, Ord a) => Expr a -> SortOrder
    Desc :: (Columnable a, Ord a) => Expr a -> SortOrder

instance Eq SortOrder where
    (==) :: SortOrder -> SortOrder -> Bool
    == :: SortOrder -> SortOrder -> Bool
(==) (Asc Expr a
_) (Asc Expr a
_) = Bool
True
    (==) (Desc Expr a
_) (Desc Expr a
_) = Bool
True
    (==) SortOrder
_ SortOrder
_ = Bool
False

sortOrderColumns :: SortOrder -> [T.Text]
sortOrderColumns :: SortOrder -> [Text]
sortOrderColumns (Asc Expr a
e) = Expr a -> [Text]
forall a. Expr a -> [Text]
getColumns Expr a
e
sortOrderColumns (Desc Expr a
e) = Expr a -> [Text]
forall a. Expr a -> [Text]
getColumns Expr a
e

mustFlipCompare :: SortOrder -> Bool
mustFlipCompare :: SortOrder -> Bool
mustFlipCompare (Asc Expr a
_) = Bool
True
mustFlipCompare (Desc Expr a
_) = Bool
False

{- | Materialize any compound sort expressions into synthetic columns on
a working dataframe, returning rewritten 'SortOrder's that reference
those columns by name.
-}
prepareSortColumns :: [SortOrder] -> DataFrame -> ([SortOrder], DataFrame)
prepareSortColumns :: [SortOrder] -> DataFrame -> ([SortOrder], DataFrame)
prepareSortColumns = Int -> [SortOrder] -> DataFrame -> ([SortOrder], DataFrame)
go Int
0
  where
    go :: Int -> [SortOrder] -> DataFrame -> ([SortOrder], DataFrame)
go Int
_ [] DataFrame
acc = ([], DataFrame
acc)
    go Int
i (SortOrder
ord : [SortOrder]
rest) DataFrame
acc =
        let (SortOrder
ord', DataFrame
acc') = Int -> SortOrder -> DataFrame -> (SortOrder, DataFrame)
materializeSortOrder Int
i SortOrder
ord DataFrame
acc
            ([SortOrder]
rest', DataFrame
acc'') = Int -> [SortOrder] -> DataFrame -> ([SortOrder], DataFrame)
go (Int
i Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1) [SortOrder]
rest DataFrame
acc'
         in (SortOrder
ord' SortOrder -> [SortOrder] -> [SortOrder]
forall a. a -> [a] -> [a]
: [SortOrder]
rest', DataFrame
acc'')

materializeSortOrder :: Int -> SortOrder -> DataFrame -> (SortOrder, DataFrame)
materializeSortOrder :: Int -> SortOrder -> DataFrame -> (SortOrder, DataFrame)
materializeSortOrder Int
_ ord :: SortOrder
ord@(Asc (Col Text
_)) DataFrame
df = (SortOrder
ord, DataFrame
df)
materializeSortOrder Int
_ ord :: SortOrder
ord@(Desc (Col Text
_)) DataFrame
df = (SortOrder
ord, DataFrame
df)
materializeSortOrder Int
i (Asc (Expr a
e :: Expr a)) DataFrame
df =
    let name :: Text
name = Int -> Text
syntheticName Int
i
     in (Expr a -> SortOrder
forall a. (Columnable a, Ord a) => Expr a -> SortOrder
Asc (Text -> Expr a
forall a. Columnable a => Text -> Expr a
Col Text
name :: Expr a), Text -> Expr a -> DataFrame -> DataFrame
forall a. Columnable a => Text -> Expr a -> DataFrame -> DataFrame
derive Text
name Expr a
e DataFrame
df)
materializeSortOrder Int
i (Desc (Expr a
e :: Expr a)) DataFrame
df =
    let name :: Text
name = Int -> Text
syntheticName Int
i
     in (Expr a -> SortOrder
forall a. (Columnable a, Ord a) => Expr a -> SortOrder
Desc (Text -> Expr a
forall a. Columnable a => Text -> Expr a
Col Text
name :: Expr a), Text -> Expr a -> DataFrame -> DataFrame
forall a. Columnable a => Text -> Expr a -> DataFrame -> DataFrame
derive Text
name Expr a
e DataFrame
df)

syntheticName :: Int -> T.Text
syntheticName :: Int -> Text
syntheticName Int
i = Text
"__sortBy_synthetic_" Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> String -> Text
T.pack (Int -> String
forall a. Show a => a -> String
show Int
i) Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> Text
"__"

{- | O(k log n) Sorts the dataframe by a given row.

> sortBy Ascending ["Age"] df
-}
sortBy ::
    [SortOrder] ->
    DataFrame ->
    DataFrame
sortBy :: [SortOrder] -> DataFrame -> DataFrame
sortBy [SortOrder]
sortOrds DataFrame
df
    | Bool -> Bool
not ([Text] -> Bool
forall a. [a] -> Bool
forall (t :: * -> *) a. Foldable t => t a -> Bool
null [Text]
missing) =
        DataFrameException -> DataFrame
forall a e. Exception e => e -> a
throw (DataFrameException -> DataFrame)
-> DataFrameException -> DataFrame
forall a b. (a -> b) -> a -> b
$
            [Text] -> Text -> [Text] -> DataFrameException
ColumnsNotFoundException
                [Text]
missing
                Text
"sortBy"
                (DataFrame -> [Text]
columnNames DataFrame
df)
    | Bool
otherwise =
        let
            ([SortOrder]
sortOrds', DataFrame
df') = [SortOrder] -> DataFrame -> ([SortOrder], DataFrame)
prepareSortColumns [SortOrder]
sortOrds DataFrame
df
            comparators :: [Int -> Int -> Ordering]
comparators = (SortOrder -> Int -> Int -> Ordering)
-> [SortOrder] -> [Int -> Int -> Ordering]
forall a b. (a -> b) -> [a] -> [b]
map (SortOrder -> DataFrame -> Int -> Int -> Ordering
`sortOrderComparator` DataFrame
df') [SortOrder]
sortOrds'
            compositeCompare :: Int -> Int -> Ordering
compositeCompare Int
i Int
j = [Ordering] -> Ordering
forall a. Monoid a => [a] -> a
mconcat [Int -> Int -> Ordering
c Int
i Int
j | Int -> Int -> Ordering
c <- [Int -> Int -> Ordering]
comparators]
            nRows :: Int
nRows = (Int, Int) -> Int
forall a b. (a, b) -> a
fst (DataFrame -> (Int, Int)
dataframeDimensions DataFrame
df')
            indexes :: Vector Int
indexes = (Int -> Int -> Ordering) -> Int -> Vector Int
sortIndices Int -> Int -> Ordering
compositeCompare Int
nRows
         in
            DataFrame
df{columns = V.map (atIndicesStable indexes) (columns df)}
  where
    referenced :: [Text]
referenced = [Text] -> [Text]
forall a. Eq a => [a] -> [a]
L.nub ((SortOrder -> [Text]) -> [SortOrder] -> [Text]
forall (t :: * -> *) a b. Foldable t => (a -> [b]) -> t a -> [b]
concatMap SortOrder -> [Text]
sortOrderColumns [SortOrder]
sortOrds)
    missing :: [Text]
missing = [Text]
referenced [Text] -> [Text] -> [Text]
forall a. Eq a => [a] -> [a] -> [a]
L.\\ DataFrame -> [Text]
columnNames DataFrame
df

{- | Build a row-index comparator from a SortOrder and a DataFrame.
The Ord dictionary is recovered from the SortOrder GADT.
-}
sortOrderComparator :: SortOrder -> DataFrame -> Int -> Int -> Ordering
sortOrderComparator :: SortOrder -> DataFrame -> Int -> Int -> Ordering
sortOrderComparator (Asc (Col Text
name :: Expr a)) DataFrame
df =
    case Text -> DataFrame -> Column
unsafeGetColumn Text
name DataFrame
df of
        BoxedColumn Maybe Bitmap
_ (Vector a
v :: V.Vector b) -> case TypeRep a -> TypeRep a -> Maybe (a :~: a)
forall a b. TypeRep a -> TypeRep b -> Maybe (a :~: b)
forall {k} (f :: k -> *) (a :: k) (b :: k).
TestEquality f =>
f a -> f b -> Maybe (a :~: b)
testEquality (forall a. Typeable a => TypeRep a
forall {k} (a :: k). Typeable a => TypeRep a
typeRep @a) (forall a. Typeable a => TypeRep a
forall {k} (a :: k). Typeable a => TypeRep a
typeRep @b) of
            Just a :~: a
Refl -> \Int
i Int
j -> a -> a -> Ordering
forall a. Ord a => a -> a -> Ordering
compare (Vector a
v Vector a -> Int -> a
forall a. Vector a -> Int -> a
`V.unsafeIndex` Int
i) (Vector a
v Vector a -> Int -> a
forall a. Vector a -> Int -> a
`V.unsafeIndex` Int
j)
            Maybe (a :~: a)
Nothing -> \Int
_ Int
_ -> Ordering
EQ
        UnboxedColumn Maybe Bitmap
_ (Vector a
v :: VU.Vector b) -> case TypeRep a -> TypeRep a -> Maybe (a :~: a)
forall a b. TypeRep a -> TypeRep b -> Maybe (a :~: b)
forall {k} (f :: k -> *) (a :: k) (b :: k).
TestEquality f =>
f a -> f b -> Maybe (a :~: b)
testEquality (forall a. Typeable a => TypeRep a
forall {k} (a :: k). Typeable a => TypeRep a
typeRep @a) (forall a. Typeable a => TypeRep a
forall {k} (a :: k). Typeable a => TypeRep a
typeRep @b) of
            Just a :~: a
Refl -> \Int
i Int
j -> a -> a -> Ordering
forall a. Ord a => a -> a -> Ordering
compare (Vector a
v Vector a -> Int -> a
forall a. Unbox a => Vector a -> Int -> a
`VU.unsafeIndex` Int
i) (Vector a
v Vector a -> Int -> a
forall a. Unbox a => Vector a -> Int -> a
`VU.unsafeIndex` Int
j)
            Maybe (a :~: a)
Nothing -> \Int
_ Int
_ -> Ordering
EQ
sortOrderComparator (Desc (Col Text
name :: Expr a)) DataFrame
df =
    case Text -> DataFrame -> Column
unsafeGetColumn Text
name DataFrame
df of
        BoxedColumn Maybe Bitmap
_ (Vector a
v :: V.Vector b) -> case TypeRep a -> TypeRep a -> Maybe (a :~: a)
forall a b. TypeRep a -> TypeRep b -> Maybe (a :~: b)
forall {k} (f :: k -> *) (a :: k) (b :: k).
TestEquality f =>
f a -> f b -> Maybe (a :~: b)
testEquality (forall a. Typeable a => TypeRep a
forall {k} (a :: k). Typeable a => TypeRep a
typeRep @a) (forall a. Typeable a => TypeRep a
forall {k} (a :: k). Typeable a => TypeRep a
typeRep @b) of
            Just a :~: a
Refl -> \Int
i Int
j -> a -> a -> Ordering
forall a. Ord a => a -> a -> Ordering
compare (Vector a
v Vector a -> Int -> a
forall a. Vector a -> Int -> a
`V.unsafeIndex` Int
j) (Vector a
v Vector a -> Int -> a
forall a. Vector a -> Int -> a
`V.unsafeIndex` Int
i)
            Maybe (a :~: a)
Nothing -> \Int
_ Int
_ -> Ordering
EQ
        UnboxedColumn Maybe Bitmap
_ (Vector a
v :: VU.Vector b) -> case TypeRep a -> TypeRep a -> Maybe (a :~: a)
forall a b. TypeRep a -> TypeRep b -> Maybe (a :~: b)
forall {k} (f :: k -> *) (a :: k) (b :: k).
TestEquality f =>
f a -> f b -> Maybe (a :~: b)
testEquality (forall a. Typeable a => TypeRep a
forall {k} (a :: k). Typeable a => TypeRep a
typeRep @a) (forall a. Typeable a => TypeRep a
forall {k} (a :: k). Typeable a => TypeRep a
typeRep @b) of
            Just a :~: a
Refl -> \Int
i Int
j -> a -> a -> Ordering
forall a. Ord a => a -> a -> Ordering
compare (Vector a
v Vector a -> Int -> a
forall a. Unbox a => Vector a -> Int -> a
`VU.unsafeIndex` Int
j) (Vector a
v Vector a -> Int -> a
forall a. Unbox a => Vector a -> Int -> a
`VU.unsafeIndex` Int
i)
            Maybe (a :~: a)
Nothing -> \Int
_ Int
_ -> Ordering
EQ
sortOrderComparator SortOrder
_ DataFrame
_ = String -> Int -> Int -> Ordering
forall a. HasCallStack => String -> a
error String
"Sorting on compound column"

-- | Sort row indices using a comparator function.
sortIndices :: (Int -> Int -> Ordering) -> Int -> VU.Vector Int
sortIndices :: (Int -> Int -> Ordering) -> Int -> Vector Int
sortIndices Int -> Int -> Ordering
cmp Int
nRows = (forall s. ST s (Vector Int)) -> Vector Int
forall a. (forall s. ST s a) -> a
runST ((forall s. ST s (Vector Int)) -> Vector Int)
-> (forall s. ST s (Vector Int)) -> Vector Int
forall a b. (a -> b) -> a -> b
$ do
    MVector s Int
withIndexes <- Vector Int -> ST s (Mutable Vector (PrimState (ST s)) Int)
forall (m :: * -> *) (v :: * -> *) a.
(PrimMonad m, Vector v a) =>
v a -> m (Mutable v (PrimState m) a)
VG.thaw (Int -> (Int -> Int) -> Vector Int
forall a. Int -> (Int -> a) -> Vector a
V.generate Int
nRows Int -> Int
forall a. a -> a
id :: V.Vector Int)
    (Int -> Int -> Ordering)
-> MVector (PrimState (ST s)) Int -> ST s ()
forall (m :: * -> *) (v :: * -> * -> *) e.
(PrimMonad m, MVector v e) =>
Comparison e -> v (PrimState m) e -> m ()
VA.sortBy Int -> Int -> Ordering
cmp MVector s Int
MVector (PrimState (ST s)) Int
withIndexes
    Vector Int
sorted <- Mutable Vector (PrimState (ST s)) Int -> ST s (Vector Int)
forall (m :: * -> *) (v :: * -> *) a.
(PrimMonad m, Vector v a) =>
Mutable v (PrimState m) a -> m (v a)
VG.unsafeFreeze MVector s Int
Mutable Vector (PrimState (ST s)) Int
withIndexes
    Vector Int -> ST s (Vector Int)
forall a. a -> ST s a
forall (m :: * -> *) a. Monad m => a -> m a
return (Vector Int -> Vector Int
forall (v :: * -> *) a (w :: * -> *).
(Vector v a, Vector w a) =>
v a -> w a
VU.convert Vector Int
sorted)

shuffle ::
    (RandomGen g) =>
    g ->
    DataFrame ->
    DataFrame
shuffle :: forall g. RandomGen g => g -> DataFrame -> DataFrame
shuffle g
pureGen DataFrame
df =
    let
        indexes :: Vector Int
indexes = g -> Int -> Vector Int
forall g. (HasCallStack, RandomGen g) => g -> Int -> Vector Int
shuffledIndices g
pureGen ((Int, Int) -> Int
forall a b. (a, b) -> a
fst (DataFrame -> (Int, Int)
dimensions DataFrame
df))
     in
        DataFrame
df{columns = V.map (atIndicesStable indexes) (columns df)}

shuffledIndices :: (HasCallStack, RandomGen g) => g -> Int -> VU.Vector Int
shuffledIndices :: forall g. (HasCallStack, RandomGen g) => g -> Int -> Vector Int
shuffledIndices g
pureGen Int
k
    | Int
k Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
0 = String -> Vector Int
forall a. HasCallStack => String -> a
error (String -> Vector Int) -> String -> Vector Int
forall a b. (a -> b) -> a -> b
$ String
"Vector index may not be a neative number: " String -> String -> String
forall a. Semigroup a => a -> a -> a
<> Int -> String
forall a. Show a => a -> String
show Int
k
    | Int
k Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
0 = Vector Int
forall a. Unbox a => Vector a
VU.empty
    | Bool
otherwise = g -> Vector Int
forall g. RandomGen g => g -> Vector Int
shuffleVec g
pureGen
  where
    shuffleVec :: (RandomGen g) => g -> VU.Vector Int
    shuffleVec :: forall g. RandomGen g => g -> Vector Int
shuffleVec g
g = (forall s. ST s (Vector Int)) -> Vector Int
forall a. (forall s. ST s a) -> a
runST ((forall s. ST s (Vector Int)) -> Vector Int)
-> (forall s. ST s (Vector Int)) -> Vector Int
forall a b. (a -> b) -> a -> b
$ do
        MVector s Int
vm <- Int -> (Int -> Int) -> ST s (MVector (PrimState (ST s)) Int)
forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
Int -> (Int -> a) -> m (MVector (PrimState m) a)
VUM.generate Int
k Int -> Int
forall a. a -> a
id
        let (Int
n, g
nGen) = (Int, Int) -> g -> (Int, g)
forall g. RandomGen g => (Int, Int) -> g -> (Int, g)
forall a g. (Random a, RandomGen g) => (a, a) -> g -> (a, g)
randomR (Int
1, Int
k Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1) g
g
        MVector (PrimState (ST s)) Int -> Int -> g -> ST s ()
forall {f :: * -> *} {t} {a}.
(RandomGen t, PrimMonad f, Unbox a) =>
MVector (PrimState f) a -> Int -> t -> f ()
go MVector s Int
MVector (PrimState (ST s)) Int
vm Int
n g
nGen
        MVector (PrimState (ST s)) Int -> ST s (Vector Int)
forall a (m :: * -> *).
(Unbox a, PrimMonad m) =>
MVector (PrimState m) a -> m (Vector a)
VU.unsafeFreeze MVector s Int
MVector (PrimState (ST s)) Int
vm

    go :: MVector (PrimState f) a -> Int -> t -> f ()
go MVector (PrimState f) a
_v (-1) t
_ = () -> f ()
forall a. a -> f a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ()
    go MVector (PrimState f) a
_v Int
0 t
_ = () -> f ()
forall a. a -> f a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ()
    go MVector (PrimState f) a
v Int
maxInd t
gen =
        let
            (Int
n, t
nextGen) = (Int, Int) -> t -> (Int, t)
forall g. RandomGen g => (Int, Int) -> g -> (Int, g)
forall a g. (Random a, RandomGen g) => (a, a) -> g -> (a, g)
randomR (Int
1, Int
maxInd) t
gen
         in
            MVector (PrimState f) a -> Int -> Int -> f ()
forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
MVector (PrimState m) a -> Int -> Int -> m ()
VUM.swap MVector (PrimState f) a
v Int
0 Int
n f () -> f () -> f ()
forall a b. f a -> f b -> f b
forall (f :: * -> *) a b. Applicative f => f a -> f b -> f b
*> MVector (PrimState f) a -> Int -> t -> f ()
go (MVector (PrimState f) a -> MVector (PrimState f) a
forall a s. Unbox a => MVector s a -> MVector s a
VUM.tail MVector (PrimState f) a
v) (Int
maxInd Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1) t
nextGen