"""Signal transform handlers for pipeline system.

This module provides handlers for signal processing transformations.
All handlers follow the standard signature: (inputs, params, step_name) -> outputs.

Available Handlers:
    - transform.resample: Resample signal to new sample rate
    - transform.offset: Add DC offset to signal
    - transform.scale: Scale signal amplitude
    - transform.invert: Invert signal polarity
    - transform.abs: Compute absolute value
    - transform.derivative: Calculate time derivative
    - transform.integral: Calculate time integral
    - transform.math: Generic math operations between traces
    - transform.window: Apply window function
    - transform.normalize: Normalize to specified range
    - transform.clip: Clip signal to value range
    - transform.rectify: Rectify signal (half-wave or full-wave)
"""

from __future__ import annotations

from typing import Any

import numpy as np

from oscura.core.config.pipeline import PipelineExecutionError
from oscura.core.types import WaveformTrace
from oscura.pipeline.handlers import register_handler

# Lazy imports to avoid circular dependencies
_interpolation = None
_windowing = None


def _get_interpolation() -> Any:
    """Lazy import interpolation module."""
    global _interpolation
    if _interpolation is None:
        from oscura.utils.math import interpolation as _interp_module

        _interpolation = _interp_module
    return _interpolation


def _get_windowing() -> Any:
    """Lazy import windowing module."""
    global _windowing
    if _windowing is None:
        from oscura.utils import windowing as _window_module

        _windowing = _window_module
    return _windowing


@register_handler("transform.resample")
def handle_transform_resample(
    inputs: dict[str, Any], params: dict[str, Any], step_name: str
) -> dict[str, Any]:
    """Resample signal to new sample rate.

    Uses high-quality polyphase resampling (scipy.signal.resample_poly) for
    integer rate changes, or linear interpolation for arbitrary rates.

    Inputs:
        trace: WaveformTrace to resample

    Parameters:
        sample_rate (float): Target sample rate in Hz
        method (str, optional): Resampling method ('auto', 'polyphase', 'interpolate')
                                (default: 'auto')

    Outputs:
        trace: Resampled WaveformTrace
        original_sample_rate: Original sample rate in Hz
        new_sample_rate: New sample rate in Hz
        num_samples: Number of samples in resampled trace

    Example:
        >>> # Downsample from 1 MHz to 100 kHz
        >>> params = {"sample_rate": 100e3}
        >>> result = handle_transform_resample({"trace": trace}, params, "downsample")
    """
    interpolation = _get_interpolation()

    trace = inputs.get("trace")
    if trace is None:
        raise PipelineExecutionError("Missing required input 'trace'", step_name=step_name)

    new_sample_rate = params.get("sample_rate")
    if new_sample_rate is None:
        raise PipelineExecutionError(
            "Missing required parameter 'sample_rate'. Add 'sample_rate: 1000000' (in Hz) to params",
            step_name=step_name,
        )

    method = params.get("method", "auto")

    try:
        # Use resample utility from interpolation module
        resampled = interpolation.resample(trace, new_sample_rate=new_sample_rate, method=method)
        original_rate = trace.metadata.sample_rate

    except Exception as e:
        raise PipelineExecutionError(f"Resampling failed: {e}", step_name=step_name) from e

    return {
        "trace": resampled,
        "original_sample_rate": original_rate,
        "new_sample_rate": new_sample_rate,
        "num_samples": len(resampled.data),
    }


@register_handler("transform.offset")
def handle_transform_offset(
    inputs: dict[str, Any], params: dict[str, Any], step_name: str
) -> dict[str, Any]:
    """Add DC offset to signal.

    Adds a constant value to all samples in the trace.

    Inputs:
        trace: WaveformTrace to offset

    Parameters:
        offset (float): DC offset value to add

    Outputs:
        trace: Offset WaveformTrace
        offset: Applied offset value

    Example:
        >>> # Remove 2.5V DC bias
        >>> params = {"offset": -2.5}
        >>> result = handle_transform_offset({"trace": trace}, params, "remove_bias")
    """
    trace = inputs.get("trace")
    if trace is None:
        raise PipelineExecutionError("Missing required input 'trace'", step_name=step_name)

    offset = params.get("offset")
    if offset is None:
        raise PipelineExecutionError(
            "Missing required parameter 'offset'. Add 'offset: 1.5' to params",
            step_name=step_name,
        )

    try:
        # Apply offset
        new_data = trace.data + offset
        new_trace = WaveformTrace(data=new_data, metadata=trace.metadata)

    except Exception as e:
        raise PipelineExecutionError(f"Offset failed: {e}", step_name=step_name) from e

    return {
        "trace": new_trace,
        "offset": offset,
    }


