"""
ATD sensor data conversion module.

Converts raw ADC values to physical units using calibration data.
Handles RL (Radial Link), LL (Load Link), and other extensometer types.
"""

import numpy as np
from typing import Tuple


def convert_radial_link_data(acceleration: np.ndarray, magnetic_field: np.ndarray,
                             temperature: np.ndarray, calibration_data: np.ndarray,
                             n_sensors: int) -> Tuple[np.ndarray, np.ndarray, np.ndarray, np.ndarray]:
    """
    Convert RL raw data to physical units.
    
    Applies calibration for acceleration (g), magnetic field (Gauss), and temperature (°C).
    
    Calibration data columns:
    0-2: caX, pIntX, iIntX (X-axis acceleration: gain, temp coeff, offset)
    3-5: caY, pIntY, iIntY (Y-axis acceleration)
    6-8: caZ, pIntZ, iIntZ (Z-axis acceleration)
    9-10: caT, intT (temperature: gain, offset)
    
    Args:
        acceleration: (n_timestamps, n_sensors*3) raw acceleration
        magnetic_field: (n_timestamps, n_sensors*3) raw magnetic field
        temperature: (n_timestamps, n_sensors) raw temperature
        calibration_data: (n_sensors, 11) calibration parameters
        n_sensors: Number of RL sensors
        
    Returns:
        Tuple of (acc_converted, mag_converted, temp_converted, err_flag)
    """
    n_timestamps = acceleration.shape[0]
    
    # Initialize output arrays
    acc_converted = np.zeros_like(acceleration)
    mag_converted = np.zeros_like(magnetic_field)
    temp_converted = np.zeros_like(temperature)
    err_flag = np.zeros_like(temperature)
    
    # Convert magnetic field from raw to Gauss (simple scaling)
    mag_converted = magnetic_field / 1000.0  # 1000 Gauss scale
    
    # Convert acceleration and temperature for each sensor
    for sensor_idx in range(n_sensors):
        # Extract calibration parameters
        caX = calibration_data[sensor_idx, 0]
        pIntX = calibration_data[sensor_idx, 1]
        iIntX = calibration_data[sensor_idx, 2]
        
        caY = calibration_data[sensor_idx, 3]
        pIntY = calibration_data[sensor_idx, 4]
        iIntY = calibration_data[sensor_idx, 5]
        
        caZ = calibration_data[sensor_idx, 6]
        pIntZ = calibration_data[sensor_idx, 7]
        iIntZ = calibration_data[sensor_idx, 8]
        
        caT = calibration_data[sensor_idx, 9]
        intT = calibration_data[sensor_idx, 10]
        
        # Convert temperature first (needed for acceleration correction)
        temp_converted[:, sensor_idx] = temperature[:, sensor_idx] * caT + intT
        
        # Convert acceleration with temperature compensation
        # Formula: acc_converted = raw * gain + (temp * temp_coeff + offset)
        temp_col = temp_converted[:, sensor_idx]
        
        # X-axis
        acc_converted[:, sensor_idx*3] = (
            acceleration[:, sensor_idx*3] * caX + 
            (temp_col * pIntX + iIntX)
        )
        
        # Y-axis
        acc_converted[:, sensor_idx*3+1] = (
            acceleration[:, sensor_idx*3+1] * caY + 
            (temp_col * pIntY + iIntY)
        )
        
        # Z-axis
        acc_converted[:, sensor_idx*3+2] = (
            acceleration[:, sensor_idx*3+2] * caZ + 
            (temp_col * pIntZ + iIntZ)
        )
    
    return acc_converted, mag_converted, temp_converted, err_flag


