"""Geomechanics: Parallel processing utilities

Pure geomechanics operations - parallel module.
Migrated from geosuite.geomech.
Layer 2: Primitives - Pure operations.
"""

from typing import TYPE_CHECKING, Any, Dict, List, Optional, Tuple, Union

import numpy as np

from geosmith.primitives.geomechanics._common import (
    NUMBA_AVAILABLE,
    PANDAS_AVAILABLE,
    njit,
    pd,
)

if TYPE_CHECKING:
    import pandas as pd

# Import prange for parallel execution (from numba if available)
try:
    from numba import prange
except ImportError:
    prange = range  # type: ignore

from geosmith.primitives.geomechanics.pressure import sv_from_density


def calculate_overburden_stress_parallel(
    depths_list: List[Union[np.ndarray, "pd.Series"]],
    rhobs_list: List[Union[np.ndarray, "pd.Series"]],
    g: float = 9.80665,
) -> List[np.ndarray]:
    """Calculate overburden stress for multiple wells in parallel.

    This function uses Numba's parallel execution to process multiple wells
    simultaneously.

    Args:
        depths_list: List of depth arrays (meters), one per well.
        rhobs_list: List of bulk density arrays (g/cc), one per well.
        g: Gravitational acceleration (m/s²), default 9.80665.

    Returns:
        List of overburden stress arrays (MPa), one per well.

    Example:
        >>> from geosmith.primitives.geomechanics import (
        ...     calculate_overburden_stress_parallel
        ... )
        >>> import numpy as np
        >>>
        >>> # Process 10 wells in parallel
        >>> depths = [np.linspace(0, 3000, 1000) for _ in range(10)]
        >>> rhobs = [np.ones(1000) * 2.5 for _ in range(10)]
        >>> sv_list = calculate_overburden_stress_parallel(depths, rhobs)
        >>> print(f"Processed {len(sv_list)} wells")
    """
    if not PANDAS_AVAILABLE:
        raise ImportError(
            "pandas is required for calculate_overburden_stress_parallel. "
            "Install with: pip install pandas"
        )

    # Convert to numpy arrays
    depths_arrays = [
        np.asarray(d, dtype=np.float64) if not isinstance(d, pd.Series) else d.values
        for d in depths_list
    ]
    rhobs_arrays = [
        np.asarray(r, dtype=np.float64) if not isinstance(r, pd.Series) else r.values
        for r in rhobs_list
    ]

    if len(depths_arrays) != len(rhobs_arrays):
        raise ValueError("depths_list and rhobs_list must have same length")

    # Note: Numba doesn't support list of arrays directly for parallel processing
    # For now, fall back to sequential processing using sv_from_density
    results = []
    for depth, rhob in zip(depths_arrays, rhobs_arrays):
        sv = sv_from_density(depth, rhob, g=g)
        results.append(sv)
    return results


def process_well_array_parallel(
    data_array: np.ndarray,
    operation: str = "mean",
) -> np.ndarray:
    """Apply statistical operation across multiple wells in parallel.

    This is a general-purpose parallel processor for well data arrays.

    Args:
        data_array: 2D array (n_wells x n_samples).
        operation: Statistical operation - 'mean', 'median', 'std', 'min',
            'max', default 'mean'.

    Returns:
        1D array of results, one per well.

    Example:
        >>> from geosmith.primitives.geomechanics import process_well_array_parallel
        >>> import numpy as np
        >>>
        >>> # Calculate mean GR for 100 wells
        >>> gr_data = np.random.normal(60, 15, (100, 1000))
        >>> means = process_well_array_parallel(gr_data, 'mean')
        >>> print(f"Mean GR per well: {means.mean():.1f} ± {means.std():.1f} API")
    """
    data_array = np.asarray(data_array, dtype=np.float64)

    if data_array.ndim != 2:
        raise ValueError("data_array must be 2D (n_wells x n_samples)")

    n_wells = data_array.shape[0]
    results = np.zeros(n_wells, dtype=np.float64)

    if NUMBA_AVAILABLE:
        # Numba-accelerated parallel version
        @njit(parallel=True, cache=True)
        def _parallel_kernel(data: np.ndarray, op: str, n: int) -> np.ndarray:
            res = np.zeros(n, dtype=np.float64)
            for i in prange(n):
                well_data = data[i]
                if op == "mean":
                    res[i] = np.mean(well_data)
                elif op == "std":
                    res[i] = np.std(well_data)
                elif op == "min":
                    res[i] = np.min(well_data)
                elif op == "max":
                    res[i] = np.max(well_data)
                else:
                    # Fallback for operations not supported in numba
                    res[i] = np.mean(well_data)
            return res

        # Convert operation to numeric code for numba
        op_codes = {"mean": "mean", "std": "std", "min": "min", "max": "max"}
        op_code = op_codes.get(operation, "mean")

        results = _parallel_kernel(data_array, op_code, n_wells)
    else:
        # Sequential fallback
        for i in range(n_wells):
            well_data = data_array[i]
            if operation == "mean":
                results[i] = np.mean(well_data)
            elif operation == "median":
                results[i] = np.median(well_data)
            elif operation == "std":
                results[i] = np.std(well_data)
            elif operation == "min":
                results[i] = np.min(well_data)
            elif operation == "max":
                results[i] = np.max(well_data)
            else:
                raise ValueError(
                    f"Unknown operation: {operation}. "
                    "Use 'mean', 'median', 'std', 'min', or 'max'"
                )

    return results


def get_parallel_info() -> dict[str, Union[bool, str, int]]:
    """Get information about parallel processing capabilities.

    Returns:
        Dictionary with Numba availability and threading info:
            - 'numba_available': Whether Numba is available
            - 'parallel_enabled': Whether parallel processing is enabled
            - 'num_threads': Number of threads (if available)
            - 'threading_layer': Threading layer type
    """
    info: dict[str, Union[bool, str, int]] = {
        "numba_available": NUMBA_AVAILABLE,
        "parallel_enabled": NUMBA_AVAILABLE,
    }

    if NUMBA_AVAILABLE:
        try:
            from numba import config as numba_config

            if numba_config is not None:
                info["num_threads"] = getattr(
                    numba_config, "NUMBA_NUM_THREADS", "unknown"
                )
                info["threading_layer"] = getattr(
                    numba_config, "THREADING_LAYER", "unknown"
                )
            else:
                info["num_threads"] = "unknown"
                info["threading_layer"] = "unknown"
        except (ImportError, AttributeError, TypeError):
            info["num_threads"] = "unknown"
            info["threading_layer"] = "unknown"
    else:
        info["num_threads"] = 1
        info["threading_layer"] = "sequential"

    return info
