# geney/mirna/expression.py — miRNA tissue expression data and filtering.
"""
Tissue-specific miRNA expression filtering.

Provides two approaches:
  1. **Bundled expression matrix**: Curated from published miRNA tissue atlases
     (Ludwig et al. 2016, miRNATissueAtlas2). Covers ~500 miRNAs across 25
     major tissue types with median RPM values.
  2. **Custom expression file**: Load from TSV (miRNA name, tissue, RPM/TPM).

Usage::

    from geney.mirna.expression import filter_by_expression, get_expressed_mirnas

    # Get miRNAs expressed in liver above 10 RPM
    expressed = get_expressed_mirnas("liver", min_rpm=10)

    # Filter a full miRBase panel to tissue-relevant miRNAs
    panel = load_mirna_fasta("mature.fa")
    filtered = filter_by_expression(panel, tissue="liver", min_rpm=10)
"""
from __future__ import annotations

import csv
import json
from pathlib import Path
from typing import Optional

# ============================================================================
# Bundled expression matrix
# ============================================================================

# Curated from miRNATissueAtlas2 (Kern et al. 2022) and Ludwig et al. 2016.
# Values are approximate median RPM (reads per million) across published datasets.
# Only miRNAs with RPM > 100 in at least one tissue are included.
# Format: {mirna_name: {tissue: rpm}}
#
# Tissues: brain, heart, liver, lung, kidney, muscle, skin, blood,
#          breast, colon, stomach, pancreas, prostate, ovary, testis,
#          thyroid, adrenal, spleen, bone_marrow, placenta, adipose,
#          esophagus, bladder, cervix, eye
#
# This is a subset — for full resolution use miRNATissueAtlas2025:
#   https://www.ccb.uni-saarland.de/tissueatlas2025

# Path resolution: try DASHR first, then fallback literature matrix
_DASHR_PATH = Path(__file__).parent / "data" / "dashr2_expression.json"
_LITERATURE_PATH = Path(__file__).parent / "expression_matrix.json"
_EXPRESSION_CACHE: dict | None = None


def _load_expression_matrix() -> dict[str, dict[str, float]]:
    """Load bundled expression matrix.

    Resolution order:
      1. DASHR v2.0 real small RNA-seq data (2,040 miRNAs × 15 tissues)
      2. Literature-curated approximation
      3. Hardcoded tissue atlas from _build_expression_from_literature()
    """
    global _EXPRESSION_CACHE
    if _EXPRESSION_CACHE is not None:
        return _EXPRESSION_CACHE

    for path in (_DASHR_PATH, _LITERATURE_PATH):
        if path.exists():
            with open(path) as f:
                _EXPRESSION_CACHE = json.load(f)
            return _EXPRESSION_CACHE

    # Final fallback: build from literature knowledge
    _EXPRESSION_CACHE = _build_expression_from_literature()
    return _EXPRESSION_CACHE


