# geney/mirna/mibsim.py — miBSIM integrative model + DS correction (Python port).
"""
Pure-Python implementation of the miBSIM prediction engine.

Port of the MATLAB codebase (Bader & Tuller, 2024) for predicting
miRNA-mediated mRNA repression. Two core components:

1. **Integrative model**: 8 elastic net regressors (2 regions × 4 seed types)
   that predict per-site nominal repression S from 152 features.

2. **DS correction**: Binding site interaction model that adjusts per-site
   repression based on neighboring sites (competition/cooperation dynamics).

The 152 features span thermodynamics (ΔG duplex/open/binding, SA, SPS),
sequence context (AU content, 3' pairing, PUM sites, RBP motifs),
conservation (phastCons/phyloP), codon usage, and biochemical (Kd, occupancy).

Parameters are loaded from mibsim_params.json (extracted from MATLAB .mat files).
"""
from __future__ import annotations

import json
import math
from pathlib import Path
from typing import Optional

import numpy as np

# Lazy-loaded parameters
_PARAMS: dict | None = None
_PARAMS_PATH = Path(__file__).parent / "mibsim_params.json"


def _load_params() -> dict:
    """Lazy-load miBSIM parameters from JSON."""
    global _PARAMS
    if _PARAMS is None:
        with open(_PARAMS_PATH) as f:
            _PARAMS = json.load(f)
    return _PARAMS


# ============================================================================
# Integrative model: elastic net per-site repression prediction
# ============================================================================

def predict_site_repression(
    feature_vector: np.ndarray | dict[str, float],
    region: str,
    seed_type: str,
) -> float:
    """Predict nominal repression S for a single binding site.

    Applies the pre-trained elastic net model for the given region/seed
    combination to the 152-feature vector.

    Args:
        feature_vector: Either a 152-element array (ordered to match
            feature_names) or a dict mapping feature names to values.
            Missing features default to 0.
        region: "orf" or "utr3".
        seed_type: "6mer", "7mer-A1", "7mer-m8", or "8mer".

    Returns:
        Nominal repression S (log2 fold-change, negative = downregulation).
        Capped at 0 (no predicted upregulation from miRNA).
    """
    params = _load_params()
    key = f"{region}_{seed_type}"
    model = params["integrative_models"].get(key)
    if model is None:
        raise ValueError(f"No model for region={region}, seed_type={seed_type}. "
                         f"Valid keys: {list(params['integrative_models'].keys())}")

    coeffs = np.array(model["coeffs"])
    intercept = model["intercept"]
    feature_names = model["feature_names"]

    # Build feature array
    if isinstance(feature_vector, dict):
        x = np.array([feature_vector.get(name, 0.0) for name in feature_names])
    else:
        x = np.asarray(feature_vector, dtype=np.float64)
        if x.shape[0] != len(coeffs):
            raise ValueError(f"Feature vector length {x.shape[0]} != {len(coeffs)}")

    # Elastic net prediction: dot(x, coeffs) + intercept, capped at 0
    prediction = float(np.dot(x, coeffs) + intercept)
    return min(prediction, 0.0)


def get_feature_names(region: str = "utr3", seed_type: str = "8mer") -> list[str]:
    """Return the ordered list of 152 feature names for a model."""
    params = _load_params()
    key = f"{region}_{seed_type}"
    return params["integrative_models"][key]["feature_names"]


# ============================================================================
# DS correction: binding site interaction model
# ============================================================================

def compute_ds_correction(
    site_repression: float,
    neighbor_repressions: list[float],
    neighbor_distances: list[int],
    site_number: int,
) -> float:
    """Compute the DS (Duplex Score) correction factor for a binding site.

    Models how neighboring miRNA binding sites interact — either competing
    for RISC/AGO resources (weakening each site's effect) or cooperating
    (enhancing repression through local stabilization).

    The DS factor scales the nominal repression:
        corrected_repression = DS * S

    Formula:
        DS = 1 - sum(|S_i| * exp(D_i / c2)) / (c1 * site_number)

    Where c1 and c2 are optimized parameters, S_i is the nominal repression
    of neighbor i, and D_i is the distance in nucleotides.

    Args:
        site_repression: Nominal repression S of this site.
        neighbor_repressions: List of nominal repressions of neighboring sites.
        neighbor_distances: List of distances (nt) to each neighbor.
        site_number: Total number of binding sites on this transcript.

    Returns:
        DS correction factor (typically 0.5-1.0).
        DS=1 means no interaction effect, DS<1 means reduced repression.
    """
    if site_number <= 1 or not neighbor_repressions:
        return 1.0

    params = _load_params()
    c1 = params["ds_correction"]["c1"]
    c2 = params["ds_correction"]["c2"]

    # Sum of neighbor contributions weighted by distance decay
    interaction_sum = 0.0
    for s_i, d_i in zip(neighbor_repressions, neighbor_distances):
        interaction_sum += abs(s_i) * math.exp(d_i / c2)

    ds = 1.0 - interaction_sum / (c1 * site_number)

    # Clamp to reasonable range
    return max(0.0, min(1.0, ds))


