# geney/transcription/model.py — Borzoi model configuration and lazy loading.
"""
Borzoi regulatory impact prediction model management.

Borzoi predicts per-tissue expression changes from 524kb raw genomic sequences.
Uses HuggingFace borzoi-pytorch models (downloads ~4GB on first use).

Reference: https://github.com/johahi/borzoi-pytorch
"""
from __future__ import annotations

import logging

import torch

logger = logging.getLogger(__name__)

# Lazy-loaded model state
_borzoi_model = None
_borzoi_model_name: str | None = None
_borzoi_device: torch.device | None = None

# Default HuggingFace model identifiers
DEFAULT_BORZOI_MODEL = "johahi/borzoi-replicate-0"
DEFAULT_FLASHZOI_MODEL = "johahi/flashzoi-replicate-0"

# Whether the loaded model is Flashzoi (requires autocast)
_is_flashzoi: bool = False


# ============================================================================
# Configuration
# ============================================================================

def set_borzoi_model(name_or_path: str) -> None:
    """Set Borzoi model (HuggingFace name or local directory).

    Resets the cached model so the next prediction triggers a fresh load.

    Args:
        name_or_path: HuggingFace model name (e.g. ``"johahi/borzoi-replicate-0"``)
            or path to a local model directory.
    """
    global _borzoi_model_name, _borzoi_model
    _borzoi_model_name = name_or_path
    _borzoi_model = None  # Reset cached model


def _default_model_name() -> str:
    """Use Flashzoi on Ampere+ CUDA GPUs (cc >= 8.0), vanilla Borzoi otherwise.

    flash-attn requires compute capability 8.0+. Turing (T4, cc 7.5) and
    older GPUs must use the standard Borzoi model.
    """
    if not torch.cuda.is_available():
        return DEFAULT_BORZOI_MODEL
    cc = torch.cuda.get_device_capability(0)
    if cc[0] < 8:
        logger.info("GPU cc %d.%d < 8.0 — using vanilla Borzoi (flash-attn needs Ampere+)", cc[0], cc[1])
        return DEFAULT_BORZOI_MODEL
    return DEFAULT_FLASHZOI_MODEL


def get_borzoi_model_name() -> str:
    """Return the Borzoi model name.

    Resolution order: setter value -> ``BORZOI_MODEL`` env var -> auto-detect
    (Flashzoi when ``flash-attn`` + CUDA available, else standard Borzoi).
    """
    global _borzoi_model_name
    if _borzoi_model_name is not None:
        return _borzoi_model_name

    import os
    env = os.environ.get("BORZOI_MODEL")
    if env:
        _borzoi_model_name = env
        return _borzoi_model_name

    _borzoi_model_name = _default_model_name()
    return _borzoi_model_name


def set_borzoi_device(device: str) -> None:
    """Override the device used by Borzoi (e.g. ``'cpu'``, ``'cuda'``).

    Resets the cached model so it reloads on the new device.
    """
    global _borzoi_device, _borzoi_model
    _borzoi_device = torch.device(device)
    _borzoi_model = None  # Force reload on next use


def get_borzoi_device() -> torch.device:
    """Return the Borzoi device (auto-detected if not overridden)."""
    return _detect_device()


# ============================================================================
# Internal helpers
# ============================================================================

def _detect_device() -> torch.device:
    """Auto-detect best available device (CUDA > MPS > CPU)."""
    import sys

    global _borzoi_device
    if _borzoi_device is not None:
        return _borzoi_device

    if torch.cuda.is_available():
        _borzoi_device = torch.device("cuda")
    elif sys.platform == "darwin" and torch.backends.mps.is_available():
        try:
            torch.tensor([1.0], device="mps")
            _borzoi_device = torch.device("mps")
        except RuntimeError:
            _borzoi_device = torch.device("cpu")
    else:
        _borzoi_device = torch.device("cpu")
    return _borzoi_device


