1
use std::collections::BTreeMap;
2
use std::fs::File;
3
use std::io::{self, BufRead};
4
use std::path::{Path, PathBuf};
5

            
6
use color_eyre::{
7
    eyre::{eyre, WrapErr},
8
    Result,
9
};
10
use polars::prelude::*;
11

            
12
use crate::structs::HapVariant;
13
use crate::core::{get_csv_reader, get_input, get_tsv_reader};
14

            
15
30
pub fn read_variable_data_file(path: PathBuf) -> Result<DataFrame> {
16
30
    let df = CsvReader::from_path(path)?
17
24
        .with_null_values(Some(NullValues::AllColumnsSingle("NA".to_string())))
18
24
        .infer_schema(Some(500))
19
24
        .has_header(true)
20
24
        .finish()?;
21

            
22
24
    Ok(df)
23
30
}
24

            
25
1518
#[derive(Debug, Clone, serde::Deserialize)]
26
struct RecombinationRow<'a> {
27
    _chr: &'a str,
28
    pos: u64,
29
    _rate: f32,
30
    cm: f32,
31
}
32

            
33
30
pub fn read_recombination_file(path: PathBuf) -> Result<BTreeMap<u64, f32>> {
34
30
    let mut rdr = get_tsv_reader(get_input(Some(path))?, false);
35
30
    let mut rates = BTreeMap::new();
36
30

            
37
30
    let mut last_cm = 0.0;
38
30
    let mut last_pos = 0;
39
1518
    for line in rdr.records() {
40
1518
        let record = line?;
41
1518
        let row: RecombinationRow = record.deserialize(None).
42
1518
            wrap_err(eyre!("Make sure no headers are present and that the recombination file is in order chr,pos,rate,cm. If the issue is not fixed, you have an invalid field in the file."))?;
43
1512
        if last_cm > row.cm {
44
            return Err(eyre!("Recombination rates file is not sorted for centimorgans"))
45
1512
        }
46
1512
        if last_pos > row.pos {
47
            return Err(eyre!("Recombination rates file is not sorted for positions"))
48
1512
        }
49
1512
        last_cm = row.cm;
50
1512
        last_pos = row.pos;
51
1512
        rates.insert(row.pos, row.cm);
52
    }
53

            
54
24
    Ok(rates)
55
30
}
56

            
57
78
pub fn read_haplotype_file(ht_path: PathBuf) -> Result<Vec<HapVariant>> {
58
78
    let input = get_input(Some(ht_path))?;
59
78
    let mut rdr = get_csv_reader(input);
60
78

            
61
78
    let mut variants: Vec<HapVariant> = vec![];
62

            
63
2502
    for line in rdr.records() {
64
2502
        let record = line?;
65
2502
        let variant: HapVariant = record.deserialize(None)?;
66
2502
        if let Some(latest) = variants.last() {
67
2424
            if latest.pos > variant.pos {
68
6
                return Err(eyre::eyre!(
69
6
                    "The haplotype file is not sorted by position. {} is larger than {}",
70
6
                    latest,
71
6
                    variant
72
6
                ));
73
2418
            }
74
78
        }
75
2496
        variants.push(variant);
76
    }
77

            
78
72
    Ok(variants)
79
78
}
80

            
81
24
pub fn read_lines<P>(filename: P) -> Result<io::Lines<io::BufReader<File>>>
82
24
where
83
24
    P: AsRef<Path>,
84
24
{
85
24
    // let file = File::open(filename)?;
86
24
    let name = filename.as_ref().display();
87
24
    let file = match File::open(&filename) {
88
18
        Ok(x) => x,
89
6
        Err(err) => {
90
6
            let msg = format!("failed to open {name}: {err}");
91
6
            return Err(std::io::Error::new(std::io::ErrorKind::NotFound, msg))?;
92
        }
93
    };
94
18
    Ok(io::BufReader::new(file).lines())
95
24
}
96

            
97
918
pub fn read_sample_ids(path: &Option<PathBuf>) -> Result<Option<Vec<String>>> {
98
918
    match path {
99
12
        Some(path) => {
100
12
            let mut samples = vec![];
101

            
102
48
            for line in read_lines(path)?.flatten() {
103
48
                let line = line.trim();
104
48
                samples.push(line.to_string());
105
48
            }
106
6
            Ok(Some(samples))
107
        }
108
906
        None => Ok(None),
109
    }
110
918
}