{-# LANGUAGE BangPatterns #-}
{-# LANGUAGE TypeFamilies #-}
module Data.Indexation.SimpleIx
( SimpleIx(..)
, Indexed(..)
, createIndexed
, withFunction
, withIndex1
, withIndex2
, withIndex3
, withIndex4
, withFilters
, insert
, deleteIndex
, delete
, size
, getPos
, getI1
, getI2
, getI3
, getI4
, getFilter
, member
, adjust
, update
, adjustWithElement
) where

import Data.HashMap.Strict (HashMap)
import qualified Data.HashMap.Strict as HM
import Data.IntSet (IntSet)
import qualified Data.IntSet as IS
import Data.IntMap.Strict (IntMap, (!))
import qualified Data.IntMap.Strict as IM
import Data.Vector (Vector, (!?))
import qualified Data.Vector as V
import Data.SafeCopy
import Data.Hashable
-- import Data.Indexation.IxData (IxData(ixDataPos, ixData))

data Indexed a i1 i2 i3 i4 
    = Indexed !(IntMap a) 
              -- ^ Where all the data is
              !(a -> Int) 
              -- ^ How to access the data 
              !(Maybe (IndexI a i1 i2 i3 i4))
              -- ^ Indexation structures
              !(Vector (a -> Bool)) 
              -- ^ Filters (i.e. >= 30, == "foo")
              !(Vector IntSet)
              -- ^ Filters results
              !IntSet
              -- ^ helper set to speed index lookup
              !Int
              -- ^ Number of items

-- TODO: the function should be (a -> b), where b is Hashable
-- this should be the "ID" function, and it should be used to check if
-- the id was changed in "adjust" and update it

data IndexI a i1 i2 i3 i4
    = IxHasI1 !(HashMap i1 IntSet) !(a -> i1)
    | IxHasI2 !(HashMap i1 IntSet) !(a -> i1)
              !(HashMap i2 IntSet) !(a -> i2)
    | IxHasI3 !(HashMap i1 IntSet) !(a -> i1)
              !(HashMap i2 IntSet) !(a -> i2)
              !(HashMap i3 IntSet) !(a -> i3)
    | IxHasI4 !(HashMap i1 IntSet) !(a -> i1)
              !(HashMap i2 IntSet) !(a -> i2)
              !(HashMap i3 IntSet) !(a -> i3)
              !(HashMap i4 IntSet) !(a -> i4)

-- * 
-- *  Creation functions
-- * 
class SimpleIx a where
    -- | Creates your indexed dataype
    type I1 a
    type I2 a
    type I3 a
    type I4 a
    create :: Indexed a (I1 a) (I2 a) (I3 a) (I4 a)

-- | createIndexed is a helper function to create the base Indexed datatype
-- it is used in the "createdUxIx" functions to avoid repetition
createIndexed :: ( Eq a, Eq i1, Hashable i1, Eq i2, Hashable i2
                 , Eq i3, Hashable i3, Eq i4, Hashable i4)
              => Indexed a i1 i2 i3 i4
createIndexed = Indexed IM.empty (const 0) Nothing V.empty V.empty IS.empty 0

withFunction :: ( Eq a, Eq i1, Hashable i1, Eq i2, Hashable i2
                , Eq i3, Hashable i3, Eq i4, Hashable i4
                )
            => Indexed a i1 i2 i3 i4 -> (a -> Int) 
            -> Indexed a i1 i2 i3 i4
withFunction (Indexed d _ iis vf vs is s) f = Indexed d f iis vf vs is s

withIndex1 :: ( Eq a 
              , Eq i1, Hashable i1, Eq i2, Hashable i2, Eq i3, Hashable i3
              , Eq i4, Hashable i4
              )
            => Indexed a i1 i2 i3 i4 -> (a -> i1) 
            -> Indexed a i1 i2 i3 i4
withIndex1 i@(Indexed d a iis vf vs is s) f = 
    case iis of
        Nothing -> Indexed d a (Just $ IxHasI1 HM.empty f) vf vs is s
        _ -> i

withIndex2 :: ( Eq a
              , Eq i1, Hashable i1, Eq i2, Hashable i2, Eq i3, Hashable i3
              , Eq i4, Hashable i4
              )
            => Indexed a i1 i2 i3 i4 -> (a -> i2) 
            -> Indexed a i1 i2 i3 i4
withIndex2 i@(Indexed d a iis vf vs is s) f = 
    case iis of
        Just (IxHasI1 m1 f1) -> 
            let niis = Just $ IxHasI2 m1 f1 HM.empty f in 
                Indexed d a niis vf vs is s
        _ -> i
        

withIndex3 :: ( Eq a
              , Eq i1, Hashable i1, Eq i2, Hashable i2, Eq i3, Hashable i3
              , Eq i4, Hashable i4
              )
            => Indexed a i1 i2 i3 i4 -> (a -> i3) 
            -> Indexed a i1 i2 i3 i4
withIndex3 i@(Indexed d a iis vf vs is s) f = 
    case iis of
        Just (IxHasI2 m1 f1 m2 f2) -> 
           let niis = Just $ IxHasI3 m1 f1 m2 f2 HM.empty f in
                Indexed d a niis vf vs is s
        _ -> i
        

withIndex4 :: ( Eq a
              , Eq i1, Hashable i1, Eq i2, Hashable i2, Eq i3, Hashable i3
              , Eq i4, Hashable i4
              )
            => Indexed a i1 i2 i3 i4 -> (a -> i4) 
            -> Indexed a i1 i2 i3 i4
withIndex4 i@(Indexed d a iis vf vs is s) f = 
    case iis of
        Just (IxHasI3 m1 f1 m2 f2 m3 f3) -> 
            let niis = Just $ IxHasI4 m1 f1 m2 f2 m3 f3 HM.empty f in
                Indexed d a niis vf vs is s
        _ -> i

withFilters :: ( Eq a
               , Eq i1, Hashable i1, Eq i2, Hashable i2, Eq i3, Hashable i3
               , Eq i4, Hashable i4
               )
            => Indexed a i1 i2 i3 i4 -> [(a -> Bool)]
            -> Indexed a i1 i2 i3 i4 
withFilters (Indexed d f iis _ _ is s) lst = 
    Indexed d f iis flst vlst is s
    where
        flst = V.fromList lst
        vlst = V.replicate (V.length flst) IS.empty

-- * 
-- *  Insertion functions
-- * 
insert :: ( Eq a
          , Eq i1, Hashable i1, Eq i2, Hashable i2, Eq i3, Hashable i3
          , Eq i4, Hashable i4
          )
       => Indexed a i1 i2 i3 i4 -> a 
       -> Indexed a i1 i2 i3 i4
insert (Indexed d f iis vf vi is s) item = 
    Indexed newd f newiis vf newvi newis (s + 1)
    where 
        pos = f item
        newd = IM.insert pos item d
        newis = IS.insert pos is
        newiis = 
            case iis of
                Nothing -> Nothing
                Just (IxHasI1 m1 f1) -> Just $
                    IxHasI1 (HM.adjust (IS.insert pos) (f1 item) m1) f1
                Just (IxHasI2 m1 f1 m2 f2) -> Just $ 
                    IxHasI2 (HM.adjust (IS.insert pos) (f1 item) m1) f1
                            (HM.adjust (IS.insert pos) (f2 item) m2) f2
                Just (IxHasI3 m1 f1 m2 f2 m3 f3) -> Just $ 
                    IxHasI3 (HM.adjust (IS.insert pos) (f1 item) m1) f1
                            (HM.adjust (IS.insert pos) (f2 item) m2) f2
                            (HM.adjust (IS.insert pos) (f3 item) m3) f3
                Just (IxHasI4 m1 f1 m2 f2 m3 f3 m4 f4) -> Just $ 
                    IxHasI4 (HM.adjust (IS.insert pos) (f1 item) m1) f1
                            (HM.adjust (IS.insert pos) (f2 item) m2) f2
                            (HM.adjust (IS.insert pos) (f3 item) m3) f3
                            (HM.adjust (IS.insert pos) (f4 item) m4) f4
        newvi = V.imap fvi vi
        -- check if the filter applies to this element, if it applies insert
        -- the element index into the intset
        fvi i e = if V.unsafeIndex vf i $ item then IS.insert pos e else e
{-# INLINABLE insert #-}

-- * 
-- *  Deletion functions
-- * 
deleteIndex :: ( Eq a
               , Eq i1, Hashable i1, Eq i2, Hashable i2, Eq i3, Hashable i3
               , Eq i4, Hashable i4
               )
            => Indexed a i1 i2 i3 i4 -> Int -> Indexed a i1 i2 i3 i4
deleteIndex i@(Indexed d a iis vf vi is s) ix = 
    case old of
        Nothing -> i
        Just _  -> Indexed newd a iis vf vi (IS.delete ix is) (s - 1)
    where
        deleteKey _ _ = Nothing
        (old, newd)   = IM.updateLookupWithKey deleteKey ix d


{-# INLINE deleteIndex #-}
    -- ^ TODO: delete from the indices and vi ?

delete :: ( Eq a
          , Eq i1, Hashable i1, Eq i2, Hashable i2, Eq i3, Hashable i3
          , Eq i4, Hashable i4
          )
         => Indexed a i1 i2 i3 i4 -> a -> Indexed a i1 i2 i3 i4
delete i@(Indexed _ f _ _ _ _ _) item = deleteIndex i (f item)
{-# INLINE delete #-}

-- * 
-- *  Accessor (getters) functions
-- * 
size :: ( Eq a
        , Eq i1, Hashable i1, Eq i2, Hashable i2, Eq i3, Hashable i3
        , Eq i4, Hashable i4
        )
     => Indexed a i1 i2 i3 i4 -> Int 
size (Indexed _ _ _ _ _ _ s) = s 
{-# INLINE size #-}

getPos :: ( Eq a
          , Eq i1, Hashable i1, Eq i2, Hashable i2, Eq i3, Hashable i3
          , Eq i4, Hashable i4
          )
      => Indexed a i1 i2 i3 i4 -> Int -> Maybe a
getPos (Indexed d _ _ _ _ _ _) pos = IM.lookup pos d
{-# INLINE getPos #-}

getSet :: IntMap a -> IntSet -> IntSet -> IntMap a
getSet im imkis is = IM.fromSet (\k -> im ! k) (IS.intersection imkis is)
{-# INLINE getSet #-}

getI1 :: ( Eq a
         , Eq i1, Hashable i1, Eq i2, Hashable i2, Eq i3, Hashable i3
         , Eq i4, Hashable i4
         )
      => Indexed a i1 i2 i3 i4 -> i1 -> Maybe (IntMap a)
getI1 (Indexed d _ iis _ _ is _) i = do 
    iis' <- iis
    set <- f iis'
    return $ getSet d is set
    where
       f (IxHasI1 m _) = HM.lookup i m
       f (IxHasI2 m _ _ _) = HM.lookup i m
       f (IxHasI3 m _ _ _ _ _) = HM.lookup i m
       f (IxHasI4 m _ _ _ _ _ _ _) = HM.lookup i m
{-# INLINABLE getI1 #-}

getI2 :: ( Eq a
         , Eq i1, Hashable i1, Eq i2, Hashable i2, Eq i3, Hashable i3
         , Eq i4, Hashable i4
         )
      => Indexed a i1 i2 i3 i4 -> i2 -> Maybe (IntMap a)
getI2 (Indexed d _ iis _ _ is _) i = do 
    iis' <- iis
    set <- f iis'
    return $ getSet d is set
    where
       f (IxHasI2 _ _ m _) = HM.lookup i m
       f (IxHasI3 _ _ m _ _ _) = HM.lookup i m
       f (IxHasI4 _ _ m _ _ _ _ _) = HM.lookup i m
       f _ = Nothing
{-# INLINABLE getI2 #-}

getI3 :: ( Eq a
         , Eq i1, Hashable i1, Eq i2, Hashable i2, Eq i3, Hashable i3
         , Eq i4, Hashable i4
         )
      => Indexed a i1 i2 i3 i4 -> i3 -> Maybe (IntMap a)
getI3 (Indexed d _ iis _ _ is _) i = do 
    iis' <- iis
    set <- f iis'
    return $ getSet d is set
    where
       f (IxHasI3 _ _ _ _ m _) = HM.lookup i m
       f (IxHasI4 _ _ _ _ m _ _ _) = HM.lookup i m
       f _ = Nothing
{-# INLINABLE getI3 #-}

getI4 :: ( Eq a
         , Eq i1, Hashable i1, Eq i2, Hashable i2, Eq i3, Hashable i3
         , Eq i4, Hashable i4
         )
      => Indexed a i1 i2 i3 i4 -> i4 -> Maybe (IntMap a)
getI4 (Indexed d _ iis _ _ is _) i = do 
    iis' <- iis
    set <- f iis'
    return $ getSet d is set
    where
       f (IxHasI4 _ _ _ _ _ _ m _) = HM.lookup i m
       f _ = Nothing
{-# INLINABLE getI4 #-}

getFilter :: ( Eq a
             , Eq i1, Hashable i1, Eq i2, Hashable i2, Eq i3, Hashable i3
             , Eq i4, Hashable i4
             )
          => Indexed a i1 i2 i3 i4 -> Int -> Maybe (IntMap a)
getFilter (Indexed d _ _ _ vi is _) filterPos = do
    set <- vi !? filterPos
    return $ getSet d is set
{-# INLINABLE getFilter #-}

member :: ( Eq a, Hashable b
          , Eq i1, Hashable i1, Eq i2, Hashable i2, Eq i3, Hashable i3
          , Eq i4, Hashable i4
          )
       => Indexed a i1 i2 i3 i4 -> b -> Bool
member (Indexed _ _ _ _ _ is _) b = IS.member (hash b) is
{-# INLINE member #-}

-- * 
-- *  Setter functions
-- * 

-- | helper function, updates the normal indices
_updateIHash :: (Eq u, Hashable u) => HashMap u IntSet -> Int -> u -> u 
             -> HashMap u IntSet
_updateIHash m val old new = 
    HM.insertWith insSet new valSet $ HM.adjust (IS.delete val) old m
    where
        insSet _ o = IS.insert val o
        valSet = IS.singleton val
-- | helper function, updates the normal indices
_updateI :: ( Eq a
            , Eq i1, Hashable i1, Eq i2, Hashable i2
            , Eq i3, Hashable i3, Eq i4, Hashable i4)
         => Maybe (IndexI a i1 i2 i3 i4) -> Int -> a -> a 
         -> Maybe (IndexI a i1 i2 i3 i4)
_updateI (Just (IxHasI1 m f)) val old new = 
        Just $ IxHasI1 (_updateIHash m val (f old) (f new)) f
_updateI (Just (IxHasI2 m1 f1 m2 f2)) val old new = 
        Just $ IxHasI2 (_updateIHash m1 val (f1 old) (f1 new)) f1
                       (_updateIHash m2 val (f2 old) (f2 new)) f2
_updateI (Just (IxHasI3 m1 f1 m2 f2 m3 f3)) val old new = 
        Just $ IxHasI3 (_updateIHash m1 val (f1 old) (f1 new)) f1
                       (_updateIHash m2 val (f2 old) (f2 new)) f2
                       (_updateIHash m3 val (f3 old) (f3 new)) f3
_updateI (Just (IxHasI4 m1 f1 m2 f2 m3 f3 m4 f4)) val old new = 
        Just $ IxHasI4 (_updateIHash m1 val (f1 old) (f1 new)) f1
                       (_updateIHash m2 val (f2 old) (f2 new)) f2
                       (_updateIHash m3 val (f3 old) (f3 new)) f3
                       (_updateIHash m4 val (f4 old) (f4 new)) f4
_updateI Nothing _ _ _ = Nothing

-- | helper function, updates the normal indices
_updateF :: Vector (a -> Bool) -> Vector IntSet -> Int -> a -> a 
         -> Vector IntSet
_updateF vf v value old new = V.imap applyF vf
    -- map the vf, apply the function to the old and the new, if diff, change 
    -- the v because it is faster to check for indexation in a vector than 
    -- to check membership in a intset
    where
        applyF pos f = let f1 = f old
                           f2 = f new
                           iset = V.unsafeIndex v pos
                       in if f1 == f2 
                            then iset
                            else if f2 then IS.insert value iset 
                                       else IS.delete value iset 

adjust :: ( Eq a
          , Eq i1, Hashable i1, Eq i2, Hashable i2, Eq i3, Hashable i3
          , Eq i4, Hashable i4
          )
       => Indexed a i1 i2 i3 i4 -> Int -> (a -> a) 
       -> Maybe (Indexed a i1 i2 i3 i4)
adjust (Indexed d a iis vf vi is s) pos f = do
    e <- me
    let newe = f e
    return $ Indexed newdata a
                     (_updateI iis pos e newe)
                     vf (_updateF vf vi pos e newe)
                     is s
    where
        (me, newdata) = IM.updateLookupWithKey (\_ v -> Just $ f v) pos d
{-# INLINABLE adjust #-}

update :: ( Eq a
          , Eq i1, Hashable i1, Eq i2, Hashable i2, Eq i3, Hashable i3
          , Eq i4, Hashable i4
          )
       => Indexed a i1 i2 i3 i4 -> Int -> a 
       -> Maybe (Indexed a i1 i2 i3 i4)
update i p a = adjust i p (const a)
{-# INLINABLE update #-}

adjustWithElement :: ( Eq a
                     , Eq i1, Hashable i1, Eq i2, Hashable i2
                     , Eq i3, Hashable i3, Eq i4, Hashable i4
                     )
                  => Indexed a i1 i2 i3 i4 -> Int -> (a -> a) 
                  -> Maybe (Indexed a i1 i2 i3 i4, a)
adjustWithElement (Indexed d a iis vf vi is s) pos f = do
    e <- me
    let newe = f e
    return (Indexed newdata a (_updateI iis pos e newe) vf 
                              (_updateF vf vi pos e newe) is s
           , newe)
    where
        (me, newdata) = IM.updateLookupWithKey (\_ v -> Just $ f v) pos d
{-# INLINABLE adjustWithElement #-}



instance ( SafeCopy a, Eq a,  SimpleIx a
         , Eq i1, Hashable i1, Eq i2, Hashable i2, Eq i3, Hashable i3
         , Eq i4, Hashable i4
         , i1 ~ I1 a, i2 ~ I2 a, i3 ~ I3 a, i4 ~ I4 a
         )  => SafeCopy (Indexed a i1 i2 i3 i4) where
    version = 0
    kind = base
    putCopy (Indexed d _ _ _ _ _ _) = contain $ safePut $ IM.elems d
    getCopy = contain $ fmap withData safeGet
        where
            withData :: (SafeCopy a, Eq a,  SimpleIx a
                        , Eq i1, Hashable i1, Eq i2, Hashable i2, Eq i3
                        , Hashable i3, Eq i4, Hashable i4
                        , i1 ~ I1 a, i2 ~ I2 a, i3 ~ I3 a
                        , i4 ~ I4 a) => IntMap a -> Indexed a i1 i2 i3 i4
            withData h = IM.foldl insert create h