-- |
-- Module      :  ELynx.Tree.Partition
-- Description :  Partitions on rose trees
-- Copyright   :  (c) Dominik Schrempf 2021
-- License     :  GPL-3.0-or-later
--
-- Maintainer  :  dominik.schrempf@gmail.com
-- Stability   :  unstable
-- Portability :  portable
--
-- Creation date: Thu Dec 12 12:58:49 2019.
module ELynx.Tree.Partition
  ( -- * Data type
    Partition (fromPartition),
    pt,
    ptUnsafe,
    bpToPt,
    ptHuman,

    -- * Work with 'Partition's
    partition,
    partitions,
    compatible,
  )
where

import Data.List hiding (partition)
import Data.Set (Set)
import qualified Data.Set as S
import ELynx.Tree.Bipartition
import ELynx.Tree.Rooted

-- | A partition of a tree is a grouping of the leaves of the tree into
-- non-overlapping, non-empty sub sets.
--
-- For unrooted trees:
--
-- - For example, each branch of an unrooted tree partitions the leaves of the
--   tree into two subsets (see 'ELynx.Tree.Bipartition').
--
-- For rooted trees:
--
-- - In a similar way, each bifurcating internal node (excluding the root node)
--   partitions the leaves into three subsets called a 'Partition'. If the tree
--   is multifurcating, and a specific node has more than two children, the
--   number of subsets induced by this node is larger than three. Partitions are
--   interesting in that we can use them for calculating incompatible splits,
--   see 'ELynx.Tree.Distance'.
--
-- The order of the subsets of a 'Partition' is meaningless. We ensure by
-- construction that the subsets are ordered, and hence, that equality checks
-- are meaningful.
newtype Partition a = Partition
  { Partition a -> Set (Set a)
fromPartition :: Set (Set a)
  }
  deriving (Partition a -> Partition a -> Bool
(Partition a -> Partition a -> Bool)
-> (Partition a -> Partition a -> Bool) -> Eq (Partition a)
forall a. Eq a => Partition a -> Partition a -> Bool
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
/= :: Partition a -> Partition a -> Bool
$c/= :: forall a. Eq a => Partition a -> Partition a -> Bool
== :: Partition a -> Partition a -> Bool
$c== :: forall a. Eq a => Partition a -> Partition a -> Bool
Eq, Eq (Partition a)
Eq (Partition a)
-> (Partition a -> Partition a -> Ordering)
-> (Partition a -> Partition a -> Bool)
-> (Partition a -> Partition a -> Bool)
-> (Partition a -> Partition a -> Bool)
-> (Partition a -> Partition a -> Bool)
-> (Partition a -> Partition a -> Partition a)
-> (Partition a -> Partition a -> Partition a)
-> Ord (Partition a)
Partition a -> Partition a -> Bool
Partition a -> Partition a -> Ordering
Partition a -> Partition a -> Partition a
forall a.
Eq a
-> (a -> a -> Ordering)
-> (a -> a -> Bool)
-> (a -> a -> Bool)
-> (a -> a -> Bool)
-> (a -> a -> Bool)
-> (a -> a -> a)
-> (a -> a -> a)
-> Ord a
forall a. Ord a => Eq (Partition a)
forall a. Ord a => Partition a -> Partition a -> Bool
forall a. Ord a => Partition a -> Partition a -> Ordering
forall a. Ord a => Partition a -> Partition a -> Partition a
min :: Partition a -> Partition a -> Partition a
$cmin :: forall a. Ord a => Partition a -> Partition a -> Partition a
max :: Partition a -> Partition a -> Partition a
$cmax :: forall a. Ord a => Partition a -> Partition a -> Partition a
>= :: Partition a -> Partition a -> Bool
$c>= :: forall a. Ord a => Partition a -> Partition a -> Bool
> :: Partition a -> Partition a -> Bool
$c> :: forall a. Ord a => Partition a -> Partition a -> Bool
<= :: Partition a -> Partition a -> Bool
$c<= :: forall a. Ord a => Partition a -> Partition a -> Bool
< :: Partition a -> Partition a -> Bool
$c< :: forall a. Ord a => Partition a -> Partition a -> Bool
compare :: Partition a -> Partition a -> Ordering
$ccompare :: forall a. Ord a => Partition a -> Partition a -> Ordering
$cp1Ord :: forall a. Ord a => Eq (Partition a)
Ord, Int -> Partition a -> ShowS
[Partition a] -> ShowS
Partition a -> String
(Int -> Partition a -> ShowS)
-> (Partition a -> String)
-> ([Partition a] -> ShowS)
-> Show (Partition a)
forall a. Show a => Int -> Partition a -> ShowS
forall a. Show a => [Partition a] -> ShowS
forall a. Show a => Partition a -> String
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
showList :: [Partition a] -> ShowS
$cshowList :: forall a. Show a => [Partition a] -> ShowS
show :: Partition a -> String
$cshow :: forall a. Show a => Partition a -> String
showsPrec :: Int -> Partition a -> ShowS
$cshowsPrec :: forall a. Show a => Int -> Partition a -> ShowS
Show, ReadPrec [Partition a]
ReadPrec (Partition a)
Int -> ReadS (Partition a)
ReadS [Partition a]
(Int -> ReadS (Partition a))
-> ReadS [Partition a]
-> ReadPrec (Partition a)
-> ReadPrec [Partition a]
-> Read (Partition a)
forall a. (Read a, Ord a) => ReadPrec [Partition a]
forall a. (Read a, Ord a) => ReadPrec (Partition a)
forall a. (Read a, Ord a) => Int -> ReadS (Partition a)
forall a. (Read a, Ord a) => ReadS [Partition a]
forall a.
(Int -> ReadS a)
-> ReadS [a] -> ReadPrec a -> ReadPrec [a] -> Read a
readListPrec :: ReadPrec [Partition a]
$creadListPrec :: forall a. (Read a, Ord a) => ReadPrec [Partition a]
readPrec :: ReadPrec (Partition a)
$creadPrec :: forall a. (Read a, Ord a) => ReadPrec (Partition a)
readList :: ReadS [Partition a]
$creadList :: forall a. (Read a, Ord a) => ReadS [Partition a]
readsPrec :: Int -> ReadS (Partition a)
$creadsPrec :: forall a. (Read a, Ord a) => Int -> ReadS (Partition a)
Read)