@register_handler("transform.scale")
def handle_transform_scale(
    inputs: dict[str, Any], params: dict[str, Any], step_name: str
) -> dict[str, Any]:
    """Scale signal amplitude.

    Multiplies all samples by a constant factor.

    Inputs:
        trace: WaveformTrace to scale

    Parameters:
        factor (float): Scale factor to apply
        unit (str, optional): Updated unit string after scaling

    Outputs:
        trace: Scaled WaveformTrace
        factor: Applied scale factor

    Example:
        >>> # Convert 10x probe to 1x
        >>> params = {"factor": 10.0, "unit": "V"}
        >>> result = handle_transform_scale({"trace": trace}, params, "probe_scale")
    """
    trace = inputs.get("trace")
    if trace is None:
        raise PipelineExecutionError("Missing required input 'trace'", step_name=step_name)

    factor = params.get("factor")
    if factor is None:
        raise PipelineExecutionError(
            "Missing required parameter 'factor'. Add 'factor: 2.0' to params",
            step_name=step_name,
        )

    try:
        # Apply scaling
        new_data = trace.data * factor
        new_metadata = trace.metadata
        if "unit" in params:
            # Update unit if provided (shallow copy to avoid mutation)
            from dataclasses import replace

            new_metadata = replace(new_metadata, unit=params["unit"])

        new_trace = WaveformTrace(data=new_data, metadata=new_metadata)

    except Exception as e:
        raise PipelineExecutionError(f"Scaling failed: {e}", step_name=step_name) from e

    return {
        "trace": new_trace,
        "factor": factor,
    }


@register_handler("transform.invert")
def handle_transform_invert(
    inputs: dict[str, Any], params: dict[str, Any], step_name: str
) -> dict[str, Any]:
    """Invert signal polarity.

    Multiplies all samples by -1.

    Inputs:
        trace: WaveformTrace or DigitalTrace to invert

    Parameters:
        None

    Outputs:
        trace: Inverted trace

    Example:
        >>> # Invert UART idle-low to idle-high
        >>> result = handle_transform_invert({"trace": trace}, {}, "invert_uart")
    """
    trace = inputs.get("trace")
    if trace is None:
        raise PipelineExecutionError("Missing required input 'trace'", step_name=step_name)

    try:
        # Invert data
        new_data = -trace.data
        new_trace = type(trace)(data=new_data, metadata=trace.metadata)

    except Exception as e:
        raise PipelineExecutionError(f"Inversion failed: {e}", step_name=step_name) from e

    return {
        "trace": new_trace,
    }


@register_handler("transform.abs")
def handle_transform_abs(
    inputs: dict[str, Any], params: dict[str, Any], step_name: str
) -> dict[str, Any]:
    """Compute absolute value of signal.

    Takes the absolute value of all samples.

    Inputs:
        trace: WaveformTrace to process

    Parameters:
        None

    Outputs:
        trace: Absolute value WaveformTrace

    Example:
        >>> # Get envelope of AC signal
        >>> result = handle_transform_abs({"trace": trace}, {}, "abs_envelope")
    """
    trace = inputs.get("trace")
    if trace is None:
        raise PipelineExecutionError("Missing required input 'trace'", step_name=step_name)

    try:
        # Compute absolute value
        new_data = np.abs(trace.data)
        new_trace = WaveformTrace(data=new_data, metadata=trace.metadata)

    except Exception as e:
        raise PipelineExecutionError(f"Absolute value failed: {e}", step_name=step_name) from e

    return {
        "trace": new_trace,
    }


@register_handler("transform.derivative")
def handle_transform_derivative(
    inputs: dict[str, Any], params: dict[str, Any], step_name: str
) -> dict[str, Any]:
    """Calculate time derivative of signal.

    Computes numerical derivative (dV/dt) using numpy.gradient for centered differences.

    Inputs:
        trace: WaveformTrace to differentiate

    Parameters:
        None

    Outputs:
        trace: Derivative WaveformTrace
        unit: Derivative unit string (e.g., "V/s")

    Example:
        >>> # Find slew rate
        >>> result = handle_transform_derivative({"trace": trace}, {}, "slew_rate")
    """
    trace = inputs.get("trace")
    if trace is None:
        raise PipelineExecutionError("Missing required input 'trace'", step_name=step_name)

    try:
        # Compute derivative using centered differences
        dt = 1.0 / trace.metadata.sample_rate
        derivative = np.gradient(trace.data, dt)

        # Update unit (V -> V/s)
        from dataclasses import replace

        original_unit = getattr(trace.metadata, "unit", "")
        new_unit = f"{original_unit}/s" if original_unit else "1/s"
        new_metadata = replace(trace.metadata, unit=new_unit)

        new_trace = WaveformTrace(data=derivative, metadata=new_metadata)

    except Exception as e:
        raise PipelineExecutionError(f"Derivative failed: {e}", step_name=step_name) from e

    return {
        "trace": new_trace,
        "unit": new_unit,
    }


