import openhtf as htf
from openhtf import Test, measures
from openhtf.plugs import plug
from openhtf.util import units
from tofupilot.openhtf import TofuPilot

from plugs.mock_dut import MockDutPlug
from utils.calibrate_sensor import calibrate_sensor
from utils.compute_noise_density import compute_noise_density
from utils.compute_r2 import compute_r2
from utils.compute_residuals import compute_residuals
from utils.compute_temp_sensitivity import compute_temp_sensitivity

@plug(dut=MockDutPlug)
def connect_dut(test: Test, dut: MockDutPlug) -> None:
    """Connect to the Device Under Test (DUT)."""
    dut.connect()

@measures(
    *(htf.Measurement("{sensor}_noise_density_{axis}")
      .doc('Noise density, normalized to √Hz')
      .in_range(0.0, {"acc": 0.003, "gyro": 0.005}.get(sensor))
      .with_units({"acc": units.METRE_PER_SECOND_SQUARED, "gyro": units.DEGREE_PER_SECOND}.get(sensor))
      .with_args(sensor=sensor, axis=axis)
      for sensor in ("acc", "gyro") for axis in ("x", "y", "z")),

    *(htf.Measurement("{sensor}_temp_sensitivity_max_{axis}")
      .doc('Max temperature sensitivity (unit/°C)')
      .with_units({"acc": units.METRE_PER_SECOND_SQUARED, "gyro": units.DEGREE_PER_SECOND}.get(sensor))
      .with_args(sensor=sensor, axis=axis)
      for sensor in ("acc", "gyro") for axis in ("x", "y", "z")),

    *(htf.Measurement("{sensor}_temp_sensitivity_ref_{axis}")
      .doc('Temperature sensitivity at 25°C (unit/°C)')
      .in_range(
        {"acc": {"x": 5e-4, "y": 5e-4, "z": 5e-4}, "gyro": {"x": 6e-5, "y": 6e-5, "z": 6e-5}}[sensor][axis], 
        {"acc": {"x": 1e-2, "y": 1e-2, "z": 1e-2}, "gyro": {"x": 1e0, "y": 1e0, "z": 1e0}}[sensor][axis])
      .with_units({"acc": units.METRE_PER_SECOND_SQUARED, "gyro": units.DEGREE_PER_SECOND}.get(sensor))
      .with_args(sensor=sensor, axis=axis)
      for sensor in ("acc", "gyro") for axis in ("x", "y", "z")),
)
@plug(dut=MockDutPlug)
def get_calibration_data(test: Test, dut: MockDutPlug) -> None:
    """Retrieve calibration data from the DUT."""
    test.state.update(dut.get_imu_data(test))

    for sensor, data_key in [("acc", "acc_data"), ("gyro", "gyro_data")]:
        sensor_data = test.state[data_key]
        temperature = sensor_data["temperature"]
        axes_data = {axis: sensor_data[f"{sensor}_{axis}"] for axis in ["x", "y", "z"]}

        for axis, axis_data in axes_data.items():
            noise_density = compute_noise_density(axis_data)
            temp_sensitivity = compute_temp_sensitivity(axis_data, temperature)

            test.measurements[f"{sensor}_noise_density_{axis}"] = noise_density
            test.measurements[f"{sensor}_temp_sensitivity_max_{axis}"] = temp_sensitivity["max_sensitivity"]
            test.measurements[f"{sensor}_temp_sensitivity_ref_{axis}"] = temp_sensitivity["sensitivity_at_ref"]


