write out ref_data.json with it's unique pattern

mandysulli · mandysulli · commit f4d5aa8a1d2a · 2025-08-05T09:10:08.000-04:00
diff --git a/src/processes/prepare_mira_reports.rs b/src/processes/prepare_mira_reports.rs
@@ -3,7 +3,8 @@ use crate::utils::{data_ingest::*, writing_outputs::*};
 use clap::Parser;
 use csv::ReaderBuilder;
 use either::Either;
-use serde::{self, Deserialize, de::DeserializeOwned};
+use serde::{self, Deserialize, Serialize, de::DeserializeOwned};
+use serde_json::json;
 use std::sync::Arc;
 use std::{
     error::Error,
@@ -115,9 +116,19 @@ pub fn prepare_mira_reports_process(args: ReportsArgs) -> Result<(), Box<dyn Err
     let vtype_data = create_vtype_data(&read_data);
     let allele_data = allele_data_collection(&args.irma_path)?;
     let indel_data = indels_data_collection(&args.irma_path)?;
-    let ref_lengths = get_reference_lens(&args.irma_path);
     //TODO: feed in organism from argument
     let seq_data = amended_consensus_data_collection(&args.irma_path, "flu");
+    let ref_lengths = match get_reference_lens(&args.irma_path) {
+        Ok(data) => data,
+        Err(e) => {
+            eprintln!("Error getting reference lengths: {e}");
+            return Err(e);
+        }
+    };
+    let (segments, segset, segcolor) =
+        return_seg_data(extract_field(coverage_data.clone(), |item| {
+            item.reference_name.clone()
+        }));
 
     //Read in DAIS-ribosome data
     let dais_ins_data = dias_insertion_data_collection(&args.irma_path);
@@ -136,7 +147,7 @@ pub fn prepare_mira_reports_process(args: ReportsArgs) -> Result<(), Box<dyn Err
     //println!("dais del data: {dais_del_data:#?}");
     //println!("dais seq data: {dais_seq_data:#?}");
     //println!("dais ref data: {dais_ref_data:#?}");
-    println!("ref length data: {ref_lengths:#?}");
+    //println!("ref length data: {ref_lengths:#?}");
 
     /////////////// Write the structs to JSON files and CSV files ///////////////
     // Writing out coverage data
@@ -290,6 +301,14 @@ pub fn prepare_mira_reports_process(args: ReportsArgs) -> Result<(), Box<dyn Err
         &indels_struct_values,
     )?;
 
+    write_ref_data_json(
+        "/home/xpa3/mira-oxide/test/ref_data.json",
+        &ref_lengths,
+        &segments,
+        &segset,
+        &segcolor,
+    )?;
+
     // Write fields to parq if flag given
     write_reads_to_parquet(&read_data, "/home/xpa3/mira-oxide/test/read_data.parquet")?;
 
