module Data.Incremental.Map (

    -- * Changes

    insert,
    delete,
    id,

    -- * Atomic changes

    AtomicChange (Insert, Delete),

    -- * Transformations

    member,
    lookup,
    filter,
    map,
    partition,
    union,
    difference,
    intersection,
    isSubmapOf,
    keysSet,
    split,
    mapKeys

) where

import           Prelude hiding (id,
                                lookup,
                                filter,
                                map)
import           Data.Map (Map)
import qualified Data.Map as Map
import           Data.DList (DList)
import qualified Data.DList as DList
import           Data.Maybe (Maybe)
import qualified Data.Maybe as Maybe
import           Data.MultiChange (MultiChange)
import qualified Data.MultiChange as MultiChange
import           Data.Incremental
import qualified Data.Incremental.Tuple as Tuple
import           Data.Set (Set)
import qualified Data.Set as Set
import qualified Data.Incremental.Set as IncSet

-- maps implementation
data AtomicChange k v = Insert k v | Delete k deriving Show

-- teeme mapi muudetavaks
instance (Ord k) => Changeable (Map k v) where

    type DefaultChange (Map k v) = MultiChange (AtomicChange k v)

insert :: (Ord k) => k -> v -> DefaultChange (Map k v)
insert k v = MultiChange.singleton (Insert k v)

delete :: (Ord k) => k -> DefaultChange (Map k v)
delete k = MultiChange.singleton (Delete k)

id :: (Ord k) => DefaultChange (Map k v)
id = MultiChange.fromList []


instance (Ord k) => Change (AtomicChange k v) where

    type Value (AtomicChange k v) = Map k v
    Insert k v $$ m = Map.insert k v m
    Delete k $$ m = Map.delete k m

append :: AtomicChange k v -> MultiChange (AtomicChange k v) -> MultiChange (AtomicChange k v)
append change changes = mappend (MultiChange.singleton change) changes

-- teeme maybe muudetavaks
instance Changeable (Maybe a)


empty :: Map k v
empty = Map.empty


singleton :: (Ord k, Changeable k, Changeable v) => (k,v) ->> Map k v
singleton = MultiChange.composeMap $ stateTrans initSingleton propSingleton

type SingletonState k v = Map k v

initSingleton :: (Ord k) => (k,v) -> (Map k v,SingletonState k v)
initSingleton (k,v) = (result,state) where
    result = Map.singleton k v
    state = result

propSingleton :: (Ord k, Changeable k, Changeable v) => Tuple.AtomicChange k v -> SingletonState k v -> (MultiChange (AtomicChange k v), SingletonState k v)
propSingleton multiChange m = 
    case multiChange of
        (Tuple.First change)  -> (mapChange, mapChange $$ m) where 
            mapChange = MultiChange.fromList [(Delete k),(Insert (change $$ k) v)]
        (Tuple.Second change) -> (mapChange, mapChange $$ m) where 
            mapChange = (insert k (change $$ v))
    where (k,v) = Map.elemAt 0 m

--member 
type MemberState k = k

member :: (Ord k) => k -> (Map k v) ->> Bool
member k = MultiChange.composeMap $ stateTrans (initMember k) propMember

initMember :: (Ord k) => k -> Map k v -> (Bool, MemberState k)
initMember k m = (isMember, state) where
    isMember = Map.member k m
    state    = k

