"""
GuideScan2 off-target scoring module for CRISPR gRNA design.

Uses GuideScan2 to calculate specificity scores for gRNA sequences
based on off-target analysis.

The GuideScan2 database is stored as a BAM file where each "read" represents
a gRNA with its off-target information. Key BAM tags:
    - cs: Specificity score (0-1, higher = more specific)
    - ds: On-target cutting efficiency score
    - k0-k4: Off-target counts at 0-4 mismatches
    - of: Encoded off-target position data
"""

import os
import logging
import subprocess
import tempfile
import numpy as np
import pandas as pd

try:
    import pysam
    PYSAM_AVAILABLE = True
except ImportError:
    PYSAM_AVAILABLE = False


logger = logging.getLogger(__name__)

SUPPORTED_GENOMES = {'hg38', 'mm10'}

# Mapping from UCSC chromosome names to RefSeq accessions for hg38 primary chromosomes.
# GuideScan2 databases built from NCBI/RefSeq references use these names instead of
# the UCSC "chr1", "chrX" convention.
_UCSC_TO_REFSEQ_HG38 = {
    'chr1': 'NC_000001.11', 'chr2': 'NC_000002.12', 'chr3': 'NC_000003.12',
    'chr4': 'NC_000004.12', 'chr5': 'NC_000005.10', 'chr6': 'NC_000006.12',
    'chr7': 'NC_000007.14', 'chr8': 'NC_000008.11', 'chr9': 'NC_000009.12',
    'chr10': 'NC_000010.11', 'chr11': 'NC_000011.10', 'chr12': 'NC_000012.12',
    'chr13': 'NC_000013.11', 'chr14': 'NC_000014.9',  'chr15': 'NC_000015.10',
    'chr16': 'NC_000016.10', 'chr17': 'NC_000017.11', 'chr18': 'NC_000018.10',
    'chr19': 'NC_000019.10', 'chr20': 'NC_000020.11', 'chr21': 'NC_000021.9',
    'chr22': 'NC_000022.11', 'chrX':  'NC_000023.11', 'chrY':  'NC_000024.10',
    'chrM':  'NC_012920.1',
}
_REFSEQ_TO_UCSC_HG38 = {v: k for k, v in _UCSC_TO_REFSEQ_HG38.items()}


def _build_chrom_translator(bam_references):
    """
    Detect BAM chromosome naming convention and return translation dicts.

    Returns (ucsc_to_bam, bam_to_ucsc). Both are empty dicts when the BAM
    already uses UCSC names (identity mapping).
    """
    bam_ref_set = set(bam_references)
    if 'chr1' in bam_ref_set or 'chr2' in bam_ref_set:
        # BAM uses UCSC names — no translation needed
        return {}, {}
    elif 'NC_000001.11' in bam_ref_set:
        logger.info("BAM uses RefSeq chromosome names — translating from UCSC")
        return _UCSC_TO_REFSEQ_HG38, _REFSEQ_TO_UCSC_HG38
    else:
        logger.warning("Could not detect BAM chromosome naming convention; attempting direct lookup")
        return {}, {}


def validate_inputs(index_path, genome):
    """
    Validate input parameters for GuideScan2 scoring.

    Args:
        index_path: Path to the GuideScan2 index
        genome: Genome identifier (hg38 or mm10)

    Raises:
        FileNotFoundError: If index_path does not exist
        ValueError: If genome is not supported
    """
    if not os.path.exists(index_path):
        raise FileNotFoundError(
            f"GuideScan2 index file not found: {index_path}\n"
            "Please download the GuideScan2 index for your genome "
            "and update the config with the correct path."
        )

    if genome not in SUPPORTED_GENOMES:
        raise ValueError(
            f"Unsupported genome: {genome}. "
            f"Supported genomes are: {', '.join(sorted(SUPPORTED_GENOMES))}"
        )