@measures(
    *(htf.Measurement("{sensor}_polynomial_coefficients_{axis}")
      .doc("Calibration polynomial coefficient matrix")
      .with_args(sensor=sensor, axis=axis)
      for sensor in ("acc", "gyro") for axis in ("x", "y", "z")),

    *(htf.Measurement("{sensor}_residual_mean_{axis}")
      .doc("Residual mean")
      .in_range(0.0, {"acc": 0.01, "gyro": 0.01}.get(sensor))
      .with_units({"acc": units.METRE_PER_SECOND_SQUARED, "gyro": units.DEGREE_PER_SECOND}.get(sensor))
      .with_args(sensor=sensor, axis=axis)
      for sensor in ("acc", "gyro") for axis in ("x", "y", "z")),

    *(htf.Measurement("{sensor}_residual_std_{axis}")
      .doc("Residual standard deviation")
      .in_range(0.0, {"acc": 5.0, "gyro": 0.3}.get(sensor))
      .with_units({"acc": units.METRE_PER_SECOND_SQUARED, "gyro": units.DEGREE_PER_SECOND}.get(sensor))
      .with_args(sensor=sensor, axis=axis)
      for sensor in ("acc", "gyro") for axis in ("x", "y", "z")),

    *(htf.Measurement("{sensor}_residual_p2p_{axis}")
      .doc("Residual peak-to-peak")
      .in_range(0.0, {"acc": {"x": 15.0, "y": 15.0, "z": 35.0}, "gyro": {"x": 2.0, "y": 2.0, "z": 2.0}}[sensor][axis])
      .with_units({"acc": units.METRE_PER_SECOND_SQUARED, "gyro": units.DEGREE_PER_SECOND}.get(sensor))
      .with_args(sensor=sensor, axis=axis)
      for sensor in ("acc", "gyro") for axis in ("x", "y", "z")),

    *(htf.Measurement("{sensor}_r2_{axis}")
      .doc("Coefficient of determination R² (unitless)")
      .in_range(0.5, 1.0)
      .with_args(sensor=sensor, axis=axis)
      for sensor in ("acc", "gyro") for axis in ("x", "y", "z")),
)
def compute_sensors_calibration(test: Test) -> None:
    """Perform calibration and metrics computation for both accelerometer and gyroscope."""
    for sensor, data_key, calibration_key in [
        ("acc", "acc_data", "acc_calibration_results"),
        ("gyro", "gyro_data", "gyro_calibration_results"),
    ]:
        sensor_data = test.state[data_key]
        temperature = sensor_data["temperature"]
        axes_data = {axis: sensor_data[f"{sensor}_{axis}"] for axis in ["x", "y", "z"]}

        test.state[calibration_key] = calibrate_sensor(
            (temperature, *axes_data.values()), sensor
        )

        for axis, axis_data in axes_data.items():
            fitted_values = test.state[calibration_key]["fitted_values"][f"{axis}_axis"]
            residuals = compute_residuals(axis_data, fitted_values)
            r2 = compute_r2(axis_data, fitted_values)

            # Update measurements
            test.measurements[f"{sensor}_polynomial_coefficients_{axis}"] = (
                test.state[calibration_key]["polynomial_coefficients"][f"{axis}_axis"].tolist()
            )
            test.measurements[f"{sensor}_residual_mean_{axis}"] = abs(residuals["mean_residual"])
            test.measurements[f"{sensor}_residual_std_{axis}"] = residuals["std_residual"]
            test.measurements[f"{sensor}_residual_p2p_{axis}"] = residuals["p2p_residual"]
            test.measurements[f"{sensor}_r2_{axis}"] = r2

        for axis, fig in zip(["x", "y", "z"], test.state[calibration_key]["figures"]):
            test.attach(f"{sensor}_calibration_figure_{axis}", fig.getvalue(), "image/png")


@plug(dut=MockDutPlug)
def save_calibration(test: Test, dut: MockDutPlug) -> None:
    """Save calibration data to the DUT."""
    dut.save_accelerometer_calibration(test.state["acc_calibration_results"])
    dut.save_gyroscope_calibration(test.state["gyro_calibration_results"])


def main():
    test = htf.Test(
        connect_dut,
        get_calibration_data,
        compute_sensors_calibration,
        save_calibration,
        procedure_id="FVT1",
        procedure_name="IMU Thermal Calibration",
        part_number="PCB01",
        part_name="Motherboard PCBA",
    )

    with TofuPilot(test):
        test.execute(lambda: "00001")  # mock operator S/N input


if __name__ == "__main__":
    main()

import time
from pathlib import Path
from typing import Dict

from openhtf import Test
from openhtf.plugs import BasePlug
from pandas import read_csv, DataFrame


class MockDutPlug(BasePlug):
    """
    A mock plug for simulating communication with a DUT.

    Provides methods to connect, disconnect, and save calibration data for the DUT.
    """

    def connect(self) -> None:
        self.logger.info("Simulated: Connecting to DUT.")
        time.sleep(1)

    def disconnect(self) -> None:
        self.logger.info("Simulated: Disconnecting from DUT.")
        time.sleep(1)

    @staticmethod
    def get_imu_data(test: Test) -> Dict:
        data = read_csv("data/imu_raw_data.csv", delimiter="\t")
        test.attach("raw_calibration_data", data.to_csv(), "text/csv")
        return {"acc_data":
                   {
                        "temperature": data["imu.temperature"],
                        "acc_x": data["imu.acc.x"],
                        "acc_y": data["imu.acc.y"],
                        "acc_z": data["imu.acc.z"] - 9.80600,  # Acceleration of freefall in Switzerland zone 2
                    },
                "gyro_data":
                   {
                       "temperature": data["imu.temperature"],
                       "gyro_x": data["imu.gyro.x"],
                       "gyro_y": data["imu.gyro.y"],
                       "gyro_z": data["imu.gyro.z"],
                   }
        }

    def save_accelerometer_calibration(self, polynomial_coefficients: dict) -> None:
        self.logger.info("Simulated: Saving IMU thermal calibration to DUT.")
        time.sleep(0.5)

    def save_gyroscope_calibration(self, polynomial_coefficients: dict) -> None:
        self.logger.info("Simulated: Saving IMU thermal calibration to DUT.")
        time.sleep(0.5)

    def tearDown(self) -> None:
        """
        OpenHTF automatically calls the tearDown method after the test phase ends.
        This ensures that any required cleanup (like disconnecting from the DUT) is performed.
        """
        self.logger.info("Simulated: Performing teardown.")
        self.disconnect()

