# geney/transcription/prediction.py — Core Borzoi inference and tissue analysis.
"""
Borzoi regulatory impact prediction: extract sequences, run model, compare
WT vs mutant on GTEx tracks, produce per-tissue fold-changes.
"""
from __future__ import annotations

import math
import logging
from pathlib import Path
from typing import Any

import numpy as np
import pandas as pd
import torch

from .encoding import BORZOI_INPUT_LEN, dna_to_onehot
from .model import _load_borzoi

logger = logging.getLogger(__name__)

HALF_LEN = BORZOI_INPUT_LEN // 2  # 262,144

# Absolute log2 fold-change threshold (~11% expression change)
LFC_THRESHOLD = 0.152

# Bundled GTEx track metadata
TARGETS_FILE = str(Path(__file__).parent / "targets_human.txt")

# Cached mapping
_gtex_mapping: dict[str, list[int]] | None = None


# ============================================================================
# GTEx track mapping
# ============================================================================

def _build_gtex_mapping(targets_path: str | None = None) -> dict[str, list[int]]:
    """Parse ``targets_human.txt`` and build a tissue -> track-index mapping.

    Returns:
        ``{"adipose_tissue": [7522, 7523, 7524], "brain": [...], ...}``
    """
    global _gtex_mapping
    if _gtex_mapping is not None:
        return _gtex_mapping

    path = targets_path or TARGETS_FILE
    df = pd.read_csv(path, delimiter="\t")
    df = df.rename(columns={"Unnamed: 0": "feat_index"})

    mapping: dict[str, list[int]] = {}
    for _, row in df.iterrows():
        identifier = str(row.get("identifier", ""))
        if not identifier.startswith("GTEX-"):
            continue
        desc = str(row.get("description", ""))
        if ":" not in desc:
            continue
        tissue = desc.split(":", 1)[1].strip()
        idx = int(row["feat_index"])
        mapping.setdefault(tissue, []).append(idx)

    logger.info("Built GTEx mapping: %d tissues, %d tracks",
                len(mapping), sum(len(v) for v in mapping.values()))
    _gtex_mapping = mapping
    return _gtex_mapping


# ============================================================================
# Sequence extraction
# ============================================================================

def _extract_sequences(chrom: str, pos: int,
                       mutations: list[tuple[int, str, str]],
                       organism: str = "hg38") -> tuple[str, str]:
    """Pull 524,288 bp WT and mutant sequences centred on *pos*.

    Uses ``SeqMat.from_fasta`` from seqmat. Handles indel padding/trimming
    to maintain exact :data:`BORZOI_INPUT_LEN` bp.  All *mutations* are
    applied together (compound mutation support).

    Args:
        chrom: Chromosome (e.g. ``"12"`` or ``"chr12"``).
        pos: Central genomic position (1-based) for the extraction window.
        mutations: List of ``(position, ref, alt)`` tuples to apply.
        organism: Genome build.

    Returns:
        ``(wt_seq, mut_seq)`` each of length :data:`BORZOI_INPUT_LEN`.
    """
    from seqmat import SeqMat

    # Ensure chr prefix for FASTA lookup
    if not chrom.startswith("chr"):
        chrom = f"chr{chrom}"

    start = max(1, pos - HALF_LEN)
    end = start + BORZOI_INPUT_LEN - 1

    wt_sm = SeqMat.from_fasta(genome=organism, chrom=chrom, start=start, end=end)
    wt_seq = wt_sm.seq

    # Pad if near chromosome boundary
    if len(wt_seq) < BORZOI_INPUT_LEN:
        wt_seq = wt_seq + ("N" * (BORZOI_INPUT_LEN - len(wt_seq)))

    # Apply all mutations
    mut_sm = wt_sm.clone()
    mut_sm.apply_mutations(mutations)
    mut_seq = mut_sm.seq

    # Fix length for indels
    len_diff = len(mut_seq) - BORZOI_INPUT_LEN
    if len_diff > 0:
        trim_left = len_diff // 2
        trim_right = len_diff - trim_left
        mut_seq = mut_seq[trim_left:len(mut_seq) - trim_right]
    elif len_diff < 0:
        mut_seq = mut_seq + ("N" * (-len_diff))

    return wt_seq, mut_seq


