#!/usr/bin/env python

import json
import time

from argaze import DataStructures
from argaze.TobiiGlassesPro2 import TobiiData

import numpy
from scipy.optimize import curve_fit

class TobiiInertialMeasureUnit():
    """Ease Tobbi IMU data handling."""

    def __init__(self):
        """Define IMU calibration data."""

        self.__gyroscope_offset = numpy.array([0., 0., 0.])
        self.__accelerometer_coefficients = numpy.array([[1., 0.], [1., 0.], [1., 0.]])

    def load_calibration_file(self, calibration_filepath):
        """Load IMU calibration from a .json  file."""

        with open(calibration_filepath) as calibration_file:

            # Deserialize .json
            # TODO find a better way
            calibration_data = json.load(calibration_file)

            # Load calibration data
            self.__gyroscope_offset = numpy.array(calibration_data['gyroscope_offset'])
            self.__accelerometer_coefficients = numpy.array(calibration_data['accelerometer_coefficients'])

    def save_calibration_file(self, calibration_filepath):
        """Save IMU calibration into .json file."""

        calibration_data = {
            'gyroscope_offset': list(self.__gyroscope_offset), 
            'accelerometer_coefficients': [list(self.__accelerometer_coefficients[0]), list(self.__accelerometer_coefficients[1]), list(self.__accelerometer_coefficients[2])]
            }

        with open(calibration_filepath, 'w', encoding='utf-8') as calibration_file:

            json.dump(calibration_data, calibration_file, ensure_ascii=False, indent=4)

    def calibrate_gyroscope_offset(self, gyroscope_ts_buffer):

        # Consider gyroscope values without timestamps
        gyroscope_values = []
        for ts, data_object in gyroscope_ts_buffer.items():
            gyroscope_values.append(data_object.value)

        # Calculate average value for each axis
        gx_offset = numpy.mean(numpy.array(gyroscope_values)[:, 0])
        gy_offset = numpy.mean(numpy.array(gyroscope_values)[:, 1])
        gz_offset = numpy.mean(numpy.array(gyroscope_values)[:, 2])

        # Store result
        self.__gyroscope_offset = numpy.array([gx_offset, gy_offset, gz_offset])

        return self.__gyroscope_offset

    def get_gyroscope_offset(self):
        """Get gyroscope offset."""

        return self.__gyroscope_offset

    def apply_gyroscope_offset(self, gyroscope_data_object):
        """Remove gyroscope offset to given gyroscope data."""

        return TobiiData.Gyroscope(gyroscope_data_object.value - self.__gyroscope_offset)

    def _accelerometer_linear_fit(self, x, a, b):

        return a * x + b

    def calibrate_accelerometer_axis_coefficients(self, axis, upward_ts_buffer, downward_ts_buffer, perpendicular_ts_buffer):
        """Calibrate one accelerometer axis using three data set (upward/+1g, downward/-1g, perpendicular/0g) for linear fit."""

        # Consider accelerometer axis values without timestamps
        accelerometer_values = []
        expected_values = []

        for (upward_ts, upward_data_object), (downward_ts, downward_data_object), (perpendicular_ts, perpendicular_data_object) in zip(upward_ts_buffer.items(), downward_ts_buffer.items(), perpendicular_ts_buffer.items()):
            
            accelerometer_values.append(upward_data_object.value[axis])
            expected_values.append(+1.0)

            accelerometer_values.append(downward_data_object.value[axis])
            expected_values.append(-1.0)

            accelerometer_values.append(perpendicular_data_object.value[axis])
            expected_values.append(0.0)

        # Find optimal coefficients according linear fit between accelerometer values and expected values
        optimal_coefficients, _ = curve_fit(self._accelerometer_linear_fit, accelerometer_values, expected_values, maxfev = 10000)

        # Store results for the given axis
        self.__accelerometer_coefficients[axis] = numpy.array(optimal_coefficients)

    def get_accelerometer_coefficients(self):
        """Get accelerometer coefficients."""

        return self.__accelerometer_coefficients

    def apply_accelerometer_coefficients(self, accelerometer_data_object):
        """Add accelerometer offset to given accelerometer data."""

        x = self._accelerometer_linear_fit(accelerometer_data_object.value[0], *self.__accelerometer_coefficients[0])
        y = self._accelerometer_linear_fit(accelerometer_data_object.value[1], *self.__accelerometer_coefficients[1])
        z = self._accelerometer_linear_fit(accelerometer_data_object.value[2], *self.__accelerometer_coefficients[2])

        return TobiiData.Accelerometer([x, y , z])