def _query_guidescan_bam(guides_df, index_path):
    """
    Query the GuideScan2 BAM database directly using pysam.

    The GuideScan2 BAM file contains pre-computed gRNAs indexed by genomic position.
    Each BAM record has a 'cs' tag containing the specificity score.

    The BAM uses 0-based coordinates (reference_start), while query_name uses 1-based.
    We match by reference_start and strand.

    If a BAM index (.bai) exists, uses efficient random access.
    Otherwise, falls back to full file scan with early termination.

    Args:
        guides_df: DataFrame with guide information including chr, start, end, strand
        index_path: Path to GuideScan2 BAM database

    Returns:
        DataFrame with specificity_score column, or None if pysam not available
    """
    if not PYSAM_AVAILABLE:
        logger.warning("pysam not available for BAM database access")
        return None

    df = guides_df.copy()

    # Check for BAM index
    index_file = index_path + '.bai'
    has_index = os.path.exists(index_file)

    try:
        # Open BAM file (ignore truncation warning for large files)
        bam = pysam.AlignmentFile(index_path, "rb", ignore_truncation=True)

        # Detect chromosome naming convention and build translators
        ucsc_to_bam, bam_to_ucsc = _build_chrom_translator(bam.references)

        # Build lookup dictionary: (ucsc_chr, start, strand) -> specificity_score
        score_lookup = {}

        # Get min/max positions per chromosome (in UCSC names)
        chr_ranges = df.groupby('chr')['start'].agg(['min', 'max']).to_dict('index')

        if has_index:
            # Use indexed access for efficient queries
            logger.info(f"Using indexed BAM access for {len(df)} guides...")

            for chrom, pos_range in chr_ranges.items():
                bam_chrom = ucsc_to_bam.get(chrom, chrom)
                try:
                    # Fetch reads in the region of interest
                    for read in bam.fetch(bam_chrom, pos_range['min'], pos_range['max'] + 1):
                        ref_start = read.reference_start
                        strand = '-' if read.is_reverse else '+'
                        # Store under UCSC name for consistent key lookup below
                        ucsc_chrom = bam_to_ucsc.get(read.reference_name, read.reference_name)

                        if read.has_tag('cs'):
                            score_lookup[(ucsc_chrom, ref_start, strand)] = read.get_tag('cs')
                except ValueError:
                    logger.debug(f"Chromosome {bam_chrom} not found in BAM index")
        else:
            # Full scan with early termination
            logger.info(f"No BAM index found. Scanning database for {len(df)} guides...")
            logger.info("Consider creating an index with: samtools index <bam_file>")

            # Track which chromosomes/positions we've seen
            found_count = 0
            total_needed = len(df)
            max_pos_needed = max(r['max'] for r in chr_ranges.values())
            chrom_order = list(bam.references)

            for read in bam.fetch(until_eof=True):
                chrom = read.reference_name
                ref_start = read.reference_start

                # Early termination: if we're past all chromosomes/positions we need
                if chrom not in chr_ranges:
                    # Check if we've passed all needed chromosomes
                    if chrom in chrom_order:
                        chrom_idx = chrom_order.index(chrom)
                        needed_chroms = set(chr_ranges.keys())
                        if all(chrom_order.index(c) < chrom_idx for c in needed_chroms if c in chrom_order):
                            logger.info("Passed all needed regions, stopping scan early")
                            break
                    continue

                pos_range = chr_ranges[chrom]

                # Skip if before our range
                if ref_start < pos_range['min'] - 10:
                    continue

                # Stop scanning this chromosome if we're past our range
                if ref_start > pos_range['max'] + 10:
                    # Remove this chromosome from our needs
                    del chr_ranges[chrom]
                    if not chr_ranges:
                        logger.info("Found all needed regions, stopping scan early")
                        break
                    continue

                # Store the specificity score under UCSC chromosome name
                strand = '-' if read.is_reverse else '+'
                if read.has_tag('cs'):
                    ucsc_chrom = bam_to_ucsc.get(chrom, chrom)
                    key = (ucsc_chrom, ref_start, strand)
                    if key not in score_lookup:
                        score_lookup[key] = read.get_tag('cs')
                        found_count += 1

        bam.close()

        # Map scores to guides
        # Coordinate offsets between our guides and GuideScan2 BAM:
        # - For + strand: GuideScan2 position is 1 less than our protospacer start
        # - For - strand: GuideScan2 stores PAM (CCN) position, which is 3bp before
        #   our protospacer start, plus the 1bp offset = 4bp total
        scores = []
        for _, row in df.iterrows():
            if row['strand'] == '+':
                lookup_pos = row['start'] - 1
            else:
                lookup_pos = row['start'] - 4  # -1 for base offset, -3 for PAM
            key = (row['chr'], lookup_pos, row['strand'])
            score = score_lookup.get(key, np.nan)
            scores.append(score)

        df['specificity_score'] = scores
        scored_count = sum(1 for s in scores if not np.isnan(s))
        logger.info(f"Retrieved specificity scores for {scored_count}/{len(df)} guides")

        return df

    except Exception as e:
        logger.warning(f"Error querying GuideScan2 BAM database: {e}")
        import traceback
        logger.warning(traceback.format_exc())
        return None