@register_handler("transform.integral")
def handle_transform_integral(
    inputs: dict[str, Any], params: dict[str, Any], step_name: str
) -> dict[str, Any]:
    """Calculate time integral of signal.

    Computes cumulative integral using trapezoidal rule.

    Inputs:
        trace: WaveformTrace to integrate

    Parameters:
        initial_value (float, optional): Initial integration constant (default: 0.0)

    Outputs:
        trace: Integrated WaveformTrace
        unit: Integral unit string (e.g., "V·s")

    Example:
        >>> # Integrate current to get charge
        >>> params = {"initial_value": 0.0}
        >>> result = handle_transform_integral({"trace": trace}, params, "charge")
    """
    trace = inputs.get("trace")
    if trace is None:
        raise PipelineExecutionError("Missing required input 'trace'", step_name=step_name)

    initial_value = params.get("initial_value", 0.0)

    try:
        # Compute integral using trapezoidal rule
        dt = 1.0 / trace.metadata.sample_rate
        integral = np.cumsum(trace.data) * dt + initial_value

        # Update unit (V -> V·s)
        from dataclasses import replace

        original_unit = getattr(trace.metadata, "unit", "")
        new_unit = f"{original_unit}·s" if original_unit else "s"
        new_metadata = replace(trace.metadata, unit=new_unit)

        new_trace = WaveformTrace(data=integral, metadata=new_metadata)

    except Exception as e:
        raise PipelineExecutionError(f"Integration failed: {e}", step_name=step_name) from e

    return {
        "trace": new_trace,
        "unit": new_unit,
    }


@register_handler("transform.math")
def handle_transform_math(
    inputs: dict[str, Any], params: dict[str, Any], step_name: str
) -> dict[str, Any]:
    """Generic math operations between two traces.

    Performs element-wise arithmetic operations (add, subtract, multiply, divide)
    between two traces. Traces must have same length or one will be interpolated.

    Inputs:
        trace1: First WaveformTrace
        trace2: Second WaveformTrace

    Parameters:
        operation (str): Operation to perform ('add', 'subtract', 'multiply', 'divide')
        auto_align (bool, optional): Auto-align traces if sample rates differ (default: True)

    Outputs:
        trace: Result WaveformTrace
        operation: Operation performed

    Example:
        >>> # Compute differential signal
        >>> params = {"operation": "subtract"}
        >>> result = handle_transform_math({"trace1": pos, "trace2": neg}, params, "diff")
    """
    trace1 = inputs.get("trace1")
    trace2 = inputs.get("trace2")

    if trace1 is None or trace2 is None:
        raise PipelineExecutionError(
            "Missing required inputs 'trace1' and 'trace2'", step_name=step_name
        )

    operation = params.get("operation")
    if operation not in {"add", "subtract", "multiply", "divide"}:
        raise PipelineExecutionError(
            f"Invalid operation '{operation}'. Must be: add, subtract, multiply, divide",
            step_name=step_name,
        )

    auto_align = params.get("auto_align", True)

    try:
        # Align traces if needed
        if auto_align and (
            trace1.metadata.sample_rate != trace2.metadata.sample_rate
            or len(trace1.data) != len(trace2.data)
        ):
            interpolation = _get_interpolation()
            # Resample trace2 to match trace1
            trace2 = interpolation.resample(trace2, new_sample_rate=trace1.metadata.sample_rate)
            # Ensure same length
            min_len = min(len(trace1.data), len(trace2.data))
            data1 = trace1.data[:min_len]
            data2 = trace2.data[:min_len]
        else:
            if len(trace1.data) != len(trace2.data):
                raise PipelineExecutionError(
                    f"Trace lengths don't match: {len(trace1.data)} vs {len(trace2.data)}. "
                    "Set 'auto_align: true' to automatically align traces",
                    step_name=step_name,
                )
            data1 = trace1.data
            data2 = trace2.data

        # Perform operation
        if operation == "add":
            result_data = data1 + data2
        elif operation == "subtract":
            result_data = data1 - data2
        elif operation == "multiply":
            result_data = data1 * data2
        elif operation == "divide":
            # Avoid division by zero
            with np.errstate(divide="ignore", invalid="ignore"):
                result_data = np.divide(data1, data2)
                result_data = np.nan_to_num(result_data, nan=0.0, posinf=0.0, neginf=0.0)
        else:
            # Should never reach here due to earlier validation
            raise ValueError(f"Unknown operation: {operation}")

        new_trace = WaveformTrace(data=result_data, metadata=trace1.metadata)

    except PipelineExecutionError:
        raise
    except Exception as e:
        raise PipelineExecutionError(f"Math operation failed: {e}", step_name=step_name) from e

    return {
        "trace": new_trace,
        "operation": operation,
    }


