# geney/transcript_tools.py — Transcript-level tool API.
"""Functions that accept a pre-mutated seqmat Transcript and run predictions."""

from __future__ import annotations

from datetime import datetime
from typing import Optional

import numpy as np
import pandas as pd

from .splicing import run_splicing_engine, adjoin_splicing_outcomes, SpliceSimulator
from .oncosplice import Oncosplice


# ---------------------------------------------------------------------------
# Internal helper: splicing prediction from a mutated SeqMat
# ---------------------------------------------------------------------------

def _predict_splicing_from_seqmat(pre_mrna, position: int, engine: str, context: int):
    """Run splicing prediction on both reference and variant from a single mutated SeqMat.

    The mutated SeqMat carries reference info in ``seq_array['ref']`` and
    the current (variant) sequence in ``.seq``.

    Returns:
        (ref_df, var_df) — prediction DataFrames with position index and
        columns ``donor_prob``, ``acceptor_prob``, ``nucleotides``.
    """
    target = pre_mrna.clone(position - context, position + context)

    # --- extract variant sequence (current / mutated) ---
    var_seq = target.seq
    var_indices = target.index

    # --- extract reference sequence ---
    # seq_array['ref'] at subidx==0 gives the original bases (including
    # deleted bases that are no longer in .seq).
    sa = target.seq_array
    ref_mask = sa['subidx'] == 0
    ref_seq = ''.join(sa['ref'][ref_mask])
    ref_indices = sa['index'][ref_mask]

    def _pad_and_predict(seq, indices):
        if len(seq) == 0:
            raise ValueError(f"No sequence data found around position {position}")

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

        if left_missing > 0 or right_missing > 0:
            step = -1 if pre_mrna.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])

        donor_probs, acceptor_probs = run_splicing_engine(seq=seq, engine=engine)

        seq = seq[5000:-5000]
        indices = indices[5000:-5000]
        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'] = pre_mrna.name
        return df

    ref_df = _pad_and_predict(ref_seq, ref_indices)
    var_df = _pad_and_predict(var_seq, var_indices)
    return ref_df, var_df


# ---------------------------------------------------------------------------
# 1. run_splicing
# ---------------------------------------------------------------------------

def run_splicing(
    transcript,
    position: Optional[int] = None,
    engine: str = "spliceai",
    context: int = 7500,
    max_distance: int = 100_000_000,
    min_increase_ratio: float = 0.2,
    tolerance: int = 500,
    min_total_prob: float = 0.5,
    oncosplice_window: int = 13,
) -> pd.DataFrame:
    """Run splicing analysis on a pre-mutated transcript.

    The transcript's ``pre_mrna`` must already carry the mutation (via
    ``apply_mutations``).  Both reference and variant splicing predictions
    are derived from the same SeqMat object.

    Args:
        transcript: A seqmat Transcript whose pre_mrna contains the mutation.
        position: Genomic position for splicing prediction window center.
            Auto-detected from ``pre_mrna.mutations`` when *None*.
        engine: Splicing engine (``"spliceai"`` or ``"pangolin"``).
        context: Half-window for sequence extraction (default 7500).
        max_distance: Max intron length in SpliceSimulator.
        min_increase_ratio: Threshold for novel splice sites.
        tolerance: Boundary matching distance in bp.
        min_total_prob: Cryptic site probability threshold for IR filtering.
        oncosplice_window: Conservation smoothing window (amino acids).

    Returns:
        DataFrame with all viable isoforms and their oncosplice scores.
        Empty DataFrame if no missplicing is detected.
    """
    pre_mrna = transcript.pre_mrna

    # --- auto-detect position ---
    if position is None:
        mutations = getattr(pre_mrna, 'mutations', None)
        if mutations is not None and len(mutations) > 0:
            positions = [m[0] if isinstance(m, (list, tuple)) else m for m in mutations]
            position = int(np.mean(positions))
        else:
            raise ValueError(
                "Cannot auto-detect position: pre_mrna has no mutations. "
                "Pass `position` explicitly."
            )

    # --- splicing predictions ---
    ref_df, var_df = _predict_splicing_from_seqmat(pre_mrna, position, engine, context)

    splicing_results = adjoin_splicing_outcomes(
        {'ref': ref_df, 'event': var_df}, transcript
    )

    # --- splice simulation ---
    ss = SpliceSimulator(
        splicing_results,
        transcript,
        feature="event",
        max_distance=max_distance,
        min_increase_ratio=min_increase_ratio,
        tolerance=tolerance,
        min_total_prob=min_total_prob,
    )

    # --- conservation vector ---
    cons_vector = getattr(transcript, 'cons_vector', None)
    if cons_vector is None:
        protein_length = len(getattr(transcript, 'protein', '') or '')
        cons_vector = np.ones(max(protein_length, 1))

    # --- build report ---
    base_report = pd.Series({
        "transcript_id": getattr(transcript, 'transcript_id', None),
        "splicing_engine": engine,
        "central_position": position,
        "time_of_execution": datetime.now().strftime("%Y-%m-%d %H:%M:%S"),
    })

    ss_metadata = ss.report(position)

    rows = []
    for variant_transcript, isoform_metadata in ss.get_viable_transcripts(metadata=True):
        onco = Oncosplice(
            transcript.protein,
            variant_transcript.protein,
            cons_vector,
            oncosplice_window,
        )
        rows.append(
            pd.concat([
                base_report,
                ss_metadata,
                isoform_metadata,
                pd.Series({
                    "reference_mrna": transcript.mature_mrna.seq,
                    "variant_mrna": variant_transcript.mature_mrna.seq,
                }),
                onco.get_analysis_series(),
            ])
        )

    return pd.DataFrame(rows)


