# geney/mirna/disruption.py — miRNA binding site disruption analysis.
"""
Analyzes how mutations alter miRNA binding sites by comparing seed matches
between reference and variant mRNA sequences. Produces a differential
binding profile showing gained, lost, and modified sites.

For full repression prediction (Kd, occupancy, miBSIM correction), use the
standalone miBSIM package. This module provides the fast, lightweight
disruption analysis suitable for pipeline integration.
"""
from __future__ import annotations

import math
from dataclasses import dataclass, field
from typing import Optional

from .seed import find_seed_matches, SeedMatch


_MIN_EXPRESSION_RPM = 10  # Minimum RPM to consider a miRNA "expressed" in a tissue


def _enrich_tissue_expression(changes: list) -> None:
    """Add tissue expression data to each SiteChange (in-place)."""
    if not changes:
        return
    try:
        from .expression import _load_expression_matrix
        matrix = _load_expression_matrix()
    except Exception:
        return

    for sc in changes:
        expr = matrix.get(sc.mirna_id, {})
        if expr:
            sc.tissue_expression = {t: round(v, 1) for t, v in expr.items() if v > 0}
            sc.top_tissues = sorted(
                [t for t, v in expr.items() if v >= _MIN_EXPRESSION_RPM],
                key=lambda t: -expr[t],
            )
        else:
            sc.tissue_expression = {}
            sc.top_tissues = []


# Strength multipliers for seed types in repression estimation
_SEED_REPRESSION = {
    "8mer": -0.31,     # Mean log2FC from TargetScan 7 (Agarwal et al.)
    "7mer-m8": -0.20,
    "7mer-A1": -0.11,
    "6mer": -0.04,
}

# Region multipliers (3'UTR sites are ~3x more effective than ORF)
_REGION_WEIGHT = {
    "utr3": 1.0,
    "orf_last_third": 0.5,
    "orf": 0.15,
}


@dataclass
class SiteChange:
    """A single miRNA binding site that was gained, lost, or modified."""
    mirna_id: str
    position: int
    change_type: str            # "gained", "lost", "modified"
    ref_seed_type: Optional[str]    # None if gained
    var_seed_type: Optional[str]    # None if lost
    region: str
    estimated_repression_delta: float  # Negative = more repression in variant
    tissue_expression: Optional[dict[str, float]] = None  # tissue → RPM
    top_tissues: Optional[list[str]] = None  # tissues where this miRNA is expressed (RPM >= 10)

    @property
    def functional_impact(self) -> str:
        """Human-readable impact description."""
        if self.change_type == "lost":
            return f"lost {self.ref_seed_type} site → reduced repression by {self.mirna_id}"
        elif self.change_type == "gained":
            return f"gained {self.var_seed_type} site → new repression by {self.mirna_id}"
        else:
            return (f"modified {self.ref_seed_type}→{self.var_seed_type} site "
                    f"for {self.mirna_id}")


@dataclass
class MiRNADisruptionResult:
    """Result of miRNA binding site disruption analysis.

    Attributes:
        status: 'completed', 'skipped', or 'error'.
        score: Disruption score (0-10). Higher = more binding site disruption.
        n_sites_ref: Total seed matches in reference.
        n_sites_var: Total seed matches in variant.
        n_gained: Sites gained in variant.
        n_lost: Sites lost in variant.
        n_modified: Sites with changed seed type.
        n_unchanged: Sites preserved.
        site_changes: List of individual site changes.
        total_repression_ref: Estimated total repression in reference.
        total_repression_var: Estimated total repression in variant.
        repression_delta: Change in estimated repression (var - ref).
        top_affected_mirnas: miRNAs with the largest binding changes.
        ref_sites: All seed matches in reference sequence.
        var_sites: All seed matches in variant sequence.
    """
    status: str
    score: float = 0.0

    n_sites_ref: int = 0
    n_sites_var: int = 0
    n_gained: int = 0
    n_lost: int = 0
    n_modified: int = 0
    n_unchanged: int = 0

    site_changes: list[SiteChange] = field(default_factory=list)

    total_repression_ref: float = 0.0
    total_repression_var: float = 0.0
    repression_delta: float = 0.0

    top_affected_mirnas: list[dict] = field(default_factory=list)

    ref_sites: list[SeedMatch] = field(default_factory=list)
    var_sites: list[SeedMatch] = field(default_factory=list)

    # Metadata
    mut_id: str = ""
    gene: str = ""
    n_mirnas_scanned: int = 0

    reason: Optional[str] = None


