# oncosplice/engines.py
from __future__ import annotations

import logging
from typing import Dict, List, Tuple, Optional, Union
import numpy as np

logger = logging.getLogger(__name__)

# Lazy-loaded model containers (loaded automatically on first use)
_pang_models = None
_pang_device = None


def _get_torch_device():
    """Get the best available device for PyTorch."""
    import sys
    import torch

    if sys.platform == 'darwin' and torch.backends.mps.is_available():
        try:
            torch.tensor([1.0], device="mps")
            return torch.device("mps")
        except RuntimeError:
            return torch.device("cpu")
    elif torch.cuda.is_available():
        return torch.device("cuda")
    return torch.device("cpu")


def _load_pangolin_models():
    """Lazy load Pangolin models."""
    global _pang_models, _pang_device

    if _pang_models is not None:
        return _pang_models

    import torch
    from pkg_resources import resource_filename
    from pangolin.model import Pangolin, L, W, AR

    _pang_device = _get_torch_device()
    logger.info("Pangolin loading to %s...", _pang_device)

    _pang_models = []
    pang_model_nums = [0, 1, 2, 3, 4, 5, 6, 7]

    for i in pang_model_nums:
        for j in range(1, 6):
            try:
                model = Pangolin(L, W, AR).to(_pang_device)
                model_path = resource_filename("pangolin", f"models/final.{j}.{i}.3")
                weights = torch.load(model_path, weights_only=True, map_location=_pang_device)
                model.load_state_dict(weights)
                model.eval()
                _pang_models.append(model)
            except Exception as e:
                logger.warning("Failed to load Pangolin model %d.%d: %s", j, i, e)

    logger.info("Pangolin loaded (%d models).", len(_pang_models))
    return _pang_models


# Cached openspliceai models to avoid reloading from disk on every call
_sai_models = None
_sai_params = None
_sai_device = None
_sai_consts = None

_OPENSPLICEAI_MODEL_DIR = None

def set_openspliceai_model_dir(path: str) -> None:
    """Set the directory containing OpenSpliceAI MANE 10000nt model weights."""
    import os
    global _OPENSPLICEAI_MODEL_DIR
    if not os.path.isdir(path):
        raise ValueError(f"OpenSpliceAI model directory does not exist: {path}")
    _OPENSPLICEAI_MODEL_DIR = path


def get_openspliceai_model_dir() -> str:
    """Return the path to the OpenSpliceAI MANE 10000nt model directory."""
    return _get_openspliceai_model_dir()


def _get_openspliceai_model_dir() -> str:
    """Return the path to the OpenSpliceAI MANE 10000nt model directory."""
    global _OPENSPLICEAI_MODEL_DIR
    if _OPENSPLICEAI_MODEL_DIR is not None:
        return _OPENSPLICEAI_MODEL_DIR

    import os

    # Models ship inside the package at geney/models/openspliceai-mane/10000nt
    pkg_dir = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
    default = os.path.join(pkg_dir, 'models', 'openspliceai-mane', '10000nt')
    default = os.path.normpath(default)

    _OPENSPLICEAI_MODEL_DIR = os.environ.get('OPENSPLICEAI_MODEL_DIR', default)
    return _OPENSPLICEAI_MODEL_DIR


def set_splicing_device(device: str) -> None:
    """Override the device used by Pangolin models (e.g. 'cpu', 'cuda', 'mps')."""
    import torch
    global _pang_device, _pang_models
    _pang_device = torch.device(device)
    _pang_models = None  # Force reload on next use


def get_splicing_device():
    """Return the current Pangolin device (auto-detected if not overridden)."""
    global _pang_device
    if _pang_device is None:
        _pang_device = _get_torch_device()
    return _pang_device


def pang_one_hot_encode(seq: str) -> np.ndarray:
    """One-hot encode DNA sequence for Pangolin model."""
    if not isinstance(seq, str):
        raise TypeError(f"Expected string, got {type(seq).__name__}")

    IN_MAP = np.asarray([[0, 0, 0, 0],  # N
                         [1, 0, 0, 0],  # A
                         [0, 1, 0, 0],  # C
                         [0, 0, 1, 0],  # G
                         [0, 0, 0, 1]]) # T

    valid_chars = set('ACGTN')
    if not all(c.upper() in valid_chars for c in seq):
        raise ValueError("Sequence contains invalid characters")

    seq = seq.upper().replace('A', '1').replace('C', '2')
    seq = seq.replace('G', '3').replace('T', '4').replace('N', '0')

    seq_array = np.asarray(list(map(int, list(seq))))
    return IN_MAP[seq_array.astype('int8')]




