# This file contains code derived from:
# https://github.com/Cimagroup/vectorization-maps
#
# Original project licensed under the GNU General Public License v3.0 (GPL-3.0).
# Modifications have been made for this project.
#
# See the LICENSE file in the project root for the full GPL-3.0 license text.

"""
Vectorization methods for converting persistence diagrams to feature vectors.

This module provides multiple methods for converting persistence diagrams (barcodes)
into fixed-length feature vectors suitable for machine learning.

Available Methods
-----------------
- betti_curve : Betti number evolution
- entropy_summary : Entropy-based summary function
- persistence_image : 2D image representation
- persistence_landscape : Piecewise-linear landscape functions
- persistence_lifespan : Normalized lifespan curve
- persistence_silhouette : Weighted silhouette representation
- persistence_stats : Comprehensive statistical features
- persistence_tropical_coordinates : Tropical polynomial features

Note
----
All methods support both snake_case names (e.g., 'persistence_stats') and 
PascalCase aliases (e.g., 'PersStats') for backward compatibility.
"""

import numpy as np
from copy import deepcopy
from gudhi import representations
from .extra_gudhi import Entropy, Lifespan


def betti_curve(barcode, resolution=100):
    """
    Compute Betti curve vectorization.
    
    The Betti curve represents the evolution of Betti numbers over the filtration.
    
    Parameters
    ----------
    barcode : numpy.ndarray
        Persistence diagram as n x 2 array (birth, death).
    resolution : int, default=100
        Number of sample points for the curve.
        
    Returns
    -------
    numpy.ndarray
        Feature vector of length resolution.
    """
    if np.size(barcode) > 0:
        betti_curve_obj = representations.vector_methods.BettiCurve(resolution=resolution)
        feature_vector = betti_curve_obj.fit_transform([barcode])[0]
    else:
        feature_vector = np.zeros(resolution)

    return feature_vector


def entropy_summary(barcode, resolution=100):
    """
    Compute entropy summary function vectorization.
    
    Based on persistence entropy, providing a statistical summary of the diagram.
    
    Parameters
    ----------
    barcode : numpy.ndarray
        Persistence diagram as n x 2 array (birth, death).
    resolution : int, default=100
        Number of sample points.
        
    Returns
    -------
    numpy.ndarray
        Feature vector of length resolution.
    """
    if hasattr(barcode, 'shape') and (barcode.shape[0]) > 1:
        ent = Entropy(mode='vector', resolution=resolution, normalized=False)
        feature_vector = ent.fit_transform([barcode]).flatten()
    else:
        feature_vector = np.zeros(resolution)

    return feature_vector


def persistence_image(barcode, bandwidth=0.2, resolution=20):
    """
    Compute persistence image vectorization.
    
    Converts the persistence diagram into a 2D image representation.
    
    Parameters
    ----------
    barcode : numpy.ndarray
        Persistence diagram as n x 2 array (birth, death).
    bandwidth : float, default=0.2
        Bandwidth for Gaussian kernel.
    resolution : int, default=20
        Resolution of the image (creates resolution x resolution grid).
        
    Returns
    -------
    numpy.ndarray
        Flattened feature vector of length resolution^2.
    """
    res = [resolution, resolution]
    if np.size(barcode) > 0:
        pers_img = representations.PersistenceImage(bandwidth=bandwidth, resolution=res)
        feature_vector = pers_img.fit_transform([barcode])[0]
    else:
        feature_vector = np.zeros(res[0] ** 2)

    return feature_vector


def persistence_landscape(barcode, resolution=100, num_landscapes=5):
    """
    Compute persistence landscape vectorization.
    
    Represents the persistence diagram as a sequence of piecewise-linear functions.
    
    Parameters
    ----------
    barcode : numpy.ndarray
        Persistence diagram as n x 2 array (birth, death).
    resolution : int, default=100
        Number of sample points per landscape.
    num_landscapes : int, default=5
        Number of landscape functions to compute.
        
    Returns
    -------
    numpy.ndarray
        Feature vector of length num_landscapes * resolution.
    """
    if np.size(barcode) > 0:
        pers_land = representations.Landscape(
            resolution=resolution,
            num_landscapes=num_landscapes
        )
        feature_vector = pers_land.fit_transform([barcode])[0]
    else:
        feature_vector = np.zeros(num_landscapes * resolution)

    return feature_vector