def convert_load_link_data(force_data: np.ndarray, temperature: np.ndarray,
                           calibration_data: np.ndarray, n_sensors: int
                           ) -> Tuple[np.ndarray, np.ndarray, np.ndarray]:
    """
    Convert LL raw data to physical units (force in kN, temperature in °C).
    
    Calibration data columns:
    0-1: caF, intF (force: gain, offset)
    2-3: caT, intT (temperature: gain, offset)
    
    Args:
        force_data: (n_timestamps, n_sensors) raw force values
        temperature: (n_timestamps, n_sensors) raw temperature
        calibration_data: (n_sensors, 4) calibration parameters
        n_sensors: Number of LL sensors
        
    Returns:
        Tuple of (force_converted, temp_converted, err_flag)
    """
    n_timestamps = force_data.shape[0]
    
    force_converted = np.zeros_like(force_data)
    temp_converted = np.zeros_like(temperature)
    err_flag = np.zeros_like(temperature)
    
    for sensor_idx in range(n_sensors):
        caF = calibration_data[sensor_idx, 0]
        intF = calibration_data[sensor_idx, 1]
        caT = calibration_data[sensor_idx, 2]
        intT = calibration_data[sensor_idx, 3]
        
        # Linear conversion: physical = raw * gain + offset
        force_converted[:, sensor_idx] = force_data[:, sensor_idx] * caF + intF
        temp_converted[:, sensor_idx] = temperature[:, sensor_idx] * caT + intT
    
    return force_converted, temp_converted, err_flag


def convert_pressure_link_data(pressure_data: np.ndarray, temperature: np.ndarray,
                               calibration_data: np.ndarray, n_sensors: int
                               ) -> Tuple[np.ndarray, np.ndarray, np.ndarray]:
    """
    Convert PL (Pressure Link) raw data to physical units.
    
    Args:
        pressure_data: (n_timestamps, n_sensors) raw pressure values
        temperature: (n_timestamps, n_sensors) raw temperature
        calibration_data: (n_sensors, 4) calibration parameters
        n_sensors: Number of PL sensors
        
    Returns:
        Tuple of (pressure_converted, temp_converted, err_flag)
    """
    pressure_converted = np.zeros_like(pressure_data)
    temp_converted = np.zeros_like(temperature)
    err_flag = np.zeros_like(temperature)
    
    for sensor_idx in range(n_sensors):
        caP = calibration_data[sensor_idx, 0]
        intP = calibration_data[sensor_idx, 1]
        caT = calibration_data[sensor_idx, 2]
        intT = calibration_data[sensor_idx, 3]
        
        pressure_converted[:, sensor_idx] = pressure_data[:, sensor_idx] * caP + intP
        temp_converted[:, sensor_idx] = temperature[:, sensor_idx] * caT + intT
    
    return pressure_converted, temp_converted, err_flag


def convert_extensometer_data(extension_data: np.ndarray, temperature: np.ndarray,
                              calibration_data: np.ndarray, n_sensors: int
                              ) -> Tuple[np.ndarray, np.ndarray, np.ndarray]:
    """
    Convert extensometer (EL, 3DEL) raw data to physical units (mm displacement).
    
    Args:
        extension_data: (n_timestamps, n_sensors) raw extension values
        temperature: (n_timestamps, n_sensors) raw temperature
        calibration_data: (n_sensors, 4) calibration parameters
        n_sensors: Number of extensometer sensors
        
    Returns:
        Tuple of (extension_converted, temp_converted, err_flag)
    """
    extension_converted = np.zeros_like(extension_data)
    temp_converted = np.zeros_like(temperature)
    err_flag = np.zeros_like(temperature)
    
    for sensor_idx in range(n_sensors):
        caE = calibration_data[sensor_idx, 0]
        intE = calibration_data[sensor_idx, 1]
        caT = calibration_data[sensor_idx, 2]
        intT = calibration_data[sensor_idx, 3]
        
        extension_converted[:, sensor_idx] = extension_data[:, sensor_idx] * caE + intE
        temp_converted[:, sensor_idx] = temperature[:, sensor_idx] * caT + intT
    
    return extension_converted, temp_converted, err_flag


