1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49

use crate::gfa::gfa::GFAtk;
use crate::load::{load_gfa, load_gfa_stdin};
use crate::utils;
use anyhow::{bail, Context, Result};
use petgraph::graph::NodeIndex;

/// Supply a sequence/segment ID from the GFA, and extract the GFA with all nodes connected to the input node.
///
/// For example:
/// ```bash
/// gfatk extract in.gfa -s 1 > out.gfa
/// ```
pub fn extract(matches: &clap::ArgMatches) -> Result<()> {
    // read in path and parse gfa
    let gfa_file = matches.value_of("GFA");
    let sequence_ids: Vec<usize> = matches.values_of_t("sequence-ids")?.into_iter().collect();
    let iterations: i32 = matches.value_of_t("iterations")?;

    let gfa: GFAtk = match gfa_file {
        Some(f) => {
            if !f.ends_with(".gfa") {
                bail!("Input file is not a GFA.")
            }
            GFAtk(load_gfa(f)?)
        }
        None => match utils::is_stdin() {
            true => GFAtk(load_gfa_stdin(std::io::stdin().lock())?),
            false => bail!("No input from STDIN. Run `gfatk extract -h` for help."),
        },
    };

    let (graph_indices, gfa_graph) = gfa.into_ungraph()?;

    // get the node index of the target sequence ID.
    let target_indices = sequence_ids
        .iter()
        .map(|e| graph_indices.seg_id_to_node_index(*e))
        .collect::<Result<Vec<NodeIndex>>>();

    let target_indices =
        target_indices.context("One of your input segment ID's does not exist in the graph.")?;

    let sequences_to_keep =
        gfa_graph.recursive_search(sequence_ids, iterations, target_indices, graph_indices)?;

    gfa.print_extract(sequences_to_keep);

    Ok(())
}