import numpy as np

def compute_noise_density(data: np.ndarray, sampling_rate: int = 100) -> float:
    """
    Computes noise density for sensor data.

    Parameters:
    data (numpy.ndarray): Raw sensor data.
    sampling_rate (int): Sensor sampling rate in Hz.

    Returns:
    float: Noise density in units/sqrt(Hz).
    """
    # Use the first 50 samples for noise analysis.
    initial_samples = data[:50]
    
    # Remove DC offset by subtracting the mean.
    detrended_data = initial_samples - np.mean(initial_samples)
    
    # Compute noise standard deviation.
    noise_std = np.std(detrended_data)
    
    # Normalize noise to the frequency domain.
    noise_density = noise_std / np.sqrt(sampling_rate)
    
    return noise_density

import numpy as np

def compute_temp_sensitivity(
    data: np.ndarray, temperatures: np.ndarray, temp_ref: float = 25
) -> dict:
    """
    Computes temperature sensitivity for sensor data.

    Parameters:
    data (numpy.ndarray): Sensor data.
    temperatures (numpy.ndarray): Corresponding temperatures for the data.
    temp_ref (float): Reference temperature for sensitivity calculations (default: 25°C).

    Returns:
    dict: Maximum sensitivity and sensitivity at reference temperature.

    Use this function to assess how sensor readings change with temperature.
    """
    # Calculate temperature changes between consecutive samples.
    d_temp = np.diff(temperatures)
    
    # Ignore negligible temperature changes to avoid division errors.
    valid_idx = np.abs(d_temp) > 1e-5

    # Compute data changes and sensitivities for valid temperature changes.
    d_data = np.diff(data)
    sensitivities = d_data[valid_idx] / d_temp[valid_idx]

    # Find the index closest to the reference temperature.
    ref_idx = np.argmin(np.abs(temperatures - temp_ref))

    # Select a range around the reference temperature for averaging.
    ref_range = slice(max(ref_idx - 10, 0), min(ref_idx + 10, len(sensitivities)))

    # Average sensitivities near the reference temperature.
    sensitivity_ref = np.mean(sensitivities[ref_range])

    return {
        "max_sensitivity": np.max(np.abs(sensitivities)),
        "sensitivity_at_ref": abs(sensitivity_ref),
    }

import numpy as np
import matplotlib
matplotlib.use("Agg")
import matplotlib.pyplot as plt
from io import BytesIO
from typing import Dict, Tuple, List