def _build_expression_from_literature() -> dict[str, dict[str, float]]:
    """Build expression matrix from well-established tissue-specific miRNAs.

    This is a curated set based on published reviews and atlases. Values
    are approximate RPM categories:
      - 10000+ : dominant in tissue
      -  1000  : highly expressed
      -   100  : moderately expressed
      -    10  : low but detectable

    Sources:
      - Ludwig et al. 2016 (miRNA Tissue Atlas, 61 tissues)
      - Landgraf et al. 2007 (mammalian miRNA expression atlas)
      - GTEx small RNA-seq summaries
      - McCall et al. 2017 (miRmine database)
    """
    # We build this as {mirna: {tissue: rpm}}
    # Using established tissue-enrichment patterns from literature
    data: dict[str, dict[str, float]] = {}

    # Helper to add expression entries
    def _add(mirna: str, tissue_rpms: dict[str, float]):
        data.setdefault(mirna, {}).update(tissue_rpms)

    # --- Ubiquitous (expressed in most tissues at moderate-high levels) ---
    _ubiq = {"brain": 500, "heart": 500, "liver": 500, "lung": 500,
             "kidney": 500, "muscle": 500, "skin": 200, "blood": 300,
             "breast": 300, "colon": 300, "stomach": 300, "pancreas": 200,
             "prostate": 200, "ovary": 200, "testis": 200, "thyroid": 300,
             "adrenal": 200, "spleen": 300, "adipose": 200, "esophagus": 200}

    for mir in ["hsa-miR-21-5p", "hsa-let-7a-5p", "hsa-let-7f-5p",
                "hsa-miR-26a-5p", "hsa-miR-30a-5p", "hsa-miR-16-5p",
                "hsa-miR-24-3p", "hsa-miR-191-5p"]:
        _add(mir, _ubiq)

    # --- Brain-enriched ---
    _brain_high = {"brain": 10000, "eye": 1000}
    _brain_mod = {"brain": 2000, "eye": 500}
    for mir in ["hsa-miR-9-5p", "hsa-miR-124-3p", "hsa-miR-128-3p"]:
        _add(mir, _brain_high)
    for mir in ["hsa-miR-132-3p", "hsa-miR-134-5p", "hsa-miR-137",
                "hsa-miR-138-5p", "hsa-miR-219a-5p"]:
        _add(mir, _brain_mod)

    # --- Liver-specific ---
    _add("hsa-miR-122-5p", {"liver": 50000, "brain": 1, "heart": 1})

    # --- Heart/muscle ---
    _card = {"heart": 10000, "muscle": 5000}
    _add("hsa-miR-1-3p", _card)
    _add("hsa-miR-133a-3p", _card)
    _add("hsa-miR-133b", _card)
    _add("hsa-miR-208a-3p", {"heart": 5000})
    _add("hsa-miR-499a-5p", {"heart": 3000})
    _add("hsa-miR-206", {"muscle": 8000})

    # --- Immune/hematopoietic ---
    _imm = {"blood": 5000, "spleen": 3000, "bone_marrow": 2000}
    _add("hsa-miR-155-5p", _imm)
    _add("hsa-miR-146a-5p", _imm)
    _add("hsa-miR-146b-5p", _imm)
    _add("hsa-miR-150-5p", _imm)
    _add("hsa-miR-142-3p", {**_imm, "thyroid": 100})
    _add("hsa-miR-223-3p", {"blood": 8000, "spleen": 3000, "bone_marrow": 5000})
    _add("hsa-miR-451a", {"blood": 20000})

    # --- Epithelial / EMT regulators ---
    _epi = {"skin": 2000, "breast": 1000, "colon": 1000, "lung": 800,
            "kidney": 1000, "bladder": 800, "cervix": 500}
    for mir in ["hsa-miR-200a-3p", "hsa-miR-200b-3p", "hsa-miR-200c-3p",
                "hsa-miR-141-3p"]:
        _add(mir, _epi)
    _add("hsa-miR-205-5p", {"skin": 5000, "breast": 1000, "cervix": 1000})
    _add("hsa-miR-203a-3p", {"skin": 8000, "cervix": 1000, "esophagus": 2000})

    # --- GI tract ---
    _gi = {"colon": 2000, "stomach": 1500, "esophagus": 1000}
    _add("hsa-miR-192-5p", {**_gi, "liver": 1000, "kidney": 2000})
    _add("hsa-miR-194-5p", {**_gi, "liver": 800, "kidney": 1500})
    _add("hsa-miR-215-5p", _gi)

    # --- Pancreas ---
    _add("hsa-miR-375-3p", {"pancreas": 10000, "stomach": 500, "adrenal": 300})
    _add("hsa-miR-217", {"pancreas": 3000})

    # --- Thyroid ---
    _add("hsa-miR-7-5p", {"brain": 2000, "pancreas": 1000, "adrenal": 500})

    # --- Placenta ---
    _add("hsa-miR-517a-3p", {"placenta": 5000})
    _add("hsa-miR-517b-3p", {"placenta": 4000})
    _add("hsa-miR-518b", {"placenta": 3000})
    _add("hsa-miR-371a-3p", {"placenta": 2000, "testis": 1000})

    # --- Testis ---
    _add("hsa-miR-449a", {"testis": 5000, "lung": 500})
    _add("hsa-miR-34c-5p", {"testis": 3000})

    # --- Kidney ---
    _add("hsa-miR-196a-5p", {"kidney": 2000, "colon": 500})

    # --- Adipose ---
    _add("hsa-miR-378a-3p", {"adipose": 3000, "muscle": 1000, "heart": 500})

    # --- let-7 family (broadly expressed, varying by tissue) ---
    for mir in ["hsa-let-7b-5p", "hsa-let-7c-5p", "hsa-let-7d-5p",
                "hsa-let-7e-5p", "hsa-let-7g-5p", "hsa-let-7i-5p"]:
        _add(mir, {k: max(100, v // 2) for k, v in _ubiq.items()})

    # --- miR-17~92 cluster (proliferative tissues) ---
    _prolif = {"blood": 1000, "bone_marrow": 800, "testis": 500,
               "breast": 300, "colon": 300, "lung": 300}
    for mir in ["hsa-miR-17-5p", "hsa-miR-18a-5p", "hsa-miR-19a-3p",
                "hsa-miR-19b-3p", "hsa-miR-20a-5p", "hsa-miR-92a-3p"]:
        _add(mir, _prolif)

    # --- p53-regulated (miR-34a, broadly induced by stress) ---
    _add("hsa-miR-34a-5p", {"brain": 1000, "liver": 500, "lung": 500,
                             "colon": 500, "kidney": 500})
    _add("hsa-miR-34b-5p", {"lung": 1000, "testis": 500})

    # --- miR-29 family (fibrosis-related) ---
    for mir in ["hsa-miR-29a-3p", "hsa-miR-29b-3p", "hsa-miR-29c-3p"]:
        _add(mir, {"liver": 1000, "lung": 800, "heart": 800, "kidney": 600,
                   "muscle": 500, "brain": 400})

    # --- miR-15/16 (cell cycle) ---
    _add("hsa-miR-15a-5p", {"blood": 1000, "spleen": 800, "liver": 500})
    _add("hsa-miR-15b-5p", {"blood": 800, "brain": 500, "liver": 400})

    # --- Angiogenesis ---
    _add("hsa-miR-126-3p", {"lung": 3000, "heart": 2000, "kidney": 1500,
                             "liver": 800, "skin": 500})

    # --- miR-181 family (immune development, broadly expressed) ---
    for mir in ["hsa-miR-181a-5p", "hsa-miR-181b-5p",
                "hsa-miR-181c-5p", "hsa-miR-181d-5p"]:
        _add(mir, {"blood": 1000, "spleen": 800, "brain": 1500,
                   "thyroid": 500, "liver": 300})

    # --- miR-221/222 (angiogenesis, cell cycle) ---
    for mir in ["hsa-miR-221-3p", "hsa-miR-222-3p"]:
        _add(mir, {"blood": 500, "skin": 300, "liver": 200,
                   "breast": 200, "thyroid": 500})

    # --- miR-148/152 (DNA methylation) ---
    for mir in ["hsa-miR-148a-3p", "hsa-miR-148b-3p", "hsa-miR-152-3p"]:
        _add(mir, {"stomach": 1000, "colon": 800, "liver": 500,
                   "kidney": 500, "pancreas": 300})

    # --- miR-486 (erythroid) ---
    _add("hsa-miR-486-5p", {"blood": 10000, "muscle": 2000, "heart": 1000})

    # --- miR-143/145 (smooth muscle) ---
    _add("hsa-miR-143-3p", {"colon": 3000, "stomach": 2000, "heart": 1000,
                             "bladder": 1500, "prostate": 800})
    _add("hsa-miR-145-5p", {"colon": 2000, "stomach": 1500, "heart": 800,
                             "bladder": 1000, "prostate": 600})

    # --- miR-210 (hypoxia-responsive, broadly expressed under stress) ---
    _add("hsa-miR-210-3p", {"placenta": 3000, "kidney": 500, "heart": 300,
                             "lung": 300, "breast": 200})

    return data


# ============================================================================
# Public API
# ============================================================================

def get_expressed_mirnas(
    tissue: str,
    min_rpm: float = 10.0,
    expression_data: dict[str, dict[str, float]] | None = None,
) -> set[str]:
    """Get the set of miRNA names expressed above threshold in a tissue.

    Args:
        tissue: Tissue name (e.g. "liver", "brain", "blood").
        min_rpm: Minimum RPM to consider expressed (default 10).
        expression_data: Custom expression matrix. If None, uses bundled data.

    Returns:
        Set of miRNA names above threshold.
    """
    data = expression_data or _load_expression_matrix()
    if not data:
        data = _build_expression_from_literature()

    expressed = set()
    tissue_lower = tissue.lower()

    for mirna, tissue_rpms in data.items():
        for t, rpm in tissue_rpms.items():
            if t.lower() == tissue_lower and rpm >= min_rpm:
                expressed.add(mirna)

    return expressed


def filter_by_expression(
    mirna_panel: dict[str, str],
    tissue: str,
    min_rpm: float = 10.0,
    expression_data: dict[str, dict[str, float]] | None = None,
) -> dict[str, str]:
    """Filter a miRNA sequence panel to tissue-expressed miRNAs.

    Args:
        mirna_panel: Full panel {name: sequence} (e.g. from load_mirna_fasta).
        tissue: Target tissue name.
        min_rpm: Minimum RPM threshold.
        expression_data: Custom expression matrix (optional).

    Returns:
        Filtered panel containing only expressed miRNAs.
    """
    expressed = get_expressed_mirnas(tissue, min_rpm, expression_data)
    return {k: v for k, v in mirna_panel.items() if k in expressed}


def load_expression_tsv(
    path: str | Path,
    mirna_col: str = "mirna",
    tissue_col: str = "tissue",
    value_col: str = "rpm",
) -> dict[str, dict[str, float]]:
    """Load miRNA expression data from a TSV file.

    Expected format: TSV with columns for miRNA name, tissue, and expression value.
    This allows loading custom expression data from miRmine, DASHR,
    miRNATissueAtlas2025, or any other source.

    Args:
        path: Path to TSV file.
        mirna_col: Column name for miRNA identifiers.
        tissue_col: Column name for tissue/cell type.
        value_col: Column name for expression value (RPM/TPM).

    Returns:
        Dict of {mirna_name: {tissue: value}}.
    """
    data: dict[str, dict[str, float]] = {}
    with open(path) as f:
        reader = csv.DictReader(f, delimiter="\t")
        for row in reader:
            mirna = row.get(mirna_col, "").strip()
            tissue = row.get(tissue_col, "").strip()
            try:
                value = float(row.get(value_col, 0))
            except ValueError:
                continue
            if mirna and tissue:
                data.setdefault(mirna, {})[tissue] = value
    return data


def list_available_tissues(
    expression_data: dict[str, dict[str, float]] | None = None,
) -> list[str]:
    """Return sorted list of tissue names in the expression matrix."""
    data = expression_data or _load_expression_matrix()
    if not data:
        data = _build_expression_from_literature()

    tissues = set()
    for tissue_rpms in data.values():
        tissues.update(tissue_rpms.keys())
    return sorted(tissues)
