{-# LANGUAGE OverloadedStrings #-}

-- |
-- Module      :  TLynx.Connect.Connect
-- Description :  Connect two phylogenies
-- Copyright   :  2021 Dominik Schrempf
-- License     :  GPL-3.0-or-later
--
-- Maintainer  :  dominik.schrempf@gmail.com
-- Stability   :  unstable
-- Portability :  portable
--
-- Creation date: Thu Sep 19 15:01:52 2019.
module TLynx.Connect.Connect
  ( connectCmd,
  )
where

import Control.Monad.IO.Class
import Control.Monad.Trans.Reader (ask)
import qualified Data.ByteString.Lazy.Char8 as BL
import Data.Default.Class
import qualified Data.Set as S
import ELynx.Tools.ELynx
import ELynx.Tools.Environment
import ELynx.Tools.Logger
import ELynx.Tree
import System.IO
import TLynx.Connect.Options
import TLynx.Parsers

-- Connect two trees with a branch in all possible ways.
--
-- Introduce a branch between two trees. If the trees have @n>2@, and @m>2@
-- nodes, respectively, there are (n-2)*(m-2) ways to connect them.
--
-- A base node label has to be given which will be used wherever the new node is
-- introduced.
--
-- Return 'Left' if one tree has a non-bifurcating root node.
connect ::
  (Semigroup e, Splittable e, Default a) =>
  e ->
  a ->
  Tree e a ->
  Tree e a ->
  Either String (Forest e a)
connect :: forall e a.
(Semigroup e, Splittable e, Default a) =>
e -> a -> Tree e a -> Tree e a -> Either String (Forest e a)
connect e
br a
lb Tree e a
l Tree e a
r = do
  Forest e a
ls <- forall e a.
(Semigroup e, Splittable e, Default a) =>
Tree e a -> Either String (Forest e a)
roots Tree e a
l
  Forest e a
rs <- forall e a.
(Semigroup e, Splittable e, Default a) =>
Tree e a -> Either String (Forest e a)
roots Tree e a
r
  forall (m :: * -> *) a. Monad m => a -> m a
return [forall e a. e -> a -> Forest e a -> Tree e a
Node e
br a
lb [Tree e a
x, Tree e a
y] | Tree e a
x <- Forest e a
ls, Tree e a
y <- Forest e a
rs]

-- | Connect two trees honoring possible constraints.
--
-- Introduce a branch between two trees. If the trees have @n>2@, and @m>2@
-- nodes, respectively, there are (n-2)*(m-2) ways to connect them.
connectCmd :: ELynx ConnectArguments ()
connectCmd :: ELynx ConnectArguments ()
connectCmd = do
  ConnectArguments
lArgs <- forall a. Environment a -> a
localArguments forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall (m :: * -> *) r. Monad m => ReaderT r m r
ask
  Handle
outH <- forall a. Reproducible a => String -> String -> ELynx a Handle
outHandle String
"results" String
".out"
  -- Do we have constraints or not?
  let cs :: Maybe String
cs = ConnectArguments -> Maybe String
constraints ConnectArguments
lArgs
      l :: String
l = ConnectArguments -> String
inFileA ConnectArguments
lArgs
      r :: String
r = ConnectArguments -> String
inFileB ConnectArguments
lArgs
  case Maybe String
cs of
    Maybe String
Nothing -> Handle -> String -> String -> ELynx ConnectArguments ()
connectOnly Handle
outH String
l String
r
    Just String
c -> Handle -> String -> String -> String -> ELynx ConnectArguments ()
connectAndFilter Handle
outH String
c String
l String
r
  forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO forall a b. (a -> b) -> a -> b
$ Handle -> IO ()
hClose Handle
outH

connectTrees ::
  Tree Length Name ->
  Tree Length Name ->
  Forest Length Name
connectTrees :: Tree Length Name -> Tree Length Name -> Forest Length Name
connectTrees Tree Length Name
t = forall a c b. (a -> c) -> (b -> c) -> Either a b -> c
either forall a. HasCallStack => String -> a
error forall a. a -> a
id forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall e a.
(Semigroup e, Splittable e, Default a) =>
e -> a -> Tree e a -> Tree e a -> Either String (Forest e a)
connect Length
0 Name
"root" Tree Length Name
t

type Constraint a = S.Set a

-- Get groups induced by multifurcations. Collect the leaves of all trees
-- induced by multifurcations.
multifurcatingGroups :: Tree e a -> [[a]]
multifurcatingGroups :: forall e a. Tree e a -> [[a]]
multifurcatingGroups (Node e
_ a
_ []) = []
multifurcatingGroups (Node e
_ a
_ [Tree e a
x]) = forall e a. Tree e a -> [[a]]
multifurcatingGroups Tree e a
x
multifurcatingGroups (Node e
_ a
_ [Tree e a
x, Tree e a
y]) = forall e a. Tree e a -> [[a]]
multifurcatingGroups Tree e a
x forall a. [a] -> [a] -> [a]
++ forall e a. Tree e a -> [[a]]
multifurcatingGroups Tree e a
y
multifurcatingGroups Tree e a
t = forall e a. Tree e a -> [a]
leaves Tree e a
t forall a. a -> [a] -> [a]
: forall (t :: * -> *) a b. Foldable t => (a -> [b]) -> t a -> [b]
concatMap forall e a. Tree e a -> [[a]]
multifurcatingGroups (forall e a. Tree e a -> Forest e a
forest Tree e a
t)

compatibleAll :: (Show a, Ord a) => Tree e a -> [Constraint a] -> Bool
compatibleAll :: forall a e. (Show a, Ord a) => Tree e a -> [Constraint a] -> Bool
compatibleAll t :: Tree e a
t@(Node e
_ a
_ [Tree e a
l, Tree e a
r]) [Constraint a]
cs =
  forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
all (forall a. (Show a, Ord a) => Partition a -> Partition a -> Bool
compatible Partition a
partitionLeft forall b c a. (b -> c) -> (a -> b) -> a -> c
. Constraint a -> Partition a
getP) [Constraint a]
cs
    Bool -> Bool -> Bool
&& forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
all (forall a. (Show a, Ord a) => Partition a -> Partition a -> Bool
compatible Partition a
partitionRight forall b c a. (b -> c) -> (a -> b) -> a -> c
. Constraint a -> Partition a
getP) [Constraint a]
cs
  where
    lvs :: Constraint a
lvs = forall a. Ord a => [a] -> Set a
S.fromList forall a b. (a -> b) -> a -> b
$ forall e a. Tree e a -> [a]
leaves Tree e a
t
    getP :: Constraint a -> Partition a
getP Constraint a
x = forall a c b. (a -> c) -> (b -> c) -> Either a b -> c
either forall a. HasCallStack => String -> a
error forall a. a -> a
id forall a b. (a -> b) -> a -> b
$ forall a. Ord a => [Set a] -> Either String (Partition a)
pt [Constraint a
x, Constraint a
lvs forall a. Ord a => Set a -> Set a -> Set a
S.\\ Constraint a
x]
    partitionLeft :: Partition a
partitionLeft = forall a c b. (a -> c) -> (b -> c) -> Either a b -> c
either forall a. HasCallStack => String -> a
error forall a. a -> a
id forall a b. (a -> b) -> a -> b
$ forall a e. Ord a => Tree e a -> Either String (Partition a)
partition Tree e a
l
    partitionRight :: Partition a
partitionRight = forall a c b. (a -> c) -> (b -> c) -> Either a b -> c
either forall a. HasCallStack => String -> a
error forall a. a -> a
id forall a b. (a -> b) -> a -> b
$ forall a e. Ord a => Tree e a -> Either String (Partition a)
partition Tree e a
r
compatibleAll Tree e a
_ [Constraint a]
_ = forall a. HasCallStack => String -> a
error String
"Tree is not bifurcating."

compatibleWith ::
  (Show b, Ord b) => (a -> b) -> [Constraint a] -> Tree e a -> Bool
compatibleWith :: forall b a e.
(Show b, Ord b) =>
(a -> b) -> [Constraint a] -> Tree e a -> Bool
compatibleWith a -> b
f [Constraint a]
cs Tree e a
t = forall a e. (Show a, Ord a) => Tree e a -> [Constraint a] -> Bool
compatibleAll (forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap a -> b
f Tree e a
t) (forall a b. (a -> b) -> [a] -> [b]
map (forall b a. Ord b => (a -> b) -> Set a -> Set b
S.map a -> b
f) [Constraint a]
cs)

parseTreeTuple ::
  FilePath ->
  FilePath ->
  ELynx
    ConnectArguments
    (Tree Length Name, Tree Length Name)
parseTreeTuple :: String
-> String
-> ELynx ConnectArguments (Tree Length Name, Tree Length Name)
parseTreeTuple String
l String
r = do
  NewickFormat
nwF <- ConnectArguments -> NewickFormat
nwFormat forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a. Environment a -> a
localArguments forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall (m :: * -> *) r. Monad m => ReaderT r m r
ask
  Tree Phylo Name
tl <- forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO forall a b. (a -> b) -> a -> b
$ NewickFormat -> String -> IO (Tree Phylo Name)
parseTree NewickFormat
nwF String
l
  Tree Phylo Name
tr <- forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO forall a b. (a -> b) -> a -> b
$ NewickFormat -> String -> IO (Tree Phylo Name)
parseTree NewickFormat
nwF String
r
  forall e.
(HasLock e, HasLogHandles e, HasVerbosity e) =>
String -> Logger e ()
logInfoS String
"Tree 1:"
  forall e.
(HasLock e, HasLogHandles e, HasVerbosity e) =>
ByteString -> Logger e ()
logInfoB forall a b. (a -> b) -> a -> b
$ forall e a.
(HasMaybeLength e, HasMaybeSupport e, HasName a) =>
Tree e a -> ByteString
toNewick Tree Phylo Name
tl
  forall e.
(HasLock e, HasLogHandles e, HasVerbosity e) =>
String -> Logger e ()
logInfoS String
"Tree 2:"
  forall e.
(HasLock e, HasLogHandles e, HasVerbosity e) =>
ByteString -> Logger e ()
logInfoB forall a b. (a -> b) -> a -> b
$ forall e a.
(HasMaybeLength e, HasMaybeSupport e, HasName a) =>
Tree e a -> ByteString
toNewick Tree Phylo Name
tr
  forall (m :: * -> *) a. Monad m => a -> m a
return (forall a c b. (a -> c) -> (b -> c) -> Either a b -> c
either forall a. HasCallStack => String -> a
error forall a. a -> a
id forall a b. (a -> b) -> a -> b
$ forall e a.
HasMaybeLength e =>
Tree e a -> Either String (Tree Length a)
toLengthTree Tree Phylo Name
tl, forall a c b. (a -> c) -> (b -> c) -> Either a b -> c
either forall a. HasCallStack => String -> a
error forall a. a -> a
id forall a b. (a -> b) -> a -> b
$ forall e a.
HasMaybeLength e =>
Tree e a -> Either String (Tree Length a)
toLengthTree Tree Phylo Name
tr)

connectOnly :: Handle -> FilePath -> FilePath -> ELynx ConnectArguments ()
connectOnly :: Handle -> String -> String -> ELynx ConnectArguments ()
connectOnly Handle
h String
l String
r = do
  (Tree Length Name
tl, Tree Length Name
tr) <- String
-> String
-> ELynx ConnectArguments (Tree Length Name, Tree Length Name)
parseTreeTuple String
l String
r
  let ts :: Forest Length Name
ts = Tree Length Name -> Tree Length Name -> Forest Length Name
connectTrees Tree Length Name
tl Tree Length Name
tr
  forall e.
(HasLock e, HasLogHandles e, HasVerbosity e) =>
String -> Logger e ()
logInfoS forall a b. (a -> b) -> a -> b
$ String
"Connected trees: " forall a. Semigroup a => a -> a -> a
<> forall a. Show a => a -> String
show (forall (t :: * -> *) a. Foldable t => t a -> Int
length Forest Length Name
ts)
  forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO forall a b. (a -> b) -> a -> b
$ Handle -> ByteString -> IO ()
BL.hPutStr Handle
h forall a b. (a -> b) -> a -> b
$ [ByteString] -> ByteString
BL.unlines forall a b. (a -> b) -> a -> b
$ forall a b. (a -> b) -> [a] -> [b]
map (forall e a.
(HasMaybeLength e, HasMaybeSupport e, HasName a) =>
Tree e a -> ByteString
toNewick forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall e a. HasMaybeLength e => Tree e a -> Tree Phylo a
lengthToPhyloTree) Forest Length Name
ts

connectAndFilter ::
  Handle -> FilePath -> FilePath -> FilePath -> ELynx ConnectArguments ()
connectAndFilter :: Handle -> String -> String -> String -> ELynx ConnectArguments ()
connectAndFilter Handle
h String
c String
l String
r = do
  NewickFormat
nwF <- ConnectArguments -> NewickFormat
nwFormat forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a. Environment a -> a
localArguments forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall (m :: * -> *) r. Monad m => ReaderT r m r
ask
  Forest Phylo Name
cts <- forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO forall a b. (a -> b) -> a -> b
$ NewickFormat -> String -> IO (Forest Phylo Name)
parseTrees NewickFormat
nwF String
c
  forall e.
(HasLock e, HasLogHandles e, HasVerbosity e) =>
String -> Logger e ()
logInfoS String
"Constraints:"
  forall e.
(HasLock e, HasLogHandles e, HasVerbosity e) =>
ByteString -> Logger e ()
logInfoB forall a b. (a -> b) -> a -> b
$ ByteString -> [ByteString] -> ByteString
BL.intercalate ByteString
"\n" forall a b. (a -> b) -> a -> b
$ forall a b. (a -> b) -> [a] -> [b]
map forall e a.
(HasMaybeLength e, HasMaybeSupport e, HasName a) =>
Tree e a -> ByteString
toNewick Forest Phylo Name
cts
  (Tree Length Name
tl, Tree Length Name
tr) <- String
-> String
-> ELynx ConnectArguments (Tree Length Name, Tree Length Name)
parseTreeTuple String
l String
r
  let ts :: Forest Length Name
ts = Tree Length Name -> Tree Length Name -> Forest Length Name
connectTrees Tree Length Name
tl Tree Length Name
tr
      cs :: [Constraint Name]
cs = forall a b. (a -> b) -> [a] -> [b]
map forall a. Ord a => [a] -> Set a
S.fromList forall a b. (a -> b) -> a -> b
$ forall (t :: * -> *) a b. Foldable t => (a -> [b]) -> t a -> [b]
concatMap forall e a. Tree e a -> [[a]]
multifurcatingGroups Forest Phylo Name
cts :: [Constraint Name]
      -- Only collect trees that are compatible with the constraints.
      ts' :: Forest Length Name
ts' = forall a. (a -> Bool) -> [a] -> [a]
filter (forall b a e.
(Show b, Ord b) =>
(a -> b) -> [Constraint a] -> Tree e a -> Bool
compatibleWith forall a. HasName a => a -> Name
getName [Constraint Name]
cs) Forest Length Name
ts
  forall e.
(HasLock e, HasLogHandles e, HasVerbosity e) =>
String -> Logger e ()
logInfoS forall a b. (a -> b) -> a -> b
$ String
"Connected  trees: " forall a. Semigroup a => a -> a -> a
<> forall a. Show a => a -> String
show (forall (t :: * -> *) a. Foldable t => t a -> Int
length Forest Length Name
ts)
  forall e.
(HasLock e, HasLogHandles e, HasVerbosity e) =>
String -> Logger e ()
logInfoS forall a b. (a -> b) -> a -> b
$ String
"Compatible trees: " forall a. Semigroup a => a -> a -> a
<> forall a. Show a => a -> String
show (forall (t :: * -> *) a. Foldable t => t a -> Int
length Forest Length Name
ts')
  forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO forall a b. (a -> b) -> a -> b
$ Handle -> ByteString -> IO ()
BL.hPutStr Handle
h forall a b. (a -> b) -> a -> b
$ [ByteString] -> ByteString
BL.unlines forall a b. (a -> b) -> a -> b
$ forall a b. (a -> b) -> [a] -> [b]
map (forall e a.
(HasMaybeLength e, HasMaybeSupport e, HasName a) =>
Tree e a -> ByteString
toNewick forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall e a. HasMaybeLength e => Tree e a -> Tree Phylo a
lengthToPhyloTree) Forest Length Name
ts'