"""
Persistent homology computation module for MedTDA.

This module provides functions for computing persistent homology of medical images
using cubical complexes with the cripser library.
"""

import numpy as np
import cripser


def clean_and_cap_ph(ph_array, cap_value):
    """
    Clean and cap a persistence homology array.
    
    Removes zero-persistence bars and caps infinite values.
    
    Parameters
    ----------
    ph_array : array-like
        Persistence homology array with columns [dimension, birth, death].
    cap_value : float
        Value to use for capping infinite death times.
        
    Returns
    -------
    numpy.ndarray
        Cleaned and capped persistence array.
    """
    ph_array = np.array(ph_array, copy=True, dtype=np.float64)

    # Remove zero-persistence bars
    persistence = ph_array[:, 2] - ph_array[:, 1]
    ph_array = ph_array[persistence != 0]

    # Cap infinity (including huge float values)
    inf_mask = np.isinf(ph_array[:, 2]) | (ph_array[:, 2] > 1e300)
    ph_array[inf_mask, 2] = cap_value

    return ph_array


def compute_ph(
    image_array,
    maxdim=None,
    construction="T",
    filtration="sublevel",
    max_value=None
):
    """
    Compute persistent homology of an image array using cubical complexes.

    Parameters
    ----------
    image_array : numpy.ndarray
        2D, 3D, or 4D image array.
    maxdim : int, optional
        Maximum homology dimension to compute.
        If None (default), automatically determined from image dimensions:
        - 2D images: maxdim=1 (compute H0, H1)
        - 3D images: maxdim=2 (compute H0, H1, H2)
        - 4D images: maxdim=3 (compute H0, H1, H2, H3)
    construction : {'T', 'V'}, default='T'
        Type of cubical complex construction:
        - 'T': T-construction (default, for general cubical complexes)
        - 'V': V-construction (for vertex-based cubical complexes)
    filtration : {'sublevel', 'superlevel'}, default='sublevel'
        Filtration type:
        - 'sublevel': Standard sublevel set filtration
        - 'superlevel': Superlevel set filtration (image is negated)
    max_value : float, optional
        Maximum value for capping infinite persistence.
        If None, uses the maximum value in the image array.

    Returns
    -------
    ph_capped : numpy.ndarray
        Cleaned and capped persistence array with columns [dimension, birth, death].
    dgms : list of numpy.ndarray
        List of persistence diagrams in GUDHI format (one per dimension).

    Raises
    ------
    ValueError
        If construction is not 'T' or 'V', or filtration is not 'sublevel' or 'superlevel'.

    Notes
    -----
    For superlevel filtration, the image is negated internally, and the max_value
    is also negated for proper capping.

    Examples
    --------
    >>> import numpy as np
    >>> from medtda.ph_computer import compute_ph
    >>> image = np.random.rand(50, 50)
    >>> ph_array, diagrams = compute_ph(image, maxdim=1, construction='T')
    """
    # Auto-determine maxdim based on image dimensions if not specified
    if maxdim is None:
        maxdim = image_array.ndim - 1

    if construction not in ["T", "V"]:
        raise ValueError(
            "Construction must be 'T' (T-construction) or 'V' (V-construction)."
        )

    if filtration not in ["sublevel", "superlevel"]:
        raise ValueError(
            "Filtration must be 'sublevel' or 'superlevel'."
        )

    # Handle superlevel filtration by negating the image
    if filtration == "superlevel":
        image_array = -image_array

    # Determine cap value
    if max_value is None:
        cap_value = np.max(image_array)
    else:
        cap_value = -max_value if filtration == "superlevel" else max_value

    # Compute persistent homology using cripser
    ph = cripser.compute_ph(image_array, maxdim=maxdim, filtration=construction)
    ph_capped = clean_and_cap_ph(ph, cap_value)

    # Convert to GUDHI diagram format
    dgms = cripser.to_gudhi_diagrams(ph_capped)

    return ph_capped, dgms
