scale_bc_parser.py

#!/usr/bin/env python3
"""
LongReader Barcode Processing Pipeline

This script processes long-read sequencing data from a FASTQ file to identify,
orient, trim, and parse barcode elements based on a YAML configuration.
It segregates reads into passing and failing files and generates a detailed log.
"""

import argparse
import gzip
import logging
import os
import re
import time
from dataclasses import dataclass
from itertools import product
from typing import Any, Dict, List, Optional, Tuple

import parasail
import pyfastx
import yaml
from Levenshtein import distance, hamming
from rich.console import Console
from rich.text import Text

#########################
# Configuration Classes
#########################


@dataclass
class SequenceElement:
  """
  Represents a single component of the expected sequence structure, such as a
  barcode, UMI, or primer.
  """
  name: str
  type: str  # "restricted" or "unrestricted"
  length: Optional[int] = None
  restriction_type: Optional[str] = None
  options: Optional[List[Tuple[Optional[str], str]]] = None
  options_file: Optional[str] = None
  output_mode: str = "skip"

  def load_options_from_fasta(self):
    """Loads barcode options from a specified FASTA file."""
    fasta = pyfastx.Fasta(self.options_file, build_index=False)
    self.options = [(name, seq) for name, seq in fasta]

  def __post_init__(self):
    """Validates the configuration of the sequence element after initialization."""
    if self.type not in ["restricted", "unrestricted"]:
      raise ValueError(
          f"Element '{self.name}' must have 'type' set to either 'restricted' or 'unrestricted'"
      )
    if self.type != "unrestricted" and (not isinstance(self.length, int) or
                                        self.length < 1):
      raise ValueError(
          f"Element '{self.name}' must have a positive integer 'length'")
    if self.restriction_type and self.restriction_type not in [
        "whitelist", "length"
    ]:
      raise ValueError(
          f"Element '{self.name}' must have 'restriction_type' set to either 'whitelist' or 'length'"
      )
    if self.restriction_type == "whitelist":
      if self.options and self.options_file:
        raise ValueError(
            f"Element '{self.name}' cannot have both 'options' and 'options_file' specified."
        )
      if self.options_file:
        if not os.path.exists(self.options_file):
          raise ValueError(
              f"Options file '{self.options_file}' for element '{self.name}' does not exist."
          )
        self.load_options_from_fasta()
        print(
            f"↳ Loaded '{self.name}' with {len(self.options)} sequence(s) from '{self.options_file}'"
        )
      elif self.options:
        self.options = [(None, opt) for opt in self.options if opt]
        print(
            f"↳ Loaded '{self.name}' with {len(self.options)} specified sequence(s)"
        )
      else:
        raise ValueError(
            f"Element '{self.name}' must have either 'options' or 'options_file' specified when using 'whitelist' restriction type."
        )
    if self.output_mode not in ["sequence", "header", "skip"]:
      raise ValueError(
          f"Element '{self.name}' must have 'output_mode' set to either 'sequence', 'header', or 'skip'"
      )
    if self.name == "transcript":
      print("↳ Loaded 'transcript' placeholder element")

  def get_query_sequence(self) -> str:
    """
    Generates a query sequence for alignment. Returns a specific sequence if
    it's a unique whitelisted option, otherwise returns a string of 'N's.
    """
    if self.restriction_type == "whitelist" and len(self.options) == 1:
      return self.options[0][1]
    return self.length * "N"