def _load_borzoi():
    """Lazy-load Borzoi model. Returns ``(model, device, dtype, use_autocast)``.

    On CUDA the model runs in fp16 except for ``human_head`` (kept fp32).
    On CPU everything is fp32.  When the loaded model is Flashzoi,
    ``use_autocast`` is ``True`` so callers can wrap the forward pass in
    ``torch.autocast("cuda")``.
    """
    global _borzoi_model, _is_flashzoi

    if _borzoi_model is not None:
        device = _detect_device()
        dtype = torch.float16 if device.type == "cuda" else torch.float32
        return _borzoi_model, device, dtype, _is_flashzoi

    from borzoi_pytorch import Borzoi

    name = get_borzoi_model_name()
    device = _detect_device()
    # fp16 only on CUDA; MPS and CPU stay fp32
    dtype = torch.float16 if device.type == "cuda" else torch.float32

    _is_flashzoi = "flashzoi" in name.lower()
    logger.info("Loading %s model '%s' on %s (%s)...",
                "Flashzoi" if _is_flashzoi else "Borzoi", name, device, dtype)
    model = Borzoi.from_pretrained(name)
    model = model.to(device, dtype=dtype).eval()

    # Keep heads in fp32 for numerical stability
    model.human_head = model.human_head.to(torch.float32)
    if hasattr(model, "mouse_head") and model.mouse_head is not None:
        model.mouse_head = model.mouse_head.to(torch.float32)

    _borzoi_model = model
    logger.info("Borzoi model loaded.")
    return _borzoi_model, device, dtype, _is_flashzoi


# ============================================================================
# Embedding extraction
# ============================================================================

class EmbeddingCapture:
    """Hook-based capture of intermediate Borzoi embeddings.

    Registers a forward hook on a target layer to capture its output
    during the forward pass. This allows access to the learned
    representation space without modifying the model architecture.

    Usage::

        model, device, dtype, use_autocast = _load_borzoi()
        with EmbeddingCapture(model) as cap:
            y = model(x)
            embeddings = cap.embeddings  # (B, hidden_dim, seq_bins)

    The default target is the transformer output (before the prediction
    heads), which captures the richest learned representation.
    """

    def __init__(self, model, target_layer: str | None = None):
        """
        Args:
            model: Loaded Borzoi model.
            target_layer: Dot-separated path to the module to hook.
                Default: auto-detect the last transformer/conv block
                before the human_head.
        """
        self.model = model
        self.target_layer = target_layer
        self._hook = None
        self.embeddings: torch.Tensor | None = None

    def _find_target_module(self):
        """Auto-detect the best layer for embedding capture.

        Strategy: find the module immediately before human_head.
        For Borzoi this is typically the final conv tower or transformer block.
        """
        if self.target_layer:
            parts = self.target_layer.split(".")
            mod = self.model
            for p in parts:
                mod = getattr(mod, p)
            return mod

        # Heuristic: look for common Borzoi architecture patterns
        # Try 'final_pointwise' first (Enformer-style), then last conv block
        for attr in ("final_pointwise", "transformer", "conv_tower"):
            if hasattr(self.model, attr):
                return getattr(self.model, attr)

        # Fallback: hook the module right before human_head
        # Walk children and return the second-to-last
        children = list(self.model.children())
        for i, child in enumerate(children):
            if child is self.model.human_head and i > 0:
                return children[i - 1]

        raise ValueError(
            "Could not auto-detect embedding layer. "
            "Specify target_layer explicitly (e.g. 'transformer' or 'conv_tower')."
        )

    def _hook_fn(self, module, input, output):
        """Forward hook that captures the layer output."""
        if isinstance(output, tuple):
            self.embeddings = output[0].detach()
        else:
            self.embeddings = output.detach()

    def __enter__(self):
        target = self._find_target_module()
        self._hook = target.register_forward_hook(self._hook_fn)
        logger.debug("Embedding hook registered on %s", type(target).__name__)
        return self

    def __exit__(self, *args):
        if self._hook is not None:
            self._hook.remove()
            self._hook = None

    def get_numpy(self) -> "np.ndarray | None":
        """Return captured embeddings as numpy array, or None."""
        if self.embeddings is None:
            return None
        return self.embeddings.cpu().float().numpy()
