# geney/tis/prediction.py — Core TIS prediction functions.
from __future__ import annotations

from typing import List, Tuple

import numpy as np
import torch

from .model import _load_titer_models_pytorch, _get_torch_device, get_codon_score
from .encoding import encode_sequences


# ============================================================================
# TIS Prediction (PyTorch)
# ============================================================================

def predict_context_score(seq: str, models=None) -> float:
    """
    Predict the ContextScore (neural network output) for a sequence.

    This is the raw neural network prediction before applying codon prior.

    Args:
        seq: 203nt sequence with codon of interest at positions 100-102.
             Pad with '$' if sequence is near the ends of the transcript.
        models: Optional pre-loaded models list (for efficiency)

    Returns:
        ContextScore (0-1) from the neural network ensemble.
    """
    if models is None:
        models = _load_titer_models_pytorch()

    device = _get_torch_device()
    encoded = encode_sequences([seq])
    x = torch.from_numpy(encoded).float().to(device)

    # Ensemble prediction across all models
    score = 0.0
    with torch.no_grad():
        for model in models:
            pred = model(x)
            score += pred[0, 0].item() / len(models)

    return float(score)


def predict_context_scores_batch(
    sequences: List[str],
    max_chunk: int = 4096,
) -> np.ndarray:
    """
    Predict ContextScores for multiple sequences efficiently.

    Sequences are processed in chunks of *max_chunk* to avoid GPU OOM on
    very large batches (e.g. thousands of mutations × ~70 ATGs each).

    Args:
        sequences: List of 203nt sequences.
        max_chunk: Maximum sequences per GPU forward pass (default 4096).

    Returns:
        Array of ContextScores.
    """
    if not sequences:
        return np.array([], dtype=np.float32)

    models = _load_titer_models_pytorch()
    device = _get_torch_device()

    scores = np.zeros(len(sequences), dtype=np.float32)

    for start in range(0, len(sequences), max_chunk):
        chunk = sequences[start : start + max_chunk]
        encoded = encode_sequences(chunk)
        x = torch.from_numpy(encoded).float().to(device)

        with torch.no_grad():
            for model in models:
                preds = model(x)
                scores[start : start + len(chunk)] += preds[:, 0].cpu().numpy() / len(models)

    return scores


def predict_tis_score(seq: str, model=None, use_prior: bool = True) -> Tuple[float, float, float]:
    """
    Predict the full TISScore for a sequence.

    From the TITER paper:
        TISScore = CodonScore x ContextScore

    Where:
        - ContextScore: Neural network prediction (sequence context features)
        - CodonScore: Prior based on codon composition preference

    Args:
        seq: 203nt sequence with codon of interest at positions 100-102
        model: Optional pre-loaded model
        use_prior: Whether to apply codon prior (default True)

    Returns:
        Tuple of (tis_score, context_score, codon_score)
    """
    context_score = predict_context_score(seq, model)

    # Extract the codon at position 100-102
    codon = seq[100:103].upper().replace('U', 'T')

    if use_prior:
        codon_score = get_codon_score(codon)
        tis_score = context_score * codon_score
    else:
        codon_score = 1.0
        tis_score = context_score

    return tis_score, context_score, codon_score


def predict_tis_probability(seq: str, model=None, use_prior: bool = False) -> float:
    """
    Predict the probability that the codon at position 100-102 is a TIS.

    For backwards compatibility, this returns ContextScore by default.
    Set use_prior=True to get the full TISScore.

    Args:
        seq: 203nt sequence with codon of interest at positions 100-102.
             Pad with '$' if sequence is near the ends of the transcript.
        model: Optional pre-loaded model (for batch efficiency)
        use_prior: Whether to apply codon prior (default False for backwards compat)

    Returns:
        Probability/score that the codon is a translation initiation site.
    """
    if use_prior:
        tis_score, _, _ = predict_tis_score(seq, model, use_prior=True)
        return tis_score
    else:
        return predict_context_score(seq, model)


def predict_tis_probabilities_batch(sequences: List[str], use_prior: bool = False) -> np.ndarray:
    """
    Predict TIS probabilities for multiple sequences efficiently.

    Args:
        sequences: List of 203nt sequences
        use_prior: Whether to apply codon prior

    Returns:
        Array of probabilities/scores
    """
    context_scores = predict_context_scores_batch(sequences)

    if not use_prior:
        return context_scores

    # Apply codon priors
    tis_scores = np.zeros_like(context_scores)
    for i, seq in enumerate(sequences):
        codon = seq[100:103].upper().replace('U', 'T')
        codon_score = get_codon_score(codon)
        tis_scores[i] = context_scores[i] * codon_score

    return tis_scores


def extract_tis_context(mrna: str, position: int, pad_char: str = '$') -> str:
    """
    Extract 203nt context around a codon position.

    Args:
        mrna: mRNA sequence
        position: 0-based position of codon start
        pad_char: Character to use for padding at sequence ends

    Returns:
        203nt string with codon at positions 100-102
    """
    # Need 100nt upstream and 100nt downstream (after the 3nt codon)
    start = position - 100
    end = position + 103  # codon (3) + 100 downstream

    left_pad = max(0, -start)
    right_pad = max(0, end - len(mrna))

    actual_start = max(0, start)
    actual_end = min(len(mrna), end)

    seq = pad_char * left_pad + mrna[actual_start:actual_end] + pad_char * right_pad

    return seq


def find_all_atg_positions(mrna: str) -> List[int]:
    """Find all ATG codon positions in an mRNA sequence."""
    positions = []
    seq_upper = mrna.upper()
    start = 0
    while True:
        pos = seq_upper.find('ATG', start)
        if pos == -1:
            break
        positions.append(pos)
        start = pos + 1
    return positions