@dataclass
class BarcodeConfig:
  """
  Holds the entire configuration for the barcode structure, loaded from a
  YAML file.
  """
  name: str
  description: str
  date: str
  author: str
  min_poly_length: int
  sequence_structure: List[SequenceElement]
  header_template: str

  @classmethod
  def from_yaml(cls, yaml_data: dict,
                header_template_override: Optional[str] = None):
    """Constructs a BarcodeConfig object from a parsed YAML dictionary."""
    name = yaml_data.get("name", "LongReaderConfig")
    description = yaml_data.get("description", "")
    date = yaml_data.get("date", "")
    author = yaml_data.get("author", "")
    min_poly_length = yaml_data.get("min_poly_length", 20)
    header_template = (header_template_override
                       if header_template_override else
                       yaml_data["output_format"]["header_template"])
    include_fields = {
        field["name"]: field["output_mode"]
        for field in yaml_data["output_format"]["include_fields"]
    }

    sequence_structure = []
    for element_data in yaml_data["sequence_structure"]:
      output_mode = include_fields.get(element_data["name"], "skip")
      sequence_structure.append(
          SequenceElement(**element_data, output_mode=output_mode))

    for field_name in include_fields:
      if not any(element.name == field_name for element in sequence_structure):
        raise ValueError(
            f"Output field '{field_name}' specified in 'include_fields' does not match any element in 'sequence_structure'."
        )

    transcript_elements = [
        element for element in sequence_structure if element.name == "transcript"
    ]
    if len(transcript_elements) != 1:
      raise ValueError(
          "There must be exactly one 'transcript' element in the sequence structure."
      )

    return cls(name, description, date, author, min_poly_length,
               sequence_structure, header_template)

  def get_5p_elements(self) -> List[SequenceElement]:
    """Returns a list of sequence elements that are 5' to the transcript."""
    transcript_idx = next(index for index, element in enumerate(
        self.sequence_structure) if element.name == "transcript")
    return self.sequence_structure[1:transcript_idx]

  def get_3p_elements(self) -> List[SequenceElement]:
    """Returns a list of sequence elements that are 3' to the transcript."""
    transcript_idx = next(index for index, element in enumerate(
        self.sequence_structure) if element.name == "transcript")
    return self.sequence_structure[transcript_idx + 1:-1]

  def get_5p_primer_sequence(self) -> str:
    """Returns the 5' primer sequence."""
    return self.sequence_structure[0].options[0][1]

  def get_3p_primer_sequence(self) -> str:
    """Returns the 3' primer sequence."""
    return self.sequence_structure[-1].options[0][1]

  def get_linker_sequence(self) -> str:
    """Returns the linker sequence from the configuration."""
    for element in self.sequence_structure:
      if element.name == "linker":
        if element.options and len(element.options) == 1:
          return element.options[0][1]
    raise ValueError(
        "A unique linker sequence was not found in the configuration.")

  def get_linker_to_3p_distance(self) -> int:
    """
    Calculates the total length of all elements between the 'linker' and
    the 3' primer in the sequence structure.
    """
    try:
      linker_index = next(i for i, el in enumerate(self.sequence_structure)
                          if el.name == "linker")
      # The 3' primer is assumed to be the last element.
      three_prime_index = len(self.sequence_structure) - 1
    except StopIteration:
      raise ValueError(
          "A 'linker' element was not found in the sequence structure.")

    distance = 0
    if linker_index < three_prime_index - 1:
      for i in range(linker_index + 1, three_prime_index):
        element = self.sequence_structure[i]
        if element.length:
          distance += element.length
    return distance