# ============================================================================
# Model inference
# ============================================================================

@torch.no_grad()
def _predict_tracks_batch(model, seqs: list[str], device: torch.device,
                          dtype: torch.dtype,
                          use_autocast: bool = False) -> list[np.ndarray]:
    """Run Borzoi forward pass on a batch of sequences.

    Args:
        model: Loaded Borzoi model.
        seqs: List of DNA strings (each :data:`BORZOI_INPUT_LEN` bp).
        device: Target torch device.
        dtype: Target torch dtype.
        use_autocast: Wrap forward pass in ``torch.autocast("cuda")``
            (required for Flashzoi).

    Returns:
        List of numpy arrays, each of shape ``(num_tracks, num_bins)``.
    """
    tensors = [dna_to_onehot(s, device, dtype) for s in seqs]
    x = torch.cat(tensors, dim=0)  # (B, 4, L)
    if use_autocast:
        with torch.autocast("cuda"):
            y = model(x)  # (B, tracks, bins)
    else:
        y = model(x)  # (B, tracks, bins)
    return [y[i].detach().cpu().float().numpy() for i in range(len(seqs))]


# ============================================================================
# Tissue-level analysis
# ============================================================================

# -- Organ system groupings for tissue clustering --
TISSUE_GROUPS: dict[str, list[str]] = {
    "nervous_system": [
        "Brain - Amygdala", "Brain - Anterior cingulate cortex (BA24)",
        "Brain - Caudate (basal ganglia)", "Brain - Cerebellar Hemisphere",
        "Brain - Cerebellum", "Brain - Cortex", "Brain - Frontal Cortex (BA9)",
        "Brain - Hippocampus", "Brain - Hypothalamus",
        "Brain - Nucleus accumbens (basal ganglia)",
        "Brain - Putamen (basal ganglia)",
        "Brain - Spinal cord (cervical c-1)",
        "Brain - Substantia nigra", "Nerve - Tibial",
    ],
    "cardiovascular": [
        "Heart - Atrial Appendage", "Heart - Left Ventricle",
        "Artery - Aorta", "Artery - Coronary", "Artery - Tibial",
    ],
    "digestive": [
        "Colon - Sigmoid", "Colon - Transverse",
        "Esophagus - Gastroesophageal Junction", "Esophagus - Mucosa",
        "Esophagus - Muscularis", "Liver", "Pancreas",
        "Small Intestine - Terminal Ileum", "Stomach",
    ],
    "reproductive": [
        "Breast - Mammary Tissue", "Cervix - Ectocervix",
        "Cervix - Endocervix", "Fallopian Tube", "Ovary",
        "Prostate", "Testis", "Uterus", "Vagina",
    ],
    "endocrine": [
        "Adrenal Gland", "Pituitary", "Thyroid",
    ],
    "immune_hematopoietic": [
        "Spleen", "Whole Blood",
        "Cells - EBV-transformed lymphocytes",
    ],
    "musculoskeletal": [
        "Muscle - Skeletal", "Adipose - Subcutaneous",
        "Adipose - Visceral (Omentum)",
    ],
    "respiratory": [
        "Lung",
    ],
    "urinary": [
        "Bladder", "Kidney - Cortex", "Kidney - Medulla",
    ],
    "integumentary": [
        "Skin - Not Sun Exposed (Suprapubic)",
        "Skin - Sun Exposed (Lower leg)",
        "Cells - Cultured fibroblasts",
    ],
    "other": [
        "Minor Salivary Gland",
    ],
}

# Inverted lookup: tissue name → group
_TISSUE_TO_GROUP: dict[str, str] = {}
for _grp, _tissues in TISSUE_GROUPS.items():
    for _t in _tissues:
        _TISSUE_TO_GROUP[_t] = _grp