def persistence_lifespan(barcode, resolution=100):
    """
    Compute persistence lifespan curve vectorization.
    
    Based on the normalized lifespan curve from persistence theory.
    
    Parameters
    ----------
    barcode : numpy.ndarray
        Persistence diagram as n x 2 array (birth, death).
    resolution : int, default=100
        Number of sample points.
        
    Returns
    -------
    numpy.ndarray
        Feature vector of length resolution.
    """
    if np.size(barcode) > 0:
        lfsp = Lifespan(resolution=resolution)
        feature_vector = lfsp.fit_transform([barcode]).flatten()
    else:
        feature_vector = np.zeros(resolution)

    return feature_vector


def persistence_silhouette(barcode, resolution=100, weight=1):
    """
    Compute persistence silhouette vectorization.
    
    A weighted summary of the persistence diagram.
    
    Parameters
    ----------
    barcode : numpy.ndarray
        Persistence diagram as n x 2 array (birth, death).
    resolution : int, default=100
        Number of sample points.
    weight : float, default=1
        Weight parameter for silhouette computation.
        
    Returns
    -------
    numpy.ndarray
        Feature vector of length resolution.
    """
    if np.size(barcode) > 0:
        pers_silhouette = representations.vector_methods.Silhouette(
            resolution=resolution,
            weight=lambda x: (x[1] - x[0]) ** weight
        )
        feature_vector = pers_silhouette.fit_transform([barcode])[0]
    else:
        feature_vector = np.zeros(resolution)

    return feature_vector


