"""Differentiable approximate convex hull algorithms."""

import jax
import jax.numpy as jnp
from jax import Array

from ..core.utils import (
    HullVertices,
    PointCloud,
    SamplingMethod,
    generate_direction_vectors,
    remove_duplicate_points,
    validate_point_cloud,
)


def approximate_convex_hull(
    points: PointCloud,
    n_directions: int = 100,
    method: SamplingMethod = "uniform",
    temperature: float = 0.1,
    random_key: Array | None = None,
    remove_duplicates: bool = True,
    tolerance: float = 1e-10,
) -> tuple[HullVertices, Array]:
    """Differentiable approximate convex hull computation.

    This function computes an approximate convex hull using direction vector
    sampling and differentiable soft selection. The approximation is suitable
    for machine learning applications where gradient computation is required.

    Args:
        points: Point cloud with shape (..., n_points, dim)
        n_directions: Number of sampling directions
        method: Sampling strategy ("uniform", "icosphere", "adaptive")
        temperature: Softmax temperature for differentiability control
                    Lower values → more sparse selection (closer to hard argmax)
                    Higher values → more uniform weighting
        random_key: JAX random key (auto-generated if None)
        remove_duplicates: Whether to remove duplicate vertices
        tolerance: Tolerance for duplicate removal

    Returns:
        Tuple of (hull_vertices, hull_indices):
            - hull_vertices: Approximate convex hull vertices
            - hull_indices: Indices of selected points in original array

    Algorithm:
        1. Generate direction vectors using specified sampling method
        2. For each direction, find the "farthest" point using soft selection:
           - Compute dot products (projection scores)
           - Apply softmax to make selection differentiable
           - Compute weighted combination of points
        3. Remove duplicate vertices if requested

    Example:
        >>> import jax.numpy as jnp
        >>> points = jnp.array([[0, 0], [1, 0], [0, 1], [1, 1]])
        >>> hull_vertices, indices = approximate_convex_hull(points, n_directions=20)
        >>> print(hull_vertices.shape)  # (n_hull_vertices, 2)
    """
    # Validate inputs
    points = validate_point_cloud(points)

    if n_directions < 1:
        raise ValueError(f"n_directions must be positive, got {n_directions}")

    if temperature <= 0:
        raise ValueError(f"temperature must be positive, got {temperature}")

    # Get dimensions
    points.shape[:-2]
    points.shape[-2]
    dim = points.shape[-1]

    # Generate random key if not provided
    if random_key is None:
        random_key = jax.random.PRNGKey(0)

    # Generate direction vectors
    directions = generate_direction_vectors(
        dimension=dim, n_directions=n_directions, method=method, random_key=random_key
    )

    # Compute projection scores for all directions
    # Shape: (..., n_points, n_directions)
    scores = jnp.dot(points, jnp.transpose(directions))

    # Apply soft selection for differentiability
    # Shape: (..., n_points, n_directions)
    weights = jax.nn.softmax(scores / temperature, axis=-2)

    # Compute soft hull points as weighted combinations
    # Shape: (..., n_directions, dim)
    soft_hull_points = jnp.sum(weights[..., :, :, None] * points[..., :, None, :], axis=-3)

    # For indices, use hard selection (non-differentiable but needed for indexing)
    hard_indices = jnp.argmax(scores, axis=-2)  # Shape: (..., n_directions)

    # Reshape to standard hull format
    hull_vertices = soft_hull_points
    hull_indices = hard_indices

    # Remove duplicates if requested
    if remove_duplicates:
        hull_vertices, unique_indices = remove_duplicate_points(hull_vertices, tolerance=tolerance)
        # Update indices to reflect unique selection
        hull_indices = hull_indices[..., unique_indices]

    return hull_vertices, hull_indices


def batched_approximate_hull(batch_points: Array, **kwargs) -> tuple[Array, Array]:
    """Batch processing version of approximate_convex_hull.

    Args:
        batch_points: Batched point clouds with shape (batch_size, n_points, dim)
        **kwargs: Arguments passed to approximate_convex_hull

    Returns:
        Tuple of batched (hull_vertices, hull_indices)

    Example:
        >>> batch_points = jnp.array([
        ...     [[0, 0], [1, 0], [0, 1]],  # Triangle 1
        ...     [[0, 0], [2, 0], [0, 2]]   # Triangle 2
        ... ])
        >>> hulls, indices = batched_approximate_hull(batch_points)
    """
    # vmap with explicit axis specification for kwargs
    # Most kwargs (like random_key, n_directions) are not batched (axis=None)
    result = jax.vmap(lambda points: approximate_convex_hull(points, **kwargs), in_axes=0, out_axes=(0, 0))(
        batch_points
    )
    return result  # type: ignore[no-any-return]


