import numpy as np
from skimage import measure
from scipy.spatial import ConvexHull, distance_matrix


def shape_2d(mask, spacing=(1.0, 1.0)):
    """
    Compute 2D shape features strictly matching the PyRadiomics 'shape2D' module.

    This function calculates shape descriptors based on the perimeter, area, 
    and principal axes of the Region of Interest (ROI). It aligns with the 
    definitions provided in the PyRadiomics documentation.

    Ref: https://pyradiomics.readthedocs.io/en/latest/features.html#radiomics-shape2d

    Parameters
    ----------
    mask : np.ndarray
        A 2D binary mask array where non-zero pixels represent the ROI.
        Format should be (Row, Col) or (Y, X).
    spacing : tuple of float, optional
        Pixel spacing in physical base units (row_spacing, col_spacing).
        Defaults to (1.0, 1.0).

    Returns
    -------
    dict
        A dictionary containing the 10 standard PyRadiomics 2D shape features:
        
        - **MeshSurface**: Area of the ROI defined by the mesh (marching squares).
        - **PixelSurface**: Area defined by the count of non-zero pixels.
        - **Perimeter**: Perimeter length of the ROI mesh.
        - **PerimeterSurfaceRatio**: Ratio of Perimeter to MeshSurface.
        - **Sphericity**: Measure of roundness (1.0 is a perfect circle).
        - **SphericalDisproportion**: Inverse of sphericity.
        - **MaximumDiameter**: Largest Euclidean distance between contour vertices.
        - **MajorAxisLength**: Principal axis length derived from image moments.
        - **MinorAxisLength**: Secondary axis length derived from image moments.
        - **Elongation**: Ratio of Minor to Major axis length.

    Raises
    ------
    ValueError
        If the input `mask` contains no non-zero pixels.

    Examples
    --------
    >>> import numpy as np
    >>> # Create a simple 2D square (10x10) inside a 20x20 grid
    >>> mask_2d = np.zeros((20, 20), dtype=int)
    >>> mask_2d[5:15, 5:15] = 1
    >>> spacing = (1.0, 1.0) # 1mm isotropic
    
    >>> # Calculate features
    >>> feats = shape_2d(mask_2d, spacing)
    
    >>> print(f"Pixel Area: {feats['PixelSurface']}")
    Pixel Area: 100.0
    >>> print(f"Elongation: {feats['Elongation']:.2f}")
    Elongation: 1.00
    """
    # 1. Setup
    mask = mask > 0
    if not np.any(mask):
        raise ValueError("Mask contains no pixels.")
    
    # 2. Pixel Surface (Area based on voxel count)
    #    PyRadiomics Definition: N_pixels * pixel_area
    pixel_area_mm2 = spacing[0] * spacing[1]
    pixel_surface = np.sum(mask) * pixel_area_mm2

    # 3. Mesh Surface & Perimeter (Marching Squares)
    #    PyRadiomics calculates perimeter and "Mesh Surface" using the
    #    contour generated by marching squares (level=0.5).
    contours = measure.find_contours(mask, 0.5)
    
    perimeter = 0.0
    mesh_surface = 0.0
    
    all_contour_points = []

    for contour in contours:
        # Scale contour to physical units (row*dy, col*dx)
        # contour is (row, col)
        scaled_c = contour * np.array(spacing)
        all_contour_points.append(scaled_c)
        
        # A. Perimeter: Sum of Euclidean distances between segments
        #    Close the loop for perimeter calculation? 
        #    measure.find_contours returns open segments or closed loops. 
        #    PyRadiomics treats them as closed for area, but sums segment lengths for perimeter.
        #    We assume the loop is closed if it's a complete shape.
        dist = np.linalg.norm(scaled_c - np.roll(scaled_c, 1, axis=0), axis=1)
        perimeter += np.sum(dist)
        
        # B. Mesh Surface: Shoelace Formula
        #    Area = 0.5 * |sum(x_i * y_{i+1} - x_{i+1} * y_i)|
        #    Note: This area is often slightly smaller than PixelSurface
        y = scaled_c[:, 0]
        x = scaled_c[:, 1]
        region_area = 0.5 * np.abs(np.dot(x, np.roll(y, 1)) - np.dot(y, np.roll(x, 1)))
        mesh_surface += region_area

    # 4. Maximum Diameter (2D)
    #    Largest Euclidean distance between mesh vertices (using Convex Hull optimization)
    if len(all_contour_points) > 0:
        pts_physical = np.vstack(all_contour_points)
        if len(pts_physical) > 2:
            try:
                hull = ConvexHull(pts_physical)
                hull_verts = pts_physical[hull.vertices]
                if len(hull_verts) > 1:
                    max_diameter = np.max(distance_matrix(hull_verts, hull_verts))
                else:
                    max_diameter = 0.0
            except:
                # Fallback if convex hull fails (e.g. collinear points)
                max_diameter = np.max(distance_matrix(pts_physical, pts_physical))
        else:
            max_diameter = 0.0
    else:
        max_diameter = 0.0

    # 5. Major & Minor Axis Length (PCA on ROI Pixels)
    #    PyRadiomics uses the SOLID pixel mask (not the contour) for axes.
    r, c = np.where(mask)
    pixel_pts_physical = np.vstack([r * spacing[0], c * spacing[1]]).T
    
    if len(pixel_pts_physical) > 1:
        cov = np.cov(pixel_pts_physical.T)
        evals, _ = np.linalg.eigh(cov)
        evals = np.sort(evals) # Sort: [Smallest, Largest]
        
        # PyRadiomics Definition: 4 * sqrt(eigenvalue)
        minor_axis = 4 * np.sqrt(max(evals[0], 0))
        major_axis = 4 * np.sqrt(max(evals[1], 0))
    else:
        minor_axis = 0.0
        major_axis = 0.0

    # 6. Derived Features
    
    # Elongation
    # PyRadiomics: sqrt(minor_eigen / major_eigen) == minor_axis / major_axis
    elongation = minor_axis / major_axis if major_axis > 0 else 1.0

    # PerimeterSurfaceRatio
    perimeter_surface_ratio = perimeter / mesh_surface if mesh_surface > 0 else 0.0

    # Sphericity
    # PyRadiomics: (2 * sqrt(pi * A)) / P
    # Range 0-1 (1 is a perfect circle)
    if perimeter > 0:
        sphericity = (2 * np.sqrt(np.pi * mesh_surface)) / perimeter
    else:
        sphericity = 0.0

    # Spherical Disproportion
    # PyRadiomics: P / (2 * sqrt(pi * A))
    # Inverse of Sphericity. Minimum 1.0.
    if mesh_surface > 0:
        spherical_disproportion = perimeter / (2 * np.sqrt(np.pi * mesh_surface))
    else:
        spherical_disproportion = 0.0

    return {
        "MeshSurface": mesh_surface,
        "PixelSurface": pixel_surface,
        "Perimeter": perimeter,
        "PerimeterSurfaceRatio": perimeter_surface_ratio,
        "Sphericity": sphericity,
        "SphericalDisproportion": spherical_disproportion,
        "MaximumDiameter": max_diameter,
        "MajorAxisLength": major_axis,
        "MinorAxisLength": minor_axis,
        "Elongation": elongation
    }


def shape_2d_units(base_unit="mm"):
    """
    Return the physical units for the 2D shape metrics.

    Parameters
    ----------
    base_unit : str, optional
        The string representation of the physical length unit (e.g., 'mm', 'cm').
        Defaults to 'mm'.

    Returns
    -------
    dict
        A dictionary mapping feature names to unit strings.

    Examples
    --------
    >>> units = shape_2d_units('mm')
    >>> print(units['MeshSurface'])
    'mm^2'
    >>> print(units['Perimeter'])
    'mm'
    """
    return {
        "Elongation": "",                         # Dimensionless (ratio)
        "MajorAxisLength": base_unit,             # Length
        "MaximumDiameter": base_unit,             # Length
        "MeshSurface": f"{base_unit}^2",          # Area
        "MinorAxisLength": base_unit,             # Length
        "Perimeter": base_unit,                   # Length
        "PerimeterSurfaceRatio": f"{base_unit}^-1", # Perimeter / Area (L / L^2 = 1/L)
        "PixelSurface": f"{base_unit}^2",         # Area
        "SphericalDisproportion": "",             # Dimensionless (ratio)
        "Sphericity": "",                         # Dimensionless (ratio)
    }