def _try_guidescan2_python_api(guides_df, index_path, genome):
    """
    Try to use the GuideScan2 Python API.

    Returns:
        DataFrame with specificity_score column, or None if API not available
    """
    try:
        import guidescan
        # Use the GuideScan2 Python API
        # Note: The exact API may vary depending on the version installed
        db = guidescan.GuideScanDB(index_path)

        scores = []
        for _, row in guides_df.iterrows():
            try:
                result = db.query(
                    row['protospacer'],
                    chrom=row['chr'],
                    start=row['start'],
                    end=row['end'],
                    strand=row['strand']
                )
                if result is not None and 'specificity' in result:
                    scores.append(result['specificity'])
                else:
                    logger.warning(
                        f"No specificity result for guide at {row['chr']}:{row['start']}-{row['end']}"
                    )
                    scores.append(np.nan)
            except Exception as e:
                logger.warning(f"GuideScan2 query failed for {row['protospacer']}: {e}")
                scores.append(np.nan)

        df = guides_df.copy()
        df['specificity_score'] = scores
        return df

    except ImportError:
        return None


def _run_guidescan2_cli(guides_df, index_path, genome):
    """
    Run GuideScan2 using the command-line interface.

    Args:
        guides_df: DataFrame with guide information
        index_path: Path to GuideScan2 index
        genome: Genome identifier

    Returns:
        DataFrame with specificity_score column
    """
    df = guides_df.copy()

    # Create temporary input file with guide sequences
    with tempfile.NamedTemporaryFile(mode='w', suffix='.txt', delete=False) as f:
        for _, row in df.iterrows():
            # GuideScan2 expects: sequence<tab>chromosome:start-end:strand
            f.write(f"{row['protospacer']}\t{row['chr']}:{row['start']}-{row['end']}:{row['strand']}\n")
        input_file = f.name

    try:
        # Create temporary output file
        with tempfile.NamedTemporaryFile(mode='w', suffix='.txt', delete=False) as f:
            output_file = f.name

        # Run GuideScan2 CLI
        # Note: Exact command may vary depending on installation
        cmd = [
            'guidescan2',
            'query',
            '-i', index_path,
            '-q', input_file,
            '-o', output_file,
            '--genome', genome
        ]

        try:
            subprocess.run(cmd, check=True, capture_output=True, text=True)

            # Parse output file
            scores = _parse_guidescan2_output(output_file, len(df))
            df['specificity_score'] = scores

        except FileNotFoundError:
            logger.warning(
                "GuideScan2 CLI not found. Returning NaN scores. "
                "Install GuideScan2 or ensure it's in your PATH."
            )
            df['specificity_score'] = np.nan

        except subprocess.CalledProcessError as e:
            logger.warning(f"GuideScan2 CLI failed: {e.stderr}")
            df['specificity_score'] = np.nan

    finally:
        # Clean up temporary files
        if os.path.exists(input_file):
            os.unlink(input_file)
        if os.path.exists(output_file):
            os.unlink(output_file)

    return df


