{-# LANGUAGE MultiParamTypeClasses, FlexibleInstances #-}

module RAExpr.RASet (RASet, fromList, fromSet, empty, attributes, RAExpr.RASet.filter,
  project, union, intersection, difference, crossProduct,
  renameColumn, renameColumns, natualJoin) where

import Control.Exception.Base (assert)
import qualified Data.Set as Set
import Data.Set (Set)
import Data.Function (on)
import RAExpr.Row(Row)
import qualified RAExpr.Row as Row
import Function(around)
import Class.Pretty
import Data.List (intercalate)

type Names names = Set names
type Rows names values = Set (Row names values)

data RASet names values = RASet (Names names) (Rows names values)
                    deriving (Eq, Ord)

-- | Instance of Show.
-- >>> RASet (Set.fromList ["a", "b"]) (Set.fromList [Row.singleton "a" "x", Row.singleton "b" "y"])
-- fromList ["a","b"] [fromList [("a","x")],fromList [("b","y")]]
instance (Show names, Show values) => Show (RASet names values)  where
  show (RASet names rows) = "fromList " ++ show (Set.toList names) ++ " " ++ show (Set.toList rows)

-- | Instance of Pretty
-- >>> pretty (empty::(RASet String String))
-- "{|}"
-- >>> pretty (fromList ["A"] ([]::[Row String String]))
-- "{A|}"
-- >>> pretty (fromList ["A"] ([Row.empty]::[Row String String]))
-- "{A|}"
-- >>> pretty (fromList ["A"] ([Row.singleton "A" "x"]))
-- "{A|(\"x\")}"
-- >>> pretty (fromList ["A","B"] ([Row.fromList [("A","x"),("B","y")]]))
-- "{A,B|(\"x\",\"y\")}"
-- >>> pretty (fromList ["A"] [Row.fromList [("A","x")], Row.fromList [("A","y")]])
-- "{A|(\"x\"),(\"y\")}"
instance Pretty (RASet String String) where
  pretty set = "{" ++ commaSeperatedNames ++ "|" ++ commaSeperatedRows ++ "}"
    where commaSeperatedNames = intercalate "," (Set.toList names)
          commaSeperatedRows = intercalate "," (Set.toList prettyRows)
          prettyRows = Set.map pretty (rows set)
          names = attributes set

-- instance (Ord names) => Foldable (RASet names) where
-- foldr f x (RASet names rows) = Set.foldr f x rows

-- | Set factory function
-- >>> empty
-- fromList [] []
empty:: (Ord names, Ord values) => RASet names values
empty = RASet Set.empty Set.empty

-- | RASet factory function
-- >>> singleton [] (Row.empty)
-- fromList [] []
-- >>> singleton [] (Row.singleton "a" 1)
-- fromList [] []
-- >>> singleton ["a"] (Row.singleton "a" 1)
-- fromList ["a"] [fromList [("a",1)]]
singleton:: (Ord names, Ord values) => [names] -> Row names values -> RASet names values
singleton attributes row = fromList attributes [row]

-- | RASet factory function
-- >>> fromList [] [Row.fromList [("a","x"),("b","y")]]
-- fromList [] []
-- >>> fromList ["a"] [Row.empty]
-- fromList ["a"] []
-- >>> fromList ["a"] [Row.singleton "a" "z", Row.fromList [("a","x"),("b","y")]]
-- fromList ["a"] [fromList [("a","x")],fromList [("a","z")]]
fromList:: (Ord names, Ord values) => [names] -> [Row names values] -> RASet names values
fromList attributes rows 
  = fromSet (Set.fromList attributes) (Set.fromList rows)

-- | Set factory function
-- >>> fromSet Set.empty Set.empty
-- fromList [] []
fromSet:: (Ord names, Ord values) => Names names -> Rows names values -> RASet names values
fromSet attributes rows = let
    mask = Set.toList attributes
    maskedRows = Set.map (Row.mask mask) rows
    filteredRows = Set.filter (not . null) maskedRows
  in RASet attributes filteredRows

-- | Get attributes of Set.
-- >>> attributes (fromList [] [Row.empty])
-- fromList []
-- >>> attributes (fromList ["a"] [Row.fromList []])
-- fromList ["a"]
attributes:: RASet names value -> Names names
attributes (RASet names _) = names

-- | Get rows of RASet.
-- >>> rows (fromList ["a"] [Row.fromList []])
-- fromList []
-- >>> rows (fromList ["a"] [Row.singleton "a" "x"])
-- fromList [fromList [("a","x")]]
rows:: RASet names values -> Rows names values
rows (RASet _ rows) = rows

filter f (RASet attributes rows) = RASet attributes (Set.filter f rows)

-- | Change the name of a column.
-- >>> renameColumn "a" "c" empty
-- fromList [] []
-- >>> renameColumn "a" "c" (fromList ["a"] [Row.singleton "a" "x"])
-- fromList ["c"] [fromList [("c","x")]]
renameColumn:: (Ord names, Ord values) => names -> names -> RASet names values -> RASet names values 
renameColumn oldName newName set = let
    oldAttributes = attributes set
    newAttributes = if oldName `Set.member` oldAttributes then
        Set.insert newName (Set.delete oldName oldAttributes)
      else
        oldAttributes
    newRows = Set.map (Row.renameColumn oldName newName) (rows set)
  in RASet newAttributes newRows

-- | Change the name of many columns.
-- >>> renameColumns [("a","c")] (fromList ["a"] [Row.singleton "a" "x"])
-- fromList ["c"] [fromList [("c","x")]]
renameColumns:: (Ord names, Ord values) => [(names, names)] -> RASet names values -> RASet names values
renameColumns changes set = 
  foldl (\set change -> (uncurry renameColumn) change set) set changes

-- | Project the column names on to a set.
-- >>> project (Set.singleton "a") (fromList ["a", "b"] [Row.singleton "a" "x", Row.singleton "b" "y"])
-- fromList ["a"] [fromList [("a","x")]]
project:: (Ord names, Ord values) => Names names -> RASet names values -> RASet names values
project attributes (RASet _ rows) = fromSet attributes rows

-- | Union of two union compatible sets.
-- union (fromList ["a"] [Row.singleton "a" "x"]]) (fromList ["a"] [Row.singleton "a" "y"]])
-- fromList ["a"] [fromList [("a","x")],fromList [("a","y")]]
union:: (Ord names, Ord values) => RASet names values -> RASet names values -> RASet names values
union = unionCompatibleOperation (Set.union `on` rows)