@register_handler("transform.window")
def handle_transform_window(
    inputs: dict[str, Any], params: dict[str, Any], step_name: str
) -> dict[str, Any]:
    """Apply window function to signal.

    Multiplies signal by a window function for spectral analysis.

    Inputs:
        trace: WaveformTrace to window

    Parameters:
        window (str): Window type ('hann', 'hamming', 'blackman', 'bartlett',
                      'rectangular', 'kaiser')
        beta (float, optional): Beta parameter for Kaiser window (default: 8.6)

    Outputs:
        trace: Windowed WaveformTrace
        window: Window type applied

    Example:
        >>> # Apply Hann window for FFT
        >>> params = {"window": "hann"}
        >>> result = handle_transform_window({"trace": trace}, params, "fft_window")
    """
    windowing = _get_windowing()

    trace = inputs.get("trace")
    if trace is None:
        raise PipelineExecutionError("Missing required input 'trace'", step_name=step_name)

    window_name = params.get("window")
    if window_name is None:
        raise PipelineExecutionError(
            "Missing required parameter 'window'. Add 'window: hann' to params",
            step_name=step_name,
        )

    try:
        # Get window coefficients
        if window_name == "kaiser":
            beta = params.get("beta", 8.6)
            window = windowing.kaiser(len(trace.data), beta)
        else:
            window = windowing.get_window(window_name, len(trace.data))

        # Apply window
        windowed_data = trace.data * window
        new_trace = WaveformTrace(data=windowed_data, metadata=trace.metadata)

    except Exception as e:
        raise PipelineExecutionError(
            f"Window application failed: {e}. "
            "Valid windows: hann, hamming, blackman, bartlett, rectangular, kaiser",
            step_name=step_name,
        ) from e

    return {
        "trace": new_trace,
        "window": window_name,
    }


@register_handler("transform.normalize")
def handle_transform_normalize(
    inputs: dict[str, Any], params: dict[str, Any], step_name: str
) -> dict[str, Any]:
    """Normalize signal to specified range.

    Scales signal to fit within a specified range [min_val, max_val].

    Inputs:
        trace: WaveformTrace to normalize

    Parameters:
        min (float, optional): Minimum value of output range (default: -1.0)
        max (float, optional): Maximum value of output range (default: 1.0)
        method (str, optional): Normalization method ('minmax', 'zscore', 'peak')
                                (default: 'minmax')

    Outputs:
        trace: Normalized WaveformTrace
        method: Normalization method used
        original_min: Original minimum value
        original_max: Original maximum value
        scale_factor: Applied scale factor

    Example:
        >>> # Normalize to [-1, 1]
        >>> params = {"min": -1.0, "max": 1.0, "method": "minmax"}
        >>> result = handle_transform_normalize({"trace": trace}, params, "norm")
    """
    trace = inputs.get("trace")
    if trace is None:
        raise PipelineExecutionError("Missing required input 'trace'", step_name=step_name)

    min_val = params.get("min", -1.0)
    max_val = params.get("max", 1.0)
    method = params.get("method", "minmax")

    if min_val >= max_val:
        raise PipelineExecutionError(
            f"min ({min_val}) must be less than max ({max_val})", step_name=step_name
        )

    try:
        original_min = float(np.min(trace.data))
        original_max = float(np.max(trace.data))

        if method == "minmax":
            # Min-max normalization
            data_range = original_max - original_min
            if data_range == 0:
                normalized = np.full_like(trace.data, (min_val + max_val) / 2)
                scale_factor = 0.0
            else:
                normalized = (trace.data - original_min) / data_range
                normalized = normalized * (max_val - min_val) + min_val
                scale_factor = (max_val - min_val) / data_range

        elif method == "zscore":
            # Z-score normalization (standardization)
            mean = np.mean(trace.data)
            std = np.std(trace.data)
            if std == 0:
                normalized = np.zeros_like(trace.data)
                scale_factor = 0.0
            else:
                normalized = (trace.data - mean) / std
                # Scale to [min, max]
                norm_min = np.min(normalized)
                norm_max = np.max(normalized)
                norm_range = norm_max - norm_min
                if norm_range > 0:
                    normalized = (normalized - norm_min) / norm_range
                    normalized = normalized * (max_val - min_val) + min_val
                scale_factor = 1.0 / std

        elif method == "peak":
            # Peak normalization (divide by absolute maximum)
            peak = np.max(np.abs(trace.data))
            if peak == 0:
                normalized = np.zeros_like(trace.data)
                scale_factor = 0.0
            else:
                normalized = trace.data / peak
                # Scale to [min, max]
                normalized = normalized * (max_val - min_val) / 2.0 + (max_val + min_val) / 2.0
                scale_factor = (max_val - min_val) / (2.0 * peak)

        else:
            raise PipelineExecutionError(
                f"Unknown normalization method '{method}'. Valid methods: minmax, zscore, peak",
                step_name=step_name,
            )

        new_trace = WaveformTrace(data=normalized, metadata=trace.metadata)

    except PipelineExecutionError:
        raise
    except Exception as e:
        raise PipelineExecutionError(f"Normalization failed: {e}", step_name=step_name) from e

    return {
        "trace": new_trace,
        "method": method,
        "original_min": original_min,
        "original_max": original_max,
        "scale_factor": scale_factor,
    }