-- | Create a partition.
pt :: Ord a => [Set a] -> Either String (Partition a)
pt :: [Set a] -> Either String (Partition a)
pt [Set a]
xs = case (Set a -> Bool) -> [Set a] -> [Set a]
forall a. (a -> Bool) -> [a] -> [a]
filter (Bool -> Bool
not (Bool -> Bool) -> (Set a -> Bool) -> Set a -> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Set a -> Bool
forall a. Set a -> Bool
S.null) [Set a]
xs of
  [] -> String -> Either String (Partition a)
forall a b. a -> Either a b
Left String
"pt: Empty list."
  [Set a]
xs' -> Partition a -> Either String (Partition a)
forall a b. b -> Either a b
Right (Partition a -> Either String (Partition a))
-> Partition a -> Either String (Partition a)
forall a b. (a -> b) -> a -> b
$ [Set a] -> Partition a
forall a. Ord a => [Set a] -> Partition a
ptUnsafe [Set a]
xs'

-- | Create a partition.
ptUnsafe :: Ord a => [Set a] -> Partition a
ptUnsafe :: [Set a] -> Partition a
ptUnsafe [Set a]
xs = Set (Set a) -> Partition a
forall a. Set (Set a) -> Partition a
Partition ([Set a] -> Set (Set a)
forall a. Ord a => [a] -> Set a
S.fromList [Set a]
xs)

-- | Convert a bipartition to a partition.
bpToPt :: Ord a => Bipartition a -> Partition a
bpToPt :: Bipartition a -> Partition a
bpToPt = [Set a] -> Partition a
forall a. Ord a => [Set a] -> Partition a
ptUnsafe ([Set a] -> Partition a)
-> (Bipartition a -> [Set a]) -> Bipartition a -> Partition a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Set a, Set a) -> [Set a]
forall a. (a, a) -> [a]
tupleToList ((Set a, Set a) -> [Set a])
-> (Bipartition a -> (Set a, Set a)) -> Bipartition a -> [Set a]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Bipartition a -> (Set a, Set a)
forall a. Bipartition a -> (Set a, Set a)
fromBipartition
  where
    -- Be careful with tuples, because 'toList' does something very weird. It only
    -- takes the second element of the tuple!
    --
    -- toList :: Foldable t => t a -> [a]
    tupleToList :: (a, a) -> [a]
tupleToList (a
x, a
y) = [a
x, a
y]

-- | Show a partition in a human readable form. Use a provided function to
-- extract the valuable information.
ptHuman :: Show a => Partition a -> String
ptHuman :: Partition a -> String
ptHuman (Partition Set (Set a)
xs) =
  String