def calculate_resultant_magnitude(acceleration: np.ndarray, magnetic_field: np.ndarray,
                                  n_sensors: int) -> Tuple[np.ndarray, np.ndarray]:
    """
    Calculate resultant magnitude vectors for acceleration and magnetic field.
    
    Args:
        acceleration: (n_timestamps, n_sensors*3) converted acceleration
        magnetic_field: (n_timestamps, n_sensors*3) converted magnetic field
        n_sensors: Number of sensors
        
    Returns:
        Tuple of (acc_magnitude, mag_magnitude)
        Each has shape (n_timestamps, n_sensors)
    """
    n_timestamps = acceleration.shape[0]
    
    acc_magnitude = np.zeros((n_timestamps, n_sensors))
    mag_magnitude = np.zeros((n_timestamps, n_sensors))
    
    for sensor_idx in range(n_sensors):
        # Acceleration magnitude: sqrt(ax^2 + ay^2 + az^2)
        ax = acceleration[:, sensor_idx*3]
        ay = acceleration[:, sensor_idx*3+1]
        az = acceleration[:, sensor_idx*3+2]
        acc_magnitude[:, sensor_idx] = np.sqrt(ax**2 + ay**2 + az**2)
        
        # Magnetic field magnitude
        mx = magnetic_field[:, sensor_idx*3]
        my = magnetic_field[:, sensor_idx*3+1]
        mz = magnetic_field[:, sensor_idx*3+2]
        mag_magnitude[:, sensor_idx] = np.sqrt(mx**2 + my**2 + mz**2)
    
    return acc_magnitude, mag_magnitude


def convert_extensometer_3d_data(displacement_data: np.ndarray, temperature: np.ndarray,
                                calibration_data: np.ndarray, n_sensors: int
                                ) -> Tuple[np.ndarray, np.ndarray, np.ndarray]:
    """
    Convert 3DEL raw data to physical units (mm displacement).
    
    Calibration data columns (per sensor):
    0-1: caX, intX (X displacement: gain, offset)
    2-3: caY, intY (Y displacement)
    4-5: caZ, intZ (Z displacement)
    6-7: caT, intT (temperature)
    
    Args:
        displacement_data: (n_timestamps, n_sensors*3) raw displacement values
        temperature: (n_timestamps, n_sensors) raw temperature
        calibration_data: (n_sensors, 8) calibration parameters
        n_sensors: Number of 3DEL sensors
        
    Returns:
        Tuple of (disp_converted, temp_converted, err_flag)
    """
    disp_converted = np.zeros_like(displacement_data)
    temp_converted = np.zeros_like(temperature)
    err_flag = np.zeros_like(temperature)
    
    for sensor_idx in range(n_sensors):
        caX = calibration_data[sensor_idx, 0]
        intX = calibration_data[sensor_idx, 1]
        caY = calibration_data[sensor_idx, 2]
        intY = calibration_data[sensor_idx, 3]
        caZ = calibration_data[sensor_idx, 4]
        intZ = calibration_data[sensor_idx, 5]
        caT = calibration_data[sensor_idx, 6]
        intT = calibration_data[sensor_idx, 7]
        
        # Convert displacements
        disp_converted[:, sensor_idx*3] = displacement_data[:, sensor_idx*3] * caX + intX
        disp_converted[:, sensor_idx*3+1] = displacement_data[:, sensor_idx*3+1] * caY + intY
        disp_converted[:, sensor_idx*3+2] = displacement_data[:, sensor_idx*3+2] * caZ + intZ
        
        # Convert temperature
        temp_converted[:, sensor_idx] = temperature[:, sensor_idx] * caT + intT
    
    return disp_converted, temp_converted, err_flag


def convert_crackmeter_data(displacement_data: np.ndarray, temperature: np.ndarray,
                           calibration_data: np.ndarray, n_sensors: int,
                           n_dimensions: int) -> Tuple[np.ndarray, np.ndarray, np.ndarray]:
    """
    Convert crackmeter raw data to physical units (mm displacement).
    
    Args:
        displacement_data: (n_timestamps, n_sensors*n_dimensions) raw values
        temperature: (n_timestamps, n_sensors) raw temperature
        calibration_data: (n_sensors, 2*n_dimensions+2) calibration parameters
        n_sensors: Number of crackmeter sensors
        n_dimensions: 1, 2, or 3 dimensions
        
    Returns:
        Tuple of (disp_converted, temp_converted, err_flag)
    """
    disp_converted = np.zeros_like(displacement_data)
    temp_converted = np.zeros_like(temperature)
    err_flag = np.zeros_like(temperature)
    
    for sensor_idx in range(n_sensors):
        # Each dimension has gain and offset
        for dim in range(n_dimensions):
            ca = calibration_data[sensor_idx, dim*2]
            offset = calibration_data[sensor_idx, dim*2+1]
            disp_converted[:, sensor_idx*n_dimensions+dim] = (
                displacement_data[:, sensor_idx*n_dimensions+dim] * ca + offset
            )
        
        # Temperature calibration
        caT = calibration_data[sensor_idx, n_dimensions*2]
        intT = calibration_data[sensor_idx, n_dimensions*2+1]
        temp_converted[:, sensor_idx] = temperature[:, sensor_idx] * caT + intT
    
    return disp_converted, temp_converted, err_flag


