# geney/splicing/branch_points.py — Branch point prediction (adapted from BPP, MIT license).
"""
Lightweight branch point scoring for intronic sequences. Detects whether a
mutation disrupts a predicted branch point, which SpliceAI/Pangolin cannot
detect.

Algorithm: 7-mer PWM scoring for branch point sequence (BPS) combined with
8-mer polypyrimidine tract (PPT) scoring. Based on BPP (Zhang & Li, 2018,
Bioinformatics 33:3166, MIT license).

Usage::

    from geney.splicing.branch_points import score_branch_points, mutation_hits_branch_point

    # Score all branch point candidates in an intron
    candidates = score_branch_points("ACGT...intron...AG")

    # Check if a mutation position hits a branch point
    hit = mutation_hits_branch_point("ACGT...intron...AG", mutation_offset=42)
"""
from __future__ import annotations

import logging
import math
import os
from functools import lru_cache
from typing import Optional

logger = logging.getLogger(__name__)

# 7-position PWM for human branch point sequence (TACTAAC consensus).
# Rows = positions 0-6, columns = A, C, G, T.
_BPS_PWM = [
    [0.2557, 0.2556, 0.2503, 0.2384],
    [0.3782, 0.1913, 0.2759, 0.1545],
    [0.1571, 0.4934, 0.1703, 0.1793],
    [0.0000, 0.0202, 0.0000, 0.9798],
    [0.4674, 0.1074, 0.3693, 0.0560],
    [1.0000, 0.0000, 0.0000, 0.0000],
    [0.1068, 0.5500, 0.0507, 0.2925],
]

_BASE_TO_IDX = {"A": 0, "C": 1, "G": 2, "T": 3}
_CONSENSUS = "TACTAAC"
_BPS_LEN = 7

# PPT scoring table path (bundled)
_PPT_FILE = os.path.join(os.path.dirname(__file__), "ppt_scores.txt")


@lru_cache(maxsize=1)
def _load_ppt_scores() -> dict[str, float]:
    """Load 8-mer PPT scoring table. Cached after first call."""
    scores = {}
    with open(_PPT_FILE) as f:
        for line in f:
            if line.startswith("#") or not line.strip():
                continue
            parts = line.split("\t")
            if len(parts) >= 5:
                scores[parts[0].upper()] = float(parts[4])
    return scores


def _bps_score(kmer: str) -> float:
    """Score a 7-mer against the branch point PWM relative to consensus."""
    if len(kmer) != _BPS_LEN:
        return 0.0
    score = 1.0
    for i, base in enumerate(kmer.upper()):
        idx = _BASE_TO_IDX.get(base)
        if idx is None:
            return 0.0
        p = _BPS_PWM[i][idx]
        # Consensus base probability
        c_idx = _BASE_TO_IDX[_CONSENSUS[i]]
        p_con = _BPS_PWM[i][c_idx]
        if p_con == 0:
            return 0.0
        score *= p / p_con
    return score


@lru_cache(maxsize=16384)
def _bps_score_cached(kmer: str) -> float:
    return _bps_score(kmer)


def _ppt_score(tract: str, ppt_dict: dict[str, float]) -> float:
    """Score a polypyrimidine tract using overlapping 8-mer lookups."""
    tract = tract.upper()
    if len(tract) < 8:
        return 0.0
    # Reference: score of all-T 8-mer
    ref = ppt_dict.get("TTTTTTTT", 1e-10)
    if ref == 0:
        ref = 1e-10
    total = 0.0
    count = 0
    for i in range(len(tract) - 7):
        octamer = tract[i:i + 8]
        if "N" in octamer:
            continue
        s = ppt_dict.get(octamer, 0.0)
        if s > 0:
            total += math.log(s / ref)
            count += 1
    return total / count if count > 0 else 0.0


def _find_agez(seq: str) -> int:
    """Find the AG exclusion zone start (last AG before 3'ss, scanning upstream)."""
    seq = seq.upper()
    # Search from position -5 upstream (relative to end)
    for i in range(len(seq) - 5, max(len(seq) - 50, -1), -1):
        if i >= 0 and seq[i:i + 2] == "AG":
            return i
    return max(0, len(seq) - 40)


