1
use std::path::PathBuf;
2

            
3
use color_eyre::{eyre::eyre, Result};
4

            
5
use crate::args::{Selection, StandardArgs};
6
use crate::read_vcf::read_vcf_to_matrix;
7
use crate::structs::{HapVariant, PhasedMatrix};
8
use crate::utils::{
9
    select_carrier_haplotypes, select_only_longest_haplotypes, shared_lengths_by_majority, push_to_output
10
};
11
use crate::core::{open_csv_writer, parse_snp_coord};
12

            
13
#[doc(hidden)]
14
18
pub fn run(args: StandardArgs, haplotype_path: PathBuf) -> Result<()> {
15
18
    let (contig, variant_pos) = parse_snp_coord(&args.coords)?;
16

            
17
18
    let mut csv_output = args.output.clone();
18
18
    push_to_output(&args, &mut csv_output, "haplotype_check", "csv");
19
18
    let mut writer = open_csv_writer(csv_output)?;
20

            
21
18
    let ht = crate::io::read_haplotype_file(haplotype_path)?;
22
18
    let start = ht.first().unwrap();
23
18
    let end = ht.last().unwrap();
24
18

            
25
18
    tracing::debug!("Reading into a sample-variant matrix");
26
18
    let vcf = match args.selection {
27
        Selection::OnlyAlts | Selection::OnlyRefs => {
28
6
            let vcf = read_vcf_to_matrix(
29
6
                &args,
30
6
                contig,
31
6
                variant_pos,
32
6
                Some((start.pos as u64, end.pos as u64)),
33
6
            )?;
34
6
            select_carrier_haplotypes(vcf, variant_pos, &args.coords, &args.selection)
35
        }
36
        Selection::OnlyLongest => {
37
6
            let vcf = read_vcf_to_matrix(&args, contig, variant_pos, None)?;
38
6
            let shared_lengths = shared_lengths_by_majority(&vcf, vcf.variant_idx());
39
6
            select_only_longest_haplotypes(&shared_lengths, vcf)
40
        }
41
6
        Selection::All => read_vcf_to_matrix(
42
6
            &args,
43
6
            contig,
44
6
            variant_pos,
45
6
            Some((start.pos as u64, end.pos as u64)),
46
6
        )?,
47
        Selection::Unphased => return Err(eyre!("Running with unphased data is not supported."))
48
    };
49

            
50
18
    tracing::debug!("Finished reading the matrix");
51

            
52
18
    let matching_indexes = identical_haplotype_count(&vcf, &ht);
53
18
    write_matches_to_csv(&matching_indexes, &mut writer, &vcf)?;
54

            
55
18
    tracing::debug!(
56
        "Haplotype found in {}/{} of alleles.",
57
        matching_indexes.len(),
58
        vcf.samples().len()
59
    );
60

            
61
18
    Ok(())
62
18
}
63

            
64
fn write_matches_to_csv(
65
    matching_indexes: &[usize],
66
    writer: &mut csv::Writer<Box<dyn std::io::Write>>,
67
    vcf: &PhasedMatrix,
68
) -> Result<()> {
69
18
    writer.write_record(vec!["id", "match"])?;
70
336
    for (idx, sample) in vcf.samples().iter().enumerate() {
71
336
        let allele = if idx >= vcf.nsamples() / 2 { 2 } else { 1 };
72
336
        writer.write_record(vec![
73
336
            format!("{sample}_{allele}"),
74
336
            format!("{}", matching_indexes.contains(&idx)),
75
336
        ])?;
76
    }
77
18
    Ok(())
78
18
}
79

            
80
54
pub fn identical_haplotype_count(vcf: &PhasedMatrix, ht: &[HapVariant]) -> Vec<usize> {
81
54
    let indexes: Vec<(usize, usize)> = ht
82
54
        .iter()
83
54
        .enumerate()
84
54
        .filter_map(|(coords_i, h)| {
85
            // Check the PartialEq implementation between HapVariant and Coord
86
15042
            if let Some(vcf_i) = vcf.coords().iter().position(|c| c == h) {
87
630
                Some((vcf_i, coords_i))
88
            } else {
89
12
                tracing::warn!("Haplotype variant {:?} not found in the vcf", h);
90
12
                None
91
            }
92
642
        })
93
54
        .collect();
94
54

            
95
54
    if indexes.is_empty() {
96
6
        return vec![];
97
48
    }
98
48

            
99
48
    let mut matching = vec![true; vcf.matrix.nrows()];
100

            
101
1176
    for (n, value) in matching.iter_mut().enumerate() {
102
8088
        'inner: for (vcf_i, coords_i) in &indexes {
103
7818
            if vcf.matrix[[n, *vcf_i]] != ht[*coords_i].gt {
104
906
                *value = false;
105
906
                break 'inner;
106
6912
            }
107
        }
108
    }
109

            
110
48
    matching
111
48
        .iter()
112
48
        .enumerate()
113
1176
        .filter_map(|(i, m)| if *m { Some(i) } else { None })
114
48
        .collect()
115
54
}
116