def convert_pcl_data(angle_data: np.ndarray, temperature: np.ndarray,
                    calibration_data: np.ndarray, n_sensors: int,
                    sensor_type: str = 'PCL') -> Tuple[np.ndarray, np.ndarray, np.ndarray]:
    """
    Convert PCL/PCLHR raw angles to physical units.
    
    Calibration data columns (per sensor):
    0-1: caX, intX (X angle: gain, offset)
    2-3: caY, intY (Y angle: gain, offset)
    4-5: caT, intT (temperature: gain, offset)
    
    Args:
        angle_data: (n_timestamps, n_sensors*2) raw angle values (ax, ay)
        temperature: (n_timestamps, n_sensors) raw temperature
        calibration_data: (n_sensors, 6) calibration parameters
        n_sensors: Number of PCL sensors
        sensor_type: 'PCL' or 'PCLHR'
        
    Returns:
        Tuple of (angles_converted, temp_converted, err_flag)
    """
    angles_converted = np.zeros_like(angle_data)
    temp_converted = np.zeros_like(temperature)
    err_flag = np.zeros_like(temperature)
    
    for sensor_idx in range(n_sensors):
        caX = calibration_data[sensor_idx, 0]
        intX = calibration_data[sensor_idx, 1]
        caY = calibration_data[sensor_idx, 2]
        intY = calibration_data[sensor_idx, 3]
        caT = calibration_data[sensor_idx, 4]
        intT = calibration_data[sensor_idx, 5]
        
        # Convert angles
        angles_converted[:, sensor_idx*2] = angle_data[:, sensor_idx*2] * caX + intX
        angles_converted[:, sensor_idx*2+1] = angle_data[:, sensor_idx*2+1] * caY + intY
        
        # Convert temperature
        temp_converted[:, sensor_idx] = temperature[:, sensor_idx] * caT + intT
    
    return angles_converted, temp_converted, err_flag


def convert_tube_link_data(angle_data: np.ndarray, temperature: np.ndarray,
                          calibration_data: np.ndarray, n_sensors: int
                          ) -> Tuple[np.ndarray, np.ndarray, np.ndarray]:
    """
    Convert TuL raw angles to physical units.
    
    Calibration data columns (per sensor):
    0-1: caX, intX (X angle: gain, offset)
    2-3: caY, intY (Y angle: gain, offset)
    4-5: caZ, intZ (Z angle/correlation: gain, offset)
    6-7: caT, intT (temperature: gain, offset)
    
    Args:
        angle_data: (n_timestamps, n_sensors*3) raw angle values (ax, ay, az)
        temperature: (n_timestamps, n_sensors) raw temperature
        calibration_data: (n_sensors, 8) calibration parameters
        n_sensors: Number of TuL sensors
        
    Returns:
        Tuple of (angles_converted, temp_converted, err_flag)
    """
    angles_converted = np.zeros_like(angle_data)
    temp_converted = np.zeros_like(temperature)
    err_flag = np.zeros_like(temperature)
    
    for sensor_idx in range(n_sensors):
        caX = calibration_data[sensor_idx, 0]
        intX = calibration_data[sensor_idx, 1]
        caY = calibration_data[sensor_idx, 2]
        intY = calibration_data[sensor_idx, 3]
        caZ = calibration_data[sensor_idx, 4]
        intZ = calibration_data[sensor_idx, 5]
        caT = calibration_data[sensor_idx, 6]
        intT = calibration_data[sensor_idx, 7]
        
        # Convert 3D angles
        angles_converted[:, sensor_idx*3] = angle_data[:, sensor_idx*3] * caX + intX
        angles_converted[:, sensor_idx*3+1] = angle_data[:, sensor_idx*3+1] * caY + intY
        angles_converted[:, sensor_idx*3+2] = angle_data[:, sensor_idx*3+2] * caZ + intZ
        
        # Convert temperature
        temp_converted[:, sensor_idx] = temperature[:, sensor_idx] * caT + intT
    
    return angles_converted, temp_converted, err_flag