# ---------------------------------------------------------------------------
# 2. run_tis
# ---------------------------------------------------------------------------

def run_tis(
    transcript,
    proximity_window: int = 500,
    probability_threshold: float = 0.1,
    use_prior: bool = False,
    skip_if_not_near_start: bool = True,
) -> dict:
    """Run TIS (Translation Initiation Site) analysis on a pre-mutated transcript.

    The transcript should already have ``mature_mrna`` generated (with the
    mutation applied at the pre-mRNA level).  Reference and variant mRNA
    are extracted from the mutated SeqMat's ``.reference_seq`` and ``.seq``.

    Args:
        transcript: A seqmat Transcript whose pre_mrna contains the mutation
            and whose mature_mrna has been generated.
        proximity_window: Max distance from 5' end to run analysis (nt).
        probability_threshold: Min probability to consider a TIS.
        use_prior: Whether to multiply by codon prior in scoring.
        skip_if_not_near_start: Skip mutations far from the start codon.

    Returns:
        Dictionary with TIS analysis results, or a dict with
        ``status='skipped'`` / ``status='error'``.
    """
    from .tis import compute_tis_shift_metrics

    # --- ensure mature mRNA exists ---
    mrna = getattr(transcript, 'mature_mrna', None)
    if mrna is None:
        return {'status': 'error', 'reason': 'no_mature_mrna'}

    ref_mrna_seq = getattr(mrna, 'reference_seq', None)
    var_mrna_seq = getattr(mrna, 'seq', None)

    if not ref_mrna_seq or not var_mrna_seq:
        return {'status': 'error', 'reason': 'missing_ref_or_var_mrna'}

    # --- find mutation positions in mRNA ---
    mutated_positions = getattr(mrna, 'mutated_positions', None)
    mutation_mrna_positions = []
    if mutated_positions is not None:
        mutation_mrna_positions = list(mutated_positions)
    elif hasattr(mrna, 'index'):
        # Fallback: compare ref and var sequences
        for i, (r, v) in enumerate(zip(ref_mrna_seq, var_mrna_seq)):
            if r != v:
                mutation_mrna_positions.append(i)

    # --- proximity check ---
    if skip_if_not_near_start and mutation_mrna_positions:
        if all(pos >= proximity_window for pos in mutation_mrna_positions):
            return {
                'status': 'skipped',
                'reason': 'mutation_far_from_start_codon',
                'transcript_id': getattr(transcript, 'transcript_id', None),
                'proximity_window': proximity_window,
            }

    # --- run TIS analysis ---
    try:
        from .pipelines import _tis_completed_dict

        result = compute_tis_shift_metrics(
            ref_mrna=ref_mrna_seq,
            var_mrna=var_mrna_seq,
            canonical_start=None,
            probability_threshold=probability_threshold,
            use_prior=use_prior,
        )

        tid = getattr(transcript, 'transcript_id', None)
        return _tis_completed_dict(
            result, mut_id=None, gene=None, transcript_id=tid,
            ref_mrna_seq=ref_mrna_seq, var_mrna_seq=var_mrna_seq,
            mrna_positions=mutation_mrna_positions,
        )
    except Exception as e:
        return {
            'status': 'error',
            'reason': str(e),
            'transcript_id': getattr(transcript, 'transcript_id', None),
        }


