1
use std::path::PathBuf;
2

            
3
use color_eyre::{eyre::eyre, Result};
4
use plotters::prelude::*;
5

            
6
use crate::args::{GraphArgs, Selection, StandardArgs};
7
use crate::read_vcf::read_vcf_to_matrix;
8
use crate::structs::{HapVariant, PhasedMatrix, SharedLength};
9
use crate::utils::{
10
    push_to_output, push_to_output_mbah, select_carrier_haplotypes, select_only_longest_alleles,
11
    shared_lengths_by_majority,
12
};
13
use crate::core::{open_csv_writer, parse_snp_coord};
14

            
15
#[doc(hidden)]
16
24
#[tracing::instrument]
17
12
pub fn run(
18
    args: StandardArgs,
19
    graph_args: GraphArgs,
20
    x_start: Option<i64>,
21
    x_stop: Option<i64>,
22
) -> Result<()> {
23
    if args.selection == Selection::Unphased {
24
        return Err(eyre!("Running with unphased data is not supported."))
25
    }
26

            
27
    let mut decay_csv = args.output.clone();
28
    push_to_output(&args, &mut decay_csv, "shares_of_mbah", "csv");
29
    let mut decay_writer = open_csv_writer(decay_csv)?;
30

            
31
    let mut ancestral_haplo_csv = args.output.clone();
32
    push_to_output_mbah(&args, &mut ancestral_haplo_csv, "mbah", "csv");
33
    let mut ancestral_writer = open_csv_writer(ancestral_haplo_csv)?;
34

            
35
    let (contig, variant_pos) = parse_snp_coord(&args.coords)?;
36
    let mut vcf = read_vcf_to_matrix(&args, contig, variant_pos, None)?;
37

            
38
    if args.selection == Selection::OnlyAlts || args.selection == Selection::OnlyRefs {
39
        vcf = select_carrier_haplotypes(vcf, variant_pos, &args.coords, &args.selection)
40
    };
41

            
42
    let shared_lengths = shared_lengths_by_majority(&vcf, vcf.variant_idx());
43
    let only_longest = args.selection == Selection::OnlyLongest;
44
    print_haplotype_lengths_csv(&shared_lengths, &vcf, only_longest, &mut decay_writer)?;
45

            
46
    let haplotype = find_ancestral_haplotype(contig, &shared_lengths, &vcf)?;
47
    print_ancestral_haplotype_csv(&haplotype, contig, &mut ancestral_writer)?;
48

            
49
    draw_plot(
50
        &args,
51
        graph_args,
52
        args.output.clone(),
53
        shared_lengths,
54
        only_longest,
55
        x_start,
56
        x_stop,
57
    )?;
58

            
59
    Ok(())
60
}
61

            
62
12
fn draw_plot(
63
12
    st_args: &StandardArgs,
64
12
    args: GraphArgs,
65
12
    mut output: PathBuf,
66
12
    sl: Vec<SharedLength>,
67
12
    only_longest: bool,
68
12
    x_start: Option<i64>,
69
12
    x_stop: Option<i64>,
70
12
) -> Result<()> {
71
12
    push_to_output(st_args, &mut output, "shares_of_mbah", "png");
72
12

            
73
12
    let root = BitMapBackend::new(&output, (args.width, args.height)).into_drawing_area();
74

            
75
12
    let mut sl = match only_longest {
76
        true => select_only_longest_alleles(&sl),
77
12
        false => sl.iter().collect(),
78
    };
79

            
80
1380
    sl.sort_by(|a, b| b.end_pos.partial_cmp(&a.end_pos).unwrap());
81
492
    sl.sort_by(|a, b| a.start_pos.partial_cmp(&b.start_pos).unwrap());
82

            
83
12
    let x_start = match x_start {
84
        Some(x_start) => x_start,
85
        None => {
86
12
            sl.iter()
87
324
                .min_by(|a, b| a.start_pos.cmp(&b.start_pos))
88
12
                .unwrap()
89
12
                .start_pos
90
        }
91
    };
92

            
93
12
    let x_stop = match x_stop {
94
        Some(x_stop) => x_stop,
95
        None => {
96
12
            sl.iter()
97
324
                .max_by(|a, b| a.end_pos.cmp(&b.end_pos))
98
12
                .unwrap()
99
12
                .end_pos
100
        }
101
    };
102

            
103
12
    root.fill(&WHITE)?;
104
12
    let mut chart = ChartBuilder::on(&root)
105
12
        .margin(10)
106
12
        .set_label_area_size(LabelAreaPosition::Left, 60)
107
12
        .set_label_area_size(LabelAreaPosition::Right, 60)
108
12
        .set_label_area_size(LabelAreaPosition::Bottom, 40)
109
12
        .build_cartesian_2d(x_start..x_stop, 0..sl.len())?;
110

            
111
12
    chart
112
12
        .configure_mesh()
113
12
        .disable_x_mesh()
114
12
        .disable_y_mesh()
115
12
        .x_labels(30)
116
12
        .max_light_lines(4)
117
12
        .y_desc("SAMPLE")
118
12
        .x_desc("POS")
119
12
        .draw()?;
120

            
121
336
    for (idx, length) in sl.iter().enumerate() {
122
336
        chart.draw_series(LineSeries::new(
123
336
            [(length.start_pos, idx), (length.end_pos, idx)],
124
336
            &BLACK,
125
336
        ))?;
126
    }
127

            
128
12
    root.present()?;
129

            
130
12
    Ok(())
131
12
}
132

            
133
#[doc(hidden)]
134
fn print_ancestral_haplotype_csv(
135
    haplotype: &Vec<HapVariant>,
136
    contig: &str,
137
    writer: &mut csv::Writer<Box<dyn std::io::Write>>,
138
) -> Result<()> {
139
12
    writer.write_record(vec!["contig", "pos", "ref", "alt", "gt"])?;
140
696
    for variant in haplotype {
141
684
        writer.write_record(vec![
142
684
            contig,
143
684
            &variant.pos.to_string(),
144
684
            &variant.reference,
145
684
            &variant.alt,
146
684
            &variant.gt.to_string(),
147
684
        ])?;
148
    }
149
12
    Ok(())
150
12
}
151

            
152
#[doc(hidden)]
153
12
fn find_ancestral_haplotype(
154
12
    contig: &str,
155
12
    shared_lengths: &[SharedLength],
156
12
    vcf: &PhasedMatrix,
157
12
) -> Result<Vec<HapVariant>> {
158
336
    let max = shared_lengths.iter().max_by_key(|l| l.end_pos).unwrap();
159
336
    let min = shared_lengths.iter().min_by_key(|l| l.start_pos).unwrap();
160
12

            
161
12
    let left = min.start()..vcf.variant_idx();
162
12
    let right = vcf.variant_idx()..max.end() + 1;
163
12

            
164
12
    let mut left = vcf.find_haplotype_for_sample(contig, left, min.index);
165
12
    let right = vcf.find_haplotype_for_sample(contig, right, max.index);
166
12

            
167
12
    left.extend(right);
168
12

            
169
12
    Ok(left)
170
12
}
171

            
172
#[doc(hidden)]
173
fn print_haplotype_lengths_csv(
174
    shared_lengths: &Vec<SharedLength>,
175
    vcf: &PhasedMatrix,
176
    only_longest: bool,
177
    writer: &mut csv::Writer<Box<dyn std::io::Write>>,
178
) -> Result<()> {
179
12
    writer.write_record(vec!["id", "start", "stop"])?;
180

            
181
12
    match only_longest {
182
        true => {
183
            for length in select_only_longest_alleles(shared_lengths) {
184
                writer.write_record(vec![
185
                    vcf.get_sample_name(length.index),
186
                    length.start_pos.to_string(),
187
                    length.end_pos.to_string(),
188
                ])?;
189
            }
190
        }
191
        false => {
192
348
            for length in shared_lengths {
193
336
                if length.index >= vcf.nsamples() / 2 {
194
168
                    writer.write_record(vec![
195
168
                        format!("{}_2", vcf.get_sample_name(length.index)),
196
168
                        length.start_pos.to_string(),
197
168
                        length.end_pos.to_string(),
198
168
                    ])?;
199
                } else {
200
168
                    writer.write_record(vec![
201
168
                        format!("{}_1", vcf.get_sample_name(length.index)),
202
168
                        length.start_pos.to_string(),
203
168
                        length.end_pos.to_string(),
204
168
                    ])?;
205
                }
206
            }
207
        }
208
    }
209
12
    Ok(())
210
12
}