def _parse_guidescan2_output(output_file, expected_count):
    """
    Parse GuideScan2 output file to extract specificity scores.

    Args:
        output_file: Path to GuideScan2 output file
        expected_count: Expected number of guides

    Returns:
        List of specificity scores
    """
    scores = [np.nan] * expected_count

    if not os.path.exists(output_file):
        return scores

    try:
        with open(output_file, 'r') as f:
            for i, line in enumerate(f):
                if i >= expected_count:
                    break
                parts = line.strip().split('\t')
                # Assuming specificity score is in a specific column
                # This may need adjustment based on actual output format
                if len(parts) >= 2:
                    try:
                        scores[i] = float(parts[-1])
                    except ValueError:
                        scores[i] = np.nan
    except Exception as e:
        logger.warning(f"Error parsing GuideScan2 output: {e}")

    return scores


def _find_bwt_index_prefix(bwt_index_dir):
    """
    Find the BWT index prefix in the extracted index directory.

    The index files have the pattern: <prefix>.index.forward, <prefix>.index.reverse, <prefix>.index.gs

    Args:
        bwt_index_dir: Path to directory containing extracted index files

    Returns:
        Path to index prefix (without .index.forward/.index.reverse/.index.gs extension)
    """
    import glob

    # Look for .index.gs file (smallest, so good for checking)
    gs_files = glob.glob(os.path.join(bwt_index_dir, "*.index.gs"))
    if gs_files:
        # Remove .index.gs suffix to get prefix
        return gs_files[0].replace(".index.gs", ".index")

    # Fallback: look for .index.forward
    forward_files = glob.glob(os.path.join(bwt_index_dir, "*.index.forward"))
    if forward_files:
        return forward_files[0].replace(".forward", "")

    raise FileNotFoundError(f"No GuideScan2 BWT index files found in {bwt_index_dir}")


def _compute_cfd_specificity(off_targets):
    """
    Compute CFD-based specificity score from off-target list.

    Specificity = 1 / (1 + sum(CFD scores for all off-targets))

    For simplicity, we use a position-independent mismatch penalty:
    - 0 mismatches: CFD = 1.0
    - 1 mismatch: CFD = 0.5
    - 2 mismatches: CFD = 0.1
    - 3 mismatches: CFD = 0.02
    - 4 mismatches: CFD = 0.005

    Args:
        off_targets: List of (mismatch_count, occurrence_count) tuples

    Returns:
        Specificity score (0-1)
    """
    cfd_weights = {0: 1.0, 1: 0.5, 2: 0.1, 3: 0.02, 4: 0.005}

    total_cfd = 0.0
    for mismatches, count in off_targets:
        if mismatches in cfd_weights:
            total_cfd += cfd_weights[mismatches] * count

    # Subtract 1 for the on-target site (0 mismatches, count=1)
    if total_cfd >= 1.0:
        total_cfd -= 1.0

    specificity = 1.0 / (1.0 + total_cfd)
    return specificity