def persistence_stats(barcode):
    """
    Compute statistical summary of persistence diagram.
    
    Computes comprehensive statistics including mean, std, median, IQR, percentiles
    for births, deaths, midpoints, and lifespans, plus total bar count and entropy.
    
    Parameters
    ----------
    barcode : numpy.ndarray
        Persistence diagram as n x 2 array (birth, death).
        
    Returns
    -------
    dict
        Feature dictionary with 38 statistical features:
        ['births_mean', 'births_std', 'births_median', 'births_iqr', 'births_range', 'births_p10',
         'births_p25', 'births_p75', 'births_p90', 'deaths_mean', 'deaths_std', 'deaths_median',
         'deaths_iqr', 'deaths_range', 'deaths_p10', 'deaths_p25', 'deaths_p75', 'deaths_p90',
         'midpoints_mean', 'midpoints_std', 'midpoints_median', 'midpoints_iqr', 'midpoints_range',
         'midpoints_p10', 'midpoints_p25', 'midpoints_p75', 'midpoints_p90', 'lifespans_mean',
         'lifespans_std', 'lifespans_median', 'lifespans_iqr', 'lifespans_range', 'lifespans_p10',
         'lifespans_p25', 'lifespans_p75', 'lifespans_p90', 'count', 'entropy']
    Notes
    -----
    Features include (in order):
    - Birth stats: mean, std, median, IQR, range, percentiles (10,25,75,90)
    - Death stats: mean, std, median, IQR, range, percentiles (10,25,75,90)
    - Midpoint stats: mean, std, median, IQR, range, percentiles (10,25,75,90)
    - Lifespan stats: mean, std, median, IQR, range, percentiles (10,25,75,90)
    - Bar count, persistence entropy
    """
    feature_dict = {
        'births_mean': 0, 'births_std': 0, 'births_median': 0, 'births_iqr': 0, 
        'births_range': 0, 'births_p10': 0, 'births_p25': 0, 'births_p75': 0, 'births_p90': 0,
        'deaths_mean': 0, 'deaths_std': 0, 'deaths_median': 0, 'deaths_iqr': 0,
        'deaths_range': 0, 'deaths_p10': 0, 'deaths_p25': 0, 'deaths_p75': 0, 'deaths_p90': 0,
        'midpoints_mean': 0, 'midpoints_std': 0, 'midpoints_median': 0, 'midpoints_iqr': 0,
        'midpoints_range': 0, 'midpoints_p10': 0, 'midpoints_p25': 0, 'midpoints_p75': 0, 'midpoints_p90': 0,
        'lifespans_mean': 0, 'lifespans_std': 0, 'lifespans_median': 0, 'lifespans_iqr': 0,
        'lifespans_range': 0, 'lifespans_p10': 0, 'lifespans_p25': 0, 'lifespans_p75': 0, 'lifespans_p90': 0,
        'count': 0, 'entropy': 0
    }
    
    if np.size(barcode) > 0:
        # Birth statistics
        births_mean, deaths_mean = np.mean(barcode, axis=0)
        births_std, deaths_std = np.std(barcode, axis=0)
        births_median, deaths_median = np.median(barcode, axis=0)
        births_iqr, deaths_iqr = np.subtract(*np.percentile(barcode, [75, 25], axis=0))
        births_range, deaths_range = np.max(barcode, axis=0) - np.min(barcode, axis=0)
        births_p10, deaths_p10 = np.percentile(barcode, 10, axis=0)
        births_p25, deaths_p25 = np.percentile(barcode, 25, axis=0)
        births_p75, deaths_p75 = np.percentile(barcode, 75, axis=0)
        births_p90, deaths_p90 = np.percentile(barcode, 90, axis=0)

        # Midpoint statistics
        avg_barcodes = (barcode[:, 1] + barcode[:, 0]) / 2
        midpoints_mean = np.mean(avg_barcodes)
        midpoints_std = np.std(avg_barcodes)
        midpoints_median = np.median(avg_barcodes)
        midpoints_iqr = np.subtract(*np.percentile(avg_barcodes, [75, 25]))
        midpoints_range = np.max(avg_barcodes) - np.min(avg_barcodes)
        midpoints_p10 = np.percentile(avg_barcodes, 10)
        midpoints_p25 = np.percentile(avg_barcodes, 25)
        midpoints_p75 = np.percentile(avg_barcodes, 75)
        midpoints_p90 = np.percentile(avg_barcodes, 90)

        # Lifespan statistics
        diff_barcode = np.subtract(
            [i[1] for i in barcode],
            [i[0] for i in barcode]
        )
        diff_barcode = np.absolute(diff_barcode)
        lifespans_mean = np.mean(diff_barcode)
        lifespans_std = np.std(diff_barcode)
        lifespans_median = np.median(diff_barcode)
        lifespans_iqr = np.subtract(*np.percentile(diff_barcode, [75, 25]))
        lifespans_range = np.max(diff_barcode) - np.min(diff_barcode)
        lifespans_p10 = np.percentile(diff_barcode, 10)
        lifespans_p25 = np.percentile(diff_barcode, 25)
        lifespans_p75 = np.percentile(diff_barcode, 75)
        lifespans_p90 = np.percentile(diff_barcode, 90)

        # Count and entropy
        count = len(diff_barcode)
        ent = Entropy()
        entropy = ent.fit_transform([barcode])[0][0]

        feature_dict.update({
            'births_mean': births_mean, 'births_std': births_std, 'births_median': births_median,
            'births_iqr': births_iqr, 'births_range': births_range, 'births_p10': births_p10,
            'births_p25': births_p25, 'births_p75': births_p75, 'births_p90': births_p90,
            'deaths_mean': deaths_mean, 'deaths_std': deaths_std, 'deaths_median': deaths_median,
            'deaths_iqr': deaths_iqr, 'deaths_range': deaths_range, 'deaths_p10': deaths_p10,
            'deaths_p25': deaths_p25, 'deaths_p75': deaths_p75, 'deaths_p90': deaths_p90,
            'midpoints_mean': midpoints_mean, 'midpoints_std': midpoints_std, 'midpoints_median': midpoints_median,
            'midpoints_iqr': midpoints_iqr, 'midpoints_range': midpoints_range, 'midpoints_p10': midpoints_p10,
            'midpoints_p25': midpoints_p25, 'midpoints_p75': midpoints_p75, 'midpoints_p90': midpoints_p90,
            'lifespans_mean': lifespans_mean, 'lifespans_std': lifespans_std, 'lifespans_median': lifespans_median,
            'lifespans_iqr': lifespans_iqr, 'lifespans_range': lifespans_range, 'lifespans_p10': lifespans_p10,
            'lifespans_p25': lifespans_p25, 'lifespans_p75': lifespans_p75, 'lifespans_p90': lifespans_p90,
            'count': count, 'entropy': entropy
        })

    # Replace non-finite values with 0
    for key in feature_dict:
        if not np.isfinite(feature_dict[key]):
            feature_dict[key] = 0

    return feature_dict