class LongReader:
  """
    A class to process long-read FASTQ files, identify barcodes, and trim
    sequences based on a provided configuration.
    """

  def __init__(self, config_path: str, input_fastq_paths: List[str],
               output_prefix: str,
               header_template_override: Optional[str] = None,
               trim_polyA: bool = False):
    """
        Initialize LongReader with paths and an optional header override.
        """
    if not os.path.exists(config_path):
      raise ValueError(f"Configuration file {config_path} not found.")
    with open(config_path, "r") as config_file:
      yaml_data = yaml.safe_load(config_file)

    self.barcode_config = BarcodeConfig.from_yaml(yaml_data,
                                                  header_template_override)
    self.input_fastq_paths = input_fastq_paths
    self.output_prefix = output_prefix
    self.orientations = {}
    self.trim_polyA = trim_polyA

  def register_orientation(self, orientation: str, element_order: list):
    """
        Register a valid orientation by specifying the ordered list of elements.
        """
    self.orientations[orientation] = element_order

  def iterate_reads(self):
    """
        Generator to iterate over reads from input FASTQ files.
        Yields tuples of (header, sequence, quality).
        """
    for fastq_path in self.input_fastq_paths:
      for name, seq, qual in pyfastx.Fastq(fastq_path, build_index=False):
        yield name, seq, qual

  def reverse_complement(self, sequence: str) -> str:
    """Return the reverse complement of a DNA sequence."""
    complement_map = str.maketrans("ACGTacgt", "TGCAtgca")
    return sequence.translate(complement_map)[::-1]

  def split_string_into_chunks(self, input_string: str,
                               lengths: List[int]) -> List[str]:
    """
        Splits the input string into a list of chunks of specified lengths.
        """
    start_index = 0
    chunks = []
    for length in lengths:
      chunks.append(input_string[start_index:start_index + length])
      start_index += length
    return chunks

  def create_custom_matrix(self):
    """
        Creates a custom Parasail scoring matrix where N-to-base matches are
        scored positively, rather than being treated as mismatches.
        """
    # Create a matrix with a match score of 5 and a mismatch score of -1.
    custom_matrix = parasail.matrix_create("ACGTN", 5, -1)

    # Modify the matrix for 'N' characters.
    matrix_size = len(custom_matrix.matrix)
    n_index = 4  # The 0-based index for 'N'
    n_match_score = 2

    for i in range(matrix_size - 1):
      custom_matrix[n_index, i] = n_match_score  # N to base match
      custom_matrix[i, n_index] = n_match_score  # Base to N match

    return custom_matrix

  def highlight_sequence(self, read_sequence: str,
                         annotations: List[Dict[str, Any]], orientation: str,
                         extra_details: str = ""):
    """
        Prints a sequence to the console with colored highlights for detected
        elements using the Rich library.
        """
    console = Console()
    text = Text(read_sequence)
    for annotation in annotations:
      text.stylize(
          annotation.get("color", "white"),
          annotation.get("start", 0),
          annotation.get("end", 0),
      )

    details_text = Text()
    details_text.append("Alignment Details:\n", style="bold")
    details_text.append(f"↳ Orientation: {orientation}\n", style="bold magenta")
    for annotation in annotations:
      details_text.append(
          f"↳ {annotation.get('label', ''):<15} at [{annotation.get('start'):>4}, {annotation.get('end'):>4})\n",
          style=annotation.get("color", "white"),
      )
    if extra_details:
      details_text.append(extra_details + "\n")

    details_text.append("\nParsed Sequence:\n", style="bold")
    details_text += text
    console.print(details_text)

  def find_ordered_alignments(
      self, alignment_sets: List[List[dict]]) -> List[Tuple[dict, ...]]:
    """
        Generates all combinations (Cartesian product) from a list of alignment
        sets and returns only those where alignments occur in sequential order.
        """

    def is_valid_sequence(sequence: Tuple[dict, ...]) -> bool:
      """Check if alignment start/end coordinates are ordered correctly."""
      return all(sequence[i]['end'] <= sequence[i + 1]['start']
                 for i in range(len(sequence) - 1))

    all_combinations = product(*alignment_sets)
    valid_combinations = [
        combination for combination in all_combinations if is_valid_sequence(combination)
    ]
    return valid_combinations

  def iterative_parasail_masking(self,
                                 query: str,
                                 reference: str,
                                 scoring_matrix,
                                 gap_open: int = 10,
                                 gap_extend: int = 1,
                                 max_hits: int = 3,
                                 mask_char: str = "X",
                                 stop_at_score_ratio: float = 0.6,
                                 require_highest_score: bool = False,
                                 match_end: bool = False,
                                 verbose: bool = False) -> List[dict]:
    """
        Performs iterative alignment and masking to find multiple non-overlapping
        hits of a query in a reference sequence.

        Returns a list of dictionaries, each containing alignment details.
        """
    alignments = []
    masked_reference = reference
    best_score = 0

    for i in range(max_hits):
      if match_end:
        # Align to the end of the reference by reversing it
        alignment = parasail.ssw(query, masked_reference[::-1], gap_open,
                                 gap_extend, scoring_matrix)
        # Convert coordinates from the reversed alignment back to original
        start = len(reference) - alignment.ref_end1 - 1
        end = len(reference) - alignment.ref_begin1 - 1
      else:
        alignment = parasail.ssw(query, masked_reference, gap_open,
                                 gap_extend, scoring_matrix)
        start = alignment.ref_begin1
        end = alignment.ref_end1

      score = alignment.score1
      if score > best_score:
        best_score = score

      # Stop if the score is too low or if no alignment is found
      if score == 0 or score < best_score * stop_at_score_ratio:
        break
      if require_highest_score and score < best_score:
        break

      alignments.append({"score": score, "start": start, "end": end})

      if verbose:
        console = Console()
        ref_text = Text(reference)
        ref_text.stylize("red3", start, end + 1)
        verbose_text = Text()
        verbose_text.append(f"Alignment {i+1} ({score}):\n", style="bold")
        verbose_text += ref_text
        console.print(verbose_text)

      # Mask the aligned region to prevent re-alignment in the next iteration
      masked_list = list(masked_reference)
      for j in range(start, end + 1):
        masked_list[j] = mask_char
      masked_reference = "".join(masked_list)

    return alignments

  def orient_and_trim_sequence(self,
                              read_sequence: str,
                              read_quality: str,
                              custom_matrix,
                              verbose: bool = False
                             ) -> Tuple[str, str, str, str]:
    """
    Orients and trims a raw read sequence and its quality string.

    Returns:
        A tuple containing:
        - trimmed_sequence (str): The extracted cDNA sequence.
        - trimmed_quality (str): The quality string for the trimmed sequence.
        - end_sequence (str): The 3' or 5' end containing barcodes.
        - orientation (str): The determined orientation (e.g., "Forward").
    """
    # Pre-compute reverse complements
    primer_5p_seq = self.barcode_config.get_5p_primer_sequence()
    primer_5p_rc_seq = self.reverse_complement(primer_5p_seq)
    primer_3p_seq = self.barcode_config.get_3p_primer_sequence()
    primer_3p_rc_seq = self.reverse_complement(primer_3p_seq)
    linker_seq = self.barcode_config.get_linker_sequence()
    linker_rc_seq = self.reverse_complement(self.barcode_config.get_linker_sequence())
    poly_length = self.barcode_config.min_poly_length
    linker_to_3p_dist = self.barcode_config.get_linker_to_3p_distance()

    # Align all expected components in both forward and reverse orientations
    alignment_3p = self.iterative_parasail_masking(primer_3p_seq, read_sequence, custom_matrix)
    alignment_5p = self.iterative_parasail_masking(primer_5p_seq, read_sequence, custom_matrix)
    alignment_3p_rc = self.iterative_parasail_masking(primer_3p_rc_seq, read_sequence, custom_matrix)
    alignment_5p_rc = self.iterative_parasail_masking(primer_5p_rc_seq, read_sequence, custom_matrix)
    alignment_linker = self.iterative_parasail_masking(linker_seq, read_sequence, custom_matrix, max_hits=5, require_highest_score=True)
    alignment_linker_rc = self.iterative_parasail_masking(linker_rc_seq, read_sequence, custom_matrix, max_hits=5, require_highest_score=True)
    alignment_polyA = self.iterative_parasail_masking("A" * poly_length, read_sequence, custom_matrix, match_end=True)
    alignment_polyT = self.iterative_parasail_masking("T" * poly_length, read_sequence, custom_matrix)

    # Find valid, ordered sets of alignments for each orientation
    base_tolerance = 5
    # Forward = 5' primer -> polyA -> linker -> 3' primer (RC)
    forward_sets = self.find_ordered_alignments(
        [alignment_5p, alignment_polyA, alignment_linker, alignment_3p_rc])
    forward_possible = [
        f_set for f_set in forward_sets if abs(
            f_set[3]['start'] - (f_set[2]['end'] + 1) - linker_to_3p_dist) <= base_tolerance
    ]

    # Reverse = 3' primer -> linker (RC) -> polyT -> 5' primer (RC)
    reverse_sets = self.find_ordered_alignments(
      [alignment_3p, alignment_linker_rc, alignment_polyT, alignment_5p_rc])
    reverse_possible = [
      r_set for r_set in reverse_sets if abs(
        r_set[1]['start'] - (r_set[0]['end'] + 1) - linker_to_3p_dist) <= base_tolerance
    ]

    if not forward_possible and not reverse_possible:
      orientation = "Artefact"
      if verbose:
        print(
            f"↳ No valid ordered alignments found. Forward possibilities: {len(forward_possible)}, Reverse possibilities: {len(reverse_possible)}"
        )
      return "", "", "", orientation

    # Initialize variables to hold the best alignment chain
    trimmed_sequence = ""
    trimmed_quality = ""
    end_sequence = ""
    orientation = "Artefact"
    best_forward_chain = None
    best_reverse_chain = None

    # Determine orientation based on the presence and score of valid alignment chains
    if forward_possible and not reverse_possible:
      orientation = "Forward"
      best_forward_chain = max(
          forward_possible,
          key=lambda chain: sum(item['score'] for item in chain))
    elif reverse_possible and not forward_possible:
      orientation = "Reverse"
      best_reverse_chain = max(
          reverse_possible,
          key=lambda chain: sum(item['score'] for item in chain))
    elif forward_possible and reverse_possible:
      # If both are possible, it's ambiguous; pick the one with the higher total alignment score
      best_forward_chain = max(
          forward_possible,
          key=lambda chain: sum(item['score'] for item in chain))
      forward_score = sum(item['score'] for item in best_forward_chain)

      best_reverse_chain = max(
          reverse_possible,
          key=lambda chain: sum(item['score'] for item in chain))
      reverse_score = sum(item['score'] for item in best_reverse_chain)

      if forward_score >= reverse_score:
        orientation = "Ambiguous_Forward"
      else:
        orientation = "Ambiguous_Reverse"

    # Process the winning orientation to extract sequences
    annotations = []
    if "Forward" in orientation and best_forward_chain:
      first_primer, poly, linker, last_primer = best_forward_chain
      trim_start = first_primer['end']
      trim_end = poly['start'] if self.trim_polyA else last_primer['start']
      annotations.extend([
        {"label": "5'_adapter", "start": first_primer['start'], "end": first_primer['end'] + 1, "color": "red3"},
        {"label": "Transcript", "start": trim_start, "end": trim_end + 1, "color": "dark_cyan"},
        {"label": "PolyA", "start": poly['start'], "end": poly['end'] + 1, "color": "dark_goldenrod"},
        {"label": "Linker", "start": linker['start'], "end": linker['end'] + 1, "color": "deep_pink4"},
        {"label": "3'_adapter_rc", "start": last_primer['start'], "end": last_primer['end'] + 1, "color": "dodger_blue2"}
      ])
      trimmed_sequence = read_sequence[trim_start:trim_end]
      trimmed_quality = read_quality[trim_start:trim_end]
      end_sequence = read_sequence[poly['start']:last_primer['end'] + 1]
    elif "Reverse" in orientation and best_reverse_chain:
      first_primer, linker, poly, last_primer = best_reverse_chain
      trim_start = poly['end'] if self.trim_polyA else first_primer['end']
      trim_end = last_primer['start']
      annotations.extend([
        {"label": "3'_adapter", "start": first_primer['start'], "end": first_primer['end'] + 1, "color": "green"},
        {"label": "Transcript", "start": trim_start, "end": trim_end + 1, "color": "dark_cyan"},
        {"label": "PolyT", "start": poly['start'], "end": poly['end'] + 1, "color": "dark_goldenrod"},
        {"label": "Linker_rc", "start": linker['start'], "end": linker['end'] + 1, "color": "deep_pink4"},
        {"label": "5'_adapter_rc", "start": last_primer['start'], "end": last_primer['end'] + 1, "color": "indian_red"}
      ])
      trimmed_sequence = self.reverse_complement(read_sequence[trim_start:trim_end])
      trimmed_quality = read_quality[trim_start:trim_end][::-1]
      end_sequence = self.reverse_complement(read_sequence[first_primer['start']:poly['end'] + 1])

    # If trim_polyA is enabled, perform a final trim of any residual polyA tail
    # (only if there are >=3 trailing 'A's from the 3' end)
    if self.trim_polyA and trimmed_sequence:
      original_len = len(trimmed_sequence)
      trimmed_sequence = re.sub(r'A{3,}$', '', trimmed_sequence)
      if len(trimmed_sequence) != original_len:
        trimmed_quality = trimmed_quality[:len(trimmed_sequence)]

    if verbose and annotations:
      self.highlight_sequence(read_sequence, annotations, orientation)

    return trimmed_sequence, trimmed_quality, end_sequence, orientation

  def find_best_match_for_element(
      self, element: SequenceElement, parsed_element_seq: str,
      min_levenshtein_threshold: int,
      min_hamming_threshold: int) -> Tuple[bool, str, str]:
    """
    Finds the best whitelist match for a parsed sequence using Levenshtein
    and Hamming-based tie-breaking.

    Returns:
      A tuple of (is_success, assignment_value, failure_reason).
    """
    name_field_idx = 1 if element.output_mode == "sequence" else 0
    get_field = lambda match: match[name_field_idx]

    # Calculate Levenshtein distances, including the special case for lig_bc_scale
    lev_distances = [(h, s, distance(parsed_element_seq, s, score_cutoff=min_levenshtein_threshold + 2)) for h, s in element.options]
    if element.name == "lig_bc_scale":
      corrected_seq = "A" + parsed_element_seq[:-1]
      lev_distances += [(h, s, distance(corrected_seq, s, score_cutoff=min_levenshtein_threshold + 2)) for h, s in element.options]

    # Determine best and passing Levenshtein matches
    min_lev_dist = min(lev_distances, key=lambda x: x[2])[2]
    best_lev_matches = [m for m in lev_distances if m[2] == min_lev_dist]
    passing_lev_matches = [m for m in best_lev_matches if m[2] <= min_levenshtein_threshold]

    # Case 1: Unique, passing Levenshtein match
    if len(passing_lev_matches) == 1:
      return True, get_field(passing_lev_matches[0]), ""

    # Case 2: Levenshtein match failed or was ambiguous
    if not passing_lev_matches:
        # Failure due to exceeding distance threshold
        reason = f"{element.name}-min_lv_dist_{min_lev_dist}_exceeds_threshold_{min_levenshtein_threshold}"
        closest_matches_str = '|'.join(map(get_field, best_lev_matches))
        assignment = f"{parsed_element_seq}({closest_matches_str})"
        return False, assignment, reason

    # Case 3: Ambiguous Levenshtein, use Hamming distance as a tie-breaker
    ham_distances = [(h, s, hamming(parsed_element_seq, s, score_cutoff=min_hamming_threshold + 2)) for h, s, _ in passing_lev_matches]
    min_ham_dist = min(ham_distances, key=lambda x: x[2])[2]
    best_ham_matches = [m for m in ham_distances if m[2] == min_ham_dist and m[2] <= min_hamming_threshold]

    if len(best_ham_matches) == 1:
      return True, get_field(best_ham_matches[0]), ""
    
    # Failure due to ambiguous match or exceeding Hamming threshold
    if not best_ham_matches:
        reason = f"{element.name}-ambiguous_lv_match_and_ham_dist_exceeds_threshold_{min_hamming_threshold}"
    else:
        reason = f"{element.name}-ambiguous_lv_match_and_ham_match_with_{len(best_ham_matches)}_sequences"
    
    closest_matches_str = '|'.join(map(get_field, best_lev_matches))
    assignment = f"{parsed_element_seq}({closest_matches_str})"
    return False, assignment, reason

  def identify_elements_from_sequence(self,
                                      end_sequence: str,
                                      elements: List[SequenceElement],
                                      custom_matrix,
                                      verbose=False,
                                      min_levenshtein_threshold=2,
                                      min_hamming_threshold=2
                                     ) -> Tuple[bool, Dict[str, str], str]:
    """
        Identifies and validates barcode elements from a given sequence segment.

        Returns:
            A tuple containing:
            - is_passing (bool): True if all elements were identified successfully.
            - element_assignments (dict): A dictionary mapping element names to sequences/headers.
            - failure_reason (str): A description of why identification failed.
        """
    is_passing = True
    failure_reason = ""
    element_assignments = {}
    element_query_seq = ''.join([element.get_query_sequence() for element in elements])

    # Find the best matching region for the combined barcode structure
    decoy_alignment = parasail.ssw(element_query_seq, end_sequence, 10, 1, custom_matrix)
    filtered_seq = end_sequence[decoy_alignment.ref_begin1:decoy_alignment.ref_end1 + 1]

    element_lengths = [element.length for element in elements if element.length]
    element_chunks = self.split_string_into_chunks(filtered_seq, element_lengths)

    for index, element in enumerate(elements):
      parsed_element_seq = element_chunks[index]
      if verbose:
        print(f"↳ Parsed sequence for '{element.name}' ({len(parsed_element_seq)} bp): {parsed_element_seq}")

      if element.restriction_type == "whitelist" and element.output_mode != "skip":
        match_is_passing, assignment, reason = self.find_best_match_for_element(
            element, parsed_element_seq, min_levenshtein_threshold,
            min_hamming_threshold)
        
        if not match_is_passing:
          is_passing = False
          failure_reason = reason
        element_assignments[element.name] = assignment

      elif element.output_mode != "skip":
        element_assignments[element.name] = parsed_element_seq
    
    if verbose:
      print(f"↳ Identification result: {is_passing}: {element_assignments}")

    # If parsing failed, populate assignments with raw chunks for logging
    if not is_passing:
      for index, element in enumerate(elements):
        if element.name not in element_assignments:
          element_assignments[element.name] = element_chunks[index]

    return (is_passing, element_assignments, failure_reason)

  def process_read(self, header: str, read_sequence: str, read_quality: str,
                   custom_matrix,
                   verbose: bool = False) -> Tuple[str, str, str, str]:
    """
        Processes a single FASTQ read, returning the trimmed sequence and orientation.
        """

    trimmed_seq, trimmed_qual, end_seq, orientation = self.orient_and_trim_sequence(read_sequence, read_quality, custom_matrix, verbose=verbose)

    if verbose:
      print(f"↳ Processed read {header} [trimmed sequence length = {len(trimmed_seq)}]: {trimmed_seq}")
    return trimmed_seq, trimmed_qual, end_seq, orientation

  def test_run(self, num_reads: int = 50, verbose: bool = True) -> None:
    """
        Runs the pipeline on a small number of reads for testing and debugging.
        """
    custom_matrix = self.create_custom_matrix()
    for i, (header, seq, qual) in enumerate(self.iterate_reads()):
      if i >= num_reads:
        break
      print(f"\n----- Processing Read {i + 1:02d} -----")
      trimmed_seq, _, _, _ = self.process_read(
          header, seq, qual, custom_matrix, verbose=verbose)
      print(f"--------------")

  def write_to_log(self, read_num: int, header: str, status: str,
                   status_detail: str, parsed_data: Dict[str, str]):
    """
        Appends a formatted entry to the run's log file.
        """
    read_id = header.split()[0]
    log_filename = self.output_prefix + ".log"

    # Format the status detail line differently for pass vs. fail
    if status.upper() == "FAIL":
      detail_line = f"↳ fail reason: {status_detail}"
    else:
      detail_line = f"↳ orientation: {status_detail.lower()}"

    log_entry = (f"({read_num}) [{status.upper()}] {read_id}\n"
                 f"{detail_line}\n")

    if parsed_data:
      log_data = {k.lower(): v for k, v in parsed_data.items()}
      data_line_parts = []
      for element in self.barcode_config.sequence_structure:
        if element.output_mode != "skip":
          element_name_lower = element.name.lower()
          if element_name_lower in log_data:
            data_line_parts.append(
                f"{element.name.upper()}={log_data[element_name_lower]}")
      log_entry += f"↳ {' '.join(data_line_parts)}\n"

    log_entry += "\n"

    with open(log_filename, "a") as log_file:
      log_file.write(log_entry)

  def log_summary(self, num_processed: int, num_passing: int, num_failing: int,
                  num_orientation_failing: int, num_parse_failing: int):
    """Writes a final summary of the run statistics to the log file."""
    summary_lines = [
      "***** SUMMARY *****",
      f"Total reads processed: {num_processed}",
      f"Passing reads: {num_passing} ({100*(num_passing/num_processed):.2f}%)",
      f"Failing reads: {num_failing} ({100*(num_failing/num_processed):.2f}%)"
    ]
    if num_failing > 0:
      summary_lines.extend([
        f"  ↳ Orientation failures: {num_orientation_failing} ({100*(num_orientation_failing/num_failing):.2f}%)",
        f"  ↳ Parsing failures: {num_parse_failing} ({100*(num_parse_failing/num_failing):.2f}%)"
      ])
    else:
      summary_lines.append("No failing reads.")

    summary_text = "\n".join(summary_lines) + "\n\n"
    log_filename = self.output_prefix + ".log"
    with open(log_filename, "a") as log_file:
      log_file.write(summary_text)

  def process_all_reads(self, verbose: bool = False):
    """
    Processes all reads from the input FASTQ, writing them to passing
    and failing output files and generating a log.
    """
    print("\nBegin processing reads...")
    start_time = time.time()
    passing_file_path = self.output_prefix + "-passing.fastq.gz"
    failing_file_path = self.output_prefix + "-failing.fastq.gz"

    passing_file_handle = gzip.open(passing_file_path, "wt")
    failing_file_handle = gzip.open(failing_file_path, "wt")

    # Initialize counters for summary statistics
    num_processed = 0
    num_passing = 0
    num_failing = 0
    num_orientation_failing = 0
    num_parse_failing = 0
    custom_matrix = self.create_custom_matrix()

    with passing_file_handle as pass_f, failing_file_handle as fail_f:
      for fastq_path in self.input_fastq_paths:
        print(f"↳ Processing {os.path.basename(fastq_path)}...")
        for header, read_sequence, read_quality in pyfastx.Fastq(fastq_path, build_index=False, full_name=True):
          num_processed += 1
          trimmed_seq, trimmed_qual, end_seq, orientation = self.process_read(
              header, read_sequence, read_quality, custom_matrix, verbose)

          # If read cannot be oriented (failure case 1)
          if not trimmed_seq:
            num_failing += 1
            num_orientation_failing += 1
            new_header = f"{header} fail_reason=ambiguous_orientation"
            fail_f.write(f"@{new_header}\n{read_sequence}\n+\n{read_quality}\n")
            self.write_to_log(num_processed, header, "FAIL", "ambiguous_orientation", {})
            continue

          # Identify barcode elements from the 5' and 3' ends
          five_prime_parsed = self.identify_elements_from_sequence( end_seq, self.barcode_config.get_5p_elements(), custom_matrix, verbose) if self.barcode_config.get_5p_elements() else (True, {}, "")
          three_prime_parsed = self.identify_elements_from_sequence(end_seq, self.barcode_config.get_3p_elements(), custom_matrix, verbose) if self.barcode_config.get_3p_elements() else (True, {}, "")
          
          five_prime_passed, five_prime_data, five_prime_fail_reason = five_prime_parsed
          three_prime_passed, three_prime_data, three_prime_fail_reason = three_prime_parsed
          parsed_data = {**five_prime_data, **three_prime_data}

          # If all elements parsed correctly (success case)
          if five_prime_passed and three_prime_passed:
            num_passing += 1
            parsed_data["existing_header"] = header
            parsed_data["existing_read_name"] = header.split(' ', 1)[0]
            new_header = self.barcode_config.header_template.format(**parsed_data).replace('\\t', '\t')
            pass_f.write(f"@{new_header}\n{trimmed_seq}\n+\n{trimmed_qual}\n")
            self.write_to_log(num_processed, header, "PASS", orientation,
                              parsed_data)
          else:
            # If orientation worked, but barcode parsing failed (failure case 2)
            num_failing += 1
            num_parse_failing += 1
            fail_reason = (five_prime_fail_reason if not five_prime_passed else three_prime_fail_reason)
            new_header = f"{header} fail_reason={fail_reason}"
            fail_f.write(f"@{new_header}\n{read_sequence}\n+\n{read_quality}\n")
            self.write_to_log(num_processed, header, "FAIL", fail_reason, parsed_data)

    print(f"Done! Total time: {time.time() - start_time:.2f} seconds")
    print(f"\nPercentage passing: {100*(num_passing / num_processed):.2f}% ({num_passing}/{num_processed})")
    print(f"Percentage failing: {100*(num_failing / num_processed):.2f}% ({num_failing}/{num_processed})")

    if num_failing > 0:
      print(f"↳ Orientation failures: {100*(num_orientation_failing / num_failing):.2f}% ({num_orientation_failing}/{num_failing})")
      print(f"↳ Parsing failures: {100*(num_parse_failing / num_failing):.2f}% ({num_parse_failing}/{num_failing})")
    else:
      print("No failing reads.")

    # Write final summary to log file
    self.log_summary(num_processed, num_passing, num_failing, num_orientation_failing, num_parse_failing)