def score_branch_points(
    intron_seq: str,
    top_n: int = 5,
    search_start: int = 12,
    search_end: int = 0,
) -> list[dict]:
    """Score branch point candidates in an intron sequence.

    The intron should end with the 3' splice site (AG dinucleotide).

    Args:
        intron_seq: Intronic sequence ending at the 3'ss.
        top_n: Return top N candidates.
        search_start: Min distance from 3'ss to search (bp).
        search_end: Max distance from 3'ss to search (bp). 0 = full intron.

    Returns:
        List of dicts sorted by combined score (descending):
            position: 0-based position in the intron
            bps_motif: 7-mer at this position
            bps_score: PWM score relative to consensus
            ppt_score: polypyrimidine tract score
            combined_score: bps_score * ppt_score
            distance_to_3ss: distance from BP to 3' splice site
    """
    seq = intron_seq.upper()
    slen = len(seq)
    if slen < search_start + _BPS_LEN:
        return []

    ppt_dict = _load_ppt_scores()

    # Search window: from search_end upstream to search_start upstream of 3'ss
    scan_from = max(0, slen - search_end) if search_end > 0 else 0
    scan_to = max(0, slen - search_start - _BPS_LEN)

    candidates = []
    for pos in range(scan_from, scan_to):
        kmer = seq[pos:pos + _BPS_LEN]
        if "N" in kmer:
            continue

        bp_sc = _bps_score_cached(kmer)
        if bp_sc < 0.001:
            continue

        # PPT is between BPS and 3'ss
        ppt_tract = seq[pos + _BPS_LEN:slen - 2]  # exclude terminal AG
        ppt_sc = _ppt_score(ppt_tract, ppt_dict)

        # Combined: BPS score weighted by PPT quality.
        # PPT score is a log-ratio (can be negative). Transform to 0-1 range.
        ppt_weight = 1.0 / (1.0 + math.exp(-ppt_sc - 8))  # sigmoid, centered at -8
        combined = bp_sc * ppt_weight

        candidates.append({
            "position": pos,
            "bps_motif": kmer,
            "bps_score": round(bp_sc, 6),
            "ppt_score": round(ppt_sc, 4),
            "combined_score": round(combined, 6),
            "distance_to_3ss": slen - pos - 4,  # distance from BP adenine (pos+5) to AG
        })

    candidates.sort(key=lambda c: c["combined_score"], reverse=True)
    return candidates[:top_n]


def mutation_hits_branch_point(
    intron_seq: str,
    mutation_offset: int,
    window: int = 3,
    min_score: float = 0.01,
) -> Optional[dict]:
    """Check if a mutation position overlaps a predicted branch point.

    Args:
        intron_seq: Intronic sequence ending at 3'ss.
        mutation_offset: 0-based position of the mutation in the intron.
        window: Positions around the BP to consider a hit (default ±3).
        min_score: Minimum combined score to consider a real BP.

    Returns:
        Dict with BP info if hit, None otherwise. Dict includes:
            disrupted: True
            branch_point: the BP candidate dict
            mutation_distance: distance from mutation to BP adenine
    """
    candidates = score_branch_points(intron_seq, top_n=10)

    for bp in candidates:
        if bp["combined_score"] < min_score:
            continue
        bp_pos = bp["position"]
        # The branch point adenine is at position bp_pos + 5 (6th base of 7-mer)
        bp_a_pos = bp_pos + 5
        distance = abs(mutation_offset - bp_a_pos)
        if distance <= window:
            return {
                "disrupted": True,
                "branch_point": bp,
                "mutation_distance": distance,
            }
        # Also check if mutation is anywhere in the 7-mer
        if bp_pos <= mutation_offset < bp_pos + _BPS_LEN:
            return {
                "disrupted": True,
                "branch_point": bp,
                "mutation_distance": mutation_offset - bp_a_pos,
            }

    return None
