# 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
# ============================================================================

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``.
    """
    tissue_fold_changes: dict[str, float] = {}
    tissue_predicted_levels: dict[str, 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)

    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)

    return {
        "tissue_fold_changes": tissue_fold_changes,
        "tissue_predicted_levels": tissue_predicted_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,
    }


# ============================================================================
# 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") -> 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"``).

    Returns:
        Dict with ``status``, ``tissue_fold_changes``, ``regulatory_score``, etc.
    """
    model, device, dtype, use_autocast = _load_borzoi()
    gtex_map = _build_gtex_mapping()

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

    wt_tracks, mut_tracks = _predict_tracks_batch(
        model, [wt_seq, mut_seq], device, dtype, use_autocast
    )

    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"

    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