def _parse_enumerate_output(output_file, guide_ids):
    """
    Parse GuideScan2 enumerate output and compute specificity scores.

    The enumerate command outputs off-target counts per mismatch level.

    Args:
        output_file: Path to enumerate output file
        guide_ids: List of guide IDs in order they were submitted

    Returns:
        Dict mapping guide_id to specificity score
    """
    scores = {}

    if not os.path.exists(output_file):
        logger.warning(f"Enumerate output file not found: {output_file}")
        return scores

    try:
        # Parse the output file
        # Format depends on GuideScan2 version, but typically:
        # guide_id<tab>sequence<tab>k0<tab>k1<tab>k2<tab>k3<tab>k4<tab>...
        with open(output_file, 'r') as f:
            for line in f:
                line = line.strip()
                if not line or line.startswith('#'):
                    continue

                parts = line.split('\t')
                if len(parts) < 3:
                    continue

                guide_id = parts[0]

                # Try to extract off-target counts (k0, k1, k2, k3, k4)
                # The exact column positions may vary
                try:
                    # Assume columns are: id, sequence, k0, k1, k2, k3, k4
                    off_targets = []
                    for i, col_idx in enumerate(range(2, min(7, len(parts)))):
                        count = int(parts[col_idx])
                        off_targets.append((i, count))

                    scores[guide_id] = _compute_cfd_specificity(off_targets)
                except (ValueError, IndexError) as e:
                    logger.debug(f"Could not parse off-target counts for {guide_id}: {e}")
                    continue

    except Exception as e:
        logger.warning(f"Error parsing enumerate output: {e}")

    return scores


def _run_guidescan2_enumerate(guides_df, bwt_index_path, genome, max_mismatches=4):
    """
    Run GuideScan2 CLI enumerate command to compute specificity scores
    for guides not found in the pre-computed BAM database.

    Args:
        guides_df: DataFrame with guides missing specificity scores
        bwt_index_path: Path to GuideScan2 BWT index directory
        genome: Genome identifier
        max_mismatches: Maximum mismatches for off-target search (default 4)

    Returns:
        List of specificity scores in same order as guides_df
    """
    if len(guides_df) == 0:
        return []

    # Find the index prefix
    try:
        index_prefix = _find_bwt_index_prefix(bwt_index_path)
        logger.info(f"Using GuideScan2 BWT index: {index_prefix}")
    except FileNotFoundError as e:
        logger.warning(str(e))
        return [np.nan] * len(guides_df)

    # Create temporary FASTA file with guide sequences
    with tempfile.NamedTemporaryFile(mode='w', suffix='.fa', delete=False) as f:
        guide_ids = []
        for idx, row in guides_df.iterrows():
            guide_id = f"guide_{idx}"
            guide_ids.append(guide_id)
            # Include PAM for enumeration (NGG for SpCas9)
            sequence = row['protospacer'] + 'NGG'
            f.write(f">{guide_id}\n{sequence}\n")
        input_file = f.name

    logger.info(f"Running GuideScan2 enumerate for {len(guides_df)} guides...")

    try:
        # Create temporary output file
        with tempfile.NamedTemporaryFile(mode='w', suffix='.tsv', delete=False) as f:
            output_file = f.name

        # Run guidescan enumerate
        cmd = [
            'guidescan', 'enumerate',
            '-i', index_prefix,
            '-o', output_file,
            '-m', str(max_mismatches),
            '-f', input_file
        ]

        logger.debug(f"Running command: {' '.join(cmd)}")

        try:
            result = subprocess.run(cmd, check=True, capture_output=True, text=True, timeout=600)
            logger.debug(f"GuideScan2 stdout: {result.stdout}")
            if result.stderr:
                logger.debug(f"GuideScan2 stderr: {result.stderr}")

            # Parse output and map back to guides
            score_dict = _parse_enumerate_output(output_file, guide_ids)

            scores = []
            for guide_id in guide_ids:
                scores.append(score_dict.get(guide_id, np.nan))

            scored_count = sum(1 for s in scores if not np.isnan(s))
            logger.info(f"GuideScan2 enumerate computed scores for {scored_count}/{len(guides_df)} guides")

            return scores

        except FileNotFoundError:
            logger.warning(
                "GuideScan2 CLI not found. Install with: conda install -c bioconda guidescan"
            )
            return [np.nan] * len(guides_df)

        except subprocess.CalledProcessError as e:
            logger.warning(f"GuideScan2 enumerate failed: {e.stderr}")
            return [np.nan] * len(guides_df)

        except subprocess.TimeoutExpired:
            logger.warning("GuideScan2 enumerate timed out after 600 seconds")
            return [np.nan] * len(guides_df)

    finally:
        # Clean up temporary files
        if os.path.exists(input_file):
            os.unlink(input_file)
        if 'output_file' in locals() and os.path.exists(output_file):
            os.unlink(output_file)


