From 48b98ab9ce78ca304df84c6a94cc5fa51ef950c3 Mon Sep 17 00:00:00 2001
From: Frederic Lemoine <fredericlemoine@users.noreply.github.com>
Date: Wed, 29 Nov 2023 16:27:07 +0100
Subject: [PATCH] Added gotree brlen add command

---
 cmd/brlenadd.go | 56 +++++++++++++++++++++++++++++++++++++++++++++++++
 test.sh         | 26 ++++++++++++++++++++++-
 tree/tree.go    | 15 +++++++++++++
 3 files changed, 96 insertions(+), 1 deletion(-)
 create mode 100644 cmd/brlenadd.go

diff --git a/cmd/brlenadd.go b/cmd/brlenadd.go
new file mode 100644
index 0000000..70e47cb
--- /dev/null
+++ b/cmd/brlenadd.go
@@ -0,0 +1,56 @@
+package cmd
+
+import (
+	goio "io"
+	"os"
+
+	"github.com/evolbioinfo/gotree/io"
+	"github.com/evolbioinfo/gotree/tree"
+	"github.com/spf13/cobra"
+)
+
+var addlengthfactor float64
+
+// addCmd represents the cut command
+var brlenAddCmd = &cobra.Command{
+	Use:   "add",
+	Short: "Add the given length to all branches of the tree.",
+	Long: `Add the given length to all branches of the tree.
+
+Example:
+
+gotree brlen add -i tree.nwk -l <length>
+`,
+	RunE: func(cmd *cobra.Command, args []string) (err error) {
+		var f *os.File
+		var treefile goio.Closer
+		var treechan <-chan tree.Trees
+
+		if f, err = openWriteFile(outtreefile); err != nil {
+			io.LogError(err)
+			return
+		}
+		defer closeWriteFile(f, outtreefile)
+
+		if treefile, treechan, err = readTrees(intreefile); err != nil {
+			io.LogError(err)
+			return
+		}
+		defer treefile.Close()
+		for tr := range treechan {
+			if tr.Err != nil {
+				io.LogError(tr.Err)
+				return tr.Err
+			}
+			tr.Tree.AddLength(addlengthfactor, brleninternal, brlenexternal)
+			f.WriteString(tr.Tree.Newick() + "\n")
+		}
+		return
+	},
+}
+
+func init() {
+	brlenCmd.AddCommand(brlenAddCmd)
+	brlenAddCmd.Flags().Float64VarP(&addlengthfactor, "addlentgh", "l", 0.0, "Length to add to all branches")
+	brlenAddCmd.PersistentFlags().StringVarP(&outtreefile, "output", "o", "stdout", "Output tree file")
+}
diff --git a/test.sh b/test.sh
index 0c35c1e..cf77a97 100755
--- a/test.sh
+++ b/test.sh
@@ -739,6 +739,31 @@ ${GOTREE} generate yuletree --seed 10 -l 10 | ${GOTREE} brlen clear | ${GOTREE}
 diff -q -b expected result
 rm -f expected result
 
+echo "->gotree brlen add"
+cat > expected1 <<EOF
+(A:11,(B:1,C:1):11);
+EOF
+
+cat > expected2 <<EOF
+(A:10,(B,C):11);
+EOF
+
+cat > expected3 <<EOF
+(A:11,(B:1,C:1):10);
+EOF
+
+echo "(A:10.0,(B,C):10);" | $GOTREE brlen add -l 1.0 > result
+diff -q -b expected1 result
+rm -f expected result
+
+echo "(A:10.0,(B,C):10);" | $GOTREE brlen add -l 1.0 --external=false > result
+diff -q -b expected2 result
+rm -f expected result
+
+echo "(A:10.0,(B,C):10);" | $GOTREE brlen add -l 1.0 --internal=false > result
+diff -q -b expected3 result
+rm -f expected result
+
 echo "->gotree prune"
 cat > expected <<EOF
 ((Tip4,(Tip7,Tip2)),((Tip8,(Tip9,Tip3)),((Tip6,Tip5),Tip1)),Tip0);
@@ -2141,4 +2166,3 @@ EOF
 echo "(((A:1,B:1):1,C:2):3,(D:4,(E:3,((F:1.5,G:1.5):0.5,H:2):1):1):1);" | $GOTREE ltt > result
 diff -q -b expected result
 rm -f expected result
-
diff --git a/tree/tree.go b/tree/tree.go
index 26c08f0..db8e679 100644
--- a/tree/tree.go
+++ b/tree/tree.go
@@ -1897,6 +1897,21 @@ func (t *Tree) ScaleLengths(factor float64, internal, external bool) {
 	}
 }
 
+// AddLength adds the given length to all branches of the tree.
+//
+// If a branch has NIL_LENGTH: then sets its length to the given length
+func (t *Tree) AddLength(brlen float64, internal, external bool) {
+	for _, e := range t.Edges() {
+		if (e.right.Tip() && external) || (!e.right.Tip() && internal) {
+			if s := e.Length(); s != NIL_LENGTH {
+				e.SetLength(s + brlen)
+			} else {
+				e.SetLength(brlen)
+			}
+		}
+	}
+}
+
 // Rounds branch lengths by a given precision. Precision is defined as 1/(power of 10)
 // Example: 6, means 10^-6.
 //