From 5639a09fb90eb1226dc184ef25be1313b1a24841 Mon Sep 17 00:00:00 2001
From: Eric Coissac <eric.coissac@metabarcoding.org>
Date: Thu, 2 Feb 2023 16:36:32 +0100
Subject: [PATCH] Add the --with-taxon-at-rank option code to obiannotate

---
 pkg/obitax/path.go                      |  2 +-
 pkg/obitax/sequence_methods.go          | 21 ++++++++++++++++-----
 pkg/obitools/obiannotate/obiannotate.go | 18 ++++++++++++++++++
 pkg/obitools/obiannotate/options.go     | 13 ++++++++++---
 4 files changed, 45 insertions(+), 9 deletions(-)

diff --git a/pkg/obitax/path.go b/pkg/obitax/path.go
index 19bec5e..ea69d96 100644
--- a/pkg/obitax/path.go
+++ b/pkg/obitax/path.go
@@ -33,7 +33,7 @@ func (taxon *TaxNode) TaxonAtRank(rank string) *TaxNode {
 		}
 	}
 
-	if taxon == taxon.pparent {
+	if taxon == taxon.pparent  && taxon.rank != rank {
 		taxon = nil
 	}
 
diff --git a/pkg/obitax/sequence_methods.go b/pkg/obitax/sequence_methods.go
index 08b8ee0..9617e4b 100644
--- a/pkg/obitax/sequence_methods.go
+++ b/pkg/obitax/sequence_methods.go
@@ -12,26 +12,37 @@ import (
 // If the taxon at the given rank doesn't exist for the taxonomy annotation
 // of the sequence, nothing happens.
 func (taxonomy *Taxonomy) SetTaxonAtRank(sequence *obiseq.BioSequence, rank string) *TaxNode {
+	var taxonAtRank *TaxNode
+
 	taxid := sequence.Taxid()
 	taxon, err := taxonomy.Taxon(taxid)
-	taxonAtRank := taxon.TaxonAtRank(rank)
-
-	if err == nil && taxonAtRank != nil {
-		sequence.SetAttribute(rank, taxonAtRank.taxid)
-		sequence.SetAttribute(rank+"_name", taxonAtRank.scientificname)
+	taxonAtRank = nil
+	if err == nil {
+		taxonAtRank = taxon.TaxonAtRank(rank)
+		if taxonAtRank != nil {
+			// log.Printf("Taxid: %d  Rank: %s --> proposed : %d (%s)", taxid, rank, taxonAtRank.taxid, *(taxonAtRank.scientificname))
+			sequence.SetAttribute(rank, taxonAtRank.taxid)
+			sequence.SetAttribute(rank+"_name", *taxonAtRank.scientificname)
+		} else {
+			sequence.SetAttribute(rank, -1)
+			sequence.SetAttribute(rank+"_name", "NA")
+		}
 	}
 
 	return taxonAtRank
 }
 
+// Setting the species of a sequence.
 func (taxonomy *Taxonomy) SetSpecies(sequence *obiseq.BioSequence) *TaxNode {
 	return taxonomy.SetTaxonAtRank(sequence, "species")
 }
 
+// Setting the genus of a sequence.
 func (taxonomy *Taxonomy) SetGenus(sequence *obiseq.BioSequence) *TaxNode {
 	return taxonomy.SetTaxonAtRank(sequence, "genus")
 }
 
+// Setting the family of a sequence.
 func (taxonomy *Taxonomy) SetFamily(sequence *obiseq.BioSequence) *TaxNode {
 	return taxonomy.SetTaxonAtRank(sequence, "family")
 }
diff --git a/pkg/obitools/obiannotate/obiannotate.go b/pkg/obitools/obiannotate/obiannotate.go
index a1f5c37..b57db51 100644
--- a/pkg/obitools/obiannotate/obiannotate.go
+++ b/pkg/obitools/obiannotate/obiannotate.go
@@ -3,6 +3,7 @@ package obiannotate
 import (
 	"git.metabarcoding.org/lecasofts/go/obitools/pkg/obiiter"
 	"git.metabarcoding.org/lecasofts/go/obitools/pkg/obiseq"
+	"git.metabarcoding.org/lecasofts/go/obitools/pkg/obitax"
 	"git.metabarcoding.org/lecasofts/go/obitools/pkg/obitools/obigrep"
 )
 
@@ -43,6 +44,17 @@ func RenameAttributeWorker(toBeRenamed map[string]string) obiseq.SeqWorker {
 	return f
 }
 
+func AddTaxonAtRankWorker(taxonomy *obitax.Taxonomy, ranks ...string) obiseq.SeqWorker {
+	f := func(s *obiseq.BioSequence) *obiseq.BioSequence {
+		for _, r := range ranks {
+			taxonomy.SetTaxonAtRank(s,r)
+		}
+		return s
+	}
+
+	return f
+}
+
 func CLIAnnotationWorker() obiseq.SeqWorker {
 	var annotator obiseq.SeqWorker
 	annotator = nil
@@ -62,6 +74,12 @@ func CLIAnnotationWorker() obiseq.SeqWorker {
 		annotator = annotator.ChainWorkers(w)
 	}
 
+	if CLIHasTaxonAtRank() {
+		taxo := obigrep.CLILoadSelectedTaxonomy()
+		w := AddTaxonAtRankWorker(taxo,CLITaxonAtRank()...)
+		annotator = annotator.ChainWorkers(w)
+	}
+
 	return annotator
 }
 
diff --git a/pkg/obitools/obiannotate/options.go b/pkg/obitools/obiannotate/options.go
index fb0018f..eb691ab 100644
--- a/pkg/obitools/obiannotate/options.go
+++ b/pkg/obitools/obiannotate/options.go
@@ -112,12 +112,19 @@ func CLIHasToBeKeptAttributes() bool {
 }
 
 func CLIToBeKeptAttributes() map[string]bool {
-	d := make(map[string]bool,len(_keepOnly))
+	d := make(map[string]bool, len(_keepOnly))
 
-	for _,v := range _keepOnly {
-		d[v]=true
+	for _, v := range _keepOnly {
+		d[v] = true
 	}
 
 	return d
 }
 
+func CLIHasTaxonAtRank() bool {
+	return len(_taxonAtRank) > 0
+}
+
+func CLITaxonAtRank() []string {
+	return _taxonAtRank
+}