1
use std::ops::Range;
2

            
3
use color_eyre::Result;
4
use ndarray::{s, ArrayBase, Axis, Dim, ViewRepr};
5

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

            
14
#[doc(hidden)]
15
132
#[tracing::instrument]
16
66
pub fn run(args: StandardArgs, missing_allowed: bool) -> Result<()> {
17
    let mut output = args.output.clone();
18
    push_to_output(&args, &mut output, "shared_haplotype", "csv");
19

            
20
    let writer = open_csv_writer(output)?;
21
    let (contig, variant_pos) = parse_snp_coord(&args.coords)?;
22

            
23
    let mut vcf = if missing_allowed {
24
        read_vcf_to_matrix_missing_allowed(&args, contig, variant_pos, None)?
25
    } else {
26
        read_vcf_to_matrix(&args, contig, variant_pos, None)?
27
    };
28

            
29
    if args.selection == Selection::OnlyAlts || args.selection == Selection::OnlyRefs {
30
        vcf = select_carrier_haplotypes(vcf, variant_pos, &args.coords, &args.selection)
31
    }
32

            
33
    if args.selection == Selection::OnlyLongest {
34
        let shared_lengths = shared_lengths_by_majority(&vcf, vcf.variant_idx());
35
        vcf = select_only_longest_haplotypes(&shared_lengths, vcf);
36
    }
37

            
38
    find_and_write_shared_haplotype(&vcf, args.selection == Selection::Unphased, writer, contig)?;
39
    Ok(())
40
}
41

            
42
#[rustfmt::skip]
43
pub fn find_and_write_shared_haplotype(
44
    vcf: &PhasedMatrix,
45
    is_rwc: bool,
46
    mut writer: csv::Writer<Box<dyn std::io::Write>>,
47
    contig: &str,
48
) -> Result<()> {
49
66
    writer.write_record(vec!["contig", "pos", "ref", "alt", "gt"])?;
50
66
    if is_rwc {
51
6
        let (left_gt, right_gt) = find_shared_rwc_run(vcf);
52
6

            
53
6
        let left_count = left_gt.len();
54
186
        for (i, gt) in left_gt.into_iter().rev().enumerate() {
55
186
            let coords = &vcf.coords().get(vcf.variant_idx - left_count + i).unwrap();
56
186
            write_haplo_record(&mut writer, contig, coords, gt)?;
57
        }
58
192
        for (i, gt) in right_gt.into_iter().enumerate() {
59
192
            let coords = &vcf.coords().get(vcf.variant_idx + i).unwrap();
60
192
            write_haplo_record(&mut writer, contig, coords, gt)?;
61
        }
62
    } else {
63
60
        let range = find_shared_haplotype_phased(vcf);
64

            
65
1860
        for (gt, index) in vcf.matrix.slice(s![0, range.clone()]).into_iter().zip(range) {
66
1860
            let coords = &vcf.coords().get(index).unwrap();
67
1860
            write_haplo_record(&mut writer, contig, coords, *gt)?;
68
        }
69

            
70
    }
71
66
    Ok(())
72
66
}
73

            
74
#[doc(hidden)]
75
fn write_haplo_record(
76
    writer: &mut csv::Writer<Box<dyn std::io::Write + 'static>>,
77
    contig: &str,
78
    coords: &crate::structs::Coord,
79
    gt: u8,
80
) -> Result<()> {
81
2238
    writer.write_record(vec![
82
2238
        contig,
83
2238
        &coords.pos.to_string(),
84
2238
        &coords.reference,
85
2238
        &coords.alt,
86
2238
        &gt.to_string(),
87
2238
    ])?;
88
2238
    Ok(())
89
2238
}
90

            
91
#[doc(hidden)]
92
6
pub fn find_shared_rwc_run(vcf: &PhasedMatrix) -> (Vec<u8>, Vec<u8>) {
93
6
    let slice = vcf.matrix.slice(s![.., vcf.variant_idx..]);
94
6
    let right_gt = find_breakpoint_unphased(slice);
95
6

            
96
6
    let mut slice = vcf.matrix.slice(s![.., ..vcf.variant_idx]);
97
6
    slice.invert_axis(Axis(1));
98
6
    let left_gt = find_breakpoint_unphased(slice);
99
6

            
100
6
    (left_gt, right_gt)
101
6
}
102

            
103
#[doc(hidden)]
104
12
pub fn find_breakpoint_unphased(matrix: ArrayBase<ViewRepr<&u8>, Dim<[usize; 2]>>) -> Vec<u8> {
105
12
    let mut gt_vec = vec![];
106

            
107
378
    for i in 0..matrix.ncols() {
108
378
        let column = matrix.slice(ndarray::s![.., i]);
109
10584
        let ones = column.into_iter().filter(|&n| *n == 1).count();
110
378

            
111
378
        let alt_af = ones as f32 / matrix.nrows() as f32;
112
378

            
113
378
        // Check no contrahomozygotes are present
114
378
        let mut alt_homozygote = false;
115
378
        let mut ref_homozygote = false;
116
5292
        for i in 0..matrix.nrows() / 2 {
117
5292
            if *column.slice(s![i]).into_scalar() == 0
118
5124
                && *column.slice(s![i + matrix.nrows() / 2]).into_scalar() == 0
119
            {
120
4872
                ref_homozygote = true
121
420
            } else if *column.slice(s![i]).into_scalar() == 1
122
168
                && *column.slice(s![i + matrix.nrows() / 2]).into_scalar() == 1
123
            {
124
                alt_homozygote = true
125
420
            }
126
        }
127

            
128
378
        if alt_homozygote && ref_homozygote {
129
            break;
130
378
        }
131
378

            
132
378
        if alt_af >= 0.5 {
133
6
            gt_vec.push(1_u8);
134
372
        } else {
135
372
            gt_vec.push(0_u8);
136
372
        }
137
    }
138
12
    gt_vec
139
12
}
140

            
141
#[doc(hidden)]
142
72
pub fn find_shared_haplotype_phased(vcf: &PhasedMatrix) -> Range<usize> {
143
72
    let variant_slice = vcf.matrix.slice(s![.., vcf.variant_idx]);
144
1092
    let zeroes = variant_slice.into_iter().filter(|&n| *n == 0).count();
145
1092
    let ones = variant_slice.into_iter().filter(|&n| *n == 1).count();
146
72

            
147
72
    if zeroes != 0 && ones != 0 {
148
6
        return vcf.variant_idx..vcf.variant_idx;
149
66
    }
150
66

            
151
66
    let slice = vcf.matrix.slice(s![.., vcf.variant_idx..]);
152
66
    let right_count = find_breakpoint(slice);
153
66

            
154
66
    let mut slice = vcf.matrix.slice(s![.., ..vcf.variant_idx]);
155
66
    slice.invert_axis(Axis(1));
156
66
    let left_count = find_breakpoint(slice);
157
66

            
158
66
    let start = vcf.variant_idx - left_count;
159
66
    let stop = vcf.variant_idx + right_count;
160
66

            
161
66
    start..stop
162
72
}
163

            
164
//
165
#[doc(hidden)]
166
#[rustfmt::skip]
167
132
pub fn find_breakpoint(matrix: ArrayBase<ViewRepr<&u8>, Dim<[usize; 2]>>) -> usize {
168
132
    let mut counter = 0;
169

            
170
2178
    for i in 0..matrix.ncols() {
171
2178
        let column = matrix.slice(ndarray::s![.., i]);
172
2178

            
173
30492
        let zeroes = column.into_iter().filter(|&n| *n == 0).count();
174
30492
        let ones = column.into_iter().filter(|&n| *n == 1).count();
175
2178

            
176
2178
        if zeroes != 0 && ones != 0 {
177
132
            break;
178
2046
        } else {
179
2046
            counter += 1;
180
2046
        }
181
    }
182
132
    counter
183
132
}