-- | Difference of two union compatible sets.
-- differance (fromList ["a"] [Row.singleton "a" "x"]]) (fromList ["a"] [Row.singleton "a" "x"]])
-- fromList ["a"] []
-- differance (fromList ["a"] [Row.singleton "a" "x"]]) (fromList ["a"] [])
-- fromList ["a"] [fromList [("a", "x")]]
difference:: (Ord names, Ord values) => RASet names values -> RASet names values -> RASet names values
difference = unionCompatibleOperation (Set.difference `on` rows)

-- | Intersection of two union compatible sets.
-- intersection (fromList ["a"] [Row.singleton "a" "x"]) (fromList ["a"] [])
-- fromList ["a"] [fromList [("a", "x")]]
-- intersection (fromList ["a"] [Row.singleton "a" "x"]]) (fromList ["a"] [Row.singleton "a" "y"]])
-- fromList ["a"] [fromList [("a","x")],fromList [("a","y")]]
intersection:: (Ord names, Ord values) => RASet names values -> RASet names values -> RASet names values
intersection = unionCompatibleOperation (Set.intersection `on` rows)

-- | CrossProduct of two anti union compatible sets.
-- >>> crossProduct (fromList ["a"] [Row.singleton "a" "x"]) (fromList ["b"] [Row.singleton "b" "y"])
-- fromList ["a","b"] [fromList [("a","x"),("b","y")]]
-- >>> crossProduct (fromList ["a"] []) (fromList ["b"] [Row.singleton "b" "y"])
-- fromList ["a","b"] []
-- >>> crossProduct (fromList ["a"] [Row.singleton "a" "x"]) (fromList ["b"] [])
-- fromList ["a","b"] []
crossProduct:: (Ord names, Ord values) => RASet names values -> RASet names values -> RASet names values
crossProduct = 
  let 
    crossProduct leftSet rightSet = Set.foldl (folder rightSet) Set.empty leftSet
    folder otherSet set row = Set.union (crossRow row otherSet) set
    crossRow row set = Set.map (Row.union row) set
  in antiUnionCompatibleOperation (crossProduct `on` rows)

