"""
Sequence filter module for CRISPR gRNA design.

Applies quality filters to candidate gRNA sequences to remove poor-quality guides.

Filters guides with:
* More than 1 T/U in the last 4 nucleotides (Pol III termination)
* More than 40% total T/U content
* 4-mer T/U repeats
* 5-mer mononucleotide repeats
* Low-complexity sequences (defined as 10 nucleotides of 2-nt repeat,
  12 nucleotides of 3- or 4-nt repeats, or 15/18 nucleotides of 5- or 6-nt repeats)
* <= 20% GC content
* >= 90% GC content
"""

import re
import pandas as pd


def _is_low_complexity(sequence, repeat_length, min_repeats):
    """
    Check if a sequence contains low-complexity regions using regex backreferences.

    Args:
        sequence (str): The DNA sequence to check.
        repeat_length (int): The length of the repeat unit.
        min_repeats (int): The minimum number of repeats to consider as low complexity.

    Returns:
        bool: True if the sequence contains a low-complexity region, False otherwise.
    """
    pattern = f'(.{{{repeat_length}}})\\1{{{min_repeats-1},}}'
    return bool(re.search(pattern, sequence))


def _pol3_termination(protospacer):
    """
    Check for Pol III termination signal: >1 T/U in the last 4 nucleotides.

    Returns True if >1 T in last 4 nt (i.e., fails the filter).
    """
    seq = protospacer.upper().replace('U', 'T')
    return seq[-4:].count('T') > 1


def _high_tu_content(protospacer):
    """
    Check for high T/U content: >40% T/U.

    Returns True if >40% T (i.e., fails the filter).
    """
    seq = protospacer.upper().replace('U', 'T')
    return seq.count('T') / len(seq) > 0.4


def _tu_homopolymer(protospacer):
    """
    Check for T/U homopolymer: 4 or more consecutive T's.

    Returns True if 4+ consecutive T's found (i.e., fails the filter).
    """
    seq = protospacer.upper().replace('U', 'T')
    return 'TTTT' in seq


def _mononucleotide_repeat(protospacer):
    """
    Check for mononucleotide repeat: 5 or more consecutive identical bases.

    Returns True if 5+ consecutive identical bases found (i.e., fails the filter).
    """
    seq = protospacer.upper()
    return any(base * 5 in seq for base in 'ATCG')


def _low_complexity(protospacer):
    """
    Check for low-complexity sequences using regex backreferences.

    Low-complexity rules:
    - 10 nucleotides of 2-nt repeat (5 repeats)
    - 12 nucleotides of 3-nt repeat (4 repeats)
    - 12 nucleotides of 4-nt repeat (3 repeats)
    - 15 nucleotides of 5-nt repeat (3 repeats)
    - 18 nucleotides of 6-nt repeat (3 repeats)

    Returns True if low-complexity pattern found (i.e., fails the filter).
    """
    seq = protospacer.upper()
    return (
        _is_low_complexity(seq, 2, 5) or   # 10 nucleotides of 2-nt repeat
        _is_low_complexity(seq, 3, 4) or   # 12 nucleotides of 3-nt repeat
        _is_low_complexity(seq, 4, 3) or   # 12 nucleotides of 4-nt repeat
        _is_low_complexity(seq, 5, 3) or   # 15 nucleotides of 5-nt repeat
        _is_low_complexity(seq, 6, 3)      # 18 nucleotides of 6-nt repeat
    )


def _low_gc(protospacer):
    """
    Check for low GC content: <=20% GC.

    Returns True if GC content is <=20% (i.e., fails the filter).
    """
    seq = protospacer.upper()
    gc_count = seq.count('G') + seq.count('C')
    gc_fraction = gc_count / len(seq)
    return gc_fraction <= 0.20


def _high_gc(protospacer):
    """
    Check for high GC content: >=90% GC.

    Returns True if GC content is >=90% (i.e., fails the filter).
    """
    seq = protospacer.upper()
    gc_count = seq.count('G') + seq.count('C')
    gc_fraction = gc_count / len(seq)
    return gc_fraction >= 0.90


def is_valid_guide(sequence):
    """
    Check if a guide sequence passes all filters.

    This is the main filter function that matches the original implementation.

    Args:
        sequence: The 20nt protospacer sequence

    Returns:
        bool: True if the guide passes all filters, False otherwise
    """
    # Convert to uppercase and replace U with T for uniformity
    seq = sequence.upper().replace('U', 'T')

    # Check for more than 1 T in the last 4 nucleotides
    if seq[-4:].count('T') > 1:
        return False

    # Check for more than 40% total T content
    if seq.count('T') / len(seq) > 0.4:
        return False

    # Check for 4-mer T repeats
    if 'TTTT' in seq:
        return False

    # Check for 5-mer mononucleotide repeats
    if any(base * 5 in seq for base in 'ATCG'):
        return False

    # Check for low-complexity sequences
    if (_is_low_complexity(seq, 2, 5) or   # 10 nucleotides of 2-nt repeat (5 repeats)
        _is_low_complexity(seq, 3, 4) or   # 12 nucleotides of 3-nt repeat (4 repeats)
        _is_low_complexity(seq, 4, 3) or   # 12 nucleotides of 4-nt repeat (3 repeats)
        _is_low_complexity(seq, 5, 3) or   # 15 nucleotides of 5-nt repeat (3 repeats)
        _is_low_complexity(seq, 6, 3)):    # 18 nucleotides of 6-nt repeat (3 repeats)
        return False

    # Check GC content
    gc_content = (seq.count('G') + seq.count('C')) / len(seq)
    if gc_content <= 0.2 or gc_content >= 0.9:
        return False

    # If all checks pass, the guide is valid
    return True


# Mapping of filter names to functions
FILTERS = {
    'pol3_termination': _pol3_termination,
    'high_tu_content': _high_tu_content,
    'tu_homopolymer': _tu_homopolymer,
    'mononucleotide_repeat': _mononucleotide_repeat,
    'low_complexity': _low_complexity,
    'low_gc': _low_gc,
    'high_gc': _high_gc
}


def apply_filters(guides_df):
    """
    Given a DataFrame of candidate guides (must have a 'protospacer' column),
    apply all sequence filters and return a filtered DataFrame with a new
    'filter_reason' column (empty string if passed, else the name of the
    failing filter).

    Args:
        guides_df: DataFrame with at least a 'protospacer' column

    Returns:
        DataFrame with added 'filter_reason' column. Guides that pass all
        filters have an empty string; guides that fail have the name of
        the first failing filter.
    """
    df = guides_df.copy()

    # Initialize filter_reason column
    df['filter_reason'] = ''

    # Apply each filter
    for filter_name, filter_func in FILTERS.items():
        # Create mask of sequences that fail this filter and haven't already failed
        mask = df.apply(
            lambda row: filter_func(row['protospacer']) if row['filter_reason'] == '' else False,
            axis=1
        )
        df.loc[mask, 'filter_reason'] = filter_name

    return df


def get_passing_guides(guides_df):
    """
    Apply filters and return only guides that pass all filters.

    Args:
        guides_df: DataFrame with at least a 'protospacer' column

    Returns:
        DataFrame containing only guides that pass all filters
    """
    filtered_df = apply_filters(guides_df)
    return filtered_df[filtered_df['filter_reason'] == ''].drop(columns=['filter_reason'])
