"""
Snakemake script for off-target scoring with GuideScan2.

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

import os
import logging
import pandas as pd

from crisprdesigner2 import score_guides_guidescan

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


def main():
    logger.info("Starting off-target scoring with GuideScan2")

    # Get inputs from snakemake
    guides_file = snakemake.input.guides
    index_file = snakemake.input.index
    output_file = snakemake.output.scored
    genome = snakemake.params.genome

    logger.info(f"Input guides: {guides_file}")
    logger.info(f"GuideScan2 index: {index_file}")
    logger.info(f"Genome: {genome}")

    # Read guides
    guides_df = pd.read_csv(guides_file, sep='\t')
    logger.info(f"Read {len(guides_df)} guides")

    # Filter to only passing guides for scoring
    if 'passed_filters' in guides_df.columns:
        passing_df = guides_df[guides_df['passed_filters'] == True].copy()
        failed_df = guides_df[guides_df['passed_filters'] == False].copy()
        logger.info(f"Scoring {len(passing_df)} guides that passed filters")
    else:
        passing_df = guides_df.copy()
        failed_df = pd.DataFrame()

    # Score passing guides
    if len(passing_df) > 0:
        try:
            scored_df = score_guides_guidescan(passing_df, index_file, genome)
            logger.info(f"Scored {len(scored_df)} guides")
        except Exception as e:
            logger.warning(f"GuideScan2 scoring failed: {e}")
            logger.warning("Assigning NaN specificity scores")
            scored_df = passing_df.copy()
            scored_df['specificity_score'] = pd.NA
    else:
        scored_df = passing_df
        scored_df['specificity_score'] = []
        logger.info("No guides to score")

    # Add specificity_score = NaN for failed guides
    if len(failed_df) > 0:
        failed_df['specificity_score'] = pd.NA
        scored_df = pd.concat([scored_df, failed_df], ignore_index=True)

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

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

    logger.info("Off-target scoring complete")


if __name__ == '__main__':
    main()
