"""
RS3 on-target scoring module for CRISPR gRNA design.

Uses the RS3 (Rule Set 3) model from the gpp-rnd/rs3 package to predict
on-target efficiency of gRNA sequences.

RS3 expects 30-nucleotide context sequences:
    4bp upstream + 20bp protospacer + 3bp PAM + 3bp downstream
"""

import numpy as np
import pandas as pd


def score_guides_rs3(guides_df, tracr='Hsu2013', n_jobs=1):
    """
    Given a DataFrame of filtered guides, compute RS3 on-target scores
    using the official rs3 library.

    Args:
        guides_df: DataFrame with a 'context_sequence' column (30nt sequences)
        tracr: tracrRNA variant to use ('Hsu2013' or 'Chen2013')
               Use 'Chen2013' if the tracrRNA doesn't have T in the 5th position
        n_jobs: Number of parallel jobs for feature calculation

    Returns:
        DataFrame with added 'RS3_score' column (float)

    Notes:
        - RS3 scores are typically in a range around -2 to 2 (log-odds style)
        - Higher scores indicate higher predicted on-target efficiency
        - Scores are NOT bounded to [0, 1] - they are raw model predictions
    """
    from rs3.seq import predict_seq

    df = guides_df.copy()

    if 'context_sequence' not in df.columns:
        raise ValueError(
            "DataFrame must contain 'context_sequence' column with 30nt sequences. "
            "Ensure guide extraction includes context extraction."
        )

    # Validate context sequences
    context_seqs = df['context_sequence'].tolist()
    for i, seq in enumerate(context_seqs):
        if len(seq) != 30:
            raise ValueError(
                f"Context sequence at index {i} is {len(seq)}nt, expected 30nt: {seq}"
            )

    # Get RS3 predictions using the official library
    # predict_seq returns numpy array of scores
    scores = predict_seq(context_seqs, sequence_tracr=tracr, n_jobs=n_jobs)

    df['RS3_score'] = scores

    return df


def score_single_guide_rs3(context_sequence, tracr='Hsu2013'):
    """
    Score a single guide sequence using RS3.

    Args:
        context_sequence: A 30-nucleotide context string
                          (4bp upstream + 20bp protospacer + 3bp PAM + 3bp downstream)
        tracr: tracrRNA variant to use ('Hsu2013' or 'Chen2013')

    Returns:
        Float RS3 score (higher = better predicted efficiency)
    """
    from rs3.seq import predict_seq

    if len(context_sequence) != 30:
        raise ValueError(
            f"Context sequence must be exactly 30 nucleotides, got {len(context_sequence)}"
        )

    scores = predict_seq([context_sequence], sequence_tracr=tracr)
    return float(scores[0])


# Legacy functions for backward compatibility with tests
def encode_guide_rs3(protospacer):
    """
    Legacy function - encode a 20nt protospacer using one-hot encoding.

    This is kept for backward compatibility with existing tests.
    The actual RS3 scoring now uses the official rs3 library with 30nt contexts.

    Args:
        protospacer: A 20-nucleotide string (ACGT only)

    Returns:
        numpy array of shape (80,) containing the one-hot encoded features
    """
    NUCLEOTIDE_MAP = {'A': 0, 'C': 1, 'G': 2, 'T': 3}
    seq = protospacer.upper()

    if len(seq) != 20:
        raise ValueError(f"Protospacer must be exactly 20 nucleotides, got {len(seq)}")

    features = np.zeros(80, dtype=np.float32)

    for i, nucleotide in enumerate(seq):
        if nucleotide not in NUCLEOTIDE_MAP:
            raise ValueError(f"Invalid nucleotide '{nucleotide}' at position {i}")
        idx = i * 4 + NUCLEOTIDE_MAP[nucleotide]
        features[idx] = 1.0

    return features


def load_rs3_model(model_path):
    """
    Legacy function - kept for backward compatibility.

    The official rs3 library loads its model internally, so this function
    is no longer needed for normal operation.
    """
    import os
    import pickle

    if not os.path.exists(model_path):
        raise FileNotFoundError(f"RS3 model file not found: {model_path}")

    with open(model_path, 'rb') as f:
        model = pickle.load(f)

    return model