def pangolin_predict_probs(seq: str, models: list = None) -> Tuple[List[float], List[float]]:
    """Predict splice site probabilities using Pangolin.

    Pangolin outputs shape (1, 12, seq_len) where:
    - 12 channels = 4 tissues × 3 prediction types
    - For each tissue: [site_usage, acceptor_gain, donor_gain] or similar

    We aggregate by taking max across tissues.
    """
    import torch

    if models is None:
        models = _load_pangolin_models()

    if not models:
        raise ValueError("No Pangolin models loaded")

    x = pang_one_hot_encode(seq)
    x = torch.tensor(x.T[None, :, :], dtype=torch.float32, device=_pang_device)

    preds = []
    with torch.no_grad():
        for model in models:
            pred = model(x)
            preds.append(pred.cpu().numpy())

    y = np.mean(preds, axis=0)  # Shape: (1, 12, seq_len)

    # Pangolin has 12 channels organized as:
    # Indices 0,3,6,9: site usage scores for 4 tissues
    # Indices 1,4,7,10: acceptor gain scores for 4 tissues
    # Indices 2,5,8,11: donor gain scores for 4 tissues
    # Take max across the 4 tissues for each type

    # Acceptor: max of channels 1, 4, 7, 10
    acceptor_channels = y[0, [1, 4, 7, 10], :]  # (4, seq_len)
    acceptor_probs = np.max(acceptor_channels, axis=0).tolist()

    # Donor: max of channels 2, 5, 8, 11
    donor_channels = y[0, [2, 5, 8, 11], :]  # (4, seq_len)
    donor_probs = np.max(donor_channels, axis=0).tolist()

    return donor_probs, acceptor_probs


def _load_openspliceai_models():
    """Cache openspliceai PyTorch models at module level.

    Avoids reloading all model weights from disk on every
    ``sai_predict_probs()`` call (~1-2s savings per call).
    """
    global _sai_models, _sai_params, _sai_device, _sai_consts
    if _sai_models is not None:
        return _sai_models, _sai_params, _sai_device, _sai_consts

    import io, sys, torch
    from openspliceai.predict.predict import load_pytorch_models, setup_device
    from openspliceai.constants import SL, CL_max

    model_dir = _get_openspliceai_model_dir()
    consts = {'SL': SL, 'CL_max': CL_max}
    device = setup_device()

    _stdout = sys.stdout
    sys.stdout = io.StringIO()
    try:
        models, params = load_pytorch_models(model_dir, device, SL, CL_max)
    finally:
        sys.stdout = _stdout

    _sai_models = models
    _sai_params = params
    _sai_device = device
    _sai_consts = consts
    return _sai_models, _sai_params, _sai_device, _sai_consts


def _sai_predict_from_cached(seq: str) -> np.ndarray:
    """Run openspliceai prediction on a single sequence using cached models.

    Returns array of shape ``(seq_len, 3)``.
    """
    import torch
    from openspliceai.predict.predict import create_datapoints

    models, params, device, consts = _load_openspliceai_models()
    sequence_length = len(seq)

    X = create_datapoints(seq, SL=consts['SL'], CL_max=consts['CL_max'])
    X = torch.tensor(X, dtype=torch.float32).permute(0, 2, 1).to(device)

    with torch.no_grad():
        y_pred = torch.mean(
            torch.stack([m(X).detach().cpu() for m in models]),
            axis=0,
        )
    y_pred = y_pred.permute(0, 2, 1).contiguous().view(-1, y_pred.shape[1])
    y_pred = y_pred[:sequence_length, :]
    return y_pred.numpy()


def sai_predict_probs(seq: str) -> Tuple[np.ndarray, np.ndarray]:
    """Predict acceptor and donor probabilities using OpenSpliceAI.

    Uses cached models to avoid reloading weights from disk each call.

    Returns (acceptor_probs, donor_probs) as numpy arrays matching the
    full input sequence length.
    """
    y = _sai_predict_from_cached(seq)  # (seq_len, 3)
    return y[:, 1], y[:, 2]  # acceptor, donor


def run_spliceai_seq(
    seq: str,
    indices: Union[List[int], np.ndarray],
    threshold: float = 0.0,
) -> Tuple[Dict[int, float], Dict[int, float]]:
    """Run SpliceAI on seq and return donor/acceptor sites above threshold."""
    if len(indices) != len(seq):
        raise ValueError(f"indices length ({len(indices)}) must match sequence length ({len(seq)})")

    acc_probs, don_probs = sai_predict_probs(seq)
    acceptor = {pos: p for pos, p in zip(indices, acc_probs) if p >= threshold}
    donor = {pos: p for pos, p in zip(indices, don_probs) if p >= threshold}
    return donor, acceptor