def soft_argmax_selection(scores: Array, temperature: float, points: Array) -> tuple[Array, Array]:
    """Differentiable soft selection of extreme points.

    This function replaces the non-differentiable argmax operation with
    a differentiable soft selection using the softmax function.

    Args:
        scores: Selection scores with shape (..., n_points)
        temperature: Softmax temperature parameter
        points: Points corresponding to scores with shape (..., n_points, dim)

    Returns:
        Tuple of (soft_selected_point, selection_weights)

    Mathematical formulation:
        Traditional (non-differentiable):
            idx = argmax(scores)
            selected_point = points[idx]

        Soft (differentiable):
            weights = softmax(scores / temperature)
            selected_point = sum(weights * points)
    """
    # Compute soft selection weights
    weights = jax.nn.softmax(scores / temperature, axis=-1)

    # Compute weighted combination
    soft_point = jnp.sum(weights[..., :, None] * points, axis=-2)

    return soft_point, weights


def adaptive_temperature_control(
    scores: Array, target_sparsity: float = 0.1, min_temperature: float = 0.01, max_temperature: float = 10.0
) -> Array:
    """Adaptive temperature control for soft selection.

    Automatically adjusts the softmax temperature to achieve a target
    sparsity level in the selection weights.

    Args:
        scores: Selection scores
        target_sparsity: Target sparsity (fraction of "active" selections)
        min_temperature: Minimum allowed temperature
        max_temperature: Maximum allowed temperature

    Returns:
        Optimal temperature value

    Note:
        This is a simplified implementation. A full implementation would
        use iterative optimization to find the optimal temperature.
    """
    # Simple heuristic: use score variance to estimate appropriate temperature
    score_std = jnp.std(scores)

    # Higher variance → lower temperature (more confident selection)
    # Lower variance → higher temperature (more uniform selection)
    temperature = jnp.clip(1.0 / (score_std + 1e-6), min_temperature, max_temperature)

    return temperature


# Removed duplicate function - using the improved version below


def multi_resolution_hull(
    points: PointCloud,
    resolution_levels: list | None = None,
    method: SamplingMethod = "uniform",
    random_key: Array | None = None,
) -> list:
    """Compute multi-resolution approximate hulls.

    This function computes multiple approximations with different numbers
    of directions, useful for hierarchical or adaptive algorithms.

    Args:
        points: Input point cloud
        resolution_levels: List of n_directions values
        method: Sampling method for all levels
        random_key: Random key for reproducibility

    Returns:
        List of (hull_vertices, hull_indices) tuples for each resolution

    Example:
        >>> points = jnp.random.normal(jax.random.PRNGKey(0), (100, 3))
        >>> hulls = multi_resolution_hull(points, [20, 50, 100])
        >>> len(hulls)  # 3 different resolutions
        3
    """
    if resolution_levels is None:
        resolution_levels = [50, 100, 200]
    if random_key is None:
        random_key = jax.random.PRNGKey(42)

    hulls = []

    for n_directions in resolution_levels:
        # Use different subkeys for each resolution
        subkey = jax.random.fold_in(random_key, n_directions)

        hull_vertices, hull_indices = approximate_convex_hull(
            points, n_directions=n_directions, method=method, random_key=subkey
        )

        hulls.append((hull_vertices, hull_indices))

    return hulls


