"""
ATD sensor data averaging module.

Applies Gaussian smoothing for noise reduction on ATD sensor data.
"""

import numpy as np
from scipy.ndimage import gaussian_filter1d
from typing import Tuple


def average_radial_link_data(acceleration: np.ndarray, magnetic_field: np.ndarray,
                             timestamps: np.ndarray, temperature: np.ndarray,
                             n_points: int) -> Tuple[np.ndarray, np.ndarray, np.ndarray, np.ndarray]:
    """
    Average RL data using Gaussian smoothing.
    
    Applies smoothing to acceleration, magnetic field, and temperature.
    Equivalent to MATLAB smoothdata(..., 'gaussian', n_points).
    
    Args:
        acceleration: (n_timestamps, n_sensors*3) converted acceleration
        magnetic_field: (n_timestamps, n_sensors*3) converted magnetic field
        timestamps: (n_timestamps,) datetime array
        temperature: (n_timestamps, n_sensors) converted temperature
        n_points: Number of points for Gaussian window
        
    Returns:
        Tuple of (acc_smoothed, mag_smoothed, temp_smoothed, err_flag)
    """
    n_timestamps = acceleration.shape[0]
    
    # Check if we have enough data
    if n_timestamps < n_points:
        n_points = n_timestamps
    
    # Calculate sigma for Gaussian filter
    # MATLAB smoothdata uses sigma = n_points / 6
    sigma = n_points / 6.0
    
    # Initialize output arrays
    acc_smoothed = np.zeros_like(acceleration)
    mag_smoothed = np.zeros_like(magnetic_field)
    temp_smoothed = np.zeros_like(temperature)
    err_flag = np.zeros_like(temperature)
    
    # Apply Gaussian filter to each column
    for col in range(acceleration.shape[1]):
        acc_smoothed[:, col] = gaussian_filter1d(acceleration[:, col], sigma=sigma)
    
    for col in range(magnetic_field.shape[1]):
        mag_smoothed[:, col] = gaussian_filter1d(magnetic_field[:, col], sigma=sigma)
    
    for col in range(temperature.shape[1]):
        temp_smoothed[:, col] = gaussian_filter1d(temperature[:, col], sigma=sigma)
    
    return acc_smoothed, mag_smoothed, temp_smoothed, err_flag


def average_load_link_data(force_data: np.ndarray, timestamps: np.ndarray,
                           temperature: np.ndarray, n_points: int
                           ) -> Tuple[np.ndarray, np.ndarray, np.ndarray]:
    """
    Average LL force data using Gaussian smoothing.
    
    Args:
        force_data: (n_timestamps, n_sensors) converted force
        timestamps: (n_timestamps,) datetime array
        temperature: (n_timestamps, n_sensors) converted temperature
        n_points: Number of points for Gaussian window
        
    Returns:
        Tuple of (force_smoothed, temp_smoothed, err_flag)
    """
    n_timestamps = force_data.shape[0]
    
    if n_timestamps < n_points:
        n_points = n_timestamps
    
    sigma = n_points / 6.0
    
    force_smoothed = np.zeros_like(force_data)
    temp_smoothed = np.zeros_like(temperature)
    err_flag = np.zeros_like(temperature)
    
    # Smooth each sensor
    for col in range(force_data.shape[1]):
        force_smoothed[:, col] = gaussian_filter1d(force_data[:, col], sigma=sigma)
        temp_smoothed[:, col] = gaussian_filter1d(temperature[:, col], sigma=sigma)
    
    return force_smoothed, temp_smoothed, err_flag


def average_pressure_link_data(pressure_data: np.ndarray, timestamps: np.ndarray,
                               temperature: np.ndarray, n_points: int
                               ) -> Tuple[np.ndarray, np.ndarray, np.ndarray]:
    """
    Average PL pressure data using Gaussian smoothing.
    
    Args:
        pressure_data: (n_timestamps, n_sensors) converted pressure
        timestamps: (n_timestamps,) datetime array
        temperature: (n_timestamps, n_sensors) converted temperature
        n_points: Number of points for Gaussian window
        
    Returns:
        Tuple of (pressure_smoothed, temp_smoothed, err_flag)
    """
    n_timestamps = pressure_data.shape[0]
    
    if n_timestamps < n_points:
        n_points = n_timestamps
    
    sigma = n_points / 6.0
    
    pressure_smoothed = np.zeros_like(pressure_data)
    temp_smoothed = np.zeros_like(temperature)
    err_flag = np.zeros_like(temperature)
    
    for col in range(pressure_data.shape[1]):
        pressure_smoothed[:, col] = gaussian_filter1d(pressure_data[:, col], sigma=sigma)
        temp_smoothed[:, col] = gaussian_filter1d(temperature[:, col], sigma=sigma)
    
    return pressure_smoothed, temp_smoothed, err_flag


