1
use std::io;
2
use std::path::PathBuf;
3
use std::thread::{self, JoinHandle};
4
use std::ffi::{OsStr, OsString};
5
use std::collections::HashMap;
6

            
7
use color_eyre::{eyre::eyre, Result};
8
use csv::{Reader, ReaderBuilder, Writer, WriterBuilder};
9
use bgzip::tabix::Tabix;
10
use crossbeam_channel::Receiver;
11
use eyre::Context;
12
use rust_htslib::bcf::{header::HeaderRecord, IndexedReader, Read};
13

            
14
30
pub fn get_tsv_reader<R: io::Read>(input: R, has_headers: bool) -> Reader<R> {
15
30
    let rdr = ReaderBuilder::new()
16
30
        .delimiter(b'\t')
17
30
        .has_headers(has_headers)
18
30
        .flexible(false)
19
30
        .from_reader(input);
20
30
    rdr
21
30
}
22

            
23
78
pub fn get_csv_reader<R: io::Read>(input: R) -> Reader<R> {
24
78
    let rdr = ReaderBuilder::new()
25
78
        .delimiter(b',')
26
78
        .has_headers(true)
27
78
        .flexible(false)
28
78
        .from_reader(input);
29
78
    rdr
30
78
}
31

            
32
204
pub fn get_csv_writer<W: io::Write>(output: W) -> Writer<W> {
33
204
    let wrtr = WriterBuilder::new()
34
204
        .delimiter(b',')
35
204
        .has_headers(false)
36
204
        .flexible(true)
37
204
        .from_writer(output);
38
204
    wrtr
39
204
}
40

            
41
108
pub fn get_input(filename: Option<PathBuf>) -> Result<Box<dyn io::Read>> {
42
108
    let input: Box<dyn io::Read> = match filename {
43
108
        Some(name) => match name.to_str() {
44
108
            Some("-") => Box::new(io::stdin()),
45
108
            Some(name) => {
46
108
                let r = match niffler::from_path(name) {
47
108
                    Ok(x) => x.0,
48
                    Err(err) => {
49
                        let msg = format!("failed to open \"{}\": {err}", name);
50
                        return Err(eyre!(msg))?;
51
                    }
52
                };
53
108
                Box::new(r)
54
            },
55
            None => return Err(eyre!("Unknown I/O error")),
56
        },
57
        None => Box::new(io::stdin()),
58
    };
59
108
    Ok(input)
60
108
}
61

            
62
204
pub fn get_output(filename: Option<PathBuf>) -> Result<Box<dyn io::Write>> {
63
204
    let output: Box<dyn io::Write> = match filename {
64
204
        Some(name) => match name.to_str() {
65
204
            Some("-") => Box::new(io::stdout()),
66
204
            Some(name) => Box::new(
67
204
                match std::fs::File::options().create(true).write(true).truncate(true).open(name) {
68
204
                    Ok(x) => x,
69
                    Err(err) => {
70
                        let msg = format!("failed to open \"{}\": {err}", name);
71
                        return Err(eyre!(msg))?;
72
                    }
73
                }
74
            ),
75
            None => return Err(eyre!("Unknown I/O error")),
76
        },
77
        None => Box::new(io::stdout()),
78
    };
79
204
    Ok(output)
80
204
}
81

            
82
198
pub fn open_csv_writer(name: PathBuf) -> Result<Writer<Box<dyn io::Write>>> {
83
198
    Ok(get_csv_writer(get_output(Some(name))?))
84
198
}
85

            
86
6
pub fn spawn_csv_collector(
87
6
    rx: Receiver<Vec<String>>,
88
6
    header: Vec<String>,
89
6
    outname: Option<PathBuf>,
90
6
    sort: bool,
91
6
) -> JoinHandle<Result<()>> {
92
6
    thread::spawn(move || -> Result<()> {
93
6
        let output = get_output(outname)?;
94
6
        let mut output = get_csv_writer(output);
95
6
        output.write_record(header)?;
96

            
97
6
        match sort {
98
            true => {
99
                let mut lines = vec![];
100

            
101
                while let Ok(line) = rx.recv() {
102
                    lines.push(line);
103
                }
104

            
105
                lines.sort_by(|a, b| alphanumeric_sort::compare_str(&a[0], &b[0]));
106

            
107
                for line in lines {
108
                    output.write_record(&line)?;
109
                }
110
            }
111
            false => {
112
12
                while let Ok(line) = rx.recv() {
113
6
                    output.write_record(&line)?;
114
                }
115
            }
116
        }
117

            
118
6
        Ok(())
119
6
    })
120
6
}
121

            
122
12
pub fn append_ext(ext: impl AsRef<OsStr>, path: &PathBuf) -> PathBuf {
123
12
    let mut os_string: OsString = path.into();
124
12
    os_string.push(".");
125
12
    os_string.push(ext.as_ref());
126
12
    os_string.into()
127
12
}
128

            
129

            
130
12
pub fn read_tabix(path: &PathBuf) -> Result<HashMap<String, (u64, usize)>> {
131
12
    let path = append_ext("tbi", path);
132
12
    let mut file = match std::fs::File::open(&path) {
133
12
        Ok(x) => x,
134
        Err(err) => {
135
            let msg = format!("failed to open \"{}\": {err}", path.display());
136
            return Err(eyre!(msg))?;
137
        }
138
    };
139
12
    let tabix = Tabix::from_reader(&mut file)?;
140

            
141
12
    let mut contigs: HashMap<String, (u64, usize)> = HashMap::new();
142

            
143
12
    let names = tabix
144
12
        .names
145
12
        .iter()
146
12
        .map(|n| Ok(std::str::from_utf8(n)?))
147
12
        .collect::<Result<Vec<&str>>>()?;
148

            
149
12
    let start = tabix
150
12
        .sequences
151
12
        .iter()
152
12
        .map(|n| {
153
12
            Ok((
154
12
                *n.intervals.first().unwrap(),
155
12
                n.number_of_distinct_bin as usize,
156
12
            ))
157
12
        })
158
12
        .collect::<Result<Vec<(u64, usize)>>>()?;
159

            
160
12
    names.iter().zip(start).for_each(|(name, (start, end))| {
161
12
        let _ = contigs.insert(name.trim_matches(char::from(0)).to_string(), (start, end));
162
12
    });
163
12

            
164
12
    Ok(contigs)
165
12
}
166

            
167
156
pub fn parse_snp_coord(coords: &str) -> Result<(&str, i64)> {
168
156
    let mut coord_split = coords.split(':');
169
156

            
170
156
    match (coord_split.next(), coord_split.next()) {
171
156
        (Some(contig), Some(value)) => {
172
156
            let value = value.parse::<i64>().wrap_err(eyre!("Failed to parse coords: {coords}"))?;
173
156
            Ok((contig, value))
174
        }
175
        _ => return Err(eyre!("Failed to parse coords: {coords}"))
176
    }
177
156
}
178

            
179
12
pub fn parse_snp_only_contig(coords: &str) -> Result<&str> {
180
12
    let mut coord_split = coords.split(':');
181
12

            
182
12
    match coord_split.next() {
183
12
        Some(contig) => Ok(contig),
184
        _ => return Err(eyre!("Failed to parse coords: {coords}"))
185
    }
186
12
}
187

            
188
6
pub fn get_htslib_bcf_contigs(path: &PathBuf) -> Result<Vec<(String, i64)>> {
189
6
    IndexedReader::from_path(path)?
190
6
        .header()
191
6
        .header_records()
192
6
        .iter()
193
36
        .filter(|r| matches!(r, HeaderRecord::Contig { .. }))
194
6
        .map(|r| match r {
195
6
            HeaderRecord::Contig { values, .. } => {
196
6
                let id = values.get("ID").expect(&format!(
197
6
                    "The input VCF has no ID for some contig record in the header",
198
6
                ));
199
6
                Ok((
200
6
                    id.clone(),
201
6
                    values
202
6
                        .get("length")
203
6
                        .expect(&format!("VCF header has no contig length for {id}"))
204
6
                        .parse::<i64>()?,
205
                ))
206
            }
207
            _ => {
208
                panic!("Header contig filtering failed for some reason. Create an issue at Github")
209
            }
210
6
        })
211
6
        .collect()
212
6
}
213

            
214

            
215
6
pub fn centromeres_hg38(chr: &str) -> (u64, u64) {
216
6
    match chr {
217
6
        "chr1" => (121700000,125100000),
218
6
        "chr2" => (91800000,96000000),
219
6
        "chr3" => (87800000,94000000),
220
6
        "chr4" => (48200000,51800000),
221
6
        "chr5" => (46100000,51400000),
222
6
        "chr6" => (58500000,62600000),
223
6
        "chr7" => (58100000,62100000),
224
6
        "chr8" => (43200000,47200000),
225
6
        "chr9" => (42200000,45500000),
226
        "chr10" => (38000000,41600000),
227
        "chr11" => (51000000,55800000),
228
        "chr12" => (33200000,37800000),
229
        "chr13" => (16500000,18900000),
230
        "chr14" => (16100000,18200000),
231
        "chr15" => (17500000,20500000),
232
        "chr16" => (35300000,38400000),
233
        "chr17" => (22700000,27400000),
234
        "chr18" => (15400000,21500000),
235
        "chr19" => (24200000,28100000),
236
        "chr20" => (25700000,30400000),
237
        "chr21" => (10900000,13000000),
238
        "chr22" => (13700000,17400000),
239
        "chrX" => (58100000,63800000),
240
        "chrY" => (10300000,10600000),
241
        _ => panic!(),
242
    }
243
6
}
244