def _generate_random_sequence(length: int) -> str:
    """Generate a random DNA sequence of given length."""
    import random
    return ''.join(random.choices('ACGT', k=length))


def run_splicing_engine(
    seq: Optional[str] = None,
    engine: str = "spliceai",
) -> Tuple[List[float], List[float]]:
    """Run specified splicing engine to predict splice site probabilities."""
    if seq is None:
        seq = _generate_random_sequence(15_001)

    if not isinstance(seq, str) or not seq:
        raise ValueError("Sequence must be a non-empty string")

    valid_chars = set("ACGTN")
    if not all(c.upper() in valid_chars for c in seq):
        raise ValueError("Sequence contains invalid nucleotides")

    match engine:
        case "spliceai":
            acc, don = sai_predict_probs(seq)
            return don.tolist(), acc.tolist()
        case "pangolin":
            return pangolin_predict_probs(seq)
        case _:
            raise ValueError(f"Engine '{engine}' not implemented. Available: 'spliceai', 'pangolin'")


# ------------------------------------------------------------------------------
# Higher-level prediction utilities (formerly in splicing_table.py)
# ------------------------------------------------------------------------------

def _prepare_splice_input(s, position: int, context: int = 7500) -> dict:
    """Extract and pad the sequence window needed for splicing prediction.

    Returns a dict with ``padded_seq`` (the full string to feed to the model),
    ``trimmed_seq``, ``trimmed_indices``, and ``name`` — everything needed by
    :func:`_assemble_splice_df` after the model runs.
    """
    if position < s.index.min() or position > s.index.max():
        raise ValueError(f"Position {position} is outside sequence bounds [{s.index.min()}, {s.index.max()}]")

    target = s.clone(position - context, position + context)

    if len(target.seq) == 0:
        raise ValueError(f"No sequence data found around position {position} with context {context}")

    seq, indices = target.seq, target.index

    if len(indices) == 0:
        raise ValueError(f"No indices found in sequence around position {position}")

    rel_pos = np.abs(indices - position).argmin()
    left_missing, right_missing = max(0, context - rel_pos), max(0, context - (len(seq) - rel_pos))

    if left_missing > 0 or right_missing > 0:
        step = -1 if s.rev else 1

        if left_missing > 0:
            left_pad = np.arange(indices[0] - step * left_missing, indices[0], step)
        else:
            left_pad = np.array([], dtype=indices.dtype)

        if right_missing > 0:
            right_pad = np.arange(indices[-1] + step, indices[-1] + step * (right_missing + 1), step)
        else:
            right_pad = np.array([], dtype=indices.dtype)

        seq = 'N' * left_missing + seq + 'N' * right_missing
        indices = np.concatenate([left_pad, indices, right_pad])

    return {
        "padded_seq": seq,
        "trimmed_seq": seq[5000:-5000],
        "trimmed_indices": indices[5000:-5000],
        "name": s.name,
    }


def _assemble_splice_df(
    donor_probs, acceptor_probs, prep: dict,
) -> "pd.DataFrame":
    """Build the splicing DataFrame from model predictions and prepared input."""
    import pandas as pd

    seq = prep["trimmed_seq"]
    indices = prep["trimmed_indices"]
    expected_len = len(seq)

    if len(donor_probs) != expected_len:
        if len(donor_probs) > expected_len:
            offset = (len(donor_probs) - expected_len) // 2
            donor_probs = donor_probs[offset:offset + expected_len]
            acceptor_probs = acceptor_probs[offset:offset + expected_len]
        else:
            pad_len = expected_len - len(donor_probs)
            donor_probs = donor_probs + [0.0] * pad_len
            acceptor_probs = acceptor_probs + [0.0] * pad_len

    df = pd.DataFrame({
        'position': indices,
        'donor_prob': donor_probs,
        'acceptor_prob': acceptor_probs,
        'nucleotides': list(seq)
    }).set_index('position').round(3)

    df.attrs['name'] = prep["name"]
    return df


def predict_splicing(s, position: int, engine: str = 'spliceai', context: int = 7500):
    """
    Predict splicing probabilities at a given position using the specified engine.

    Args:
        s: Sequence object with .seq, .index, .clone(), .rev attributes
        position: The genomic position to predict splicing probabilities for.
        engine: The prediction engine to use. Supported: 'spliceai', 'pangolin'.
        context: The length of the target central region (default: 7500).

    Returns:
        pd.DataFrame with position index and columns: donor_prob, acceptor_prob, nucleotides
    """
    prep = _prepare_splice_input(s, position, context)
    donor_probs, acceptor_probs = run_splicing_engine(seq=prep["padded_seq"], engine=engine)
    return _assemble_splice_df(donor_probs, acceptor_probs, prep)