def average_extensometer_data(extension_data: np.ndarray, timestamps: np.ndarray,
                              temperature: np.ndarray, n_points: int
                              ) -> Tuple[np.ndarray, np.ndarray, np.ndarray]:
    """
    Average extensometer data using Gaussian smoothing.
    
    Args:
        extension_data: (n_timestamps, n_sensors) converted extension
        timestamps: (n_timestamps,) datetime array
        temperature: (n_timestamps, n_sensors) converted temperature
        n_points: Number of points for Gaussian window
        
    Returns:
        Tuple of (extension_smoothed, temp_smoothed, err_flag)
    """
    n_timestamps = extension_data.shape[0]
    
    if n_timestamps < n_points:
        n_points = n_timestamps
    
    sigma = n_points / 6.0
    
    extension_smoothed = np.zeros_like(extension_data)
    temp_smoothed = np.zeros_like(temperature)
    err_flag = np.zeros_like(temperature)
    
    for col in range(extension_data.shape[1]):
        extension_smoothed[:, col] = gaussian_filter1d(extension_data[:, col], sigma=sigma)
        temp_smoothed[:, col] = gaussian_filter1d(temperature[:, col], sigma=sigma)
    
    return extension_smoothed, temp_smoothed, err_flag


def average_resultant_vectors(acc_magnitude: np.ndarray, mag_magnitude: np.ndarray,
                              n_points: int) -> Tuple[np.ndarray, np.ndarray]:
    """
    Average resultant magnitude vectors.
    
    Args:
        acc_magnitude: (n_timestamps, n_sensors) acceleration magnitude
        mag_magnitude: (n_timestamps, n_sensors) magnetic field magnitude
        n_points: Number of points for Gaussian window
        
    Returns:
        Tuple of (acc_mag_smoothed, mag_mag_smoothed)
    """
    n_timestamps = acc_magnitude.shape[0]
    
    if n_timestamps < n_points:
        n_points = n_timestamps
    
    sigma = n_points / 6.0
    
    acc_mag_smoothed = np.zeros_like(acc_magnitude)
    mag_mag_smoothed = np.zeros_like(mag_magnitude)
    
    for col in range(acc_magnitude.shape[1]):
        acc_mag_smoothed[:, col] = gaussian_filter1d(acc_magnitude[:, col], sigma=sigma)
        mag_mag_smoothed[:, col] = gaussian_filter1d(mag_magnitude[:, col], sigma=sigma)
    
    return acc_mag_smoothed, mag_mag_smoothed


def average_extensometer_3d_data(displacement_data: np.ndarray, timestamps: np.ndarray,
                                temperature: np.ndarray, n_points: int
                                ) -> Tuple[np.ndarray, np.ndarray, np.ndarray]:
    """
    Average 3DEL data using Gaussian smoothing.
    
    Args:
        displacement_data: (n_timestamps, n_sensors*3) converted 3D displacement
        timestamps: (n_timestamps,) datetime array
        temperature: (n_timestamps, n_sensors) converted temperature
        n_points: Number of points for Gaussian window
        
    Returns:
        Tuple of (disp_smoothed, temp_smoothed, err_flag)
    """
    n_timestamps = displacement_data.shape[0]
    
    if n_timestamps < n_points:
        n_points = n_timestamps
    
    sigma = n_points / 6.0
    
    disp_smoothed = np.zeros_like(displacement_data)
    temp_smoothed = np.zeros_like(temperature)
    err_flag = np.zeros_like(temperature)
    
    for col in range(displacement_data.shape[1]):
        disp_smoothed[:, col] = gaussian_filter1d(displacement_data[:, col], sigma=sigma)
    
    for col in range(temperature.shape[1]):
        temp_smoothed[:, col] = gaussian_filter1d(temperature[:, col], sigma=sigma)
    
    return disp_smoothed, temp_smoothed, err_flag