def progressive_hull_refinement(
    points: PointCloud,
    initial_directions: int = 20,
    max_directions: int = 200,
    refinement_steps: int = 3,
    convergence_threshold: float = 1e-4,
    random_key: Array | None = None,
) -> tuple[HullVertices, Array, dict]:
    """Progressive refinement of approximate hull.

    Starts with a coarse approximation and progressively refines it
    until convergence or maximum resolution is reached.

    Args:
        points: Input point cloud
        initial_directions: Starting number of directions
        max_directions: Maximum number of directions
        refinement_steps: Number of refinement iterations
        convergence_threshold: Threshold for convergence detection
        random_key: Random key for reproducibility

    Returns:
        Tuple of (final_hull_vertices, final_hull_indices, refinement_info)
    """
    if random_key is None:
        random_key = jax.random.PRNGKey(123)

    current_directions = initial_directions
    previous_hull = None
    refinement_info = {"iterations": [], "converged": False, "final_directions": current_directions}
    from typing import cast

    iterations_list = cast(list, refinement_info["iterations"])

    for step in range(refinement_steps):
        subkey = jax.random.fold_in(random_key, step)

        hull_vertices, hull_indices = approximate_convex_hull(
            points, n_directions=current_directions, random_key=subkey
        )

        # Check convergence if we have a previous hull
        if previous_hull is not None:
            # Simple convergence check: compare hull centers
            current_center = jnp.mean(hull_vertices, axis=-2)
            previous_center = jnp.mean(previous_hull, axis=-2)
            center_change = jnp.linalg.norm(current_center - previous_center)

            iterations_list.append(
                {
                    "step": step,
                    "directions": current_directions,
                    "center_change": center_change,
                    "n_vertices": hull_vertices.shape[-2],
                }
            )

            if center_change < convergence_threshold:
                refinement_info["converged"] = True
                break

        previous_hull = hull_vertices

        # Increase resolution for next iteration
        current_directions = min(current_directions * 2, max_directions)
        if current_directions >= max_directions:
            break

    refinement_info["final_directions"] = current_directions

    return hull_vertices, hull_indices, refinement_info


# JIT-compiled versions for performance
approximate_convex_hull_jit = jax.jit(approximate_convex_hull, static_argnames=["method", "remove_duplicates"])


# =============================================================================
# PHASE 2: IMPROVED ALGORITHM WITH VERTEX COUNT CONSTRAINT
# =============================================================================


def improved_approximate_convex_hull(
    points: PointCloud,
    n_directions: int = 100,
    method: SamplingMethod = "uniform",
    temperature: float = 0.1,
    random_key: Array | None = None,
    max_vertices: int | None = None,
) -> tuple[HullVertices, Array]:
    """Improved differentiable convex hull with vertex count constraint.

    This algorithm addresses the mathematical issues in the original implementation
    by using a staged selection method that ensures output vertices ≤ input vertices.

    Stages:
    1. Coarse boundary detection to identify candidate vertices
    2. Differentiable refinement with constraints
    3. Hard vertex limit enforcement

    Args:
        points: Point cloud with shape (..., n_points, dim)
        n_directions: Number of direction vectors for sampling
        method: Direction sampling method
        temperature: Temperature for soft selection (lower = more selective)
        random_key: JAX random key for reproducibility
        max_vertices: Maximum output vertices (default: input vertex count)

    Returns:
        Tuple of (hull_vertices, hull_indices) where hull_vertices.shape[0] <= points.shape[0]

    Mathematical guarantee:
        len(output_vertices) <= len(input_points)
    """
    points = validate_point_cloud(points)
    n_points, _dim = points.shape[-2], points.shape[-1]

    max_vertices_int: int = n_points if max_vertices is None else min(max_vertices, n_points)

    if random_key is None:
        random_key = jax.random.PRNGKey(0)

    # Stage 1: Coarse boundary detection
    boundary_candidates, candidate_scores = _stage1_coarse_boundary_detection(points, n_directions, method, random_key)

    # Stage 2: Differentiable refinement with constraints
    refined_vertices, refined_indices = _stage2_constrained_refinement(
        points, boundary_candidates, candidate_scores, temperature, max_vertices_int
    )

    # Stage 3: Hard vertex limit enforcement
    final_vertices, final_indices = _stage3_enforce_vertex_limit(refined_vertices, refined_indices, max_vertices_int)

    return final_vertices, final_indices


def _stage1_coarse_boundary_detection(
    points: Array, n_directions: int, method: SamplingMethod, random_key: Array
) -> tuple[Array, Array]:
    """Stage 1: Coarse boundary detection using direction-based scoring.

    Identifies points that are likely to be on the convex hull boundary
    by checking how often they are extreme in various directions.

    Args:
        points: Input point cloud
        n_directions: Number of directions to test
        method: Direction sampling method
        random_key: Random key

    Returns:
        Tuple of (candidate_points, boundary_scores)
    """
    n_points, dim = points.shape

    # Generate direction vectors
    directions = generate_direction_vectors(dim, n_directions, method, random_key)

    # For each direction, count how often each point is "most extreme"
    boundary_scores = jnp.zeros(n_points)

    for i in range(directions.shape[0]):
        direction = directions[i]
        # Project points onto direction
        projections = jnp.dot(points, direction)

        # Find the point(s) with maximum projection
        max_projection = jnp.max(projections)

        # Give points that are exactly at the maximum a score of 1
        # This is more precise than soft scoring for boundary detection
        is_extreme = jnp.abs(projections - max_projection) < 1e-6
        boundary_scores = boundary_scores + is_extreme.astype(jnp.float32)

    # Normalize by number of directions tested
    boundary_scores = boundary_scores / n_directions

    return points, boundary_scores


