"""
Snakemake script for guide extraction and filtering.

This script is called by the extract_and_filter_guides rule.
"""

import os
import logging
import pandas as pd

from crisprdesigner2 import (
    extract_guides_from_bed,
    apply_filters,
    get_passing_guides,
    load_predesigned_guides,
    filter_predesigned_regions,
)

# Set up logging
logging.basicConfig(
    filename=snakemake.log[0],
    level=logging.INFO,
    format='%(asctime)s - %(levelname)s - %(message)s'
)
logger = logging.getLogger(__name__)


def read_regions_bed(bed_path):
    """Read BED file with regions."""
    columns = ['chr', 'start', 'end', 'name']
    try:
        df = pd.read_csv(bed_path, sep='\t', header=None, usecols=[0, 1, 2, 3])
        df.columns = columns[:len(df.columns)]
    except Exception:
        # Try reading with only 3 columns
        df = pd.read_csv(bed_path, sep='\t', header=None, usecols=[0, 1, 2])
        df.columns = ['chr', 'start', 'end']
        df['name'] = df['chr'] + ':' + df['start'].astype(str) + '-' + df['end'].astype(str)
    return df


def main():
    logger.info("Starting guide extraction and filtering")

    # Get inputs from snakemake
    regions_file = snakemake.input.regions
    genome_fasta = snakemake.input.genome_fasta
    output_file = snakemake.output.guides
    predesigned_path = snakemake.params.predesigned

    logger.info(f"Input regions: {regions_file}")
    logger.info(f"Genome FASTA: {genome_fasta}")
    logger.info(f"Pre-designed guides: {predesigned_path or 'None'}")

    # Load pre-designed guides if available
    predesigned_df = load_predesigned_guides(predesigned_path)
    logger.info(f"Loaded {len(predesigned_df)} pre-designed guides")

    # Read regions
    regions_df = read_regions_bed(regions_file)
    logger.info(f"Read {len(regions_df)} regions from input BED file")

    # Filter out regions that already have pre-designed guides
    if not predesigned_df.empty:
        regions_df = filter_predesigned_regions(regions_df, predesigned_df)
        logger.info(f"After filtering pre-designed regions: {len(regions_df)} regions remain")

    if regions_df.empty:
        logger.info("All regions already have pre-designed guides. Creating empty output.")
        # Create empty DataFrame with expected columns
        columns = ['chr', 'start', 'end', 'locus', 'strand',
                   'GuideSequenceWithPAM', 'guideSet', 'protospacer', 'filter_reason']
        empty_df = pd.DataFrame(columns=columns)
        empty_df.to_csv(output_file, sep='\t', index=False)
        return

    # Extract guides from all regions
    logger.info("Extracting guides from regions...")
    all_guides = extract_guides_from_bed(regions_file, genome_fasta)
    logger.info(f"Extracted {len(all_guides)} candidate guides")

    # Convert to DataFrame
    if all_guides:
        guides_df = pd.DataFrame(all_guides)
    else:
        guides_df = pd.DataFrame(columns=['chr', 'start', 'end', 'locus', 'strand',
                                          'GuideSequenceWithPAM', 'guideSet', 'protospacer'])

    # Apply sequence filters
    logger.info("Applying sequence filters...")
    filtered_df = apply_filters(guides_df)

    # Get only passing guides
    passing_df = get_passing_guides(guides_df)
    logger.info(f"Guides passing all filters: {len(passing_df)}")

    # Keep all guides with filter_reason for downstream analysis
    # but mark which ones passed
    filtered_df['passed_filters'] = filtered_df['filter_reason'] == ''

    # Create output directory if needed
    os.makedirs(os.path.dirname(output_file), exist_ok=True)

    # Write output
    filtered_df.to_csv(output_file, sep='\t', index=False)
    logger.info(f"Wrote {len(filtered_df)} guides to {output_file}")

    logger.info("Guide extraction and filtering complete")


if __name__ == '__main__':
    main()
