"""Functions for reading and processing FASTQ files"""

# AUTOGENERATED! DO NOT EDIT! File to edit: ../nbs/01_fastq.ipynb.

# %% auto 0
__all__ = ['FastqHandler', 'OligoCodonProcessor']

# %% ../nbs/01_fastq.ipynb 4
import os
import gzip
from collections import Counter, defaultdict
from typing import List, Dict, Tuple, Iterator, Optional
from pathlib import Path
from Bio.SeqIO.QualityIO import FastqGeneralIterator
from tqdm.auto import tqdm
from .core import OligoRegex, OligoCounter

# %% ../nbs/01_fastq.ipynb 6
class FastqHandler:
    """Handles FASTQ file operations."""
    
    @staticmethod
    def find_fastq_pairs(directory: str) -> Tuple[str, str]:
        """Find paired FASTQ files in a directory.
        
        Args:
            directory: Directory to search for FASTQ files
            
        Returns:
            Tuple of (read1_path, read2_path)
            
        Raises:
            ValueError: If paired FASTQ files cannot be found
        """
        files = os.listdir(directory)
        # Common naming patterns for paired FASTQ files
        r1_patterns = ('_1.fq.gz', '_R1.fastq.gz', '_R1_001.fastq.gz')
        r2_patterns = ('_2.fq.gz', '_R2.fastq.gz', '_R2_001.fastq.gz')
        
        # Find files matching the patterns
        r1_file = None
        r2_file = None
        
        for f in files:
            if any(f.endswith(pattern) for pattern in r1_patterns):
                r1_file = f
            elif any(f.endswith(pattern) for pattern in r2_patterns):
                r2_file = f
        
        if not (r1_file and r2_file):
            raise ValueError(f"Could not find paired FASTQ files in {directory}")
            
        return os.path.join(directory, r1_file), os.path.join(directory, r2_file)
    
    @staticmethod
    def open_fastq(filepath: str) -> gzip.GzipFile:
        """Open a FASTQ file, handling gzipped and non-gzipped files.
        
        Args:
            filepath: Path to the FASTQ file
            
        Returns:
            An opened file handle
        """
        if filepath.endswith(('.gz', '.gzip')):
            return gzip.open(filepath, "rt")
        else:
            return open(filepath, "rt")
    
    @staticmethod
    def fastq_pair_iterator(r1_path: str, r2_path: str) -> Iterator[Tuple[Tuple, Tuple]]:
        """Create an iterator over paired FASTQ reads.
        
        Args:
            r1_path: Path to the first read file
            r2_path: Path to the second read file
            
        Returns:
            Iterator yielding pairs of (read1, read2) tuples where each read is 
            a tuple of (title, sequence, quality)
        """
        with FastqHandler.open_fastq(r1_path) as f1, FastqHandler.open_fastq(r2_path) as f2:
            for read1, read2 in zip(FastqGeneralIterator(f1), FastqGeneralIterator(f2)):
                yield read1, read2

# %% ../nbs/01_fastq.ipynb 8
class OligoCodonProcessor:
    """Processes FASTQ files to count oligo codons."""
    
    def __init__(self, oligos: List[str]):
        """Initialize the processor with a list of oligos.
        
        Args:
            oligos: List of oligo sequences to search for
        """
        self.oligos = oligos
        self.compiled_oligos = [OligoRegex(oligo) for oligo in oligos]
    
    def process_fastq_pair(self, r1_path: str, r2_path: str,) -> OligoCounter:
        """Process a pair of FASTQ files to count oligo codons.
        
        Args:
            r1_path: Path to the first read file
            r2_path: Path to the second read file
            progress_callback: Optional callback function for progress reporting
            
        Returns:
            Dictionary mapping oligo indices to codon counters
        """
        results = defaultdict(Counter)
        
        processed_reads = 0
        for read_1, read_2 in tqdm(FastqHandler.fastq_pair_iterator(r1_path, r2_path), ):
            for i, compiled in enumerate(self.compiled_oligos):
                codon = compiled.find_codon(read_1[1], read_2[1])
                results[i][codon] += 1
  
        return results