natualJoin:: (Ord names, Ord values) => RASet names values -> RASet names values -> RASet names values
natualJoin set1 set2 = undefined

-- | Test if two RASets are union compatible.
-- >>> unionCompatible (fromList ["a"] [Row.empty]) (fromList ["a"] [Row.empty])
-- True
-- >>> unionCompatible (fromList ["a"] [Row.empty]) (fromList ["b"] [Row.empty])
-- False
-- >>> unionCompatible (fromList ["a", "b"] [Row.empty]) (fromList ["b", "c"] [Row.empty])
-- False
unionCompatible:: (Ord names, Ord values) => RASet names values -> RASet names values -> Bool
unionCompatible = (==) `on` attributes

-- | Test if two RASets are not union compatible.
-- >>> notUnionCompatible (fromList ["a"] [Row.empty]) (fromList ["a"] [Row.empty])
-- False
-- >>> notUnionCompatible (fromList ["a"] [Row.empty]) (fromList ["b"] [Row.empty])
-- True
-- >>> notUnionCompatible (fromList ["a", "b"] [Row.empty]) (fromList ["b", "c"] [Row.empty])
-- True
notUnionCompatible:: (Ord names, Ord values) => RASet names values -> RASet names values -> Bool
notUnionCompatible = not `around` unionCompatible

-- | Test if two RASets are anti union compatible.
-- >>> antiUnionCompatible (fromList ["a"] [Row.empty]) (fromList ["a"] [Row.empty])
-- False
-- >>> antiUnionCompatible (fromList ["a"] [Row.empty]) (fromList ["b"] [Row.empty])
-- True
-- >>> antiUnionCompatible (fromList ["a", "b"] [Row.empty]) (fromList ["b", "c"] [Row.empty])
-- False
antiUnionCompatible:: (Ord names, Ord values) => RASet names values -> RASet names values -> Bool
antiUnionCompatible = (Set.empty==) `around` (Set.intersection `on` attributes)

-- | Assert two sets are union compatible.
-- >>> assertUnionCompatible empty empty True
-- True
assertUnionCompatible:: (Ord names, Ord values) => RASet names values -> RASet names values -> a -> a
assertUnionCompatible = assert `around` unionCompatible

-- | Assert two sets are anti Union compatible.
-- >>> assertAntiUnionCompatible (fromList ["a"] []) (fromList ["b"] []) True
-- True
assertAntiUnionCompatible:: (Ord names, Ord values) => RASet names values -> RASet names values -> a -> a
assertAntiUnionCompatible = assert `around` antiUnionCompatible

-- | Preform an opperation if the two sets are union compatible.
-- >>> unionCompatibleOperation (\x y -> Set.empty) (fromList ["a"] []) (fromList ["a"] [])
-- fromList ["a"] []
unionCompatibleOperation:: (Ord n, Ord v) => (RASet n v -> RASet n v -> Rows n v) -> RASet n v-> RASet n v -> RASet n v
unionCompatibleOperation operator set1 set2 = 
  let newRows = assertUnionCompatible set1 set2 (set1 `operator` set2)
      newAttributes = attributes set1
  in RASet newAttributes newRows

-- | Preform an opperation if the two sets are anti union compatible.
-- >>> antiUnionCompatibleOperation (\x y -> Set.empty) (fromList ["a"] []) (fromList ["b"] [])
-- fromList ["a","b"] []
antiUnionCompatibleOperation:: (Ord n, Ord v) => (RASet n v -> RASet n v-> Rows n v) -> RASet n v-> RASet n v -> RASet n v
antiUnionCompatibleOperation operator set1 set2 = 
  let newRows = assertAntiUnionCompatible set1 set2 (set1 `operator` set2)
      newAttributes = (Set.union `on` attributes) set1 set2
  in RASet newAttributes newRows

