"""API details."""

# AUTOGENERATED! DO NOT EDIT! File to edit: ../nbs/riann.ipynb.

# %% auto 0
__all__ = ['RIANN']

# %% ../nbs/riann.ipynb 3
import onnxruntime as rt
import numpy as np
import pkg_resources

# %% ../nbs/riann.ipynb 4
class RIANN:
    """
    Robust IMU-based Attitude Neural Network for orientation estimation using IMU data.
    
    This class implements efficient quaternion-based attitude estimation from 
    accelerometer and gyroscope data using a neural network approach. The implementation
    is optimized for step-by-step processing in real-time applications.
    
    Parameters
    ----------
    onnx_path : str, optional, default: ``None``
        Path to the RIANN onnx file. Defaults to the file provided with the package.

    Examples
    --------
    >>> import numpy as np
    >>> from riann.riann import RIANN
    >>> 
    >>> # Initialize RIANN and set sampling rate
    >>> riann = RIANN()
    >>> riann.set_sampling_rate(100)  # 100 Hz
    >>> 
    >>> # Process a single IMU measurement
    >>> acc = np.array([0.0, 0.0, 9.81])  # Acceleration in m/s²
    >>> gyr = np.array([0.0, 0.0, 0.0])   # Gyroscope in rad/s
    >>> quaternion = riann.predict_step(acc, gyr)
    >>> 
    >>> # Process a sequence of measurements
    >>> acc_seq = np.ones((100, 3))   # 100 measurements
    >>> gyr_seq = np.zeros((100, 3))  # 100 measurements
    >>> quaternions, states = riann.predict_sequence(acc_seq, gyr_seq, 100)
    """

    def __init__(self, onnx_path=None):
        """
        Initialize RIANN with optimized buffers for step-by-step processing.
        
        Parameters
        ----------
        onnx_path : str, optional, default: ``None``
            Path to the RIANN onnx file. Defaults to the file provided with the package.
        
        Notes
        -----
        Initializes with a default sampling rate of 200 Hz.
        """
        if onnx_path is None:
            try:
                onnx_path = pkg_resources.resource_string(__name__, "riann.onnx")
            except: #pkg_resources fails if the code is executed in the jupyter-notebook
                onnx_path = 'riann/riann.onnx'
        
        self.session = rt.InferenceSession(onnx_path)
        
        # Get input and output names for faster reference
        self.input_name = 'input'
        self.h0_name = 'h0'
        self.output_names = []  # Empty list for backward compatibility with run([])
        
        # Pre-allocated buffer for a single step (direct buffer assignment optimization)
        self.input_buffer = np.zeros((1, 1, 7), dtype=np.float32)
        
        self.set_sampling_rate(200.0)
        self.reset_state()
    
    def reset_state(self):
        """Reset the hidden state to initial values."""
        self.hidden_state = np.zeros((2, 1, 1, 200), dtype=np.float32)

    def set_sampling_rate(self, fs):
        """
        Set sampling rate for subsequent predict_step calls.
        
        Parameters
        ----------
        fs : float
            Sampling rate in Hz
        """
        self.fs = fs
        self.dt = 1.0 / fs
        self.input_buffer[0, 0, 6] = self.dt
    
    def predict_step(self, acc, gyr, fs=None):
        """
        Process a single step of IMU data, preserving the hidden state.
        
        Parameters
        ----------
        acc: numpy-array [3]
            Acceleration data of the IMU. The axis order is x,y,z.
        gyr: numpy-array [3]
            Gyroscope data of the IMU. The axis order is x,y,z.
        fs: float, optional
            Sampling rate of the provided IMU data. If None, uses the previously set rate.
            
        Returns
        -------
        attitude unit-quaternion [4]
        """
        if acc.shape[0] != 3 or gyr.shape[0] != 3:
            raise ValueError(f"Expected acc and gyr to have shape [3], got {acc.shape} and {gyr.shape}")
    

        self.input_buffer[0, 0, :3] = acc  # Acceleration components
        self.input_buffer[0, 0, 3:6] = gyr  # Gyroscope components
        if fs is not None: self.set_sampling_rate(fs)
        
        # Run inference with current hidden state
        outputs = self.session.run([], {self.input_name: self.input_buffer, self.h0_name: self.hidden_state})
        
        # Extract attitude
        attitude = outputs[0][0][0]
        
        # Update hidden state if model returns it
        if len(outputs) > 1:
            self.hidden_state = outputs[1]
        
        return attitude
    
    def predict(self, acc, gyr, fs):
        """
        Process a full sequence of IMU data at once (original behavior).
        
        Parameters
        ----------
        acc: numpy-array [sequence_length x 3]
            Acceleration data of the IMU. The axis order is x,y,z.
        gyr: numpy-array [sequence_length x 3]
            Gyroscope data of the IMU. The axis order is x,y,z.
        fs: float
            Sampling rate of the provided IMU data
            
        Returns
        -------
        attitude unit-quaternions [sequence_length x 4]
        """
        # For full-sequence processing, use the original implementation
        np_inp = np.concatenate([acc, gyr, np.tile(1/fs, (acc.shape[0], 1))], axis=-1).astype(np.float32)[None,...]
        
        # Save current hidden state
        old_state = self.hidden_state.copy()
        self.reset_state()
        
        outputs = self.session.run([], {self.input_name: np_inp, self.h0_name: self.hidden_state})
        
        # Restore the previous hidden state
        self.hidden_state = old_state
        
        return outputs[0][0]
    
    def predict_sequence(self, acc, gyr, fs, reset_state=True):
        """
        Process a sequence of IMU data step by step, preserving hidden state between steps.
        
        Parameters
        ----------
        acc: numpy-array [sequence_length x 3]
            Acceleration data of the IMU. The axis order is x,y,z.
        gyr: numpy-array [sequence_length x 3]
            Gyroscope data of the IMU. The axis order is x,y,z.
        fs: float
            Sampling rate of the provided IMU data
        reset_state: bool, default=True
            Whether to reset the hidden state before processing
            
        Returns
        -------
        tuple:
            - attitude unit-quaternions [sequence_length x 4]
            - hidden states [sequence_length x hidden_state_size]
        """
        if reset_state:
            self.reset_state()
        
        sequence_length = acc.shape[0]
        results = []
        hidden_states = []
        self.set_sampling_rate(fs)
        
        for i in range(sequence_length):
            quat = self.predict_step(acc[i], gyr[i])
            results.append(quat)
            hidden_states.append(self.hidden_state.copy())
        
        return np.array(results), np.array(hidden_states)