def _estimate_repression(sites: list[SeedMatch]) -> float:
    """Estimate total miRNA-mediated repression from seed matches.

    Uses TargetScan-derived mean repression per seed type, weighted by
    region. This is a first-order approximation; for full prediction
    use the miBSIM integrative model.

    Returns:
        Estimated log2 fold-change (negative = downregulation).
    """
    total = 0.0
    for site in sites:
        base = _SEED_REPRESSION.get(site.seed_type, -0.04)
        weight = _REGION_WEIGHT.get(site.region, 0.15)
        total += base * weight
    return total


def _match_key(site: SeedMatch) -> tuple:
    """Key for matching sites between ref and var: (mirna_id, position)."""
    return (site.mirna_id, site.position)


def analyze_mirna_disruption(
    ref_mrna: str,
    var_mrna: str,
    mirna_seeds: dict[str, str],
    utr3_start: Optional[int] = None,
    orf_start: int = 0,
    orf_end: Optional[int] = None,
    position_tolerance: int = 3,
    mut_id: str = "",
    gene: str = "",
) -> MiRNADisruptionResult:
    """Analyze how a mutation disrupts miRNA binding sites.

    Compares seed matches between reference and variant mRNA to identify
    gained, lost, and modified binding sites. Estimates the change in
    miRNA-mediated repression.

    Args:
        ref_mrna: Reference mRNA sequence.
        var_mrna: Variant mRNA sequence.
        mirna_seeds: Dict of {mirna_id: seed_or_mature_sequence}.
            For genome-wide scanning, use a pre-filtered set of miRNAs
            expressed in the relevant tissue.
        utr3_start: 0-based position of 3'UTR start.
        orf_start: 0-based position of ORF start.
        orf_end: 0-based position of ORF end.
        position_tolerance: Max positional shift to consider a site "modified"
            rather than gained+lost (accounts for indel shifts).
        mut_id: Mutation identifier for metadata.
        gene: Gene name for metadata.

    Returns:
        MiRNADisruptionResult with full disruption analysis.
    """
    if not mirna_seeds:
        return MiRNADisruptionResult(
            status="skipped", reason="no_mirna_seeds_provided",
            mut_id=mut_id, gene=gene,
        )

    # Find seed matches in both sequences
    ref_sites = find_seed_matches(
        ref_mrna, mirna_seeds, utr3_start, orf_start, orf_end)
    var_sites = find_seed_matches(
        var_mrna, mirna_seeds, utr3_start, orf_start, orf_end)

    # Build position-indexed lookups
    # Key: (mirna_id, position) → SeedMatch
    ref_by_key = {}
    for s in ref_sites:
        ref_by_key[_match_key(s)] = s

    var_by_key = {}
    for s in var_sites:
        var_by_key[_match_key(s)] = s

    # Also build mirna_id → [positions] for fuzzy matching
    ref_by_mirna: dict[str, list[SeedMatch]] = {}
    for s in ref_sites:
        ref_by_mirna.setdefault(s.mirna_id, []).append(s)

    var_by_mirna: dict[str, list[SeedMatch]] = {}
    for s in var_sites:
        var_by_mirna.setdefault(s.mirna_id, []).append(s)

    # Classify changes
    changes: list[SiteChange] = []
    matched_ref: set[tuple] = set()
    matched_var: set[tuple] = set()

    # Pass 1: exact position matches → unchanged or modified
    for key, ref_s in ref_by_key.items():
        if key in var_by_key:
            var_s = var_by_key[key]
            matched_ref.add(key)
            matched_var.add(key)

            if ref_s.seed_type != var_s.seed_type:
                ref_rep = _SEED_REPRESSION.get(ref_s.seed_type, 0) * _REGION_WEIGHT.get(ref_s.region, 0.15)
                var_rep = _SEED_REPRESSION.get(var_s.seed_type, 0) * _REGION_WEIGHT.get(var_s.region, 0.15)
                changes.append(SiteChange(
                    mirna_id=ref_s.mirna_id, position=ref_s.position,
                    change_type="modified",
                    ref_seed_type=ref_s.seed_type, var_seed_type=var_s.seed_type,
                    region=ref_s.region,
                    estimated_repression_delta=var_rep - ref_rep,
                ))

    # Pass 2: fuzzy match within tolerance (handles indel shifts)
    for key, ref_s in ref_by_key.items():
        if key in matched_ref:
            continue
        mirna = ref_s.mirna_id
        # Look for nearby var site for same miRNA
        best_match = None
        best_dist = position_tolerance + 1
        for var_s in var_by_mirna.get(mirna, []):
            vk = _match_key(var_s)
            if vk in matched_var:
                continue
            dist = abs(ref_s.position - var_s.position)
            if dist <= position_tolerance and dist < best_dist:
                best_match = var_s
                best_dist = dist

        if best_match is not None:
            matched_ref.add(key)
            matched_var.add(_match_key(best_match))

            if ref_s.seed_type != best_match.seed_type:
                ref_rep = _SEED_REPRESSION.get(ref_s.seed_type, 0) * _REGION_WEIGHT.get(ref_s.region, 0.15)
                var_rep = _SEED_REPRESSION.get(best_match.seed_type, 0) * _REGION_WEIGHT.get(best_match.region, 0.15)
                changes.append(SiteChange(
                    mirna_id=mirna, position=ref_s.position,
                    change_type="modified",
                    ref_seed_type=ref_s.seed_type, var_seed_type=best_match.seed_type,
                    region=ref_s.region,
                    estimated_repression_delta=var_rep - ref_rep,
                ))
            # else: unchanged, no change entry needed

    # Pass 3: unmatched ref sites → lost
    for key, ref_s in ref_by_key.items():
        if key in matched_ref:
            continue
        ref_rep = _SEED_REPRESSION.get(ref_s.seed_type, 0) * _REGION_WEIGHT.get(ref_s.region, 0.15)
        changes.append(SiteChange(
            mirna_id=ref_s.mirna_id, position=ref_s.position,
            change_type="lost",
            ref_seed_type=ref_s.seed_type, var_seed_type=None,
            region=ref_s.region,
            estimated_repression_delta=-ref_rep,  # Positive = less repression
        ))

    # Pass 4: unmatched var sites → gained
    for key, var_s in var_by_key.items():
        if key in matched_var:
            continue
        var_rep = _SEED_REPRESSION.get(var_s.seed_type, 0) * _REGION_WEIGHT.get(var_s.region, 0.15)
        changes.append(SiteChange(
            mirna_id=var_s.mirna_id, position=var_s.position,
            change_type="gained",
            ref_seed_type=None, var_seed_type=var_s.seed_type,
            region=var_s.region,
            estimated_repression_delta=var_rep,  # Negative = more repression
        ))

    # Enrich changed sites with tissue expression
    _enrich_tissue_expression(changes)

    # Counts
    n_gained = sum(1 for c in changes if c.change_type == "gained")
    n_lost = sum(1 for c in changes if c.change_type == "lost")
    n_modified = sum(1 for c in changes if c.change_type == "modified")
    n_unchanged = len(ref_sites) - n_lost - n_modified

    # Repression estimates
    total_ref = _estimate_repression(ref_sites)
    total_var = _estimate_repression(var_sites)
    repression_delta = total_var - total_ref

    # Score: based on magnitude of repression change and number of strong site changes
    abs_delta = abs(repression_delta)
    site_disruption_count = n_gained + n_lost + n_modified
    # Nonlinear: each strong site change (~0.3 log2FC) → ~3 points
    raw_score = 10.0 * (1.0 - math.exp(-3.0 * abs_delta))
    # Bonus for multiple disruptions
    if site_disruption_count > 1:
        raw_score = min(10.0, raw_score * (1.0 + 0.1 * (site_disruption_count - 1)))
    score = min(10.0, raw_score)

    # Top affected miRNAs
    mirna_impact: dict[str, float] = {}
    for c in changes:
        mirna_impact[c.mirna_id] = mirna_impact.get(c.mirna_id, 0.0) + abs(c.estimated_repression_delta)
    top_mirnas = sorted(mirna_impact.items(), key=lambda x: x[1], reverse=True)[:5]
    top_affected = [
        {"mirna_id": m, "total_abs_impact": round(v, 4),
         "changes": [c for c in changes if c.mirna_id == m]}
        for m, v in top_mirnas
    ]

    return MiRNADisruptionResult(
        status="completed",
        score=round(score, 2),
        n_sites_ref=len(ref_sites),
        n_sites_var=len(var_sites),
        n_gained=n_gained,
        n_lost=n_lost,
        n_modified=n_modified,
        n_unchanged=n_unchanged,
        site_changes=changes,
        total_repression_ref=round(total_ref, 4),
        total_repression_var=round(total_var, 4),
        repression_delta=round(repression_delta, 4),
        top_affected_mirnas=top_affected,
        ref_sites=ref_sites,
        var_sites=var_sites,
        mut_id=mut_id,
        gene=gene,
        n_mirnas_scanned=len(mirna_seeds),
    )