def average_crackmeter_data(displacement_data: np.ndarray, timestamps: np.ndarray,
                           temperature: np.ndarray, n_points: int
                           ) -> Tuple[np.ndarray, np.ndarray, np.ndarray]:
    """
    Average crackmeter data using Gaussian smoothing.
    
    Args:
        displacement_data: (n_timestamps, n_sensors*n_dimensions) converted displacement
        timestamps: (n_timestamps,) datetime array
        temperature: (n_timestamps, n_sensors) converted temperature
        n_points: Number of points for Gaussian window
        
    Returns:
        Tuple of (disp_smoothed, temp_smoothed, err_flag)
    """
    n_timestamps = displacement_data.shape[0]
    
    if n_timestamps < n_points:
        n_points = n_timestamps
    
    sigma = n_points / 6.0
    
    disp_smoothed = np.zeros_like(displacement_data)
    temp_smoothed = np.zeros_like(temperature)
    err_flag = np.zeros_like(temperature)
    
    for col in range(displacement_data.shape[1]):
        disp_smoothed[:, col] = gaussian_filter1d(displacement_data[:, col], sigma=sigma)
    
    for col in range(temperature.shape[1]):
        temp_smoothed[:, col] = gaussian_filter1d(temperature[:, col], sigma=sigma)
    
    return disp_smoothed, temp_smoothed, err_flag


def average_pcl_data(angle_data: np.ndarray, timestamps: np.ndarray,
                    temperature: np.ndarray, n_points: int
                    ) -> Tuple[np.ndarray, np.ndarray, np.ndarray]:
    """
    Average PCL angle data using Gaussian smoothing.
    
    Args:
        angle_data: (n_timestamps, n_sensors*2) converted angles (ax, ay)
        timestamps: (n_timestamps,) datetime array
        temperature: (n_timestamps, n_sensors) converted temperature
        n_points: Number of points for Gaussian window
        
    Returns:
        Tuple of (angles_smoothed, temp_smoothed, err_flag)
    """
    n_timestamps = angle_data.shape[0]
    
    if n_timestamps < n_points:
        n_points = n_timestamps
    
    sigma = n_points / 6.0
    
    angles_smoothed = np.zeros_like(angle_data)
    temp_smoothed = np.zeros_like(temperature)
    err_flag = np.zeros_like(temperature)
    
    for col in range(angle_data.shape[1]):
        angles_smoothed[:, col] = gaussian_filter1d(angle_data[:, col], sigma=sigma)
    
    for col in range(temperature.shape[1]):
        temp_smoothed[:, col] = gaussian_filter1d(temperature[:, col], sigma=sigma)
    
    return angles_smoothed, temp_smoothed, err_flag


def average_tube_link_data(angle_data: np.ndarray, timestamps: np.ndarray,
                          temperature: np.ndarray, n_points: int
                          ) -> Tuple[np.ndarray, np.ndarray, np.ndarray]:
    """
    Average TuL angle data using Gaussian smoothing.
    
    Args:
        angle_data: (n_timestamps, n_sensors*3) converted angles (ax, ay, az)
        timestamps: (n_timestamps,) datetime array
        temperature: (n_timestamps, n_sensors) converted temperature
        n_points: Number of points for Gaussian window
        
    Returns:
        Tuple of (angles_smoothed, temp_smoothed, err_flag)
    """
    n_timestamps = angle_data.shape[0]
    
    if n_timestamps < n_points:
        n_points = n_timestamps
    
    sigma = n_points / 6.0
    
    angles_smoothed = np.zeros_like(angle_data)
    temp_smoothed = np.zeros_like(temperature)
    err_flag = np.zeros_like(temperature)
    
    for col in range(angle_data.shape[1]):
        angles_smoothed[:, col] = gaussian_filter1d(angle_data[:, col], sigma=sigma)
    
    for col in range(temperature.shape[1]):
        temp_smoothed[:, col] = gaussian_filter1d(temperature[:, col], sigma=sigma)
    
    return angles_smoothed, temp_smoothed, err_flag
