{-# OPTIONS_GHC -Wno-orphans #-}

-- |
-- Module      :  ELynx.Tree.Export.Newick
-- Description :  Export tree objects to Newick format
-- Copyright   :  (c) Dominik Schrempf 2021
-- License     :  GPL-3.0-or-later
--
-- Maintainer  :  dominik.schrempf@gmail.com
-- Stability   :  unstable
-- Portability :  portable
--
-- Creation date: Thu Jan 17 13:51:47 2019.
--
-- Some functions are inspired by
-- [Biobase.Newick.Import](https://hackage.haskell.org/package/BiobaseNewick).
--
-- See nomenclature in 'ELynx.Tree'.
module ELynx.Tree.Export.Newick
  ( toNewick,
    toNewickBuilder,
  )
where

import qualified Data.ByteString.Builder as BB
import qualified Data.ByteString.Lazy.Char8 as BL
import Data.List (intersperse)
import ELynx.Tree.Length
import ELynx.Tree.Name
import ELynx.Tree.Phylogeny
import ELynx.Tree.Rooted
import ELynx.Tree.Support

-- Allow export of trees having branches with lengths only.
instance HasMaybeSupport Length where
  getMaybeSupport :: Length -> Maybe Support
getMaybeSupport = Maybe Support -> Length -> Maybe Support
forall a b. a -> b -> a
const Maybe Support
forall a. Maybe a
Nothing

-- Allow export of trees having branches with support values only.
instance HasMaybeLength Support where
  getMaybeLength :: Support -> Maybe Length
getMaybeLength = Maybe Length -> Support -> Maybe Length
forall a b. a -> b -> a
const Maybe Length
forall a. Maybe a
Nothing

buildBrLen :: Length -> BB.Builder
buildBrLen :: Length -> Builder
buildBrLen Length
bl = Char -> Builder
BB.char8 Char
':' Builder -> Builder -> Builder
forall a. Semigroup a => a -> a -> a
<> Double -> Builder
BB.doubleDec (Length -> Double
fromLength Length
bl)

buildBrSup :: Support -> BB.Builder
buildBrSup :: Support -> Builder
buildBrSup Support
bs = Char -> Builder
BB.char8 Char
'[' Builder -> Builder -> Builder
forall a. Semigroup a => a -> a -> a
<> Double -> Builder
BB.doubleDec (Support -> Double
fromSupport Support
bs) Builder -> Builder -> Builder
forall a. Semigroup a => a -> a -> a
<> Char -> Builder
BB.char8 Char
']'

-- | See 'toNewick'.
toNewickBuilder :: (HasMaybeLength e, HasMaybeSupport e, HasName a) => Tree e a -> BB.Builder
toNewickBuilder :: Tree e a -> Builder
toNewickBuilder Tree e a
t = Tree e a -> Builder
forall a e.
(HasName a, HasMaybeLength e, HasMaybeSupport e) =>
Tree e a -> Builder
go Tree e a
t Builder -> Builder -> Builder
forall a. Semigroup a => a -> a -> a
<> Char -> Builder
BB.char8 Char
';'
  where
    go :: Tree e a -> Builder
go (Node e
b a
l []) = e -> a -> Builder
forall a e.
(HasName a, HasMaybeLength e, HasMaybeSupport e) =>
e -> a -> Builder
lbl e
b a
l
    go (Node e
b a
l [Tree e a]
ts) =
      Char -> Builder
BB.char8 Char
'('
        Builder -> Builder -> Builder
forall a. Semigroup a => a -> a -> a
<> [Builder] -> Builder
forall a. Monoid a => [a] -> a
mconcat (Builder -> [Builder] -> [Builder]
forall a. a -> [a] -> [a]
intersperse (Char -> Builder
BB.char8 Char
',') ([Builder] -> [Builder]) -> [Builder] -> [Builder]
forall a b. (a -> b) -> a -> b
$ (Tree e a -> Builder) -> [Tree e a] -> [Builder]
forall a b. (a -> b) -> [a] -> [b]
map Tree e a -> Builder
go [Tree e a]
ts)
        Builder -> Builder -> Builder
forall a. Semigroup a => a -> a -> a
<> Char -> Builder
BB.char8 Char
')'
        Builder -> Builder -> Builder
forall a. Semigroup a => a -> a -> a
<> e -> a -> Builder
forall a e.
(HasName a, HasMaybeLength e, HasMaybeSupport e) =>
e -> a -> Builder
lbl e
b a
l
    mBrSupBuilder :: e -> Builder
mBrSupBuilder e
x = Builder -> (Support -> Builder) -> Maybe Support -> Builder
forall b a. b -> (a -> b) -> Maybe a -> b
maybe Builder
forall a. Monoid a => a
mempty Support -> Builder
buildBrSup (e -> Maybe Support
forall e. HasMaybeSupport e => e -> Maybe Support
getMaybeSupport e
x)
    mBrLenBuilder :: e -> Builder
mBrLenBuilder e
x = Builder -> (Length -> Builder) -> Maybe Length -> Builder
forall b a. b -> (a -> b) -> Maybe a -> b
maybe Builder
forall a. Monoid a => a
mempty Length -> Builder
buildBrLen (e -> Maybe Length
forall e. HasMaybeLength e => e -> Maybe Length
getMaybeLength e
x)
    lbl :: e -> a -> Builder
lbl e
x a
y =
      ByteString -> Builder
BB.lazyByteString (Name -> ByteString
fromName (Name -> ByteString) -> Name -> ByteString
forall a b. (a -> b) -> a -> b
$ a -> Name
forall a. HasName a => a -> Name
getName a
y)
        Builder -> Builder -> Builder
forall a. Semigroup a => a -> a -> a
<> e -> Builder
forall e. HasMaybeLength e => e -> Builder
mBrLenBuilder e
x
        -- After reading several discussions, I go for the "more semantical
        -- form" with branch support values in square brackets.
        Builder -> Builder -> Builder
forall a. Semigroup a => a -> a -> a
<> e -> Builder
forall e. HasMaybeSupport e => e -> Builder
mBrSupBuilder e
x
{-# SPECIALIZE toNewickBuilder :: Tree Length Name -> BB.Builder #-}
{-# SPECIALIZE toNewickBuilder :: Tree Length Int -> BB.Builder #-}
{-# SPECIALIZE toNewickBuilder :: Tree Phylo Name -> BB.Builder #-}
{-# SPECIALIZE toNewickBuilder :: Tree Phylo Int -> BB.Builder #-}

-- | General conversion of a tree into a Newick 'BL.ByteString'.
--
-- Functions to write key value pairs for nodes are not provided. Those can just
-- be set as node labels. For example, the posterior density and the confidence
-- interval of a node can be encoded by setting the node label to a
-- 'BL.ByteString':
--
-- @
-- "ACTUALNAME[posterior=-2839.2,age_95%_HPD={4.80804,31.6041}]"
-- @
toNewick :: (HasMaybeLength e, HasMaybeSupport e, HasName a) => Tree e a -> BL.ByteString
toNewick :: Tree e a -> ByteString
toNewick = Builder -> ByteString
BB.toLazyByteString (Builder -> ByteString)
-> (Tree e a -> Builder) -> Tree e a -> ByteString
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Tree e a -> Builder
forall e a.
(HasMaybeLength e, HasMaybeSupport e, HasName a) =>
Tree e a -> Builder
toNewickBuilder
{-# SPECIALIZE toNewick :: Tree Length Name -> BL.ByteString #-}
{-# SPECIALIZE toNewick :: Tree Length Int -> BL.ByteString #-}
{-# SPECIALIZE toNewick :: Tree Phylo Name -> BL.ByteString #-}
{-# SPECIALIZE toNewick :: Tree Phylo Int -> BL.ByteString #-}