propMember :: (Ord k) => (AtomicChange k v) -> MemberState k -> (PrimitiveChange Bool, MemberState k)
propMember (Insert k v) k'  
        | k == k'   = (ReplaceBy True,k')
        | otherwise = (Keep,k')
propMember (Delete k) k' 
        | k == k'   =(ReplaceBy False,k')
        | otherwise = (Keep,k')

--lookup 
type LookupState k = k

lookup :: Ord k => k -> (Map k v) ->> Maybe v
lookup k = MultiChange.composeMap $ stateTrans (initLookup k) propLookup

initLookup :: (Ord k) => k -> Map k v -> (Maybe v, LookupState k)
initLookup k m = (result, state)
    where result = Map.lookup k m
          state  = k

propLookup :: Ord k => AtomicChange k v -> LookupState k -> (PrimitiveChange (Maybe v), LookupState k)
propLookup (Insert k v) k'  
        | k == k'   = (ReplaceBy (Just v),k')
        | otherwise = (Keep,k')
propLookup (Delete k) k' 
        | k == k'   =(ReplaceBy Nothing,k')
        | otherwise = (Keep,k')



--filter
type FilterState v = (v -> Bool)

filter :: (Ord k) => (v -> Bool) -> Map k v ->> Map k v
filter f = MultiChange.composeMap $ stateTrans (initFilter f) propFilter

initFilter :: (Ord k) => (v -> Bool) -> Map k v -> (Map k v, FilterState v)
initFilter f m = (result, state) 
    where result = Map.filter f m
          state  = f

propFilter :: (Ord k) =>  AtomicChange k v -> FilterState v -> (MultiChange (AtomicChange k v), FilterState v)
propFilter (Insert k v) f
        | f v           = (insert k v, f)
        | otherwise     = (delete k, f)
propFilter (Delete k) f = (delete k, f)


--map
type MapState v a = (v -> a)

map :: (Ord k) => (v -> a) -> Map k v ->> Map k a
map f = MultiChange.composeMap $ stateTrans (initMap f) propMap

initMap :: (Ord k) => (v -> a) -> Map k v -> (Map k a, MapState v a)
initMap f m = (result,state)
    where result = Map.map f m
          state  = f

propMap :: (Ord k) => AtomicChange k v -> MapState v a -> (MultiChange (AtomicChange k a), MapState v a)
propMap (Insert k v) f = (insert k (f v),f)
propMap (Delete k) f = (delete k,f)


--partition
type PartitionState v = ((v -> Bool))

partition :: (Ord k) => (v -> Bool) -> Map k v ->> (Map k v, Map k v)
partition f = MultiChange.composeMap $ stateTrans (initPartition f) propPartition

initPartition :: (Ord k) => (v -> Bool) -> Map k v -> ((Map k v, Map k v), PartitionState v)
initPartition f m = (result, state) 
    where result = Map.partition f m
          state  = f 

propPartition :: (Ord k) => AtomicChange k v -> PartitionState v -> (MultiChange (Tuple.AtomicChange (Map k v) (Map k v)), PartitionState v)
propPartition mapChange f = 
    case mapChange of 
        (Insert k v)            
            | f v       -> (MultiChange.fromList [Tuple.First (insert k v), Tuple.Second (delete k)],f)
            | otherwise -> (MultiChange.fromList [Tuple.First (delete k), Tuple.Second (insert k v)],f)
        (Delete k)      -> (MultiChange.fromList [Tuple.First (delete k), Tuple.Second (delete k)],f)

--union
type UnionState k v = (Map k v, Map k v)

union :: (Ord k) => (Map k v, Map k v) ->> Map k v
union = MultiChange.composeMap $ stateTrans initUnion propUnion

initUnion :: Ord k => (Map k v, Map k v) -> (Map k v, UnionState k v)
initUnion (l,r) = (result, state)
    where result = Map.union l r
          state = (l,r)

propUnion :: (Ord k) => (Tuple.AtomicChange (Map k v) (Map k v)) -> UnionState k v -> (MultiChange (AtomicChange k v), UnionState k v)
propUnion multiChange state = 
    case multiChange of
        (Tuple.First changes)  -> foldl applyUnionLeft (id,state) changes
        (Tuple.Second changes) -> foldl applyUnionRight (id,state) changes

--hetkel ei näe võimalust neid kokku tõsta, sest need ei ole päris samasugused meetodid, mis tuleneb unioni asümmeetrilisusest.
applyUnionLeft :: Ord k => (MultiChange (AtomicChange k v), UnionState k v) -> AtomicChange k v -> (MultiChange (AtomicChange k v), UnionState k v)
applyUnionLeft (multiChange,(l,r)) change = prop change where
    prop (Delete k)
        | inFirst && inSecond == False = (append change multiChange, (change $$ l, r))
        | inFirst && inSecond          = case secondElem of (Just v) -> (append (Insert k v) multiChange, (change $$ l, r))
        | otherwise                    = (multiChange, (l, r))  
        where   secondElem = Map.lookup k r
                inFirst = Map.member k l
                inSecond = Maybe.isJust secondElem
    prop (Insert k v) = (append change multiChange, (change $$ l, r))  

applyUnionRight :: Ord k => (MultiChange (AtomicChange k v), UnionState k v) -> AtomicChange k v -> (MultiChange (AtomicChange k v), UnionState k v)
applyUnionRight (multiChange,(l,r)) change = prop change where
    prop (Delete k)
        | inFirst = (multiChange, (l, change $$ r))
        | otherwise = (append change multiChange, (l, change $$ r))      
        where   inFirst = Map.member k l                           
    prop (Insert k v)
        | inFirst = (multiChange, (l, change $$ r))
        | otherwise = (append change multiChange, (l, change $$ r))
        where inFirst = Map.member k l 

--difference
type DifferenceState k v = (Map k v, Map k v)

difference :: (Ord k) => (Map k v, Map k v) ->> Map k v
difference = MultiChange.composeMap $ stateTrans initDifference propDifference

initDifference :: Ord k => (Map k v, Map k v) -> (Map k v, DifferenceState k v)
initDifference (l,r) = (result, state)
    where result = Map.difference l r
          state = (l,r)

propDifference :: (Ord k) => (Tuple.AtomicChange (Map k v) (Map k v)) -> DifferenceState k v -> (MultiChange (AtomicChange k v), DifferenceState k v)
propDifference multiChange state = 
    case multiChange of
        (Tuple.First changes)  -> foldl applyDifferenceLeft (id,state) changes
        (Tuple.Second changes) -> foldl applyDifferenceRight (id,state) changes

applyDifferenceLeft :: Ord k => (MultiChange (AtomicChange k v), (Map k v, Map k v)) -> AtomicChange k v -> (MultiChange (AtomicChange k v), (Map k v, Map k v))
applyDifferenceLeft (multiChange,(l,r)) change =  
    case change of
        (Insert k v)
            | inSecond == False  -> (append change multiChange, (change $$ l, r))
            | otherwise          -> (multiChange, (change $$ l, r))   
            where inSecond = Map.member k r
        (Delete k) -> (append change multiChange, (change $$ l, r)) 

applyDifferenceRight :: Ord k => (MultiChange (AtomicChange k v), (Map k v, Map k v)) -> AtomicChange k v -> (MultiChange (AtomicChange k v), (Map k v, Map k v))
applyDifferenceRight (multiChange,(l,r)) change = 
    case change of
        (Insert k v)
            | inFirst   -> (append (Delete k) multiChange, (l, change $$ r))
            | otherwise -> (multiChange, (l, change $$ r))   
            where inFirst = Map.member k l                          
        (Delete k) -> case Map.lookup k l of 
                (Just v) -> (append (Insert k v) multiChange, (l, change $$ r))
                Nothing  -> (multiChange, (l, change $$ r))

--intersection
type IntersectionState k v = (Map k v, Map k v)

intersection :: (Ord k) => (Map k v, Map k v) ->> Map k v
intersection = MultiChange.composeMap $ stateTrans initIntersection propIntersection

initIntersection :: Ord k => (Map k v, Map k v) -> (Map k v, IntersectionState k v)
initIntersection (l,r) = (result, state)
    where result = Map.intersection l r
          state = (l,r)

propIntersection :: (Ord k) => (Tuple.AtomicChange (Map k v) (Map k v)) -> IntersectionState k v -> (MultiChange (AtomicChange k v), IntersectionState k v)
propIntersection multiChange state = 
    case multiChange of
        (Tuple.First changes)  -> foldl applyIntersectionLeft (id,state) changes
        (Tuple.Second changes) -> foldl applyIntersectionRight (id,state) changes

applyIntersectionLeft :: Ord k => (MultiChange (AtomicChange k v), (Map k v, Map k v)) -> AtomicChange k v -> (MultiChange (AtomicChange k v), IntersectionState k v)
applyIntersectionLeft (multiChange,(l,r)) change =  
    case change of
        (Insert k v)
            | inSecond  -> (append change multiChange, (change $$ l, r))
            | otherwise -> (multiChange, (change $$ l, r))   
            where inSecond = Map.member k r
        (Delete k)
            | inSecond  -> (append change multiChange, (change $$ l, r))
            | otherwise -> (multiChange, (change $$ l, r))   
            where inSecond = Map.member k r
            

applyIntersectionRight :: Ord k => (MultiChange (AtomicChange k v), (Map k v, Map k v)) -> AtomicChange k v -> (MultiChange (AtomicChange k v), IntersectionState k v)
applyIntersectionRight (multiChange,(l,r)) change = 
    case change of
        (Insert k v) ->
            case Map.lookup k l of
                (Just v) -> (append (Insert k v) multiChange, (l, change $$ r)) --testimine tuvastas vea mis sai parandatud
                Nothing  -> (multiChange, (l, change $$ r))   
        (Delete k)
            | inFirst   -> (append change multiChange, (l, change $$ r))
            | otherwise -> (multiChange, (l, change $$ r))   
            where inFirst = Map.member k l

--submap
--state: left, right, difference
--efektiivne ainult sellisel juhul, kui hulgad on piisavalt sarnased
--originaalalgoritmi puhul piisab ühest vastunäitest, et anda vastus False, seega vähim
--keerukus on O(1)! Samas allpoololevas tehakse igal juhul kolm O(log n) keerukusega operatsiooni
type SubmapOfState k v = (Map k v, Map k v, Map k v)

isSubmapOf :: (Ord k, Eq v) => (Map k v, Map k v) ->> Bool
isSubmapOf = MultiChange.composeMap $ stateTrans initSubmapOf propSubmapOf

initSubmapOf :: (Ord k, Eq v) => (Map k v, Map k v) -> (Bool, SubmapOfState k v)
initSubmapOf (l,r) = (result, state)
    where result = Map.isSubmapOf l r
          state = (l, r, Map.differenceWith discardEqualValues l r) 

discardEqualValues :: (Eq a) => (a -> a -> Maybe a)
discardEqualValues a b  
    | a == b    = Nothing
    | otherwise = Just a

propSubmapOf :: (Ord k, Eq v) => (Tuple.AtomicChange (Map k v) (Map k v)) -> SubmapOfState k v -> (PrimitiveChange Bool, SubmapOfState k v)
propSubmapOf multiChange state = 
    case multiChange of
        (Tuple.First changes)  -> foldl applySubmapOfLeft (Keep,state) changes
        (Tuple.Second changes) -> foldl applySubmapOfRight (Keep,state) changes

applySubmapOfLeft :: (Ord k, Eq v) =>  (PrimitiveChange Bool, SubmapOfState k v) -> AtomicChange k v -> (PrimitiveChange Bool, SubmapOfState k v )
applySubmapOfLeft (boolChange,(l,r,lrDiff)) change = prop change where 
    prop (Insert k v) = right k v (Map.lookup k r) where
        right k v (Just v') = if (v == v') 
                                then (mappend (Keep) boolChange, (insert k v $$ l, r, delete k $$ lrDiff)) 
                                else (mappend (ReplaceBy False) boolChange, (insert k v $$ l, r, insert k v $$ lrDiff)) 
        right k v (Nothing) = (mappend (ReplaceBy False) boolChange, (insert k v $$ l, r, insert k v $$ lrDiff))
    prop (Delete k) = diff k (Map.lookup k lrDiff) where
        diff k (Just _) = if (length lrDiff == 1) 
                                then (mappend (ReplaceBy True) boolChange, (delete k $$ l, r, delete k $$ lrDiff)) 
                                else (mappend Keep boolChange, (delete k $$ l, r, delete k $$ lrDiff)) 
        diff k (Nothing) = (mappend Keep boolChange, (delete k $$ l, r, delete k $$ lrDiff)) 
       
applySubmapOfRight :: (Ord k, Eq v) =>  (PrimitiveChange Bool, SubmapOfState k v) -> AtomicChange k v -> (PrimitiveChange Bool, SubmapOfState k v )
applySubmapOfRight (boolChange,(l,r,lrDiff)) change = prop change where
    prop (Insert k v) = diff k v (Map.lookup k lrDiff) where
        diff k v (Just v') = if (v == v' && length lrDiff == 1) then (mappend (ReplaceBy True) boolChange, (l,insert k v' $$ r, delete k $$ lrDiff)) 
                                else if (v == v') then (mappend (Keep) boolChange, (l,insert k v $$ r, delete k $$ lrDiff)) 
                                else (mappend Keep boolChange, (l, insert k v $$ r, lrDiff)) 
        diff k v (Nothing) = left k v (Map.lookup k l) where
            left k v (Just v') = if (v == v') then (mappend Keep boolChange, (l, insert k v $$ r, lrDiff)) 
                                else (mappend (ReplaceBy False) boolChange, (l, insert k v $$ r, insert k v' $$ lrDiff))
            left k v (Nothing) = (mappend Keep boolChange, (l, insert k v $$ r, lrDiff))
    prop (Delete k) = left k (Map.lookup k l) where
        left k (Just v') = (mappend (ReplaceBy False) boolChange, (l, delete k $$ r, insert k v' $$ lrDiff)) 
        left k (Nothing) = (mappend (Keep) boolChange, (l, delete k $$ r, lrDiff))


--kas siin võiks kasutada (Map k ()) ? Aga sellisel juhul ei oleks väljund võrreldav Data.Map omaga
--keysset
keysSet :: (Ord k) => Map k v ->> (Set k)
keysSet = MultiChange.bind $ simpleTrans initKeysSet propKeysSet

initKeysSet :: (Ord k) => Map k v -> (Set k)
initKeysSet m = Map.keysSet m

propKeysSet :: (Ord k) => AtomicChange k v -> MultiChange (IncSet.AtomicChange k)
propKeysSet mapChange = 
    case mapChange of 
        (Insert k v) -> MultiChange.singleton (IncSet.Insert k)
        (Delete k)   -> MultiChange.singleton (IncSet.Delete k)

--split
--O(logn ) -> O(log n), mõttetu
type SplitState k = k

split :: (Ord k) => k -> Map k v ->> (Map k v, Map k v)
split k = MultiChange.composeMap $ stateTrans (initSplit k) propSplit

initSplit :: (Ord k) => k -> Map k v -> ((Map k v, Map k v), SplitState k)
initSplit k m = (result,state)
    where result = Map.split k m 
          state  = k

propSplit :: (Ord k) => AtomicChange k v -> SplitState k -> (MultiChange (Tuple.AtomicChange (Map k v) (Map k v)), SplitState k)
propSplit change state = 
    case change of
        (Insert k v)            
            | k < state  -> (MultiChange.fromList [Tuple.First (insert k v)], state)
            | k > state  -> (MultiChange.fromList [Tuple.Second (insert k v)], state)
            | otherwise  -> (MultiChange.fromList [], state)
        (Delete k)
            | k < state  -> (MultiChange.fromList [Tuple.First (delete k)], state)
            | k > state  -> (MultiChange.fromList [Tuple.Second (delete k)], state)
            | otherwise  -> (MultiChange.fromList [], state)


--mapkeys 
--values of the new keys are collapsed
--to the value of largest key (before f is applied)
type MapKeysState a b v = ((a -> b), Map b (Map a v))

mapKeys :: (Ord a, Ord b) => (a -> b) -> Map a v ->> Map b v
mapKeys fun = MultiChange.composeMap $ stateTrans (initMapKeys fun) propMapKeys

initMapKeys :: (Ord a, Ord b) => (a -> b) -> Map a v -> (Map b v, MapKeysState a b v)
initMapKeys fun m = (result,state) where
    result = Map.mapKeys fun m
    state  = (fun, (Map.foldrWithKey (groupKeys fun) Map.empty m))

--creates a map of keys and list of values
groupKeys :: (Ord a, Ord b) => (a -> b) -> a -> v -> Map b (Map a v) -> Map b (Map a v)
groupKeys fun key value bMap = 
    let fKey        = fun key
        entry       = Map.lookup fKey bMap
        addTo Nothing = Map.insert fKey (Map.singleton key value) bMap
        addTo (Just aMap) = Map.insert fKey (Map.insert key value aMap) bMap
    in (addTo entry)

propMapKeys :: (Ord a, Ord b) => AtomicChange a v -> MapKeysState a b v -> (MultiChange (AtomicChange b v), MapKeysState a b v)
propMapKeys change state = prop change state where
    prop (Insert k v) (f,bMap) = (chg', (f,bMap')) where
            bMap' = updateBmap $ Map.lookup (f k) bMap
            chg' = chg $ initVal $ Map.lookup (f k) bMap'
            updateBmap Nothing  = Map.insert (f k) (Map.singleton k v) bMap   
            updateBmap (Just aMap) = Map.insert (f k) (Map.insert k v aMap) bMap 
            chg Nothing = insert (f k) v
            chg (Just (_,v')) = insert (f k) v'
            initVal (Nothing) = Nothing
            initVal (Just m) = if (length m > 0) 
                                then Just (Map.findMax m) 
                                else Nothing
            --NOTE have to make sure, that the order of evaluation is correct
            --this can be achieved, by giving bMap' as an argument to chg'

                       --loogika: 
               --1. uuendan bMap
               --2. otsin uue key järgi bMapist kõiki vanu vasteid
               --3. kui ei ole, siis võib teha chg Delete
               --4. kui on, siis tuleb teha suurimale vanale vastele Insert
    prop (Delete k) (f,bMap) = (chg', (f,bMap')) where  
            bMap' = updateBmap $ Map.lookup (f k) bMap 
            chg' = chg $ initVal $ Map.lookup (f k) bMap'                  
            updateBmap Nothing  = bMap                                            
            updateBmap (Just aMap) = Map.insert (f k) (Map.delete k aMap) bMap   
            chg Nothing = delete (f k)  
            chg (Just (k',v')) = insert (f k) v'
            initVal (Nothing) = Nothing
            initVal (Just m) = if (length m > 0) 
                                then Just (Map.findMax m) 
                                else Nothing
            --loogika: 
               --1. uuendan bMap
               --2. otsin uue key järgi bMapist kõiki vanu vasteid
               --3. kui ei ole, siis võib teha chg Delete
               --4. kui on, siis tuleb teha suurimale vanale vastele Insert