def calibrate_sensor(
        data: Tuple[List[float], List[float], List[float], List[float]],
        sensor: str, polynomial_order: int = 3
        ) -> Dict[str, Dict[str, np.ndarray]]:
    """
    Fit a polynomial model to the sensor data against temperature.

    Args:
    data (tuple): A tuple containing:
    - temperature data (numpy array)
    - sensor data for each axis (numpy arrays for x, y, z)
    polynomial_order (int): The degree of the polynomial to fit.

    Returns:
    dict: A dictionary containing:
    - polynomial_coefficients: Coefficients of the fitted polynomial for each axis.
    - fitted_values: Fitted values for each axis at the given temperature points.
    - figures: Matplotlib figure objects as in-memory images for each axis.
    """

    # Define color names
    colors = {
        "zinc": "#09090B",
        "white": "#ffffff",
        "lime": "#bef264",
        "pink": "#f9a8d4",
    }

    # Convert the tuple elements to NumPy arrays
    temp, *sensor_data = (np.array(arr) for arr in data)

    poly_coeffs: Dict[str, np.ndarray] = {}
    fitted_values: Dict[str, np.ndarray] = {}
    figures: List[BytesIO] = []
    axis_list = ('x', 'y', 'z')

    for i, axis_data in enumerate(sensor_data):
        if sensor == "acc":
            sensor_name = "Accelerometer"
            unit = "m/s²"
        else:
            sensor_name = "Gyroscope"
            unit = "°/s"
        axis_name = f"{axis_list[i]}_axis"

        # Fit polynomial to the data
        coeffs = np.polyfit(temp, axis_data, polynomial_order)
        poly_coeffs[axis_name] = coeffs

        # Compute fitted values
        fitted = np.polyval(coeffs, temp)
        fitted_values[axis_name] = fitted

        # Calculate residuals
        residuals = axis_data - fitted

        # Generate plot and store figure
        fig, axs = plt.subplots(2, 1, figsize=(8, 10), gridspec_kw={'height_ratios': [3, 1]})
        fig.patch.set_facecolor(colors["zinc"])
        axs[0].set_facecolor(colors["zinc"])
        axs[1].set_facecolor(colors["zinc"])

        # Plot sensor data with the fitted curve
        axs[0].plot(temp, axis_data, ".", color=colors["lime"], label=f"{sensor_name} data")
        axs[0].plot(temp, fitted, "-", color=colors["pink"], label="Fitted Curve")
        axs[0].set_title(f"{sensor_name} {axis_name[0].capitalize()} axis calibration", color=colors["white"])
        axs[0].set_xlabel("Temperature (°C)", color=colors["white"])
        axs[0].set_ylabel(f"{sensor_name} value ({unit})", color=colors["white"])
        axs[0].tick_params(colors=colors["white"])
        axs[0].spines['bottom'].set_color(colors["white"])
        axs[0].spines['left'].set_color(colors["white"])
        axs[0].spines['top'].set_color(colors["white"])
        axs[0].spines['right'].set_color(colors["white"])
        axs[0].legend(facecolor=colors["zinc"], edgecolor=colors["white"], labelcolor=colors["white"])

        # Plot histogram of residuals
        axs[1].hist(residuals, bins=30, color=colors["lime"], edgecolor=colors["white"], alpha=0.8)
        axs[1].set_title(f"Residual distribution ({sensor_name.lower()} {axis_name[0].capitalize()} axis)", color=colors["white"])
        axs[1].set_xlabel(f"Residual value ({unit})", color=colors["white"])
        axs[1].set_ylabel("Occurences", color=colors["white"])
        axs[1].tick_params(colors=colors["white"])
        axs[1].spines['bottom'].set_color(colors["white"])
        axs[1].spines['left'].set_color(colors["white"])
        axs[1].spines['top'].set_color(colors["white"])
        axs[1].spines['right'].set_color(colors["white"])

        # Add spacing between subplots
        plt.subplots_adjust(hspace=0.3)

        # Convert the figure to an in-memory image
        buffer = BytesIO()
        plt.savefig(buffer, format="png")
        buffer.seek(0)
        figures.append(buffer)
        plt.close(fig)

    return {
        "polynomial_coefficients": poly_coeffs,
        "fitted_values": fitted_values,
        "figures": figures,
    }

import numpy as np

def compute_residuals(data: np.ndarray, fit_model: np.ndarray) -> dict:
    """
    Computes residuals between sensor data and a fitted model.

    Parameters:
    data (numpy.ndarray): Sensor data.
    fit_model (numpy.ndarray): Fitted model values for the same data.

    Returns:
    dict: Mean, standard deviation, and peak-to-peak residuals.

    Use this function to evaluate calibration quality by analyzing deviations.
    """
    # Residuals: differences between data and model.
    residuals = data - fit_model
    
    return {
        # Mean residual: measures systematic error.
        "mean_residual": np.mean(residuals),

        # Standard deviation: indicates scatter.
        "std_residual": np.std(residuals),
        
        # Peak-to-peak: captures extreme errors.
        "p2p_residual": np.ptp(residuals),
    }

import numpy as np

def compute_r2(data: np.ndarray, fit_model: np.ndarray) -> float:
    """
    Computes R-squared to evaluate how well a model fits the data.

    Parameters:
    data (numpy.ndarray): Sensor data.
    fit_model (numpy.ndarray): Fitted model values for the data.

    Returns:
    float: R-squared value.

    Use this function to quantify the goodness-of-fit of the calibration model.
    """
    # Residuals: differences between data and model.
    residuals = data - fit_model

    # Total variation in the data.
    total_variation = np.sum((data - np.mean(data)) ** 2)
    
    # Variation not explained by the model.
    residual_variation = np.sum(residuals**2)
    
    # R²: fraction of variance explained by the model.
    r2 = 1 - residual_variation / total_variation
    return r2