# ---------------------------------------------------------------------------
# 3. run_folding
# ---------------------------------------------------------------------------

def run_folding(
    transcript,
    window_size: int = 39,
) -> dict:
    """Compute delta-MFE for a pre-mutated transcript with rich per-window output.

    The transcript should already have ``mature_mrna`` generated (with the
    mutation applied at the pre-mRNA level).  Reference and variant mRNA
    are extracted from the mutated SeqMat.

    Returns per-window structures, sequences, individual MFEs, and a 0-10
    normalized score alongside the mean delta-MFE.

    Args:
        transcript: A seqmat Transcript whose pre_mrna contains the mutation
            and whose mature_mrna has been generated.
        window_size: Local folding window in nt (default 39).

    Returns:
        Dictionary with ``status``, per-mutation ``mutations`` list (each with
        per-window detail), mean ``delta_mfe``, and overall ``score``.
    """
    from .folding.delta_mfe import compute_folding_profile
    from .folding.vienna import fold_available

    if not fold_available():
        return {
            'status': 'error',
            'reason': 'ViennaRNA not installed (pip install ViennaRNA)',
        }

    # --- ensure mature mRNA exists ---
    mrna = getattr(transcript, 'mature_mrna', None)
    if mrna is None or not hasattr(mrna, 'seq'):
        return {'status': 'error', 'reason': 'no_mature_mrna'}

    ref_mrna = getattr(mrna, 'reference_seq', None)
    var_mrna = getattr(mrna, 'seq', None)

    if not ref_mrna or not var_mrna:
        return {'status': 'error', 'reason': 'missing_ref_or_var_mrna'}

    ref_mrna = ref_mrna.upper()
    var_mrna = var_mrna.upper()

    # --- early exit if mRNA sequences are identical ---
    if ref_mrna == var_mrna:
        return {
            'status': 'skipped',
            'reason': 'mrna_identical',
            'transcript_id': getattr(transcript, 'transcript_id', None),
        }

    # --- find mutation positions in mRNA ---
    mutated_positions = getattr(mrna, 'mutated_positions', None)
    if mutated_positions is not None:
        mrna_positions = list(mutated_positions)
    else:
        # Fallback: compare ref and var sequences character by character
        mrna_positions = [
            i for i, (r, v) in enumerate(zip(ref_mrna, var_mrna)) if r != v
        ]

    if not mrna_positions:
        return {
            'status': 'skipped',
            'reason': 'no_mutations_in_mrna',
            'transcript_id': getattr(transcript, 'transcript_id', None),
        }

    # --- compute folding profile for each mutation position ---
    per_mut: list[dict] = []
    for mrna_pos in mrna_positions:
        try:
            profile = compute_folding_profile(ref_mrna, var_mrna, mrna_pos, window_size)
            per_mut.append({
                'mrna_position': mrna_pos,
                'status': 'ok',
                'ref_base': profile.ref_base,
                'var_base': profile.var_base,
                'delta_mfe': profile.delta_mfe,
                'score': profile.score,
                'num_windows': len(profile.windows),
                'windows': [
                    {
                        'start': w.start,
                        'end': w.end,
                        'ref_seq': w.ref_seq,
                        'var_seq': w.var_seq,
                        'ref_structure': w.ref_structure,
                        'var_structure': w.var_structure,
                        'ref_mfe': w.ref_mfe,
                        'var_mfe': w.var_mfe,
                        'delta_mfe': w.delta_mfe,
                        'mutation_offset': w.mutation_offset,
                    }
                    for w in profile.windows
                ],
            })
        except Exception as e:
            per_mut.append({
                'mrna_position': mrna_pos,
                'status': 'fold_error',
                'delta_mfe': None,
                'score': None,
                'error': str(e),
            })

    ok_deltas = [m['delta_mfe'] for m in per_mut if m['status'] == 'ok' and m['delta_mfe'] is not None]
    mean_delta = float(np.mean(ok_deltas)) if ok_deltas else None

    return {
        'status': 'completed',
        'transcript_id': getattr(transcript, 'transcript_id', None),
        'ref_mrna_length': len(ref_mrna),
        'var_mrna_length': len(var_mrna),
        'window_size': window_size,
        'mutations': per_mut,
        'delta_mfe': mean_delta,
        'score': min(10.0, abs(mean_delta)) if mean_delta is not None else None,
    }


# ---------------------------------------------------------------------------
# 4. MutationContextLibrary
# ---------------------------------------------------------------------------