"(" String -> ShowS
forall a. [a] -> [a] -> [a]
++ String -> [String] -> String
forall a. [a] -> [[a]] -> [a]
intercalate String
"|" ((Set a -> String) -> [Set a] -> [String]
forall a b. (a -> b) -> [a] -> [b]
map Set a -> String
forall a. Show a => Set a -> String
setShow (Set (Set a) -> [Set a]
forall a. Set a -> [a]
S.toList Set (Set a)
xs)) String -> ShowS
forall a. [a] -> [a] -> [a]
++ String
")"

-- Show the elements of a set in a human readable format.
setShow :: Show a => Set a -> String
setShow :: Set a -> String
setShow = String -> [String] -> String
forall a. [a] -> [[a]] -> [a]
intercalate String
"," ([String] -> String) -> (Set a -> [String]) -> Set a -> String
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (a -> String) -> [a] -> [String]
forall a b. (a -> b) -> [a] -> [b]
map a -> String
forall a. Show a => a -> String
show ([a] -> [String]) -> (Set a -> [a]) -> Set a -> [String]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Set a -> [a]
forall a. Set a -> [a]
S.toList

-- | Get partition defined by the root of the tree.
--
-- Return 'Left' if:
-- - the tree is a leaf;
-- - the tree contains duplicate leaves.
partition :: Ord a => Tree e a -> Either String (Partition a)
partition :: Tree e a -> Either String (Partition a)
partition (Node e
_ a
_ []) = String -> Either String (Partition a)
forall a b. a -> Either a b
Left String
"partition: Encountered a leaf."
partition t :: Tree e a
t@(Node e
_ a
_ [Tree e a]
ts)
  | Tree e a -> Bool
forall a e. Ord a => Tree e a -> Bool
duplicateLeaves Tree e a
t = String -> Either String (Partition a)
forall a b. a -> Either a b
Left String
"partition: Tree contains duplicate leaves."
  | Bool
otherwise = [Set a] -> Either String (Partition a)
forall a. Ord a => [Set a] -> Either String (Partition a)
pt ([Set a] -> Either String (Partition a))
-> [Set a] -> Either String (Partition a)
forall a b. (a -> b) -> a -> b
$ (Tree e a -> Set a) -> [Tree e a] -> [Set a]
forall a b. (a -> b) -> [a] -> [b]
map ([a] -> Set a
forall a. Ord a => [a] -> Set a
S.fromList ([a] -> Set a) -> (Tree e a -> [a]) -> Tree e a -> Set a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Tree e a -> [a]
forall e a. Tree e a -> [a]
leaves) [Tree e a]
ts

-- | Get all 'Partition's of a tree.
--
-- Return 'Left' if tree contains duplicate leaves.
partitions :: Ord a => Tree e a -> Either String (Set (Partition a))
partitions :: Tree e a -> Either String (Set (Partition a))
partitions Tree e a
t
  | Tree e a -> Bool
forall a e. Ord a => Tree e a -> Bool
duplicateLeaves Tree e a
t = String -> Either String (Set (Partition a))
forall a b. a -> Either a b
Left String
"partitions: Tree contains duplicate leaves."
  | Bool
otherwise = Set (Partition a) -> Either String (Set (Partition a))
forall a b. b -> Either a b
Right (Set (Partition a) -> Either String (Set (Partition a)))
-> Set (Partition a) -> Either String (Set (Partition a))
forall a b. (a -> b) -> a -> b
$ Set a -> Tree e (Set a) -> Set (Partition a)
forall a e. Ord a => Set a -> Tree e (Set a) -> Set (Partition a)
partitions' Set a
forall a. Set a
S.empty (Tree e (Set a) -> Set (Partition a))
-> Tree e (Set a) -> Set (Partition a)
forall a b. (a -> b) -> a -> b
$ [a] -> Set a
forall a. Ord a => [a] -> Set a
S.fromList ([a] -> Set a) -> Tree e [a] -> Tree e (Set a)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Tree e a -> Tree e [a]
forall e a. Tree e a -> Tree e [a]
groups Tree e a
t

-- See 'partitions', but do not check if leaves are unique.
partitions' :: Ord a => Set a -> Tree e (Set a) -> Set (Partition a)
partitions' :: Set a -> Tree e (Set a) -> Set (Partition a)
partitions' Set a
_ (Node e
_ Set a
_ []) = Set (Partition a)
forall a. Set a
S.empty
partitions' Set a
p t :: Tree e (Set a)
t@(Node e
_ Set a
_ [Tree e (Set a)]
ts) =
  [Set (Partition a)] -> Set (Partition a)
