"""Inference utilities for encoder/decoder models."""

import numpy as np
import scipy.sparse as sp
import torch
import torch.nn as nn


def get_corrected_expression(
    encoder: nn.Module,
    decoder: nn.Module,
    X_input: np.ndarray,
    use_batch_conditioning: bool = True,
    target_batch: int = 0,
    device: str = "cpu",
    batch_size: int = 2048,
) -> np.ndarray:
    """Get batch-corrected expression from encoder/decoder.

    Returns expression in log2(CPM+1) scale, matching the input scale.

    Parameters
    ----------
    encoder : nn.Module
        Trained encoder model.
    decoder : nn.Module
        Trained decoder model.
    X_input : np.ndarray or sparse matrix
        Normalized expression matrix for encoder input (n_cells, n_genes).
        Expected to be in log2(CPM+1) scale.
    use_batch_conditioning : bool, default=True
        Whether decoder uses batch conditioning.
    target_batch : int, default=0
        Target batch index for correction. All cells will be decoded
        as if they belong to this batch, effectively removing batch effects.
    device : str, default="cpu"
        Device for model inference.
    batch_size : int, default=2048
        Batch size for inference.

    Returns
    -------
    X_corrected : np.ndarray
        Batch-corrected expression in log2(CPM+1) scale (n_cells, n_genes).
    """
    # Move models to device and set eval mode
    encoder.to(device)
    decoder.to(device)
    encoder.eval()
    decoder.eval()

    n_cells = X_input.shape[0]
    n_genes = X_input.shape[1]
    use_cuda = device != "cpu"

    # Pre-allocate result array
    result = np.empty((n_cells, n_genes), dtype=np.float32)

    with torch.inference_mode():
        for i in range(0, n_cells, batch_size):
            end_idx = min(i + batch_size, n_cells)
            X_batch = X_input[i:end_idx]

            # Convert sparse to dense if needed
            if sp.issparse(X_batch):
                X_batch = X_batch.toarray()

            # Transfer to device (non_blocking for GPU)
            X_tensor = torch.from_numpy(np.asarray(X_batch, dtype=np.float32))
            X_tensor = X_tensor.to(device, non_blocking=use_cuda)

            # Create batch override tensor directly on device
            if use_batch_conditioning:
                batch_override = torch.full(
                    (X_tensor.shape[0],), target_batch, dtype=torch.long, device=device
                )
            else:
                batch_override = None

            z = encoder(X_tensor)
            rho = decoder.get_rho(z, batch=batch_override)

            # Store result directly (log2(CPM+1) scale)
            result[i:end_idx] = np.log2(rho.cpu().numpy() * 1e6 + 1)

    return result