def parse_arguments():
  parser = argparse.ArgumentParser(description='Barcode parsing script for FASTQ files.')
  parser.add_argument('-i', '--input', required=True,
                      help='Comma-separated list of input FASTQ files (can be gzipped).')
  parser.add_argument('-o', '--output_prefix', required=True,
                      help='Output FASTQ prefix (fastq outputs will be gzipped).')
  parser.add_argument('-c', '--config', required=True,
                      help='YAML configuration file.')
  parser.add_argument('--header_template', required=False, default=None,
                      help='Override header template from YAML configuration.')
  parser.add_argument('--trim_polyA', action="store_true", default=False,
                      help='Enable trimming of polyA/T tails from the transcript sequence.')
  parser.add_argument('--test_run', action="store_true", default=False,
                      help='Boolean flag to process and print 50 reads with colorful outputs')
  parser.add_argument('--verbose', action="store_true", default=False,
                      help='Turn on verbose logging for interactive runs (default: False).')
  return parser.parse_args()

def main():
  """Main function to execute the barcode parsing pipeline."""
  args = parse_arguments()
  input_files = [f.strip() for f in args.input.split(',')]
  long_reader = LongReader(args.config, input_files, args.output_prefix,
                           args.header_template, trim_polyA=args.trim_polyA)
  if args.test_run:
    long_reader.test_run(num_reads=50, verbose=True)
  else:
    long_reader.process_all_reads(verbose=args.verbose)

if __name__ == '__main__':
  main()