# geney/epistasis.py — Post-hoc epistasis detection from concurrent results.
"""
Compares compound (double) variant results against their constituent
single-variant results to classify interaction effects per mechanism.

Usage::

    from geney import analyze_concurrent
    from geney.epistasis import analyze_epistasis

    results = analyze_concurrent([
        "GENE:1:100:A:G",           # variant A
        "GENE:1:200:C:T",           # variant B
        "GENE:1:100:A:G|GENE:1:200:C:T",  # compound A|B
    ])

    epistasis = analyze_epistasis(
        compound_id="GENE:1:100:A:G|GENE:1:200:C:T",
        compound_result=results[2],
        individual_ids=["GENE:1:100:A:G", "GENE:1:200:C:T"],
        individual_results=[results[0], results[1]],
    )
"""
from __future__ import annotations

from dataclasses import dataclass, field
from typing import Optional


# -- Classification thresholds ------------------------------------------------
# These are on the 0-10 normalized score scale.
# Scores are nonlinear transforms of raw measurements, so we use simple
# heuristics rather than arithmetic expectations.

_SYNERGY_DELTA = 1.5       # compound must exceed max individual by this much
_ANTAGONISM_DELTA = -1.5   # compound must fall below max individual by this much
_NOISE_FLOOR = 0.5         # ignore mechanisms where all scores are below this


@dataclass
class MechanismInteraction:
    """Epistasis classification for a single mechanism."""
    mechanism: str
    compound_score: float
    individual_scores: list[float]           # one per constituent variant
    individual_ids: list[str]

    # Derived
    max_individual: float = 0.0
    sum_individual: float = 0.0
    delta_vs_max: float = 0.0                # compound - max(individuals)
    ratio_vs_max: Optional[float] = None     # compound / max(individuals)

    # Classification
    interaction_type: str = "neutral"        # synergistic | antagonistic | additive | neutral
    magnitude: float = 0.0                   # abs(delta_vs_max), for bar height
    direction: int = 0                       # +1 synergistic, -1 antagonistic, 0 neutral

    # For visualization: normalized bar positions (-1 to +1 scale)
    bar_position: float = 0.0

    def __post_init__(self):
        self.max_individual = max(self.individual_scores) if self.individual_scores else 0.0
        self.sum_individual = sum(self.individual_scores)
        self.delta_vs_max = self.compound_score - self.max_individual

        if self.max_individual > 0:
            self.ratio_vs_max = self.compound_score / self.max_individual
        else:
            self.ratio_vs_max = None

        self._classify()

    def _classify(self):
        # If all scores are below noise floor, nothing to compare
        if self.compound_score < _NOISE_FLOOR and self.max_individual < _NOISE_FLOOR:
            self.interaction_type = "neutral"
            self.direction = 0
            self.magnitude = 0.0
            self.bar_position = 0.0
            return

        if self.delta_vs_max >= _SYNERGY_DELTA:
            self.interaction_type = "synergistic"
            self.direction = 1
        elif self.delta_vs_max <= _ANTAGONISM_DELTA:
            self.interaction_type = "antagonistic"
            self.direction = -1
        elif self.compound_score > self.sum_individual + 0.5:
            # Exceeds sum — supralinear even if doesn't hit synergy threshold
            self.interaction_type = "synergistic"
            self.direction = 1
        elif self.compound_score < self.max_individual - 0.5:
            # Below max individual — corrective
            self.interaction_type = "antagonistic"
            self.direction = -1
        else:
            self.interaction_type = "additive"
            self.direction = 0

        self.magnitude = abs(self.delta_vs_max)

        # Bar position: normalized to -1..+1 for visualization
        # Clamp to [-1, 1] using a soft scale
        if self.max_individual > 0:
            self.bar_position = max(-1.0, min(1.0, self.delta_vs_max / max(self.max_individual, 2.0)))
        else:
            self.bar_position = 0.0

    def to_dict(self) -> dict:
        return {
            "mechanism": self.mechanism,
            "compound_score": round(self.compound_score, 1),
            "individual_scores": [round(s, 1) for s in self.individual_scores],
            "individual_ids": self.individual_ids,
            "max_individual": round(self.max_individual, 1),
            "sum_individual": round(self.sum_individual, 1),
            "delta_vs_max": round(self.delta_vs_max, 2),
            "ratio_vs_max": round(self.ratio_vs_max, 2) if self.ratio_vs_max is not None else None,
            "interaction_type": self.interaction_type,
            "magnitude": round(self.magnitude, 2),
            "direction": self.direction,
            "bar_position": round(self.bar_position, 3),
        }