forall (f :: * -> *) a. (Foldable f, Ord a) => f (Set a) -> Set a
S.unions ([Set (Partition a)] -> Set (Partition a))
-> [Set (Partition a)] -> Set (Partition a)
forall a b. (a -> b) -> a -> b
$
    (String -> Set (Partition a))
-> (Partition a -> Set (Partition a))
-> Either String (Partition a)
-> Set (Partition a)
forall a c b. (a -> c) -> (b -> c) -> Either a b -> c
either (Set (Partition a) -> String -> Set (Partition a)
forall a b. a -> b -> a
const Set (Partition a)
forall a. Set a
S.empty) Partition a -> Set (Partition a)
forall a. a -> Set a
S.singleton ([Set a] -> Either String (Partition a)
forall a. Ord a => [Set a] -> Either String (Partition a)
pt (Set a
p Set a -> [Set a] -> [Set a]
forall a. a -> [a] -> [a]
: (Tree e (Set a) -> Set a) -> [Tree e (Set a)] -> [Set a]
forall a b. (a -> b) -> [a] -> [b]
map Tree e (Set a) -> Set a
forall e a. Tree e a -> a
label [Tree e (Set a)]
ts)) Set (Partition a) -> [Set (Partition a)] -> [Set (Partition a)]
forall a. a -> [a] -> [a]
:
    (Set a -> Tree e (Set a) -> Set (Partition a))
-> [Set a] -> [Tree e (Set a)] -> [Set (Partition a)]
forall a b c. (a -> b -> c) -> [a] -> [b] -> [c]
zipWith Set a -> Tree e (Set a) -> Set (Partition a)
forall a e. Ord a => Set a -> Tree e (Set a) -> Set (Partition a)
partitions' [Set a]
cs [Tree e (Set a)]
ts
  where
    cs :: [Set a]
cs = Set a -> Tree e (Set a) -> [Set a]
forall a e. Ord a => Set a -> Tree e (Set a) -> [Set a]
getComplementaryLeaves Set a
p Tree e (Set a)
t

-- | 'Partition's are compatible if they do not contain conflicting
-- information. This function checks if two partitions are compatible with
-- each other. Thereby, a variation of the following algorithm is used:
--
-- @
-- mp1 `compatible` mp2
-- for set1 in mp1:
--   for set2 in mp2:
--     if set1 `S.isSubSetOf` set2:
--       remove set1 from mp1
--     if set2 `S.isSubSetOf` set1:
--       remove set2 from mp2
-- if either mp2 or mp2 is empty, they are compatible
-- @
compatible :: (Show a, Ord a) => Partition a -> Partition a -> Bool
compatible :: Partition a -> Partition a -> Bool
compatible Partition a
l Partition a
r = Set (Set a) -> Bool
forall a. Set a -> Bool
S.null ((Set a -> Bool) -> Set (Set a) -> Set (Set a)
forall a. (a -> Bool) -> Set a -> Set a
S.filter (Set a -> Set (Set a) -> Bool
forall a. Ord a => Set a -> Set (Set a) -> Bool
`remove` Set (Set a)
rs) Set (Set a)
ls) Bool -> Bool -> Bool
|| Set (Set a) -> Bool
forall a. Set a -> Bool
S.null ((Set a -> Bool) -> Set (Set a) -> Set (Set a)
forall a. (a -> Bool) -> Set a -> Set a
S.filter (Set a -> Set (Set a) -> Bool
forall a. Ord a => Set a -> Set (Set a) -> Bool
`remove` Set (Set a)
ls) Set (Set a)
rs)
  where
    ls :: Set (Set a)
ls = Partition a -> Set (Set a)
forall a. Partition a -> Set (Set a)
fromPartition Partition a
l
    rs :: Set (Set a)
rs = Partition a -> Set (Set a)
forall a. Partition a -> Set (Set a)
fromPartition Partition a
r

remove :: Ord a => Set a -> Set (Set a) -> Bool
remove :: Set a -> Set (Set a) -> Bool
remove Set a
s = Bool -> Bool
not (Bool -> Bool) -> (Set (Set a) -> Bool) -> Set (Set a) -> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Set a -> Bool) -> Set (Set a) -> Bool
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
any (Set a
s Set a -> Set a -> Bool
forall a. Ord a => Set a -> Set a -> Bool
`S.isSubsetOf`)