{-# LANGUAGE GADTs #-}
{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE RankNTypes #-}
{-# LANGUAGE CPP #-}

module Data.Store.Internal.Function
where

--------------------------------------------------------------------------------
import           Control.Applicative hiding (empty)
--------------------------------------------------------------------------------
import           Data.Monoid ((<>))
import           Data.Functor.Identity 
import qualified Data.List
#if MIN_VERSION_containers(0,5,0)
import qualified Data.IntMap.Strict as Data.IntMap
import qualified Data.Map.Strict as Data.Map
#else
import qualified Data.IntMap
import qualified Data.Map
#endif
import qualified Data.IntSet
import qualified Data.IntSet.Extra
--------------------------------------------------------------------------------
import qualified Data.Store.Internal.Type as I
--------------------------------------------------------------------------------

moduleName :: String
moduleName = "Data.Store.Internal.Function"

genericSubset :: I.Empty (I.Index irs ts)
              => Data.IntSet.IntSet
              -> I.Store tag krs irs ts v
              -> I.Store tag krs irs ts v
genericSubset ids (I.Store vs _ _) =
  Data.IntSet.foldr (\i acc ->
    case Data.IntMap.lookup i vs of
      Just (ik, e) -> runIdentity $ genericInsert indexInsertID'' ik e acc
      _ -> acc
    ) I.empty ids
{-# INLINE genericSubset #-}

genericLookup :: Data.IntSet.IntSet
              -> I.Store tag krs irs ts v
              -> [(I.RawKey krs ts, v)]
genericLookup ids (I.Store vs _ _) = {-# SCC "genericLookup" #-} 
  Data.IntSet.foldr (\i acc ->
    case Data.IntMap.lookup i vs of
      Just (ik, v) -> (keyInternalToRaw ik, v) : acc
      _ -> acc
    ) [] ids
{-# INLINE genericLookup #-}

genericLookupAsc :: I.GetDimension n (I.IKey krs ts)
                 => Data.IntSet.IntSet
                 -> n
                 -> I.Store tag krs irs ts v
                 -> [(I.RawKey krs ts, v)]
genericLookupAsc ids n (I.Store vs _ _) =
  map (\(ik, v) -> (keyInternalToRaw ik, v)) . Data.List.sortBy (lookupSortComparator n) $
  Data.IntSet.foldr (\i acc ->
    case Data.IntMap.lookup i vs of
      Just (ik, v) -> (ik, v) : acc
      _ -> acc
    ) [] ids
{-# INLINE genericLookupAsc #-}

genericLookupDesc :: I.GetDimension n (I.IKey krs ts)
                  => Data.IntSet.IntSet
                  -> n
                  -> I.Store tag krs irs ts v
                  -> [(I.RawKey krs ts, v)]
genericLookupDesc ids n (I.Store vs _ _) =
  map (\(ik, v) -> (keyInternalToRaw ik, v)) . Data.List.sortBy (flip $ lookupSortComparator n) $
  Data.IntSet.foldr (\i acc ->
    case Data.IntMap.lookup i vs of
      Just (ik, v) -> (ik, v) : acc
      _ -> acc
    ) [] ids
{-# INLINE genericLookupDesc #-}

lookupSortComparator :: I.GetDimension n (I.IKey krs ts) => n -> (I.IKey krs ts, e) -> (I.IKey krs ts, e) -> Ordering
lookupSortComparator n (k1, _) (k2, _) = case (I.getDimension n k1, I.getDimension n k2) of
  (I.IKeyDimensionO d1, I.IKeyDimensionO d2) -> d1 `compare` d2
  (I.IKeyDimensionM  _