def _compute_tissue_changes(wt_tracks: np.ndarray, mut_tracks: np.ndarray,
                            gtex_map: dict[str, list[int]]) -> dict[str, Any]:
    """Compute per-tissue log2 fold-changes from WT and mutant track arrays.

    Returns dict with ``tissue_fold_changes``, ``tissue_predicted_levels``,
    ``max_lfc_tissue``, ``max_lfc_value``, ``max_abs_lfc``,
    ``num_affected_tissues``, and enriched summary statistics.
    """
    tissue_fold_changes: dict[str, float] = {}
    tissue_predicted_levels: dict[str, float] = {}
    tissue_wt_levels: dict[str, float] = {}

    # Per-tissue: also collect per-track LFCs for statistical testing
    tissue_track_lfcs: dict[str, list[float]] = {}

    for tissue, track_indices in gtex_map.items():
        wt_vals = np.array([wt_tracks[idx, :].sum() for idx in track_indices])
        mut_vals = np.array([mut_tracks[idx, :].sum() for idx in track_indices])

        wt_mean = float(wt_vals.mean())
        mut_mean = float(mut_vals.mean())

        lfc = math.log2((mut_mean + 1.0) / (wt_mean + 1.0))
        ratio = (mut_mean + 1.0) / (wt_mean + 1.0)

        tissue_fold_changes[tissue] = round(lfc, 4)
        tissue_predicted_levels[tissue] = round(ratio, 4)
        tissue_wt_levels[tissue] = round(wt_mean, 4)

        # Per-track LFCs for within-tissue variance estimation
        per_track = []
        for idx in track_indices:
            wt_s = float(wt_tracks[idx, :].sum())
            mut_s = float(mut_tracks[idx, :].sum())
            per_track.append(math.log2((mut_s + 1.0) / (wt_s + 1.0)))
        tissue_track_lfcs[tissue] = per_track

    abs_lfcs = {t: abs(v) for t, v in tissue_fold_changes.items()}
    max_tissue = max(abs_lfcs, key=abs_lfcs.get) if abs_lfcs else ""
    max_lfc = abs_lfcs.get(max_tissue, 0.0)
    num_affected = sum(1 for v in abs_lfcs.values() if v > LFC_THRESHOLD)

    # --- Enriched statistics ---

    # Median absolute LFC across all tissues
    all_abs_lfcs = list(abs_lfcs.values())
    median_abs_lfc = float(np.median(all_abs_lfcs)) if all_abs_lfcs else 0.0

    # Tissue specificity index (tau): measures how concentrated the effect is
    # tau = 0 → uniform across tissues, tau = 1 → perfectly specific
    if all_abs_lfcs and max_lfc > 0:
        n_tissues = len(all_abs_lfcs)
        tau = sum(1.0 - (v / max_lfc) for v in all_abs_lfcs) / (n_tissues - 1) if n_tissues > 1 else 0.0
        tau = round(max(0.0, min(1.0, tau)), 4)
    else:
        tau = 0.0

    # Top affected organ system (group-level aggregation)
    group_abs_lfcs: dict[str, list[float]] = {}
    for tissue, lfc_val in tissue_fold_changes.items():
        grp = _TISSUE_TO_GROUP.get(tissue, "other")
        group_abs_lfcs.setdefault(grp, []).append(abs(lfc_val))

    group_mean_abs_lfc: dict[str, float] = {
        grp: round(float(np.mean(vals)), 4) for grp, vals in group_abs_lfcs.items()
    }
    top_group = max(group_mean_abs_lfc, key=group_mean_abs_lfc.get) if group_mean_abs_lfc else ""
    n_affected_groups = sum(
        1 for v in group_mean_abs_lfc.values() if v > LFC_THRESHOLD
    )

    # Effect direction consistency: are most tissues going the same direction?
    n_up = sum(1 for v in tissue_fold_changes.values() if v > LFC_THRESHOLD)
    n_down = sum(1 for v in tissue_fold_changes.values() if v < -LFC_THRESHOLD)
    if n_up + n_down > 0:
        direction_consistency = round(max(n_up, n_down) / (n_up + n_down), 4)
        predominant_direction = "upregulated" if n_up > n_down else "downregulated"
    else:
        direction_consistency = 0.0
        predominant_direction = "unchanged"

    # Top 5 most affected tissues (for quick summary)
    sorted_tissues = sorted(abs_lfcs.items(), key=lambda x: x[1], reverse=True)
    top_affected = [
        {"tissue": t, "lfc": round(tissue_fold_changes[t], 4), "abs_lfc": round(a, 4),
         "wt_level": tissue_wt_levels.get(t, 0.0),
         "group": _TISSUE_TO_GROUP.get(t, "other")}
        for t, a in sorted_tissues[:5]
    ]

    return {
        "tissue_fold_changes": tissue_fold_changes,
        "tissue_predicted_levels": tissue_predicted_levels,
        "tissue_wt_levels": tissue_wt_levels,
        "max_lfc_tissue": max_tissue,
        "max_lfc_value": round(tissue_fold_changes.get(max_tissue, 0.0), 4),
        "max_abs_lfc": round(max_lfc, 4),
        "num_affected_tissues": num_affected,
        # Enriched fields
        "median_abs_lfc": round(median_abs_lfc, 4),
        "tissue_specificity_tau": tau,
        "top_affected_group": top_group,
        "group_mean_abs_lfc": group_mean_abs_lfc,
        "n_affected_groups": n_affected_groups,
        "direction_consistency": direction_consistency,
        "predominant_direction": predominant_direction,
        "top_affected_tissues": top_affected,
        "n_upregulated": n_up,
        "n_downregulated": n_down,
    }