diff --git a/src/utils/data_ingest.rs b/src/utils/data_ingest.rs
@@ -3,7 +3,7 @@ use either::Either;
 use glob::glob;
 use serde::{self, Deserialize, Serialize, de::DeserializeOwned};
 use std::{
-    collections::HashMap,
+    collections::{HashMap, HashSet},
     error::Error,
     fs::{self, File, OpenOptions},
     io::{self, BufRead, BufReader, Read, Stdin, stdin},
@@ -142,12 +142,6 @@ pub struct SeqData {
     sequence: String,
 }
 
-#[derive(Debug)]
-pub struct RefLengthData {
-    name: String,
-    length: usize,
-}
-
 /////////////// Structs to hold dais-ribosome data ///////////////
 /// Insertion Data
 #[derive(Serialize, Deserialize, Debug)]
@@ -602,61 +596,88 @@ pub fn create_vtype_data(reads_data: &Vec<ReadsData>) -> Vec<ProcessedRecord> {
     processed_records
 }
 
-/// Collect the reference lengths from the IRMA outputs
+// Function to collect reference lengths from IRMA outputs
 pub fn get_reference_lens(
     irma_path: &PathBuf,
-) -> Result<Vec<RefLengthData>, Box<dyn std::error::Error>> {
-    // Determine the glob pattern based on the organism
+) -> Result<HashMap<String, usize>, Box<dyn std::error::Error>> {
     let pattern = format!(
         "{}/*/IRMA/*/intermediate/0-ITERATIVE-REFERENCES/R0*ref",
         irma_path.to_string_lossy()
     );
 
-    let mut ref_len_data: Vec<RefLengthData> = Vec::new();
+    let mut ref_len_map: HashMap<String, usize> = HashMap::new();
 
-    // Iterate over all files matching the pattern
-    for entry in glob(&pattern).expect("Failed to read glob pattern") {
+    for entry in glob(&pattern)? {
         match entry {
             Ok(path) => {
                 let file = File::open(&path)?;
                 let reader = BufReader::new(file);
 
-                // Parse the file line by line (assuming FASTA format)
                 let mut ref_name = String::new();
                 let mut current_sequence = String::new();
 
                 for line in reader.lines() {
                     let line = line?;
                     if line.starts_with('>') {
-                        // If there's an existing sequence, save its length
                         if !ref_name.is_empty() {
-                            ref_len_data.push(RefLengthData {
-                                name: ref_name.clone(),
-                                length: current_sequence.len(),
-                            });
+                            ref_len_map.insert(ref_name.clone(), current_sequence.len());
                         }
-                        // Start a new sequence
-                        ref_name = line[1..].to_string(); // Remove '>'
+                        ref_name = line[1..].to_string();
                         current_sequence.clear();
                     } else {
-                        // Append to the current sequence
                         current_sequence.push_str(&line);
                     }
                 }
 
-                // Save the last sequence's length
                 if !ref_name.is_empty() {
-                    ref_len_data.push(RefLengthData {
-                        name: ref_name,
-                        length: current_sequence.len(),
-                    });
+                    // Insert the last reference name and its length into the HashMap
+                    ref_len_map.insert(ref_name, current_sequence.len());
                 }
             }
             Err(e) => println!("Error reading file: {e}"),
         }
     }
 
-    Ok(ref_len_data)
+    Ok(ref_len_map)
+}
+
+// Function to process reference names and generate segments, segset, and segcolor
+pub fn return_seg_data(
+    reference_names: Vec<String>,
+) -> (Vec<String>, Vec<String>, HashMap<String, &'static str>) {
+    let mut segments: Vec<String> = reference_names.into_iter().collect();
+    segments.sort();
+    segments.dedup();
+
+    let color_palette = vec![
+        "#1f77b4", "#ff7f0e", "#2ca02c", "#d62728", "#9467bd", "#8c564b", "#e377c2", "#7f7f7f",
+        "#bcbd22", "#17becf",
+    ];
+
+    let mut segset: Vec<String> = Vec::new();
+    for segment in &segments {
+        let parts: Vec<&str> = segment.split('_').collect();
+        if parts.len() > 1 {
+            segset.push(parts[1].to_string());
+        } else {
+            segset.push(segment.clone());
+        }
+    }
+
+    let segset: Vec<String> = segset
+        .into_iter()
+        .collect::<std::collections::HashSet<_>>()
+        .into_iter()
+        .collect();
+
+    let mut segcolor: HashMap<String, &str> = HashMap::new();
+    for (i, seg) in segset.iter().enumerate() {
+        if i < color_palette.len() {
+            segcolor.insert(seg.clone(), color_palette[i]);
+        }
+    }
+
+    (segments, segset, segcolor)
 }
 
 /////////////// Data reading functions for DAIS-ribosome ///////////////
diff --git a/src/utils/writing_outputs.rs b/src/utils/writing_outputs.rs
@@ -6,6 +6,8 @@ use csv::Writer;
 use parquet::arrow::ArrowWriter;
 use serde::Serialize;
 use serde_json::{Value, json};
+use std::collections::HashMap;
+use std::io::Write;
 use std::sync::Arc;
 use std::{error::Error, fs::File};
 
@@ -81,6 +83,29 @@ pub fn write_structs_to_csv_file<T: Serialize>(
     Ok(())
 }
 
+/// make ref_data.json - has unique set up
+pub fn write_ref_data_json(
+    file_path: &str,
+    ref_lens: &HashMap<String, usize>,
+    segments: &Vec<String>,
+    segset: &Vec<String>,
+    segcolor: &HashMap<String, &str>,
+) -> Result<(), Box<dyn Error>> {
+    let json_data = json!({
+        "ref_lens": ref_lens,
+        "segments": segments,
+        "segset": segset,
+        "segcolor": segcolor,
+    });
+
+    // Write JSON to a file
+    let mut file = File::create(file_path)?;
+    file.write_all(serde_json::to_string_pretty(&json_data)?.as_bytes())?;
+
+    println!("Data written to ref_data.json");
+
+    Ok(())
+}
 /////////////// Functions to write parquet files out ///////////////
 /// Functions to convert values in a vecxtor of structs to vector
 /// Some perform type converions