def _stage2_constrained_refinement(
    points: Array, boundary_candidates: Array, candidate_scores: Array, temperature: float, max_vertices: int
) -> tuple[Array, Array]:
    """Stage 2: Differentiable refinement with vertex count constraints.

    Selects points with high boundary scores, but only those that actually
    contribute to the hull boundary.

    Args:
        points: Original input points
        boundary_candidates: Candidate points from stage 1
        candidate_scores: Boundary likelihood scores
        temperature: Soft selection temperature (currently unused to maintain exactness)
        max_vertices: Maximum number of vertices to select

    Returns:
        Tuple of (refined_vertices, refined_indices)
    """
    n_points = points.shape[0]

    # Only select points that have non-zero boundary scores
    # This filters out interior points that never appear as extreme points
    has_boundary_score = candidate_scores > 1e-6
    boundary_point_indices = jnp.where(has_boundary_score, size=n_points, fill_value=-1)[0]

    # Filter out the fill values (-1)
    valid_boundary_indices = boundary_point_indices[boundary_point_indices >= 0]

    # Limit to max_vertices
    n_boundary_points = jnp.sum(has_boundary_score).astype(int)
    n_selected = min(int(n_boundary_points), max_vertices)

    if n_selected == 0:
        # Fallback: if no boundary points found, take the first point
        selected_indices = jnp.array([0])
        selected_vertices = points[selected_indices]
    else:
        # Take the top n_selected boundary points
        selected_indices = valid_boundary_indices[:n_selected]
        selected_vertices = points[selected_indices]

    return selected_vertices, selected_indices


def _stage3_enforce_vertex_limit(vertices: Array, indices: Array, max_vertices: int) -> tuple[Array, Array]:
    """Stage 3: Hard enforcement of vertex count constraint.

    Final safety check to ensure we never exceed the vertex limit.

    Args:
        vertices: Candidate vertices from stage 2
        indices: Corresponding indices
        max_vertices: Hard limit on vertex count

    Returns:
        Tuple of (final_vertices, final_indices) with guaranteed constraint satisfaction
    """
    n_vertices = vertices.shape[0]

    if n_vertices <= max_vertices:
        return vertices, indices

    # If we somehow exceeded the limit, take the first max_vertices
    # In practice, this shouldn't happen due to stage 2 constraints
    return vertices[:max_vertices], indices[:max_vertices]


def compute_hull_quality_metrics(hull_vertices: Array, original_points: Array) -> dict[str, float]:
    """Compute quality metrics for hull approximation.

    Args:
        hull_vertices: Computed hull vertices
        original_points: Original input points

    Returns:
        Dictionary of quality metrics
    """
    n_hull = hull_vertices.shape[0]
    n_original = original_points.shape[0]

    # Vertex count ratio (should be ≤ 1.0)
    vertex_ratio = n_hull / n_original

    # Coverage: fraction of original points that are hull vertices
    # (This is a simplified metric - true coverage would check containment)
    exact_matches = 0
    for hull_vertex in hull_vertices:
        for original_point in original_points:
            if jnp.allclose(hull_vertex, original_point, atol=1e-6):
                exact_matches += 1
                break

    coverage = exact_matches / n_original

    # Boundary efficiency: how well we use our vertex budget
    boundary_efficiency = exact_matches / n_hull if n_hull > 0 else 0.0

    return {
        "vertex_count_ratio": float(vertex_ratio),
        "coverage": float(coverage),
        "boundary_efficiency": float(boundary_efficiency),
        "n_hull_vertices": int(n_hull),
        "n_original_points": int(n_original),
        "constraint_satisfied": bool(vertex_ratio <= 1.0),
    }


# JIT-compiled version of improved algorithm
improved_approximate_convex_hull_jit = jax.jit(improved_approximate_convex_hull, static_argnames=["method"])