@dataclass
class EpistasisResult:
    """Epistasis analysis for a compound variant vs its constituents."""
    compound_id: str
    individual_ids: list[str]
    mechanisms: dict[str, MechanismInteraction] = field(default_factory=dict)

    # Summary
    has_epistasis: bool = False
    dominant_interaction: str = "neutral"     # overall classification
    most_epistatic_mechanism: str = ""
    max_magnitude: float = 0.0

    def to_dict(self) -> dict:
        return {
            "compound_id": self.compound_id,
            "individual_ids": self.individual_ids,
            "has_epistasis": self.has_epistasis,
            "dominant_interaction": self.dominant_interaction,
            "most_epistatic_mechanism": self.most_epistatic_mechanism,
            "max_magnitude": round(self.max_magnitude, 2),
            "mechanisms": {k: v.to_dict() for k, v in self.mechanisms.items()},
        }

    def summary_line(self) -> str:
        """One-line human-readable summary."""
        if not self.has_epistasis:
            return "No significant epistasis detected"
        m = self.most_epistatic_mechanism
        mi = self.mechanisms[m]
        if mi.direction > 0:
            return (f"Synergistic in {m}: compound={mi.compound_score:.1f} "
                    f"vs max individual={mi.max_individual:.1f} (+{mi.delta_vs_max:.1f})")
        else:
            return (f"Antagonistic in {m}: compound={mi.compound_score:.1f} "
                    f"vs max individual={mi.max_individual:.1f} ({mi.delta_vs_max:.1f})")


# -- Annotation helpers -------------------------------------------------------

@dataclass
class EpistasisAnnotation:
    """Lightweight annotation attached to an individual mutation's result."""
    compound_id: str
    partner_ids: list[str]
    interaction_type: str        # overall dominant interaction
    most_epistatic_mechanism: str
    compound_scores: dict[str, float]    # mechanism → compound score
    bar_positions: dict[str, float]      # mechanism → bar position


def _annotate_individual(epistasis: EpistasisResult, individual_id: str) -> EpistasisAnnotation:
    """Build an annotation for one constituent variant."""
    partners = [i for i in epistasis.individual_ids if i != individual_id]
    return EpistasisAnnotation(
        compound_id=epistasis.compound_id,
        partner_ids=partners,
        interaction_type=epistasis.dominant_interaction,
        most_epistatic_mechanism=epistasis.most_epistatic_mechanism,
        compound_scores={k: round(v.compound_score, 1) for k, v in epistasis.mechanisms.items()},
        bar_positions={k: round(v.bar_position, 3) for k, v in epistasis.mechanisms.items()},
    )


# -- Main entry point ---------------------------------------------------------

_MECHANISM_NAMES = ("splicing", "protein", "tis", "folding", "transcription", "mirna")


def analyze_epistasis(
    compound_id: str,
    compound_result,
    individual_ids: list[str],
    individual_results: list,
) -> EpistasisResult:
    """Compare a compound variant's scores against its individual constituents.

    Args:
        compound_id: Mutation ID of the compound variant (pipe-separated).
        compound_result: MutationResult from the compound variant.
        individual_ids: Mutation IDs of each constituent variant.
        individual_results: MutationResults for each constituent, same order.

    Returns:
        EpistasisResult with per-mechanism interaction classifications.
    """
    compound_scores = compound_result.scores
    individual_score_sets = [r.scores for r in individual_results]

    mechanisms = {}
    for mech in _MECHANISM_NAMES:
        c_score = compound_scores[mech].score if mech in compound_scores else 0.0
        i_scores = [
            (s[mech].score if mech in s else 0.0)
            for s in individual_score_sets
        ]

        mi = MechanismInteraction(
            mechanism=mech,
            compound_score=c_score,
            individual_scores=i_scores,
            individual_ids=individual_ids,
        )
        mechanisms[mech] = mi

    # Summary: find the most epistatic mechanism
    epistatic = [m for m in mechanisms.values() if m.interaction_type != "neutral" and m.interaction_type != "additive"]
    has_epistasis = len(epistatic) > 0

    dominant = "neutral"
    most_epistatic = ""
    max_mag = 0.0
    if epistatic:
        top = max(epistatic, key=lambda m: m.magnitude)
        most_epistatic = top.mechanism
        max_mag = top.magnitude
        dominant = top.interaction_type

    result = EpistasisResult(
        compound_id=compound_id,
        individual_ids=individual_ids,
        mechanisms=mechanisms,
        has_epistasis=has_epistasis,
        dominant_interaction=dominant,
        most_epistatic_mechanism=most_epistatic,
        max_magnitude=max_mag,
    )
    return result