@register_handler("transform.clip")
def handle_transform_clip(
    inputs: dict[str, Any], params: dict[str, Any], step_name: str
) -> dict[str, Any]:
    """Clip signal to value range.

    Limits signal values to [min_val, max_val] range.

    Inputs:
        trace: WaveformTrace to clip

    Parameters:
        min (float, optional): Minimum value (default: -inf)
        max (float, optional): Maximum value (default: +inf)

    Outputs:
        trace: Clipped WaveformTrace
        clipped_samples: Number of samples clipped

    Example:
        >>> # Clip to ADC range
        >>> params = {"min": 0.0, "max": 3.3}
        >>> result = handle_transform_clip({"trace": trace}, params, "adc_clip")
    """
    trace = inputs.get("trace")
    if trace is None:
        raise PipelineExecutionError("Missing required input 'trace'", step_name=step_name)

    min_val = params.get("min", -np.inf)
    max_val = params.get("max", np.inf)

    if min_val >= max_val:
        raise PipelineExecutionError(
            f"min ({min_val}) must be less than max ({max_val})", step_name=step_name
        )

    try:
        original_data = trace.data.copy()
        clipped_data = np.clip(trace.data, min_val, max_val)

        # Count clipped samples
        clipped_samples = int(np.sum(clipped_data != original_data))

        new_trace = WaveformTrace(data=clipped_data, metadata=trace.metadata)

    except Exception as e:
        raise PipelineExecutionError(f"Clipping failed: {e}", step_name=step_name) from e

    return {
        "trace": new_trace,
        "clipped_samples": clipped_samples,
    }


@register_handler("transform.rectify")
def handle_transform_rectify(
    inputs: dict[str, Any], params: dict[str, Any], step_name: str
) -> dict[str, Any]:
    """Rectify signal (half-wave or full-wave).

    Performs rectification for envelope detection or power analysis.

    Inputs:
        trace: WaveformTrace to rectify

    Parameters:
        mode (str): Rectification mode ('half', 'full') (default: 'full')

    Outputs:
        trace: Rectified WaveformTrace
        mode: Rectification mode applied

    Example:
        >>> # Full-wave rectification for power measurement
        >>> params = {"mode": "full"}
        >>> result = handle_transform_rectify({"trace": trace}, params, "rectify")
    """
    trace = inputs.get("trace")
    if trace is None:
        raise PipelineExecutionError("Missing required input 'trace'", step_name=step_name)

    mode = params.get("mode", "full")

    if mode not in {"half", "full"}:
        raise PipelineExecutionError(
            f"Invalid rectification mode '{mode}'. Valid modes: half, full",
            step_name=step_name,
        )

    try:
        if mode == "full":
            # Full-wave rectification (absolute value)
            rectified_data = np.abs(trace.data)
        else:
            # Half-wave rectification (zero negative values)
            rectified_data = np.maximum(trace.data, 0.0)

        new_trace = WaveformTrace(data=rectified_data, metadata=trace.metadata)

    except Exception as e:
        raise PipelineExecutionError(f"Rectification failed: {e}", step_name=step_name) from e

    return {
        "trace": new_trace,
        "mode": mode,
    }