def persistence_tropical_coordinates(barcode, r=28):
    """
    Compute tropical coordinate vectorization.
    
    Uses tropical polynomial invariants to create a feature vector.
    
    Parameters
    ----------
    barcode : numpy.ndarray
        Persistence diagram as n x 2 array (birth, death).
    r : float, default=28
        Parameter for tropical coordinate computation.
        
    Returns
    -------
    numpy.ndarray
        Feature vector with 7 tropical coordinate features.
    """
    feature_vector = np.zeros(7)
    if np.size(barcode) > 0:
        # Change the deaths by the lifetime
        new_barcode = deepcopy(barcode)
        new_barcode[:, 1] = new_barcode[:, 1] - new_barcode[:, 0]
        # Sort them so the bars with the longest lifetime appears first
        new_barcode = new_barcode[np.argsort(-new_barcode[:, 1])]
        # Write the output of the selected polynomials
        feature_vector[0] = new_barcode[0, 1]
        if barcode.shape[0] > 1:
            feature_vector[1] = new_barcode[0, 1] + new_barcode[1, 1]
            if barcode.shape[0] > 2:
                feature_vector[2] = new_barcode[0, 1] + new_barcode[1, 1] + new_barcode[2, 1]
                if barcode.shape[0] > 3:
                    feature_vector[3] = new_barcode[0, 1] + new_barcode[1, 1] + new_barcode[2, 1] + new_barcode[3, 1]
        feature_vector[4] = sum(new_barcode[:, 1])
        # In each row, take the minimum between the birth time and r*lifetime
        aux_array = np.array(list(map(lambda x: min(r * x[1], x[0]), new_barcode)))
        feature_vector[5] = sum(aux_array)
        M = max(aux_array + new_barcode[:, 1])
        feature_vector[6] = sum(M - (aux_array + new_barcode[:, 1]))

    return feature_vector


# Mapping of vectorization method names to functions
VECTORIZATION_METHODS = {
    # Preferred snake_case names
    'betti_curve': betti_curve,
    'entropy_summary': entropy_summary,
    'persistence_image': persistence_image,
    'persistence_landscape': persistence_landscape,
    'persistence_lifespan': persistence_lifespan,
    'persistence_silhouette': persistence_silhouette,
    'persistence_stats': persistence_stats,
    'persistence_tropical_coordinates': persistence_tropical_coordinates,
    # Backward compatibility aliases (PascalCase/original names)
    'BettiCurve': betti_curve,
    'EntropySummary': entropy_summary,
    'PersImage': persistence_image,
    'PersLandscape': persistence_landscape,
    'PersLifespan': persistence_lifespan,
    'PersSilhouette': persistence_silhouette,
    'PersStats': persistence_stats,
    'PersTropicalCoordinates': persistence_tropical_coordinates,
}
