# 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:
    """Always use Flashzoi on CUDA (requires flash-attn), fall back to Borzoi on CPU."""
    if not torch.cuda.is_available():
        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