def adjoin_splicing_outcomes(
    splicing_predictions: Dict[str, 'pd.DataFrame'],
    transcript: Optional[object] = None,
) -> 'pd.DataFrame':
    """
    Combine splicing predictions for multiple mutations into a multi-index DataFrame.

    Args:
        splicing_predictions: {label -> DF with 'donor_prob','acceptor_prob','nucleotides'}
        transcript: optional transcript (must have .acceptors, .donors, .rev)
    """
    import pandas as pd

    if not splicing_predictions:
        raise ValueError("splicing_predictions cannot be empty")

    dfs = []
    for label, df in splicing_predictions.items():
        if not isinstance(df, pd.DataFrame):
            raise TypeError(f"Expected DataFrame for '{label}', got {type(df).__name__}")

        required_cols = ["donor_prob", "acceptor_prob", "nucleotides"]
        missing = [c for c in required_cols if c not in df.columns]
        if missing:
            raise ValueError(f"DataFrame for '{label}' missing required columns: {missing}")

        # Handle duplicate indices by keeping the first occurrence
        if df.index.duplicated().any():
            df = df[~df.index.duplicated(keep='first')]

        var_df = df.rename(
            columns={
                "donor_prob": ("donors", f"{label}_prob"),
                "acceptor_prob": ("acceptors", f"{label}_prob"),
                "nucleotides": ("nts", f"{label}"),
            }
        )
        dfs.append(var_df)

    # Use outer join to handle index mismatches between ref and variant
    try:
        full_df = pd.concat(dfs, axis=1, join='outer')
    except Exception as e:
        # Fallback: align on intersection of indices
        common_idx = dfs[0].index
        for df in dfs[1:]:
            common_idx = common_idx.intersection(df.index)
        if len(common_idx) == 0:
            raise ValueError(f"Failed to concatenate DataFrames: no common indices found") from e
        aligned_dfs = [df.loc[df.index.intersection(common_idx)] for df in dfs]
        full_df = pd.concat(aligned_dfs, axis=1)

    if not isinstance(full_df.columns, pd.MultiIndex):
        full_df.columns = pd.MultiIndex.from_tuples(full_df.columns)

    if transcript is not None:
        full_df[("acceptors", "annotated")] = full_df.apply(
            lambda row: row.name in transcript.acceptors, axis=1
        )
        full_df[("donors", "annotated")] = full_df.apply(
            lambda row: row.name in transcript.donors, axis=1
        )
        full_df.sort_index(axis=1, level=0, inplace=True)
        full_df.sort_index(ascending=not transcript.rev, inplace=True)
    else:
        full_df.sort_index(axis=1, level=0, inplace=True)

    return full_df


# ============================================================================
# Cross-mutation batched prediction
# ============================================================================