def score_guides_guidescan(guides_df, index_path, genome, bwt_index_path=None):
    """
    Given a DataFrame of filtered guides, query the GuideScan2 index and
    return the DataFrame with a new 'specificity_score' column.

    The specificity score ranges from 0 to 1, where:
    - 1.0 = unique in the genome (no off-targets)
    - 0.0 = many off-targets
    - NaN = guide could not be scored (e.g., not in index)

    This function first tries to look up scores in the pre-computed BAM database.
    If some guides are not found and bwt_index_path is provided, it runs the
    GuideScan2 CLI enumerate command to compute scores for the missing guides.

    Args:
        guides_df: DataFrame with 'protospacer', 'chr', 'start', 'end', 'strand' columns
        index_path: Path to the GuideScan2 index file (BAM format)
        genome: Genome identifier ('hg38' or 'mm10')
        bwt_index_path: Optional path to GuideScan2 BWT index directory for CLI fallback

    Returns:
        DataFrame with added 'specificity_score' column (float, 0-1)

    Raises:
        FileNotFoundError: If index_path does not exist
        ValueError: If genome is not supported
    """
    # Validate inputs
    validate_inputs(index_path, genome)

    result = None

    # Try reading BAM database directly with pysam (preferred method)
    if index_path.endswith('.bam') or index_path.endswith('.bam.sorted'):
        logger.info("Querying GuideScan2 BAM database directly")
        result = _query_guidescan_bam(guides_df, index_path)

    # If BAM lookup didn't work, try Python API
    if result is None:
        result = _try_guidescan2_python_api(guides_df, index_path, genome)

    # If Python API didn't work, fall back to CLI
    if result is None:
        result = _run_guidescan2_cli(guides_df, index_path, genome)

    # Check for guides without scores and try CLI enumerate if BWT index is available
    if result is not None and bwt_index_path and os.path.exists(bwt_index_path):
        missing_mask = result['specificity_score'].isna()
        missing_count = missing_mask.sum()

        if missing_count > 0:
            logger.info(f"Found {missing_count} guides without specificity scores")
            logger.info("Running GuideScan2 CLI enumerate for missing guides...")

            missing_guides = result[missing_mask].copy()
            cli_scores = _run_guidescan2_enumerate(missing_guides, bwt_index_path, genome)

            # Update the missing scores
            result.loc[missing_mask, 'specificity_score'] = cli_scores

            # Log final coverage
            final_missing = result['specificity_score'].isna().sum()
            total_scored = len(result) - final_missing
            logger.info(f"Final specificity score coverage: {total_scored}/{len(result)} guides")

    return result


def compute_specificity_mock(guides_df):
    """
    Compute mock specificity scores for testing purposes.

    This function generates plausible specificity scores based on
    guide sequence characteristics, without actually running GuideScan2.

    Args:
        guides_df: DataFrame with 'protospacer' column

    Returns:
        DataFrame with 'specificity_score' column
    """
    df = guides_df.copy()

    # Generate mock scores based on sequence properties
    scores = []
    for protospacer in df['protospacer']:
        # Higher GC content generally means more potential off-targets
        gc_content = (protospacer.count('G') + protospacer.count('C')) / len(protospacer)

        # Generate a plausible score
        # Real scores would be calculated by GuideScan2 algorithm
        base_score = 0.7 + 0.3 * (1 - gc_content)  # Lower GC -> higher specificity
        noise = np.random.normal(0, 0.05)
        score = np.clip(base_score + noise, 0.0, 1.0)
        scores.append(score)

    df['specificity_score'] = scores
    return df
