-- |
-- Module      :  ELynx.Tree.Partition
-- Description :  Partitions on rose trees
-- Copyright   :  2021 Dominik Schrempf
-- 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
  { forall a. Partition a -> Set (Set a)
fromPartition :: Set (Set a)
  }
  deriving (Partition a -> Partition a -> Bool
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, 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
Ord, Int -> Partition a -> ShowS
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)
ReadS [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 :: forall a. Ord a => [Set a] -> Either String (Partition a)
pt [Set a]
xs = case forall a. (a -> Bool) -> [a] -> [a]
filter (Bool -> Bool
not forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a. Set a -> Bool
S.null) [Set a]
xs of
  [] -> forall a b. a -> Either a b
Left String
"pt: Empty list."
  [Set a]
xs' -> forall a b. b -> Either a b
Right forall a b. (a -> b) -> a -> b
$ forall a. Ord a => [Set a] -> Partition a
ptUnsafe [Set a]
xs'

-- | Create a partition.
ptUnsafe :: Ord a => [Set a] -> Partition a
ptUnsafe :: forall a. Ord a => [Set a] -> Partition a
ptUnsafe [Set a]
xs = forall a. Set (Set a) -> Partition a
Partition (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 :: forall a. Ord a => Bipartition a -> Partition a
bpToPt = forall a. Ord a => [Set a] -> Partition a
ptUnsafe forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall {a}. (a, a) -> [a]
tupleToList forall b c a. (b -> c) -> (a -> b) -> a -> c
. 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 :: forall a. Show a => Partition a -> String
ptHuman (Partition Set (Set a)
xs) =
  String
"(" forall a. [a] -> [a] -> [a]
++ forall a. [a] -> [[a]] -> [a]
intercalate String
"|" (forall a b. (a -> b) -> [a] -> [b]
map forall a. Show a => Set a -> String
setShow (forall a. Set a -> [a]
S.toList Set (Set a)
xs)) forall a. [a] -> [a] -> [a]
++ String
")"

-- Show the elements of a set in a human readable format.
setShow :: Show a => Set a -> String
setShow :: forall a. Show a => Set a -> String
setShow = forall a. [a] -> [[a]] -> [a]
intercalate String
"," forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a b. (a -> b) -> [a] -> [b]
map forall a. Show a => a -> String
show forall b c a. (b -> c) -> (a -> b) -> a -> c
. 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 :: forall a e. Ord a => Tree e a -> Either String (Partition a)
partition (Node e
_ 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)
  | forall a e. Ord a => Tree e a -> Bool
duplicateLeaves Tree e a
t = forall a b. a -> Either a b
Left String
"partition: Tree contains duplicate leaves."
  | Bool
otherwise = forall a. Ord a => [Set a] -> Either String (Partition a)
pt forall a b. (a -> b) -> a -> b
$ forall a b. (a -> b) -> [a] -> [b]
map (forall a. Ord a => [a] -> Set a
S.fromList forall b c a. (b -> c) -> (a -> b) -> a -> c
. 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 :: forall a e. Ord a => Tree e a -> Either String (Set (Partition a))
partitions Tree e a
t
  | forall a e. Ord a => Tree e a -> Bool
duplicateLeaves Tree e a
t = forall a b. a -> Either a b
Left String
"partitions: Tree contains duplicate leaves."
  | Bool
otherwise = forall a b. b -> Either a b
Right forall a b. (a -> b) -> a -> b
$ forall a e. Ord a => Set a -> Tree e (Set a) -> Set (Partition a)
partitions' forall a. Set a
S.empty forall a b. (a -> b) -> a -> b
$ forall a. Ord a => [a] -> Set a
S.fromList forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> 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' :: forall a e. Ord a => Set a -> Tree e (Set a) -> Set (Partition a)
partitions' Set a
_ (Node e
_ Set 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) =
  forall (f :: * -> *) a. (Foldable f, Ord a) => f (Set a) -> Set a
S.unions forall a b. (a -> b) -> a -> b
$
    forall a c b. (a -> c) -> (b -> c) -> Either a b -> c
either (forall a b. a -> b -> a
const forall a. Set a
S.empty) forall a. a -> Set a
S.singleton (forall a. Ord a => [Set a] -> Either String (Partition a)
pt (Set a
p forall a. a -> [a] -> [a]
: forall a b. (a -> b) -> [a] -> [b]
map forall e a. Tree e a -> a
label [Tree e (Set a)]
ts))
      forall a. a -> [a] -> [a]
: forall a b c. (a -> b -> c) -> [a] -> [b] -> [c]
zipWith 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 = 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 :: forall a. (Show a, Ord a) => Partition a -> Partition a -> Bool
compatible Partition a
l Partition a
r = forall a. Set a -> Bool
S.null (forall a. (a -> Bool) -> Set a -> Set a
S.filter (forall a. Ord a => Set a -> Set (Set a) -> Bool
`remove` Set (Set a)
rs) Set (Set a)
ls) Bool -> Bool -> Bool
|| forall a. Set a -> Bool
S.null (forall a. (a -> Bool) -> Set a -> Set a
S.filter (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 = forall a. Partition a -> Set (Set a)
fromPartition Partition a
l
    rs :: Set (Set a)
rs = forall a. Partition a -> Set (Set a)
fromPartition Partition a
r

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