def batch_sai_predict_probs(
    sequences: List[str],
    max_chunk: int = 64,
) -> List[Tuple[np.ndarray, np.ndarray]]:
    """Predict acceptor/donor probabilities for multiple sequences using cached
    openspliceai models in a single batched forward pass.

    All sequences' windows are concatenated into one tensor and passed through
    each model once, then split back per sequence.  This avoids N separate GPU
    round-trips.

    Args:
        sequences: List of DNA sequences (typically ~15 000 nt each).
        max_chunk: Max windows per GPU forward pass to avoid OOM.

    Returns:
        List of ``(acceptor, donor)`` tuples, one per input sequence.
    """
    if not sequences:
        return []

    import torch
    from openspliceai.predict.predict import create_datapoints

    models, params, device, consts = _load_openspliceai_models()
    SL, CL_max = consts['SL'], consts['CL_max']

    # 1. Create datapoints for all sequences and track offsets
    all_windows = []       # list of tensors, each (num_windows_i, 4, window_size)
    window_counts = []     # num_windows per sequence
    seq_lengths = []       # original sequence lengths

    for seq in sequences:
        X = create_datapoints(seq, SL=SL, CL_max=CL_max)
        X = torch.tensor(X, dtype=torch.float32).permute(0, 2, 1)  # (nw, 4, ws)
        all_windows.append(X)
        window_counts.append(X.shape[0])
        seq_lengths.append(len(seq))

    # 2. Concatenate all windows
    X_all = torch.cat(all_windows, dim=0).to(device)  # (total_windows, 4, ws)
    total_windows = X_all.shape[0]

    # 3. Forward pass through each model, chunked if needed
    # Model input: (N, 4, SL+CL_max) → output: (N, 3, SL)
    n_models = len(models)
    y_accum = None  # initialized on first prediction to get correct output size

    with torch.no_grad():
        for model in models:
            for start in range(0, total_windows, max_chunk):
                end = min(start + max_chunk, total_windows)
                chunk = X_all[start:end]
                pred = model(chunk).detach().cpu()  # (chunk_size, 3, SL)
                if y_accum is None:
                    out_sl = pred.shape[2]
                    y_accum = torch.zeros(total_windows, 3, out_sl, dtype=torch.float32)
                y_accum[start:end] += pred / n_models

    # 4. Split back per sequence and extract acceptor/donor
    # Each sequence's windows produce SL positions each; concatenated they
    # cover the original sequence.  Permute to (nw, SL, 3) then flatten.
    results: List[Tuple[np.ndarray, np.ndarray]] = []
    offset = 0
    for i, (nw, seq_len) in enumerate(zip(window_counts, seq_lengths)):
        y_seq = y_accum[offset:offset + nw]  # (nw, 3, SL)
        y_seq = y_seq.permute(0, 2, 1).contiguous().view(-1, 3)  # (nw * SL, 3)
        y_flat = y_seq[:seq_len, :].numpy()
        results.append((y_flat[:, 1], y_flat[:, 2]))  # acceptor, donor
        offset += nw

    return results


def batch_pangolin_predict_probs(
    sequences: List[str],
    max_chunk: int = 16,
) -> List[Tuple[List[float], List[float]]]:
    """Predict donor/acceptor probabilities for multiple sequences using
    Pangolin in batched forward passes.

    Sequences are one-hot encoded, batched together, and passed through
    each Pangolin model once.

    Args:
        sequences: List of DNA sequences.
        max_chunk: Max sequences per GPU forward pass to avoid OOM.

    Returns:
        List of ``(donor_probs, acceptor_probs)`` tuples, one per input sequence.
    """
    if not sequences:
        return []

    import torch

    models = _load_pangolin_models()
    if not models:
        raise ValueError("No Pangolin models loaded")

    # All sequences must be the same length for batched tensor
    # (Pangolin expects fixed-length input per batch)
    seq_len = len(sequences[0])
    same_length = all(len(s) == seq_len for s in sequences)

    if not same_length:
        # Fallback: process individually when lengths differ
        return [pangolin_predict_probs(seq, models) for seq in sequences]

    # Batch encode
    encoded = [pang_one_hot_encode(seq) for seq in sequences]
    X = np.stack(encoded, axis=0)  # (N, seq_len, 4)
    X = torch.tensor(np.transpose(X, (0, 2, 1)), dtype=torch.float32, device=_pang_device)

    # Batched forward pass through each model, chunked
    N = len(sequences)
    y_accum = np.zeros((N, 12, seq_len), dtype=np.float32)
    n_models = len(models)

    with torch.no_grad():
        for model in models:
            for start in range(0, N, max_chunk):
                end = min(start + max_chunk, N)
                pred = model(X[start:end]).cpu().numpy()  # (chunk, 12, seq_len)
                y_accum[start:end] += pred / n_models

    # Extract donor/acceptor per sequence
    results: List[Tuple[List[float], List[float]]] = []
    for i in range(N):
        y = y_accum[i]  # (12, seq_len)
        acceptor_probs = np.max(y[[1, 4, 7, 10], :], axis=0).tolist()
        donor_probs = np.max(y[[2, 5, 8, 11], :], axis=0).tolist()
        results.append((donor_probs, acceptor_probs))

    return results


def batch_run_splicing_engine(
    sequences: List[str],
    engine: str = "spliceai",
) -> List[Tuple[List[float], List[float]]]:
    """Batch version of :func:`run_splicing_engine`.

    Returns a list of ``(donor_probs, acceptor_probs)`` tuples.  Each element
    matches what ``run_splicing_engine()`` would return for a single sequence.
    """
    if not sequences:
        return []

    match engine:
        case "spliceai":
            sai_results = batch_sai_predict_probs(sequences)
            # sai returns (acceptor, donor); convert to (donor, acceptor) lists
            return [(don.tolist(), acc.tolist()) for acc, don in sai_results]
        case "pangolin":
            return batch_pangolin_predict_probs(sequences)
        case _:
            raise ValueError(f"Engine '{engine}' not implemented. Available: 'spliceai', 'pangolin'")
