{-# LANGUAGE OverloadedStrings #-}

module Funcons.Operations.Maps where

import Funcons.Operations.Booleans
import Funcons.Operations.Internal
import Funcons.Operations.Sets
import qualified Data.Map as M
import qualified Data.Set as S 

library :: (HasValues t, Ord t) => Library t
library = libFromList [
    ("map-empty", NullaryExpr map_empty)
  , ("map-singleton", BinaryExpr map_singleton)
  , ("is-map-empty", UnaryExpr is_map_empty)
  , ("map-insert", TernaryExpr map_insert)
  , ("map-lookup", BinaryExpr map_lookup)
  , ("lookup", BinaryExpr map_lookup)
  , ("map-domain", UnaryExpr domain)
  , ("domain", UnaryExpr domain)
  , ("map-delete", BinaryExpr map_delete)
  , ("is-in-domain", BinaryExpr is_in_domain)
  , ("map-unite", NaryExpr map_unite)
  , ("map-override", BinaryExpr map_override)
  , ("maps", BinaryExpr maps)
  , ("map", NaryExpr map_)
  , ("map-elements", UnaryExpr map_elements)
  , ("map-points", UnaryExpr map_points)
  ]

map_ :: (Ord t, HasValues t) => [OpExpr t] -> OpExpr t
map_ = vNaryOp "map" (Normal . inject . Map . M.fromList . mkPairs)

mkPairs :: [a] -> [(a,a)]
mkPairs []        = []
mkPairs [x]       = []
mkPairs (x:y:ys)  = (x,y) : mkPairs ys

maps_ :: HasValues t => [OpExpr t] -> OpExpr t
maps_ = binaryOp maps
maps :: HasValues t => OpExpr t -> OpExpr t -> OpExpr t
maps = vBinaryOp "maps" op
  where op (ComputationType (Type t1)) (ComputationType (Type t2)) = Normal $ injectT (Maps t1 t2)
        op _ _ = SortErr "maps not applied to two types"

map_empty_ :: HasValues t => [OpExpr t] -> OpExpr t 
map_empty_ = nullaryOp map_empty
map_empty :: HasValues t => OpExpr t 
map_empty = NullaryOp "map-empty" (Normal $ inject (Map M.empty))

map_singleton_ :: (HasValues t,Ord t) => [OpExpr t] -> OpExpr t
map_singleton_ = binaryOp map_singleton
map_singleton :: (HasValues t, Ord t) => OpExpr t -> OpExpr t -> OpExpr t 
map_singleton k v = RewritesTo "map-insert" (map_insert map_empty k v) [k,v]

is_map_empty_ :: HasValues t => [OpExpr t] -> OpExpr t
is_map_empty_ = unaryOp is_map_empty
is_map_empty :: HasValues t => OpExpr t -> OpExpr t 
is_map_empty = vUnaryOp "is-map-empty" op 
  where op (Map m)  = Normal $ inject $ tobool (null m)
        op _        = SortErr "is-map-empty(M) not applied to a map"

map_insert_ :: (HasValues t, Ord t) => [OpExpr t] -> OpExpr t
map_insert_ = ternaryOp map_insert 
map_insert :: (HasValues t, Ord t) => OpExpr t -> OpExpr t -> OpExpr t -> OpExpr t
map_insert = vTernaryOp "map-insert" op
  where op xv k v = case xv of 
              Map m -> Normal $ inject $ Map (M.insert k v m)
              _     -> SortErr "map-insert(M,K,V) not applied to a map (first argument)"

map_lookup_ :: (HasValues t, Ord t) => [OpExpr t] -> OpExpr t
map_lookup_ = binaryOp map_lookup
map_lookup :: (HasValues t, Ord t) => OpExpr t -> OpExpr t -> OpExpr t
map_lookup = vBinaryOp "map-lookup" op
  where op xv k = case xv of 
                    Map m -> case M.lookup k m of 
                        Nothing -> DomErr "key not in domain"
                        Just v  -> Normal $ inject v
                    _ -> SortErr "map-lookup(M,V) not applied to a map and a value"

map_delete_ :: (HasValues t, Ord t) => [OpExpr t] -> OpExpr t
map_delete_ = binaryOp map_delete
map_delete :: (HasValues t, Ord t) => OpExpr t -> OpExpr t -> OpExpr t
map_delete = vBinaryOp "map-delete" op
  where op (Map m) (Set s) = Normal $ inject $ Map (foldr M.delete m s)
        op _ _ = SortErr "map-delete(M,S) not applied to a map and a set"

is_in_domain_ :: (Ord t, HasValues t) => [OpExpr t] -> OpExpr t 
is_in_domain_ = binaryOp is_in_domain
is_in_domain :: (Ord t, HasValues t) => OpExpr t -> OpExpr t -> OpExpr t 
is_in_domain x y = RewritesTo "is-in-domain" (is_in_set x (domain y)) [x,y] 

domain_ :: (HasValues t, Ord t) => [OpExpr t] -> OpExpr t
domain_ = unaryOp domain
domain :: (HasValues t, Ord t) => OpExpr t -> OpExpr t
domain = vUnaryOp "domain" op
  where op (Map m)  = Normal $ inject $ Set $ S.fromList $ M.keys m
        op _        = SortErr "domain(M) not applied to a map"

map_override_ :: (HasValues t, Ord t) => [OpExpr t] -> OpExpr t
map_override_ = binaryOp map_override
map_override :: (HasValues t, Ord t) => OpExpr t -> OpExpr t -> OpExpr t
map_override = vBinaryOp "map-override" op
  where op (Map m1) (Map m2) = Normal $ inject $ Map (M.union m1 m2)
        op _ _ = SortErr "map-override(M,M) not applied tOpExpr two maps"

map_unite_ :: (HasValues t, Ord t) => [OpExpr t] -> OpExpr t
map_unite_ = map_unite 
map_unite :: (HasValues t, Ord t) => [OpExpr t] -> OpExpr t
map_unite = vNaryOp "map-unite" op
  where op args 
          | all isMap args =
              let maps = map toMap args
                  domains = map (M.keysSet) maps
              in if all (null . uncurry S.intersection) (allDomainPairs domains)
                  then Normal $ inject $ Map $ M.unions maps
                  else DomErr "union with domain intersection"
          | otherwise     = SortErr "map-unite(M1,...,Mn) not applied to maps"
          where isMap (Map _) = True
                isMap _       = False
                toMap (Map m) = m
                toMap _       = error "map_unite"

allDomainPairs :: [a] -> [(a,a)] 
allDomainPairs (x:xs) = [ (x,y)  | y <- xs ] ++ allDomainPairs xs
allDomainPairs [] = []

map_elements_ :: (Ord t, HasValues t) => [OpExpr t] -> OpExpr t
map_elements_ = unaryOp map_elements
map_elements :: (Ord t, HasValues t) => OpExpr t -> OpExpr t
map_elements = vUnaryOp "map-elements" op
  where op (Map m) = Normal $ inject $ ADTVal "list" (map inject $ M.foldrWithKey combine [] m)
          where combine k v ls = k:v:ls
        op _ = SortErr "map-elements not applied to a map"

map_points_ :: (Ord t, HasValues t) => [OpExpr t] -> OpExpr t
map_points_ = unaryOp map_points
map_points :: (Ord t, HasValues t) => OpExpr t -> OpExpr t
map_points = vUnaryOp "map-points" op
  where op (Map m) = Normal $ inject $ ADTVal "list" (map inject $ M.foldrWithKey combine [] m)
          where combine k v ls = ADTVal "tuple" [inject k, inject v]:ls
        op _ = SortErr "map-points not applied to a map"