"""
Log Feature Engineering Module

Multi-well feature engineering for machine learning on well logs.
Creates features that capture both pointwise log values and contextual patterns.

Implements feature extraction strategies described in automated well log
interpretation workflows.
"""

from collections.abc import Callable
from dataclasses import dataclass

import numpy as np
import pandas as pd
from scipy.ndimage import gaussian_filter1d
from scipy.spatial.distance import cdist

try:
    from numba import njit

    NUMBA_AVAILABLE = True
except ImportError:
    NUMBA_AVAILABLE = False

    def njit(*args, **kwargs):
        def decorator(func):
            return func

        return decorator if not args else decorator(args[0])


@njit(cache=True, fastmath=True)
def _compute_weighted_average_fast(
    offset_array: np.ndarray,
    weights: np.ndarray,
    null_value: float,
) -> tuple[np.ndarray, np.ndarray]:
    """Numba-accelerated inverse-distance weighted averaging."""
    n_wells = offset_array.shape[0]
    n_depths = offset_array.shape[1]

    weighted_avg = np.full(n_depths, null_value)
    weighted_std = np.full(n_depths, null_value)

    for i in range(n_depths):
        # Find valid (non-null) values at this depth
        valid_count = 0
        valid_sum = 0.0
        weight_sum = 0.0

        for j in range(n_wells):
            if offset_array[j, i] != null_value:
                valid_count += 1
                valid_sum += weights[j] * offset_array[j, i]
                weight_sum += weights[j]

        if valid_count > 0 and weight_sum > 0:
            weighted_avg[i] = valid_sum / weight_sum

            # Compute weighted standard deviation
            if valid_count > 1:
                var_sum = 0.0
                for j in range(n_wells):
                    if offset_array[j, i] != null_value:
                        diff = offset_array[j, i] - weighted_avg[i]
                        var_sum += weights[j] * diff * diff
                weighted_std[i] = np.sqrt(var_sum / weight_sum)

    return weighted_avg, weighted_std


@dataclass
class FeatureSet:
    """Container for engineered features with metadata."""

    features: pd.DataFrame
    feature_names: list[str]
    feature_groups: dict[str, list[str]]  # Group name -> feature list
    description: str

    def __str__(self) -> str:
        return (
            f"Feature Set: {self.description}\n"
            f"  Total features: {len(self.feature_names)}\n"
            f"  Feature groups: {list(self.feature_groups.keys())}\n"
            f"  Samples: {len(self.features)}"
        )