def _compute_track_deltas(
    wt_tracks: np.ndarray,
    mut_tracks: np.ndarray,
    gtex_map: dict[str, list[int]],
    n_top_tracks: int = 10,
    n_top_bins: int = 5,
) -> dict[str, Any]:
    """Analyze the raw Borzoi track-level and spatial-bin-level deltas.

    Instead of collapsing tracks to per-tissue sums, this examines the full
    (tracks, bins) delta matrix to identify:
      1. Individual tracks with the largest absolute change
      2. Spatial bins where the biggest regulatory shifts occur
      3. Whether changes are concentrated in specific genomic regions

    This catches effects invisible to tissue-level means: enhancer
    repositioning, promoter-proximal shifts, localized regulatory element
    disruption.

    Args:
        wt_tracks: WT Borzoi output, shape (num_tracks, num_bins).
        mut_tracks: Mutant Borzoi output, same shape.
        gtex_map: Tissue -> track indices mapping.
        n_top_tracks: Number of top-changed tracks to report.
        n_top_bins: Number of top-changed spatial bins to report.

    Returns:
        Dict with track_deltas, spatial_profile, and summary statistics.
    """
    # Full delta matrix
    delta = mut_tracks - wt_tracks  # (tracks, bins)
    n_tracks, n_bins = delta.shape

    # --- Per-track analysis: which individual tracks changed most? ---
    track_abs_sums = np.abs(delta).sum(axis=1)  # (tracks,)
    track_signed_sums = delta.sum(axis=1)        # Preserve direction
    track_max_abs_bin = np.abs(delta).max(axis=1) # Peak change in any bin

    top_track_indices = np.argsort(track_abs_sums)[::-1][:n_top_tracks]

    # Build reverse map: track_idx -> tissue
    idx_to_tissue: dict[int, str] = {}
    for tissue, indices in gtex_map.items():
        for idx in indices:
            idx_to_tissue[idx] = tissue

    top_tracks = []
    for idx in top_track_indices:
        idx = int(idx)
        tissue = idx_to_tissue.get(idx, "unknown")
        top_tracks.append({
            "track_index": idx,
            "tissue": tissue,
            "abs_delta_sum": round(float(track_abs_sums[idx]), 4),
            "signed_delta_sum": round(float(track_signed_sums[idx]), 4),
            "peak_bin_delta": round(float(track_max_abs_bin[idx]), 4),
            "direction": "up" if track_signed_sums[idx] > 0 else "down",
        })

    # --- Per-bin (spatial) analysis: where in the 524kb window? ---
    # Sum absolute deltas across all tracks per bin → spatial disruption profile
    bin_disruption = np.abs(delta).sum(axis=0)   # (bins,)
    bin_signed = delta.sum(axis=0)               # Net direction per bin

    top_bin_indices = np.argsort(bin_disruption)[::-1][:n_top_bins]

    # Borzoi bins map to genomic positions: each bin = 128bp (Borzoi resolution)
    # The center bin corresponds to the mutation position
    bin_resolution = 128  # bp per bin
    center_bin = n_bins // 2

    top_bins = []
    for b in top_bin_indices:
        b = int(b)
        offset_from_center = (b - center_bin) * bin_resolution
        top_bins.append({
            "bin_index": b,
            "offset_bp": offset_from_center,  # Negative = upstream, positive = downstream
            "abs_disruption": round(float(bin_disruption[b]), 4),
            "net_direction": round(float(bin_signed[b]), 4),
            "relative_position": (
                "at_mutation" if abs(offset_from_center) < bin_resolution
                else f"{'upstream' if offset_from_center < 0 else 'downstream'}_{abs(offset_from_center)//1000}kb"
            ),
        })

    # --- Global spatial summary ---
    # Is the disruption concentrated or diffuse?
    if float(bin_disruption.sum()) > 0:
        # Gini coefficient of spatial disruption (0=uniform, 1=concentrated)
        sorted_bins = np.sort(bin_disruption)
        n = len(sorted_bins)
        cumulative = np.cumsum(sorted_bins)
        gini = (2 * np.sum((np.arange(1, n + 1) * sorted_bins)) / (n * sorted_bins.sum()) - (n + 1) / n)
        spatial_gini = round(float(max(0, min(1, gini))), 4)
    else:
        spatial_gini = 0.0

    # Fraction of total disruption within ±10kb of mutation
    center_range = slice(max(0, center_bin - 10000 // bin_resolution),
                         min(n_bins, center_bin + 10000 // bin_resolution + 1))
    local_disruption = float(bin_disruption[center_range].sum())
    total_disruption = float(bin_disruption.sum())
    local_fraction = round(local_disruption / total_disruption, 4) if total_disruption > 0 else 0.0

    # Track-level summary
    n_tracks_changed = int(np.sum(track_abs_sums > np.percentile(track_abs_sums, 90)))

    return {
        "top_tracks": top_tracks,
        "top_spatial_bins": top_bins,
        "spatial_gini": spatial_gini,
        "local_disruption_fraction": local_fraction,
        "n_tracks_with_large_change": n_tracks_changed,
        "total_track_disruption": round(float(track_abs_sums.sum()), 4),
        "mean_track_disruption": round(float(track_abs_sums.mean()), 4),
        "n_bins": n_bins,
        "bin_resolution_bp": bin_resolution,
    }


# ============================================================================
# Score normalization
# ============================================================================

def _normalize_score(max_abs_lfc: float) -> float:
    """Normalize max absolute LFC to a 0-10 score.

    ``10 * (1 - exp(-2 * |max_lfc|))``
    """
    k = 2.0
    return 10.0 * (1.0 - math.exp(-k * abs(max_abs_lfc)))


# ============================================================================
# Full prediction entry point
# ============================================================================

def predict_transcription_effect(chrom: str, pos: int,
                                 mutations: list[tuple[int, str, str]],
                                 organism: str = "hg38",
                                 capture_embeddings: bool = False) -> dict:
    """Full Borzoi pipeline: extract sequences, predict, compare, score.

    Args:
        chrom: Chromosome (e.g. ``"12"`` or ``"chr12"``).
        pos: Central genomic position for the extraction window.
        mutations: List of ``(position, ref, alt)`` tuples.
        organism: Genome build (default ``"hg38"``).
        capture_embeddings: If True, also capture the intermediate
            embedding representations for both WT and mutant. These
            can be used for downstream similarity analyses, clustering,
            or more granular effect characterization. Stored under
            ``"wt_embedding"`` and ``"mut_embedding"`` keys as numpy arrays.

    Returns:
        Dict with ``status``, ``tissue_fold_changes``, ``regulatory_score``, etc.
        If ``capture_embeddings=True``, also includes ``wt_embedding`` and
        ``mut_embedding`` numpy arrays.
    """
    from .model import EmbeddingCapture

    model, device, dtype, use_autocast = _load_borzoi()
    gtex_map = _build_gtex_mapping()

    wt_seq, mut_seq = _extract_sequences(chrom, pos, mutations, organism)

    if capture_embeddings:
        # Run with embedding hooks
        with EmbeddingCapture(model) as cap:
            wt_tracks_list = _predict_tracks_batch(
                model, [wt_seq], device, dtype, use_autocast
            )
            wt_embedding = cap.get_numpy()

        with EmbeddingCapture(model) as cap:
            mut_tracks_list = _predict_tracks_batch(
                model, [mut_seq], device, dtype, use_autocast
            )
            mut_embedding = cap.get_numpy()

        wt_tracks = wt_tracks_list[0]
        mut_tracks = mut_tracks_list[0]
    else:
        wt_tracks, mut_tracks = _predict_tracks_batch(
            model, [wt_seq, mut_seq], device, dtype, use_autocast
        )
        wt_embedding = None
        mut_embedding = None

    tissue_data = _compute_tissue_changes(wt_tracks, mut_tracks, gtex_map)
    score = _normalize_score(tissue_data["max_abs_lfc"])
    tissue_data["regulatory_score"] = round(score, 4)
    tissue_data["status"] = "completed"

    # Raw track-level delta analysis
    track_deltas = _compute_track_deltas(wt_tracks, mut_tracks, gtex_map)
    tissue_data["track_deltas"] = track_deltas

    if capture_embeddings:
        tissue_data["wt_embedding"] = wt_embedding
        tissue_data["mut_embedding"] = mut_embedding
        # Compute embedding distance as a summary metric
        if wt_embedding is not None and mut_embedding is not None:
            # Cosine distance between mean-pooled embeddings
            wt_pooled = wt_embedding.mean(axis=-1).flatten()
            mut_pooled = mut_embedding.mean(axis=-1).flatten()
            cos_sim = float(np.dot(wt_pooled, mut_pooled) /
                          (np.linalg.norm(wt_pooled) * np.linalg.norm(mut_pooled) + 1e-8))
            tissue_data["embedding_cosine_distance"] = round(1.0 - cos_sim, 6)
            # L2 distance
            tissue_data["embedding_l2_distance"] = round(
                float(np.linalg.norm(wt_pooled - mut_pooled)), 4)

    return tissue_data


# ============================================================================
# Cross-mutation batched prediction
# ============================================================================

def predict_transcription_effects_batch(
    specs: list[dict],
    max_batch_size: int = 8,
    n_io_workers: int = 8,
) -> list[dict]:
    """Batched Borzoi transcription-effect prediction for multiple mutations.

    Extracts WT/mutant sequences in parallel (I/O bound), then batches all
    sequences through the model in GPU-memory-sized chunks.

    Args:
        specs: List of dicts, each with keys ``chrom``, ``pos``,
            ``mutations``, and optionally ``organism``.
        max_batch_size: Maximum sequences per GPU forward pass (controls
            GPU memory).  Each mutation contributes 2 sequences (WT + mut).
        n_io_workers: Thread pool size for parallel sequence extraction.

    Returns:
        List of result dicts (same format as
        :func:`predict_transcription_effect`), one per input spec.
    """
    from concurrent.futures import ThreadPoolExecutor

    if not specs:
        return []

    model, device, dtype, use_autocast = _load_borzoi()
    gtex_map = _build_gtex_mapping()

    # -- Parallel sequence extraction (I/O bound) --
    def _extract(spec):
        return _extract_sequences(
            spec["chrom"], spec["pos"], spec["mutations"],
            spec.get("organism", "hg38"),
        )

    with ThreadPoolExecutor(max_workers=min(n_io_workers, len(specs))) as pool:
        seq_pairs = list(pool.map(_extract, specs))

    # -- Flatten all sequences for batched inference --
    all_seqs: list[str] = []
    for wt, mut in seq_pairs:
        all_seqs.append(wt)
        all_seqs.append(mut)

    # -- Chunked GPU forward passes --
    all_tracks: list[np.ndarray] = []
    for i in range(0, len(all_seqs), max_batch_size):
        chunk = all_seqs[i : i + max_batch_size]
        all_tracks.extend(
            _predict_tracks_batch(model, chunk, device, dtype, use_autocast)
        )

    # -- Compute per-mutation tissue changes --
    results: list[dict] = []
    for idx, spec in enumerate(specs):
        wt_tracks = all_tracks[idx * 2]
        mut_tracks = all_tracks[idx * 2 + 1]
        tissue_data = _compute_tissue_changes(wt_tracks, mut_tracks, gtex_map)
        score = _normalize_score(tissue_data["max_abs_lfc"])
        tissue_data["regulatory_score"] = round(score, 4)
        tissue_data["status"] = "completed"
        results.append(tissue_data)

    return results