def predict_transcript_repression(
    sites: list[dict],
) -> dict:
    """Predict total miRNA-mediated repression for a transcript.

    Takes a list of binding sites (each with features, position, region,
    seed_type) and computes:
    1. Per-site nominal repression S (via integrative model)
    2. DS correction for site interactions
    3. Total corrected repression

    Args:
        sites: List of dicts, each with:
            - "features": dict[str, float] or 152-element array
            - "position": int (nt position on mRNA)
            - "region": "orf" or "utr3"
            - "seed_type": "6mer", "7mer-A1", "7mer-m8", "8mer"

    Returns:
        Dict with per_site_results (list) and total_repression.
    """
    if not sites:
        return {"per_site_results": [], "total_repression": 0.0}

    # Step 1: predict nominal repression for each site
    per_site = []
    for site in sites:
        S = predict_site_repression(
            site["features"], site["region"], site["seed_type"])
        per_site.append({
            "position": site["position"],
            "region": site["region"],
            "seed_type": site["seed_type"],
            "nominal_repression": S,
        })

    # Step 2: compute DS correction for each site
    n_sites = len(per_site)
    for i, site_i in enumerate(per_site):
        # Only apply DS to UTR3 sites with seed_type != 6mer
        if site_i["region"] != "utr3" or site_i["seed_type"] == "6mer":
            site_i["ds"] = 1.0
            site_i["corrected_repression"] = site_i["nominal_repression"]
            continue

        # Gather neighbors
        neighbor_S = []
        neighbor_D = []
        for j, site_j in enumerate(per_site):
            if i == j:
                continue
            neighbor_S.append(site_j["nominal_repression"])
            neighbor_D.append(abs(site_i["position"] - site_j["position"]))

        ds = compute_ds_correction(
            site_i["nominal_repression"],
            neighbor_S, neighbor_D,
            n_sites,
        )
        site_i["ds"] = round(ds, 4)
        site_i["corrected_repression"] = round(ds * site_i["nominal_repression"], 4)

    # Step 3: total repression
    total = sum(s["corrected_repression"] for s in per_site)

    return {
        "per_site_results": per_site,
        "total_repression": round(total, 4),
        "n_sites": n_sites,
    }


# ============================================================================
# Biochemical model: Kd → occupancy (logistic dose-response)
# ============================================================================

def predict_occupancy(
    log_kd: float,
    in_orf: bool = False,
    log_sa_diff: float = 0.0,
    threep_canon: float = 0.0,
    pct: float = 0.0,
    free_ago: float = -7.0,
) -> dict:
    """Predict miRNA-mRNA binding occupancy from Kd and site features.

    Implements the Biochemical+ model (logistic dose-response with
    correction terms for 3' pairing, structural accessibility, etc.).

    Args:
        log_kd: log10(Kd) for this site (from CNN or lookup).
        in_orf: Whether site is in ORF (vs 3'UTR).
        log_sa_diff: log(structural accessibility) difference from background.
        threep_canon: 3' canonical pairing score.
        pct: Probability of conserved targeting.
        free_ago: Free AGO concentration parameter (default -7).

    Returns:
        Dict with occ, occ_init, log_occ, log_occ_init, pred.
    """
    params = _load_params()
    bio = params["biochemical_model"]

    log_decay = bio["log_decay"]
    coefs = bio["feature_coefs"]  # [log_KA, in_ORF, logSA_diff, Threep_canon]
    decay = math.exp(log_decay)

    log_ka = -log_kd  # Affinity = 1/Kd

    # Feature contributions
    offset = coefs[3] * threep_canon  # + PCT coef if available
    offset_all = coefs[1] * float(in_orf) + coefs[2] * log_sa_diff

    # Occupancy with site
    z = free_ago + coefs[0] * log_ka + offset + offset_all
    occ = 1.0 / (1.0 + math.exp(-z))

    # Occupancy without site (baseline)
    z_init = free_ago + offset_all
    occ_init = 1.0 / (1.0 + math.exp(-z_init))

    # Predicted repression contribution
    pred = math.log(1 + decay * occ_init) - math.log(1 + decay * occ)

    return {
        "occ": round(occ, 6),
        "occ_init": round(occ_init, 6),
        "log_occ": round(math.log(occ + 1e-10), 4),
        "log_occ_init": round(math.log(occ_init + 1e-10), 4),
        "pred": round(pred, 6),
    }
