1
use std::collections::BTreeMap;
2
use std::path::PathBuf;
3

            
4
use color_eyre::Result;
5
use termion::color;
6

            
7
use crate::{io::read_haplotype_file, structs::HapVariant};
8
use crate::core::open_csv_writer;
9

            
10
6
pub fn run(
11
6
    haplotypes: Vec<PathBuf>,
12
6
    mut output: PathBuf,
13
6
    prefix: Option<String>,
14
6
    csv: bool,
15
6
    hide_missing: bool,
16
6
    tag_rows: bool,
17
6
) -> Result<()> {
18
6
    match prefix {
19
6
        None => output.push(format!("ht_comparison.csv")),
20
        Some(prefix) => output.push(format!("{prefix}_ht_comparison.csv")),
21
    }
22

            
23
6
    let mut names = vec![];
24
6
    let mut hts = vec![];
25
18
    for path in haplotypes {
26
12
        names.push(path.display().to_string());
27
12
        hts.push(read_haplotype_file(path)?);
28
    }
29

            
30
6
    let aligner = HaplotypeAligner::new(hts, names);
31
6

            
32
6
    match csv {
33
6
        true => aligner.to_csv(output)?,
34
        false => aligner.print(hide_missing, tag_rows),
35
    }
36
6
    Ok(())
37
6
}
38

            
39
pub struct HaplotypeAligner {
40
    names: Vec<String>,
41
    alignment: BTreeMap<i64, Vec<Option<String>>>,
42
}
43

            
44
impl HaplotypeAligner {
45
7
    pub fn new(haplotypes: Vec<Vec<HapVariant>>, names: Vec<String>) -> Self {
46
7
        let mut last_contig = &haplotypes[0][0].contig;
47
7

            
48
7
        // Register all positions first
49
7
        let mut positions = vec![];
50
22
        for ht in &haplotypes {
51
15
            let curr_contig = &ht[0].contig;
52
15
            if last_contig != curr_contig {
53
                panic!("Haplotypes are from different contigs {last_contig} vs {curr_contig}")
54
15
            }
55
15
            last_contig = curr_contig;
56
653
            for variant in ht {
57
638
                positions.push((variant.pos, vec![None; names.len()]))
58
            }
59
        }
60
7
        positions.sort();
61
7
        positions.dedup();
62
7

            
63
7
        // Register present or missing genotypes for each position
64
7
        let mut alignment = BTreeMap::from_iter(positions);
65
391
        for (k, v) in alignment.iter_mut() {
66
831
            for (idx, ht) in haplotypes.iter().enumerate() {
67
36317
                for var in ht {
68
35486
                    if k == &var.pos {
69
638
                        v[idx] = Some(var.genotype().to_owned());
70
34848
                    }
71
                }
72
            }
73
        }
74

            
75
7
        Self { names, alignment }
76
7
    }
77

            
78
    pub fn print(&self, hide_missing: bool, tag_rows: bool) {
79
        // Prepare header
80
        let names = self
81
            .names
82
            .iter()
83
            .fold(" pos".to_string(), |s, r| format!("{s} {r}"));
84

            
85
        println!("{}{names} status", color::Fg(color::AnsiValue::rgb(3, 3, 5)));
86

            
87
        for (k, v) in &self.alignment {
88
            // If missing genotypes present, yellow is true
89
            // If contradictory genotypes present, red is also true
90
            let (yellow, red) = self.find_mismatch(v);
91
            let color = if red {
92
                color::Fg(color::AnsiValue::rgb(5, 0, 1))
93
            } else if yellow {
94
                color::Fg(color::AnsiValue::rgb(5, 5, 1))
95
            } else {
96
                color::Fg(color::AnsiValue::rgb(0, 5, 1))
97
            };
98

            
99
            // Switch Nones to - and fold the vec into a string
100
            let line = v
101
                .iter()
102
                .map(|v| match v {
103
                    None => "-".to_string(),
104
                    Some(v) => v.to_owned(),
105
                })
106
                .fold(format!(" {k}"), |s, r| format!("{s} {r}"));
107

            
108
            if !hide_missing || !yellow {
109
                if tag_rows {
110
                    match (yellow, red) {
111
                        (_, true) => println!("{color}{line} err"),
112
                        (true, false) => println!("{color}{line} mis"),
113
                        (false, false) => println!("{color}{line} ok"),
114
                    }
115
                } else {
116
                    println!("{color}{line}");
117
                }
118
            }
119

            
120
        }
121
    }
122

            
123
15
    fn find_mismatch(&self, v: &[Option<String>]) -> (bool, bool) {
124
15
        let mut values = v.iter().flatten().collect::<Vec<&String>>();
125
15
        let yellow = values.len() != v.len();
126
15
        values.sort();
127
15
        values.dedup();
128
15
        let red = values.len() > 1;
129
15
        (yellow, red)
130
15
    }
131

            
132
6
    pub fn to_csv(&self, path: PathBuf) -> Result<()> {
133
6
        let mut wrtr = open_csv_writer(path)?;
134
6
        let mut header = vec!["pos".to_string()];
135
6
        header.extend(self.names.clone());
136
6

            
137
6
        wrtr.write_record(&header)?;
138
348
        for (k, v) in &self.alignment {
139
342
            let mut record = vec![format!("{k}")];
140
684
            v.iter().for_each(|v| match v {
141
156
                None => record.push("-".to_string()),
142
528
                Some(v) => record.push(v.to_owned()),
143
684
            });
144
342
            wrtr.write_record(&record)?;
145
        }
146
6
        Ok(())
147
6
    }
148
}
149

            
150
#[cfg(test)]
151
#[rustfmt::skip]
152
mod tests {
153
    use super::*;
154

            
155
110
    fn get_ref_alt() -> (String, String) {
156
110
        ("A".to_string(), "T".to_string())
157
110
    }
158

            
159
3
    fn create_haplotype(missing: i64, gt: i64) -> Vec<HapVariant> {
160
3
        let mut haps = vec![];
161
153
        for i in 0..50 {
162
150
            if i % missing != 0 {
163
110
                let (reference, alt) = get_ref_alt();
164
110
                let gt = match i % gt == 0 {
165
9
                    true => 1,
166
101
                    false => 0,
167
                };
168
110
                let hap = HapVariant {
169
110
                    contig: "".to_string(),
170
110
                    pos: i,
171
110
                    reference,
172
110
                    alt,
173
110
                    gt,
174
110
                };
175
110
                haps.push(hap);
176
40
            }
177
        }
178
3
        haps
179
3
    }
180

            
181
1
    #[test]
182
1
    fn test_aligner_alignment() {
183
1
        let hap1 = create_haplotype(11, 9);
184
1
        let hap2 = create_haplotype(2, 10);
185
1
        let hap3 = create_haplotype(5, 11);
186
1
        let aligner = HaplotypeAligner::new(
187
1
            vec![hap1, hap2, hap3],
188
1
            vec!["hap1".to_string(), "hap2".to_string(), "hap3".to_string()],
189
1
        );
190
1
        let pos: Vec<i64> = (1..50).collect();
191
1
        let colors = vec![
192
1
            (false, false),
193
1
            (true, false),
194
1
            (false, false),
195
1
            (true, false),
196
1
            (true, false),
197
1
            (true, false),
198
1
            (false, false),
199
1
            (true, false),
200
1
            (false, true),
201
1
            (true, false),
202
1
        ];
203

            
204
10
        for (i, (k, v)) in aligner.alignment.iter().enumerate() {
205
10
            assert_eq!(&pos[i], k);
206
10
            assert_eq!(colors[i], aligner.find_mismatch(v));
207

            
208
10
            if i == 9 {
209
1
                break;
210
9
            }
211
        }
212
1
    }
213

            
214
1
    #[test]
215
1
    fn test_aligner_mismatch() {
216
1
        let aligner = HaplotypeAligner { names : vec![], alignment: BTreeMap::new() };
217
1

            
218
1
        let foo: Vec<Option<String>> = vec![Some("A".into()), None, Some("T".into())];
219
1
        let (yellow, red) = aligner.find_mismatch(&foo);
220
1
        assert_eq!((yellow, red), (true, true));
221

            
222
1
        let foo: Vec<Option<String>> = vec![Some("A".into()), Some("A".into()), Some("T".into())];
223
1
        let (yellow, red) = aligner.find_mismatch(&foo);
224
1
        assert_eq!((yellow, red), (false, true));
225

            
226
1
        let foo: Vec<Option<String>> = vec![Some("A".into()), None, None];
227
1
        let (yellow, red) = aligner.find_mismatch(&foo);
228
1
        assert_eq!((yellow, red), (true, false));
229

            
230
1
        let foo: Vec<Option<String>> = vec![None, None, None];
231
1
        let (yellow, red) = aligner.find_mismatch(&foo);
232
1
        assert_eq!((yellow, red), (true, false));
233

            
234
1
        let foo: Vec<Option<String>> = vec![Some("A".into()), Some("A".into()), Some("A".into())];
235
1
        let (yellow, red) = aligner.find_mismatch(&foo);
236
1
        assert_eq!((yellow, red), (false, false));
237
1
    }
238
}