class MutationContextLibrary:
    """Library of (transcript, context_seqmat) pairs for every transcript of a gene.

    Given a mutation ID, loads the gene and builds one entry per transcript.
    Each entry is a ``(transcript, mutation_context)`` tuple where:

    * **transcript** — a cloned, mutated Transcript with ``mature_mrna`` and
      ``protein`` already generated.
    * **mutation_context** — a SeqMat clone of the mutated ``pre_mrna`` covering
      *context* bp on each side of the mutation position.  Because the
      mutation is applied before cloning, the SeqMat carries both the variant
      sequence (``.seq``) and the original reference (``seq_array['ref']``).

    Usage::

        lib = MutationContextLibrary("KRAS:12:25227343:G:T", context=1000)

        # Access by transcript ID
        transcript, ctx = lib["ENST00000256078"]

        # Iterate all entries
        for tid, (transcript, ctx) in lib:
            print(tid, len(ctx.seq))

    Args:
        mut_id: Mutation ID (``GENE:CHR:POS:REF:ALT``; pipe-separated for compound).
        organism: Genome build (default ``"hg38"``).
        context: Number of base pairs on each side of the mutation to extract.
        protein_coding_only: If *True*, skip non-protein-coding transcripts.
        permissive: If *True*, ignore reference-allele mismatches when applying
            mutations.
    """

    def __init__(
        self,
        mut_id: str,
        organism: str = "hg38",
        context: int = 1000,
        protein_coding_only: bool = True,
        permissive: bool = True,
    ):
        from seqmat import Gene
        from .variants import MutationalEvent
        from .utils import _apply_mutation_safe, _is_protein_coding

        m = MutationalEvent(mut_id)
        if not m.compatible():
            raise ValueError(f"Mutations in event are incompatible: {mut_id}")

        gene = Gene.from_file(m.gene, organism=organism)

        self.gene_name = m.gene
        self.mut_id = mut_id
        self.context = context
        self._entries: dict[str, tuple] = {}
        self._errors: dict[str, str] = {}

        # --- collect raw transcripts from gene ---
        transcripts_map = getattr(gene, 'transcripts', None) or getattr(gene, '_transcripts', {})
        if hasattr(transcripts_map, 'items'):
            items = list(transcripts_map.items())
        elif hasattr(transcripts_map, '__iter__'):
            items = [
                (getattr(t, 'transcript_id', str(i)), t)
                for i, t in enumerate(transcripts_map)
            ]
        else:
            items = []

        mutation_args = m.mutation_args()

        for tid, raw_transcript in items:
            try:
                if protein_coding_only and not _is_protein_coding(raw_transcript):
                    continue

                t = raw_transcript.clone().generate_mature_mrna()

                # Apply mutations directly to mature mRNA (genomic coords preserved)
                for pos, ref, alt in mutation_args:
                    try:
                        _apply_mutation_safe(t.mature_mrna, pos, ref, alt, permissive=permissive)
                    except Exception:
                        pass  # mutation outside this transcript's exonic range

                # Scoped pre-mRNA read for mutation context (intronic context)
                central_pos = m.central_position
                ctx = raw_transcript.clone()
                ctx.generate_pre_mrna(
                    region_start=central_pos - context,
                    region_end=central_pos + context,
                )
                for pos, ref, alt in mutation_args:
                    try:
                        _apply_mutation_safe(ctx.pre_mrna, pos, ref, alt, permissive=permissive)
                    except Exception:
                        pass
                context_seqmat = ctx.pre_mrna

                # Generate protein from mutated mature mRNA
                t.generate_protein()

                self._entries[tid] = (t, context_seqmat)
            except Exception as e:
                self._errors[tid] = str(e)

    # --- dict-like access ---

    def __getitem__(self, transcript_id: str):
        """Return ``(transcript, context_seqmat)`` for a given transcript ID."""
        return self._entries[transcript_id]

    def __contains__(self, transcript_id: str) -> bool:
        return transcript_id in self._entries

    def __iter__(self):
        """Yield ``(transcript_id, (transcript, context_seqmat))`` tuples."""
        return iter(self._entries.items())

    def __len__(self) -> int:
        return len(self._entries)

    def __repr__(self) -> str:
        return (
            f"MutationContextLibrary({self.gene_name!r}, "
            f"{len(self._entries)} transcripts, "
            f"context={self.context}bp)"
        )

    @property
    def transcript_ids(self) -> list[str]:
        """All transcript IDs that were successfully processed."""
        return list(self._entries.keys())

    @property
    def errors(self) -> dict[str, str]:
        """Transcript IDs that failed, mapped to error messages."""
        return dict(self._errors)
