"""Core utilities and classes for OligoSeeker"""

# AUTOGENERATED! DO NOT EDIT! File to edit: ../nbs/00_core.ipynb.

# %% auto 0
__all__ = ['OligoCounter', 'FastqPair', 'COMPLEMENT', 'DNAUtils', 'OligoRegex', 'OligoLoader']

# %% ../nbs/00_core.ipynb 4
from typing import Dict, List, Tuple, Set
from collections import Counter
import re
from Bio.Seq import Seq

# %% ../nbs/00_core.ipynb 6
# Type aliases
OligoCounter = Dict[int, Counter]
FastqPair = Tuple[str, str]

# %% ../nbs/00_core.ipynb 8
# Constants for complement mapping
COMPLEMENT = {
    'A': 'T', 'C': 'G', 'G': 'C', 'T': 'A',
    '[': ']', ']': '[', '(': ')', ')': '(',
    '.': '.', 'N': 'N'
}

# %% ../nbs/00_core.ipynb 9
class DNAUtils:
    """Utility class for DNA sequence operations."""
    
    @staticmethod
    def reverse_complement(sequence: str) -> str:
        """Generate reverse complement of a DNA sequence.
        
        Args:
            sequence: DNA sequence string
            
        Returns:
            Reverse complement of the input sequence
        """
        return ''.join(COMPLEMENT.get(char, char) for char in reversed(sequence))
    
    @staticmethod
    def validate_oligos(oligos: List[str]) -> Tuple[bool, Set[str]]:
        """Validate a list of oligo sequences.
        
        Args:
            oligos: List of oligo sequences to validate
            
        Returns:
            Tuple of (is_valid, invalid_characters)
        """
        invalid_chars = set()
        for oligo in oligos:
            for b in oligo:
                if b not in COMPLEMENT:
                    invalid_chars.add(b)
        return len(invalid_chars) == 0, invalid_chars

# %% ../nbs/00_core.ipynb 11
class OligoRegex:
    """Compiles and manages regex patterns for oligo searching."""
    
    def __init__(self, oligo: str):
        """Initialize regex patterns for forward and reverse oligo sequences.
        
        Args:
            oligo: Oligo sequence with potential NNN codon sites
        """
        self.oligo = oligo
        self.forward = re.compile(oligo.replace('NNN', '(...)'))  # Captures the codon in group 1
        rev_oligo = DNAUtils.reverse_complement(oligo)
        self.reverse = re.compile(rev_oligo.replace('NNN', '(...)'))  # Captures the codon in group 1
    
    
    def find_codon(self, read_1: str, read_2: str) -> str:
        """Find codon in paired reads.
        
        Args:
            read_1: Sequence from the first read
            read_2: Sequence from the second read
            
        Returns:
            The codon found or 'none' if not found
        """
        # Combine reads with a unique separator for pattern matching
        combined_read = f"{read_1}@{read_2}"
        
        # Try matching in forward direction
        match = self.forward.search(combined_read)
        if match:
            return match.group(1)
        
        # Try matching in reverse direction
        match = self.reverse.search(combined_read)
        if match:
            return str(Seq(match.group(1)).reverse_complement())
        
        # No match found
        return 'none'

# %% ../nbs/00_core.ipynb 16
class OligoLoader:
    """Loads and validates oligo sequences from different sources."""
    
    @staticmethod
    def from_file(filepath: str) -> List[str]:
        """Load oligos from a file.
        
        Args:
            filepath: Path to the file containing oligo sequences (one per line)
            
        Returns:
            List of oligo sequences
            
        Raises:
            ValueError: If the file contains duplicate oligos or invalid characters
        """
        with open(filepath) as f:
            oligos = [line.strip() for line in f if line.strip()]
        return OligoLoader.validate_oligos(oligos)
    
    @staticmethod
    def from_string(oligo_string: str, delimiter: str = ',') -> List[str]:
        """Load oligos from a delimited string.
        
        Args:
            oligo_string: String containing oligo sequences
            delimiter: Delimiter character (default: comma)
            
        Returns:
            List of oligo sequences
            
        Raises:
            ValueError: If the string contains duplicate oligos or invalid characters
        """
        oligos = [o.strip() for o in oligo_string.split(delimiter) if o.strip()]
        return OligoLoader.validate_oligos(oligos)
    
    @staticmethod
    def validate_oligos(oligos: List[str]) -> List[str]:
        """Validate a list of oligo sequences.

        Args:
            oligos: List of oligo sequences to validate

        Returns:
            The validated list of oligos

        Raises:
            ValueError: If duplicate oligos found, invalid characters detected, or NNN pattern issues
        """
        # Check for duplicates
        if len(oligos) != len(set(oligos)):
            raise ValueError("Duplicate oligos found")

        # Check for valid characters
        is_valid, invalid_chars = DNAUtils.validate_oligos(oligos)
        if not is_valid:
            invalid_chars_str = ','.join(invalid_chars)
            raise ValueError(f"Invalid characters found in oligos: {invalid_chars_str}")

        # Check for NNN pattern - each oligo must contain exactly one NNN codon
        invalid_oligos = []
        for oligo in oligos:
            nnn_count = oligo.count('NNN')
            if nnn_count != 1:
                raise ValueError(f"{oligo} (contains {nnn_count} NNN patterns, must have exactly 1)")
                

        if invalid_oligos:
            error_msg = "Invalid NNN pattern in oligos:\n" + "\n".join(invalid_oligos)
            raise ValueError(error_msg)

        return oligos
