{-# LANGUAGE OverloadedStrings #-}

-- |
-- Description :  Analyze trees
-- Copyright   :  2021 Dominik Schrempf
-- License     :  GPL-3.0-or-later
--
-- Maintainer  :  dominik.schrempf@gmail.com
-- Stability   :  unstable
-- Portability :  portable
--
-- Creation date: Fri May 24 13:47:56 2019.
module TLynx.Examine.Examine
  ( examine,
  )
where

import Control.Comonad
import Control.Monad (unless)
import Control.Monad.IO.Class
import Control.Monad.Trans.Reader (ask)
import qualified Data.ByteString.Lazy.Char8 as BL
import Data.Containers.ListUtils (nubOrd)
import Data.List (foldl', (\\))
import qualified Data.Map as M
import ELynx.Tools.ByteString
import ELynx.Tools.ELynx
import ELynx.Tools.Environment
import ELynx.Tools.Logger
import ELynx.Tree
import System.IO
  ( Handle,
    hPutStrLn,
  )
import TLynx.Examine.Options
import TLynx.Parsers
import Text.Printf

pretty :: Length -> String
pretty :: Length -> String
pretty = forall r. PrintfType r => String -> r
printf String
"%.5f" forall b c a. (b -> c) -> (a -> b) -> a -> c
. Length -> Double
fromLength

prettyRow :: String -> String -> BL.ByteString
prettyRow :: String -> String -> ByteString
prettyRow String
name String
val = Int -> ByteString -> ByteString
alignLeft Int
33 ByteString
n forall a. Semigroup a => a -> a -> a
<> Int -> ByteString -> ByteString
alignRight Int
8 ByteString
v
  where
    n :: ByteString
n = String -> ByteString
BL.pack String
name
    v :: ByteString
v = String -> ByteString
BL.pack String
val

-- | Examine branches of a tree.
summarizeLengths :: HasLength e => Tree e a -> BL.ByteString
summarizeLengths :: forall e a. HasLength e => Tree e a -> ByteString
summarizeLengths Tree e a
t =
  ByteString -> [ByteString] -> ByteString
BL.intercalate
    ByteString
"\n"
    [ String -> String -> ByteString
prettyRow String
"Origin height: " forall a b. (a -> b) -> a -> b
$ Length -> String
pretty Length
h,
      String -> String -> ByteString
prettyRow String
"Mean distance origin leaves: " forall a b. (a -> b) -> a -> b
$ Length -> String
pretty Length
h',
      String -> String -> ByteString
prettyRow String
"Total branch length: " forall a b. (a -> b) -> a -> b
$ Length -> String
pretty Length
b
    ]
  where
    n :: Int
n = forall (t :: * -> *) a. Foldable t => t a -> Int
length forall a b. (a -> b) -> a -> b
$ forall e a. Tree e a -> [a]
leaves Tree e a
t
    h :: Length
h = forall e a. HasLength e => Tree e a -> Length
height Tree e a
t
    h' :: Length
h' = forall (t :: * -> *) a. (Foldable t, Num a) => t a -> a
sum (forall e a. HasLength e => Tree e a -> [Length]
distancesOriginLeaves Tree e a
t) forall a. Fractional a => a -> a -> a
/ forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
n
    b :: Length
b = forall e a. HasLength e => Tree e a -> Length
totalBranchLength Tree e a
t

readTrees :: FilePath -> ELynx ExamineArguments (Forest Phylo Name)
readTrees :: String -> ELynx ExamineArguments (Forest Phylo Name)
readTrees String
fp = do
  forall e.
(HasLock e, HasLogHandles e, HasVerbosity e) =>
String -> Logger e ()
logInfoS forall a b. (a -> b) -> a -> b
$ String
"Read tree(s) from file " forall a. Semigroup a => a -> a -> a
<> String
fp forall a. Semigroup a => a -> a -> a
<> String
"."
  NewickFormat
nf <- ExamineArguments -> NewickFormat
argsNewickFormat 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
  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
nf String
fp

countElements :: (Ord a, Foldable f) => f a -> M.Map a Int
countElements :: forall a (f :: * -> *). (Ord a, Foldable f) => f a -> Map a Int
countElements = forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl' forall {k} {a}. (Ord k, Num a) => Map k a -> k -> Map k a
f forall k a. Map k a
M.empty
  where
    f :: Map k a -> k -> Map k a
f Map k a
m k
x = forall k a.
Ord k =>
(Maybe a -> Maybe a) -> k -> Map k a -> Map k a
M.alter forall {a}. Num a => Maybe a -> Maybe a
g k
x Map k a
m
    g :: Maybe a -> Maybe a
g Maybe a
Nothing = forall a. a -> Maybe a
Just a
1
    g (Just a
x) = forall a. a -> Maybe a
Just forall a b. (a -> b) -> a -> b
$ a
x forall a. Num a => a -> a -> a
+ a
1

examineTree :: HasName a => Handle -> Tree Phylo a -> IO ()
examineTree :: forall a. HasName a => Handle -> Tree Phylo a -> IO ()
examineTree Handle
h Tree Phylo a
t = do
  Handle -> String -> IO ()
hPutStrLn Handle
h forall a b. (a -> b) -> a -> b
$ String
"Number of leaves: " forall a. [a] -> [a] -> [a]
++ forall a. Show a => a -> String
show (forall (t :: * -> *) a. Foldable t => t a -> Int
length [ByteString]
lvs)
  Handle -> String -> IO ()
hPutStrLn Handle
h forall a b. (a -> b) -> a -> b
$ String
"Degree of root node: " forall a. [a] -> [a] -> [a]
++ forall a. Show a => a -> String
show (forall e a. Tree e a -> Int
degree Tree Phylo a
t)
  if forall e a. Tree e a -> Bool
bifurcating Tree Phylo a
t
    then Handle -> String -> IO ()
hPutStrLn Handle
h String
"Tree is bifurcating."
    else
      let degrees :: Tree Phylo Int
degrees = forall (w :: * -> *) a b. Comonad w => (w a -> b) -> w a -> w b
extend forall e a. Tree e a -> Int
degree Tree Phylo a
t
          degreeMax :: Int
degreeMax = forall (t :: * -> *) a. (Foldable t, Ord a) => t a -> a
maximum Tree Phylo Int
degrees
       in do
            if Int
degreeMax forall a. Ord a => a -> a -> Bool
> Int
2
              then Handle -> String -> IO ()
hPutStrLn Handle
h String
"Tree is multifurcating."
              else Handle -> String -> IO ()
hPutStrLn Handle
h String
"Tree is bifurcating but has degree two nodes."
            Handle -> String -> IO ()
hPutStrLn Handle
h forall a b. (a -> b) -> a -> b
$ String
"List of degrees with counts: " forall a. Semigroup a => a -> a -> a
<> forall a. Show a => a -> String
show (forall k a. Map k a -> [(k, a)]
M.toList forall a b. (a -> b) -> a -> b
$ forall a (f :: * -> *). (Ord a, Foldable f) => f a -> Map a Int
countElements Tree Phylo Int
degrees)
  let l :: Either String (Tree Length a)
l = forall e a.
HasMaybeLength e =>
Tree e a -> Either String (Tree Length a)
toLengthTree Tree Phylo a
t
  case Either String (Tree Length a)
l of
    Left String
_ -> Handle -> String -> IO ()
hPutStrLn Handle
h String
"Branch lengths not available."
    Right Tree Length a
t' -> Handle -> ByteString -> IO ()
BL.hPutStrLn Handle
h forall a b. (a -> b) -> a -> b
$ forall e a. HasLength e => Tree e a -> ByteString
summarizeLengths Tree Length a
t'
  forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
unless (forall (t :: * -> *) a. Foldable t => t a -> Bool
null [ByteString]
dups) forall a b. (a -> b) -> a -> b
$ do
    Handle -> String -> IO ()
hPutStrLn Handle
h String
""
    Handle -> String -> IO ()
hPutStrLn Handle
h (String
"Duplicate leaves: " forall a. [a] -> [a] -> [a]
++ forall a. Show a => a -> String
show [ByteString]
dups)
  Handle -> ByteString -> IO ()
BL.hPutStrLn Handle
h forall a b. (a -> b) -> a -> b
$ ByteString
"Leave names: " forall a. Semigroup a => a -> a -> a
<> ByteString -> [ByteString] -> ByteString
BL.intercalate ByteString
" " [ByteString]
lvs
  where
    lvs :: [ByteString]
lvs = forall a b. (a -> b) -> [a] -> [b]
map (Name -> ByteString
fromName forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a. HasName a => a -> Name
getName) forall a b. (a -> b) -> a -> b
$ forall e a. Tree e a -> [a]
leaves Tree Phylo a
t
    dups :: [ByteString]
dups = [ByteString]
lvs forall a. Eq a => [a] -> [a] -> [a]
\\ forall a. Ord a => [a] -> [a]
nubOrd [ByteString]
lvs

-- | Examine phylogenetic trees.
examine :: ELynx ExamineArguments ()
examine :: Logger (Environment ExamineArguments) ()
examine = do
  ExamineArguments
l <- 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
  let inFn :: String
inFn = ExamineArguments -> String
argsInFile ExamineArguments
l
  Forest Phylo Name
trs <- String -> ELynx ExamineArguments (Forest Phylo Name)
readTrees String
inFn
  Handle
outH <- forall a. Reproducible a => String -> String -> ELynx a Handle
outHandle String
"results" String
".out"
  forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO forall a b. (a -> b) -> a -> b
$ forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
(a -> m b) -> t a -> m ()
mapM_ (forall a. HasName a => Handle -> Tree Phylo a -> IO ()
examineTree Handle
outH) Forest Phylo Name
trs