class LogFeatureEngineer:
    """
    Feature engineering for well log machine learning.

    Creates derived features from raw log curves including:
    - Curve derivatives (rate of change)
    - Cross-curve ratios (petrophysical indicators)
    - Rolling window statistics (local context)
    - Spatial features from offset wells
    """

    def __init__(self, null_value: float = -999.25):
        """
        Initialize feature engineer.

        Args:
            null_value: Value indicating missing data
        """
        self.null_value = null_value
        self.feature_groups = {}

    def compute_derivatives(
        self,
        data: pd.DataFrame,
        curves: list[str] | None = None,
        method: str = "gradient",
        smooth_sigma: float = 2.0,
    ) -> pd.DataFrame:
        """
        Compute derivatives (rate of change) for log curves.

        Derivatives are useful for detecting boundaries and transitions.

        Args:
            data: DataFrame with log curves
            curves: List of curves to process (None = all)
            method: 'gradient' (numerical) or 'savgol' (Savitzky-Golay filter)
            smooth_sigma: Smoothing parameter for Gaussian filter

        Returns:
            DataFrame with derivative curves (suffix '_DERIV')
        """
        if curves is None:
            curves = [
                col
                for col in data.columns
                if not col.endswith(("_DERIV", "_RATIO", "_ROLL"))
            ]

        derivatives = pd.DataFrame(index=data.index)
        depth_step = np.median(np.diff(data.index))

        for curve in curves:
            if curve not in data.columns:
                continue

            values = data[curve].values.copy()
            valid_mask = values != self.null_value

            if valid_mask.sum() < 5:
                derivatives[f"{curve}_DERIV"] = self.null_value
                continue

            # Apply Gaussian smoothing first
            smoothed = values.copy()
            smoothed[~valid_mask] = np.nan
            smoothed = pd.Series(smoothed).interpolate(limit=5).values
            smoothed = gaussian_filter1d(smoothed, sigma=smooth_sigma)

            # Compute derivative
            if method == "gradient":
                deriv = np.gradient(smoothed, depth_step)
            elif method == "savgol":
                from scipy.signal import savgol_filter

                # window_length must be odd
                window_length = (
                    int(5 * smooth_sigma)
                    if int(5 * smooth_sigma) % 2 == 1
                    else int(5 * smooth_sigma) + 1
                )
                window_length = max(5, min(window_length, len(smoothed) // 2))
                deriv = savgol_filter(
                    smoothed,
                    window_length=window_length,
                    polyorder=2,
                    deriv=1,
                    delta=depth_step,
                )
            else:
                raise ValueError(
                    f"Unknown derivative method: {method}",
                    "Valid methods: 'gradient', 'savgol'",
                )

            # Restore null values
            deriv[~valid_mask] = self.null_value
            derivatives[f"{curve}_DERIV"] = deriv

        self.feature_groups["derivatives"] = list(derivatives.columns)
        return derivatives

    def compute_ratios(
        self,
        data: pd.DataFrame,
        ratio_definitions: dict[str, tuple[str, str, Callable]] | None = None,
    ) -> pd.DataFrame:
        """
        Compute cross-curve ratios (petrophysical indicators).

        Args:
            data: DataFrame with log curves
            ratio_definitions: Dict of {name: (numerator, denominator, transform_func)}
                              If None, uses standard petrophysical ratios

        Returns:
            DataFrame with ratio features
        """
        if ratio_definitions is None:
            ratio_definitions = self._get_standard_ratios()

        ratios = pd.DataFrame(index=data.index)

        for ratio_name, (num_curve, den_curve, transform) in ratio_definitions.items():
            if num_curve not in data.columns or den_curve not in data.columns:
                continue

            num_values = data[num_curve].values.copy()
            den_values = data[den_curve].values.copy()

            # Handle null values
            valid_mask = (
                (num_values != self.null_value)
                & (den_values != self.null_value)
                & (den_values != 0)
            )

            ratio_values = np.full(len(data), self.null_value)

            if valid_mask.sum() > 0:
                # Compute ratio with transform function
                ratio_values[valid_mask] = transform(
                    num_values[valid_mask], den_values[valid_mask]
                )

            ratios[ratio_name] = ratio_values

        self.feature_groups["ratios"] = list(ratios.columns)
        return ratios

    def _get_standard_ratios(self) -> dict[str, tuple[str, str, Callable]]:
        """Standard petrophysical ratio definitions."""
        return {
            "VSH_GR": ("GR", "GR", lambda gr, _: self._gamma_ray_to_vshale(gr)),
            "PORO_DENSITY": (
                "RHOB",
                "RHOB",
                lambda rhob, _: self._density_to_porosity(rhob),
            ),
            "RESISTIVITY_RATIO": ("ILD", "ILM", lambda deep, med: np.log10(deep / med)),
            "NPHI_RHOB_XPLOT": ("NPHI", "RHOB", lambda nphi, rhob: nphi * rhob),
            "DT_DENSITY": ("DT", "RHOB", lambda dt, rhob: dt / rhob),
        }

    def _gamma_ray_to_vshale(
        self, gr: np.ndarray, gr_clean: float = 20, gr_shale: float = 150
    ) -> np.ndarray:
        """Convert gamma ray to volume of shale using linear method."""
        vsh = (gr - gr_clean) / (gr_shale - gr_clean)
        return np.clip(vsh, 0, 1)

    def _density_to_porosity(
        self,
        rhob: np.ndarray,
        rho_matrix: float = 2.65,
        rho_fluid: float = 1.0,
    ) -> np.ndarray:
        """Convert bulk density to porosity."""
        phi = (rho_matrix - rhob) / (rho_matrix - rho_fluid)
        return np.clip(phi, 0, 1)

    def compute_rolling_statistics(
        self,
        data: pd.DataFrame,
        curves: list[str] | None = None,
        window_sizes: list[int] = [5, 10, 20],
        statistics: list[str] = ["mean", "std", "min", "max"],
    ) -> pd.DataFrame:
        """
        Compute rolling window statistics.

        Captures local context and variability patterns.

        Args:
            data: DataFrame with log curves
            curves: List of curves to process (None = all)
            window_sizes: List of window sizes (in depth samples)
            statistics: List of statistics: 'mean', 'std', 'min', 'max', 'median'

        Returns:
            DataFrame with rolling statistics features
        """
        if curves is None:
            curves = [
                col
                for col in data.columns
                if not col.endswith(("_DERIV", "_RATIO", "_ROLL"))
            ]

        rolling_stats = pd.DataFrame(index=data.index)

        for curve in curves:
            if curve not in data.columns:
                continue

            # Replace null with NaN for rolling calculations
            series = data[curve].replace(self.null_value, np.nan)

            for window in window_sizes:
                rolling = series.rolling(window=window, center=True, min_periods=1)

                for stat in statistics:
                    if stat == "mean":
                        values = rolling.mean()
                    elif stat == "std":
                        values = rolling.std()
                    elif stat == "min":
                        values = rolling.min()
                    elif stat == "max":
                        values = rolling.max()
                    elif stat == "median":
                        values = rolling.median()
                    else:
                        continue

                    col_name = f"{curve}_ROLL{window}_{stat.upper()}"
                    rolling_stats[col_name] = values.fillna(self.null_value)

        self.feature_groups["rolling_statistics"] = list(rolling_stats.columns)
        return rolling_stats

    def compute_spatial_features(
        self,
        target_well: pd.DataFrame,
        offset_wells: dict[str, pd.DataFrame],
        well_locations: dict[str, tuple[float, float]],
        target_location: tuple[float, float],
        max_distance: float = 5000,  # meters
        curves: list[str] | None = None,
    ) -> pd.DataFrame:
        """
        Compute spatial features from offset wells.

        Uses inverse-distance weighting to estimate expected values at target
        location based on nearby wells.

        Args:
            target_well: Target well DataFrame with depth index
            offset_wells: Dict of offset well DataFrames
            well_locations: Dict of {well_name: (x, y)} coordinates
            target_location: (x, y) coordinates of target well
            max_distance: Maximum distance for offset well influence (meters)
            curves: List of curves to use (None = all common curves)

        Returns:
            DataFrame with spatial features (suffix '_SPATIAL')
        """
        if curves is None:
            # Find common curves across all wells
            all_curves = set(target_well.columns)
            for well_data in offset_wells.values():
                all_curves &= set(well_data.columns)
            curves = [
                c
                for c in all_curves
                if not c.endswith(("_DERIV", "_RATIO", "_ROLL", "_SPATIAL"))
            ]

        # Compute distances
        target_xy = np.array(target_location).reshape(1, -1)
        offset_names = []
        offset_coords = []

        for well_name, (x, y) in well_locations.items():
            if well_name in offset_wells:
                offset_names.append(well_name)
                offset_coords.append([x, y])

        if not offset_coords:
            # No offset wells, return empty DataFrame
            return pd.DataFrame(index=target_well.index)

        distances = cdist(target_xy, np.array(offset_coords))[0]

        # Filter by max distance
        valid_idx = distances <= max_distance
        if not valid_idx.any():
            return pd.DataFrame(index=target_well.index)

        distances = distances[valid_idx]
        offset_names = [name for name, valid in zip(offset_names, valid_idx) if valid]

        # Compute inverse-distance weights
        weights = 1 / (distances + 1)  # +1 to avoid division by zero
        weights = weights / weights.sum()

        # Compute weighted averages
        spatial_features = pd.DataFrame(index=target_well.index)

        for curve in curves:
            # Interpolate offset well curves to target well depths
            offset_values = []

            for well_name in offset_names:
                well_data = offset_wells[well_name]
                if curve not in well_data.columns:
                    continue

                # Interpolate to target depths
                interp_values = np.interp(
                    target_well.index,
                    well_data.index,
                    well_data[curve],
                    left=self.null_value,
                    right=self.null_value,
                )
                offset_values.append(interp_values)

            if not offset_values:
                continue

            # Compute weighted average using Numba acceleration
            offset_array = np.array(offset_values)

            if NUMBA_AVAILABLE:
                weighted_avg, spatial_std = _compute_weighted_average_fast(
                    offset_array, weights, self.null_value
                )
            else:
                # Fallback to pure Python
                valid_mask = offset_array != self.null_value
                weighted_avg = np.full(len(target_well), self.null_value)
                spatial_std = np.full(len(target_well), self.null_value)

                for i in range(len(target_well)):
                    col_valid = valid_mask[:, i]
                    if col_valid.sum() > 0:
                        col_weights = weights[col_valid] / weights[col_valid].sum()
                        weighted_avg[i] = np.average(
                            offset_array[col_valid, i], weights=col_weights
                        )
                    if col_valid.sum() > 1:
                        spatial_std[i] = np.std(offset_array[col_valid, i])

            spatial_features[f"{curve}_SPATIAL"] = weighted_avg
            spatial_features[f"{curve}_SPATIAL_STD"] = spatial_std

        self.feature_groups["spatial"] = list(spatial_features.columns)
        return spatial_features

    def create_feature_set(
        self,
        data: pd.DataFrame,
        include_derivatives: bool = True,
        include_ratios: bool = True,
        include_rolling_stats: bool = True,
        include_spatial: bool = False,
        offset_wells: dict[str, pd.DataFrame] | None = None,
        well_locations: dict[str, tuple[float, float]] | None = None,
        target_location: tuple[float, float] | None = None,
        **kwargs,
    ) -> FeatureSet:
        """
        Create complete feature set for machine learning.

        Args:
            data: Raw log data
            include_derivatives: Add derivative features
            include_ratios: Add ratio features
            include_rolling_stats: Add rolling statistics
            include_spatial: Add spatial features from offset wells
            offset_wells: Dict of offset well data (required if include_spatial=True)
            well_locations: Dict of well locations (required if include_spatial=True)
            target_location: Target well location (required if include_spatial=True)
            **kwargs: Additional arguments for specific feature functions

        Returns:
            FeatureSet with all requested features
        """
        features = data.copy()
        self.feature_groups = {"original": list(data.columns)}

        if include_derivatives:
            derivatives = self.compute_derivatives(
                data, **kwargs.get("derivative_params", {})
            )
            features = pd.concat([features, derivatives], axis=1)

        if include_ratios:
            ratios = self.compute_ratios(data, **kwargs.get("ratio_params", {}))
            features = pd.concat([features, ratios], axis=1)

        if include_rolling_stats:
            rolling = self.compute_rolling_statistics(
                data, **kwargs.get("rolling_params", {})
            )
            features = pd.concat([features, rolling], axis=1)

        if include_spatial:
            if (
                offset_wells is None
                or well_locations is None
                or target_location is None
            ):
                raise ValueError(
                    "Spatial features require offset_wells, well_locations, "
                    "and target_location. "
                    "Provide offset well data and coordinates",
                )

            spatial = self.compute_spatial_features(
                data,
                offset_wells,
                well_locations,
                target_location,
                **kwargs.get("spatial_params", {}),
            )
            features = pd.concat([features, spatial], axis=1)

        return FeatureSet(
            features=features,
            feature_names=list(features.columns),
            feature_groups=self.feature_groups.copy(),
            description=f"Feature set with {len(features.columns)} features",
        )

    def select_features(
        self,
        feature_set: FeatureSet,
        target: pd.Series,
        method: str = "correlation",
        n_features: int = 20,
    ) -> list[str]:
        """
        Select most informative features for modeling.

        Args:
            feature_set: FeatureSet from create_feature_set
            target: Target variable (e.g., facies labels)
            method: 'correlation', 'mutual_info', or 'variance'
            n_features: Number of features to select

        Returns:
            List of selected feature names
        """
        from sklearn.feature_selection import f_classif, mutual_info_classif

        features = feature_set.features

        # Remove null values for feature selection
        valid_mask = target != self.null_value
        for col in features.columns:
            valid_mask &= features[col] != self.null_value

        if valid_mask.sum() < 10:
            raise ValueError(
                "Insufficient valid data for feature selection",
                "Need at least 10 samples with all features and target",
            )

        X = features.loc[valid_mask].values
        y = target.loc[valid_mask].values

        if method == "correlation":
            # Use F-statistic
            scores, _ = f_classif(X, y)
        elif method == "mutual_info":
            scores = mutual_info_classif(X, y, random_state=42)
        elif method == "variance":
            # Simple variance-based selection
            scores = np.var(X, axis=0)
        else:
            raise ValueError(
                f"Unknown feature selection method: {method}",
                "Valid methods: 'correlation', 'mutual_info', 'variance'",
            )

        # Handle NaN scores
        scores = np.nan_to_num(scores, nan=0.0)

        # Select top features
        top_indices = np.argsort(scores)[::-1][:n_features]
        selected_features = [feature_set.feature_names[i] for i in top_indices]

        return selected_features


def prepare_ml_dataset(
    wells: dict[str, pd.DataFrame],
    target_column: str,
    feature_engineer: LogFeatureEngineer | None = None,
    test_well: str | None = None,
    **feature_kwargs,
) -> tuple[pd.DataFrame, pd.Series, pd.DataFrame | None, pd.Series | None]:
    """
    Prepare complete ML dataset from multiple wells.

    Args:
        wells: Dict of {well_name: DataFrame} with log data and targets
        target_column: Name of target column (e.g., 'Facies')
        feature_engineer: Optional LogFeatureEngineer instance
        test_well: Optional well name to hold out for testing
        **feature_kwargs: Arguments passed to create_feature_set

    Returns:
        Tuple of (X_train, y_train, X_test, y_test)
        X_test and y_test are None if no test_well specified
    """
    if feature_engineer is None:
        feature_engineer = LogFeatureEngineer()

    # Separate test well if specified
    if test_well:
        if test_well not in wells:
            raise ValueError(
                f"Test well '{test_well}' not found in wells dictionary",
                f"Available wells: {list(wells.keys())}",
            )
        test_data = wells[test_well]
        train_wells = {k: v for k, v in wells.items() if k != test_well}
    else:
        train_wells = wells
        test_data = None

    # Create features for each training well
    train_features_list = []
    train_targets_list = []

    for well_name, well_data in train_wells.items():
        if target_column not in well_data.columns:
            continue

        feature_set = feature_engineer.create_feature_set(
            well_data.drop(columns=[target_column]), **feature_kwargs
        )

        # Remove null values
        valid_mask = well_data[target_column] != feature_engineer.null_value
        for col in feature_set.features.columns:
            valid_mask &= feature_set.features[col] != feature_engineer.null_value

        if valid_mask.sum() > 0:
            train_features_list.append(feature_set.features.loc[valid_mask])
            train_targets_list.append(well_data.loc[valid_mask, target_column])

    X_train = pd.concat(train_features_list, axis=0)
    y_train = pd.concat(train_targets_list, axis=0)

    # Process test well if provided
    if test_data is not None:
        test_feature_set = feature_engineer.create_feature_set(
            test_data.drop(columns=[target_column]), **feature_kwargs
        )

        valid_mask = test_data[target_column] != feature_engineer.null_value
        for col in test_feature_set.features.columns:
            valid_mask &= test_feature_set.features[col] != feature_engineer.null_value

        X_test = test_feature_set.features.loc[valid_mask]
        y_test = test_data.loc[valid_mask, target_column]
    else:
        X_test = None
        y_test = None

    return X_train, y_train, X_test, y_test
