"""Shared utility functions for h5seurat conversion operations."""

from __future__ import annotations

from collections.abc import Iterable
from typing import Any

import h5py
import numpy as np
import pandas as pd
from scipy import sparse

from readseurat.convert.constants import (
    LOGICAL_FALSE,
    LOGICAL_NA,
    LOGICAL_TRUE,
    R_FACTOR_BASE_INDEX,
    R_FACTOR_NA_VALUE,
)
from readseurat.convert.seurat.model import Assay


def _choose_primary_matrix(assay: Assay, use_counts: bool) -> sparse.spmatrix | np.ndarray:
    """Choose which matrix to use as the primary X matrix.

    Args:
        assay: The assay containing the matrices
        use_counts: If True and counts are available, use counts; otherwise use data

    Returns:
        The selected matrix (sparse or dense)
    """
    if use_counts and assay.counts is not None:
        return assay.counts
    return assay.data


def _ensure_cells_by_features(matrix: Any, n_cells: int, n_features: int) -> sparse.spmatrix | np.ndarray:
    """Ensure matrix is in cells x features orientation.

    Args:
        matrix: Input matrix (sparse or dense, from various sources)
        n_cells: Expected number of cells
        n_features: Expected number of features

    Returns:
        Matrix in cells x features orientation (transposed if necessary)
    """
    if sparse.issparse(matrix):
        m = matrix  # type: sparse.spmatrix
        shape = m.shape
        if shape == (n_cells, n_features):
            return m
        if shape == (n_features, n_cells):
            return m.T.tocsr()
        return m

    arr = np.asarray(matrix)
    if arr.shape == (n_cells, n_features):
        return arr
    if arr.shape == (n_features, n_cells):
        return arr.T
    return arr


def _ensure_features_by_cells(matrix: Any, n_cells: int, n_features: int) -> sparse.spmatrix | np.ndarray:
    """Ensure matrix is in features x cells orientation.

    Args:
        matrix: Input matrix (sparse or dense, from various sources)
        n_cells: Expected number of cells
        n_features: Expected number of features

    Returns:
        Matrix in features x cells orientation (transposed if necessary)
    """
    if sparse.issparse(matrix):
        m = matrix  # type: sparse.spmatrix
        shape = m.shape
        if shape == (n_features, n_cells):
            return m
        if shape == (n_cells, n_features):
            return m.T.tocsc()
        return m

    arr = np.asarray(matrix)
    if arr.shape == (n_features, n_cells):
        return arr
    if arr.shape == (n_cells, n_features):
        return arr.T
    return arr


def _build_var_frame(assay: Assay) -> pd.DataFrame:
    """Build the var DataFrame from assay feature information.

    Args:
        assay: The assay containing feature metadata

    Returns:
        DataFrame with features as index and metadata as columns
    """
    features = assay.features

    # Handle if features is 2D (take first column if so)
    if isinstance(features, np.ndarray) and features.ndim == 2:
        features = features[:, 0]

    if assay.meta_features is None:
        var = pd.DataFrame(index=features)
    else:
        var = assay.meta_features.copy()
        if len(var.index) != len(features):
            var.index = features
        else:
            # Safely compare indices - handles potential edge cases
            try:
                if not np.array_equal(np.asarray(var.index.values), features):
                    var.index = features
            except (ValueError, AttributeError):
                var.index = features

    # Add variable features indicator
    if assay.variable_features is not None:
        var_features = set(assay.variable_features.tolist())
        var["highly_variable"] = [f in var_features for f in features]

    return var


def _logical_to_int(values: Iterable[Any]) -> np.ndarray:
    """Convert logical/boolean values to integer encoding for R compatibility.

    R logical vectors are represented as integers:
    - FALSE -> 0
    - TRUE -> 1
    - NA -> 2

    Args:
        values: Iterable of boolean or mixed values

    Returns:
        Integer array with R-compatible encoding
    """
    mapped = []
    for value in values:
        if value is pd.NA or value is None or (isinstance(value, float) and np.isnan(value)):
            mapped.append(LOGICAL_NA)
        elif bool(value):
            mapped.append(LOGICAL_TRUE)
        else:
            mapped.append(LOGICAL_FALSE)
    return np.asarray(mapped, dtype=np.int8)


def _write_factor_group(group: h5py.Group, series: pd.Series) -> None:
    """Write a factor (categorical variable) to an HDF5 group.

    Factors are encoded with R-compatible 1-based indexing.
    Missing values are encoded as 0.

    Args:
        group: HDF5 group to write to
        series: Pandas Series with categorical data
    """
    cat = series.astype("category")
    levels = np.asarray(cat.cat.categories.astype(str), dtype=object)
    codes = cat.cat.codes.to_numpy()

    # Convert to 1-based indexing for R compatibility
    values = (codes + R_FACTOR_BASE_INDEX).astype(np.int32)
    values[codes == -1] = R_FACTOR_NA_VALUE

    # Write string dataset for levels
    _write_string_dataset(group, "levels", levels)
    group.create_dataset("values", data=values, dtype=np.int32)


def _write_string_dataset(group: h5py.Group, name: str, data: Iterable[Any]) -> None:
    """Write string data to HDF5 dataset with proper encoding.

    Args:
        group: HDF5 group to write to
        name: Name of the dataset
        data: Iterable of string data
    """
    dtype = h5py.string_dtype("utf-8")
    group.create_dataset(name, data=np.asarray(list(data), dtype=object), dtype=dtype)