def analyze_epistasis_batch(
    compound_ids: list[str],
    compound_results: list,
    individual_ids_per_compound: list[list[str]],
    all_results: dict[str, object],
) -> list[EpistasisResult]:
    """Batch epistasis analysis for multiple compound variants.

    Args:
        compound_ids: List of compound mutation IDs.
        compound_results: MutationResults for each compound.
        individual_ids_per_compound: For each compound, list of constituent mut IDs.
        all_results: Dict mapping mut_id → MutationResult for all individuals.

    Returns:
        List of EpistasisResult, one per compound.
    """
    results = []
    for comp_id, comp_result, ind_ids in zip(compound_ids, compound_results, individual_ids_per_compound):
        ind_results = [all_results[i] for i in ind_ids if i in all_results]
        if len(ind_results) != len(ind_ids):
            continue
        results.append(analyze_epistasis(comp_id, comp_result, ind_ids, ind_results))
    return results


def annotate_results(
    epistasis_results: list[EpistasisResult],
    all_results: dict[str, object],
) -> dict[str, list[EpistasisAnnotation]]:
    """Attach epistasis annotations to individual mutation results.

    Returns a dict mapping individual mut_id → list of EpistasisAnnotation
    (one per compound variant that individual participates in).
    """
    annotations: dict[str, list[EpistasisAnnotation]] = {}
    for ep in epistasis_results:
        if not ep.has_epistasis:
            continue
        for ind_id in ep.individual_ids:
            if ind_id not in annotations:
                annotations[ind_id] = []
            annotations[ind_id].append(_annotate_individual(ep, ind_id))
    return annotations


# -- Top-level convenience function -------------------------------------------

def _split_compound(compound_id: str) -> list[str]:
    """Split a pipe-separated compound ID into individual mutation IDs."""
    import re
    return [p.strip() for p in re.split(r"[|,]", compound_id) if p.strip()]


def analyze_with_epistasis(
    mut_ids: list[str],
    **concurrent_kwargs,
) -> list:
    """Run concurrent analysis and epistasis in one call.

    Accepts a mixed list of single and compound (pipe-separated) mutation IDs.
    Automatically detects compounds, ensures all constituents are included,
    runs analyze_concurrent on the full set, then populates the ``epistasis``
    field on each MutationResult.

    Returns the same ``list[MutationResult]`` as ``analyze_concurrent``, but
    with ``result.epistasis`` populated:

    - Single variant, no compound involvement → ``epistasis = None``
    - Single variant implicated in compound(s) → ``epistasis`` = dict with
      ``status="constituent"`` and links to compound variant IDs
    - Compound variant → ``epistasis`` = EpistasisResult with per-mechanism
      interaction bars

    Any auto-added constituent IDs (not in the original ``mut_ids``) are
    appended to the end of the returned list.

    Args:
        mut_ids: Mutation IDs — singles and compounds mixed freely.
        **concurrent_kwargs: Forwarded to analyze_concurrent
            (run_splicing, run_mirna, n_workers, etc.)

    Returns:
        list[MutationResult] — one per input mutation (plus auto-added
        constituents), with ``epistasis`` field populated where relevant.
    """
    from .concurrent import analyze_concurrent

    # 1. Separate compounds from singles, ensure all constituents are queued
    all_ids = list(dict.fromkeys(mut_ids))  # preserve order, dedupe
    compound_ids = []
    compound_parts: dict[str, list[str]] = {}

    for mid in list(all_ids):
        parts = _split_compound(mid)
        if len(parts) > 1:
            compound_ids.append(mid)
            compound_parts[mid] = parts
            for part in parts:
                if part not in all_ids:
                    all_ids.append(part)

    # 2. Run everything through analyze_concurrent
    results_list = analyze_concurrent(all_ids, **concurrent_kwargs)
    results_map = dict(zip(all_ids, results_list))

    # 3. Run epistasis on compounds and populate results
    epistasis_results = []
    for comp_id in compound_ids:
        parts = compound_parts[comp_id]
        ind_results = [results_map[p] for p in parts if p in results_map]
        if len(ind_results) == len(parts) and comp_id in results_map:
            ep = analyze_epistasis(comp_id, results_map[comp_id], parts, ind_results)
            epistasis_results.append(ep)
            # Attach full EpistasisResult to the compound variant
            results_map[comp_id].epistasis = ep

    # 4. Annotate individual variants with back-references
    for ep in epistasis_results:
        for ind_id in ep.individual_ids:
            if ind_id in results_map:
                existing = results_map[ind_id].epistasis
                if existing is None:
                    existing = []
                if not isinstance(existing, list):
                    existing = [existing]
                existing.append(_annotate_individual(ep, ind_id))
                results_map[ind_id].epistasis = existing

    return results_list
