#!/usr/bin/env python3
# -*- coding: utf-8 -*-

"""
File: mag_compensation.py
Author: Tanja Baumann
Email: tanja@auterion.com
Github: https://github.com/baumanta
Description:
    Computes linear coefficients for mag compensation from thrust and current
Usage:
    python mag_compensation.py /path/to/log/logfile.ulg current --instance 1

Remark:
    If your logfile does not contain some of the topics, e.g.battery_status/current_a
    you will have to comment out the corresponding parts in the script
"""


import matplotlib.pylab as plt
from mpl_toolkits.mplot3d import Axes3D
from pyulog import ULog
from pyulog.px4 import PX4ULog
from pylab import *
import numpy as np
import textwrap as tw
import argparse

#arguments
parser = argparse.ArgumentParser(description='Calculate compensation parameters from ulog')

parser.add_argument('logfile', type=str, nargs='?', default=[],
                    help='full path to ulog file')
parser.add_argument('type', type=str, nargs='?', choices=['current', 'thrust'], default=[],
                    help='Power signal used for compensation, supported is "current" or "thrust".')
parser.add_argument('--instance', type=int, nargs='?', default=0,
                    help='instance of the current or thrust signal to use (0 or 1)')

args = parser.parse_args()
log_name = args.logfile
comp_type = args.type
comp_instance = args.instance

#Load the log data (produced by pyulog)
log = ULog(log_name)
pxlog = PX4ULog(log)

def get_data(topic_name, variable_name, index):
   try:
      dataset = log.get_dataset(topic_name, index)
      return dataset.data[variable_name]
   except:
      return []

def ms2s_list(time_ms_list):
    if len(time_ms_list) > 0:
        return 1e-6 * time_ms_list
    else:
        return time_ms_list

# Select msgs and copy into arrays
armed = get_data('vehicle_status', 'arming_state', 0)
t_armed =  ms2s_list(get_data('vehicle_status', 'timestamp', 0))

if comp_type == "thrust":
	power = get_data('vehicle_rates_setpoint', 'thrust_body[2]', comp_instance)
	power_t = ms2s_list(get_data('vehicle_rates_setpoint', 'timestamp', comp_instance))
	comp_type_param = 1
	factor = 1
	unit = "[G]"
elif comp_type == "current":
	power = get_data('battery_status', 'current_a', comp_instance)
	power = np.true_divide(power, 1000) #kA
	power_t = ms2s_list(get_data('battery_status', 'timestamp', comp_instance))
	comp_type_param = 2 + comp_instance
	factor = -1
	unit = "[G/kA]"
else:
	print("unknown compensation type {}. Supported is either 'thrust' or 'current'.".format(comp_type))
	sys.exit(1)

if len(power) == 0:
	print("could not retrieve power signal from log, zero data points")
	sys.exit(1)

mag0X_body = get_data('sensor_mag', 'x', 0)
mag0Y_body = get_data('sensor_mag', 'y', 0)
mag0Z_body = get_data('sensor_mag', 'z', 0)
t_mag0 =  ms2s_list(get_data('sensor_mag', 'timestamp', 0))
mag0_ID = get_data('sensor_mag', 'device_id', 0)

mag1X_body = get_data('sensor_mag', 'x', 1)
mag1Y_body = get_data('sensor_mag', 'y', 1)
mag1Z_body = get_data('sensor_mag', 'z', 1)
t_mag1 = ms2s_list(get_data('sensor_mag', 'timestamp', 1))
mag1_ID = get_data('sensor_mag', 'device_id', 1)

mag2X_body = get_data('sensor_mag', 'x', 2)
mag2Y_body = get_data('sensor_mag', 'y', 2)
mag2Z_body = get_data('sensor_mag', 'z', 2)
t_mag2 = ms2s_list(get_data('sensor_mag', 'timestamp', 2))
mag2_ID = get_data('sensor_mag', 'device_id', 2)

mag3X_body = get_data('sensor_mag', 'x', 3)
mag3Y_body = get_data('sensor_mag', 'y', 3)
mag3Z_body = get_data('sensor_mag', 'z', 3)
t_mag3 = ms2s_list(get_data('sensor_mag', 'timestamp', 3))
mag3_ID = get_data('sensor_mag', 'device_id', 3)

magX_body = []
magY_body = []
magZ_body = []
mag_id = []
t_mag = []

if len(mag0X_body) > 0:
    magX_body.append(mag0X_body)
    magY_body.append(mag0Y_body)
    magZ_body.append(mag0Z_body)
    t_mag.append(t_mag0)
    mag_id.append(mag0_ID[0])

if len(mag1X_body) > 0:
    magX_body.append(mag1X_body)
    magY_body.append(mag1Y_body)
    magZ_body.append(mag1Z_body)
    t_mag.append(t_mag1)
    mag_id.append(mag1_ID[0])

if len(mag2X_body) > 0:
    magX_body.append(mag2X_body)
    magY_body.append(mag2Y_body)
    magZ_body.append(mag2Z_body)
    t_mag.append(t_mag2)
    mag_id.append(mag2_ID[0])

if len(mag3X_body) > 0:
    magX_body.append(mag3X_body)
    magY_body.append(mag3Y_body)
    magZ_body.append(mag3Z_body)
    t_mag.append(t_mag3)
    mag_id.append(mag3_ID[0])

n_mag = len(magX_body)

#log index does not necessarily match mag calibration instance number
calibration_instance = []
instance_found = False
for idx in range(n_mag):
    instance_found = False
    for j in range(4):
        if mag_id[idx] == log.initial_parameters["CAL_MAG{}_ID".format(j)]:
                calibration_instance.append(j)
                instance_found = True
    if not instance_found:
        print('Mag {} calibration instance not found, run compass calibration first.'.format(mag_id[idx]))

#get first arming sequence from data
start_time = 0
stop_time = 0
for i in range(len(armed)-1):
    if armed[i] == 1 and armed[i+1] == 2:
        start_time = t_armed[i+1]
    if armed[i] == 2 and armed[i+1] == 1:
        stop_time = t_armed[i+1]
        break

#cut unarmed sequences from mag data
index_start = 0
index_stop = 0

for idx in range(n_mag):
    for i in range(len(t_mag[idx])):
        if t_mag[idx][i] > start_time:
            index_start = i
            break

    for i in range(len(t_mag[idx])):
        if t_mag[idx][i] > stop_time:
            index_stop = i -1
            break

    t_mag[idx] = t_mag[idx][index_start:index_stop]
    magX_body[idx] = magX_body[idx][index_start:index_stop]
    magY_body[idx] = magY_body[idx][index_start:index_stop]
    magZ_body[idx] = magZ_body[idx][index_start:index_stop]


#resample data
power_resampled = []

for idx in range(n_mag):
    power_resampled.append(interp(t_mag[idx], power_t, power))

#fit linear to get coefficients
px = []
py = []
pz = []

for idx in range(n_mag):
    px_temp, res_x, _, _, _ = polyfit(power_resampled[idx], magX_body[idx], 1,full = True)
    py_temp, res_y, _, _, _ = polyfit(power_resampled[idx], magY_body[idx], 1,full = True)
    pz_temp, res_z, _, _, _ = polyfit(power_resampled[idx], magZ_body[idx], 1, full = True)

    px.append(px_temp)
    py.append(py_temp)
    pz.append(pz_temp)

#print to console
for idx in range(n_mag):
    print('Mag{} device ID {} (calibration instance {})'.format(idx, mag_id[idx], calibration_instance[idx]))
print('\033[91m \n{}-based compensation: \033[0m'.format(comp_type))
print('\nparam set CAL_MAG_COMP_TYP {}'.format(comp_type_param))
for idx in range(n_mag):
    print('\nparam set CAL_MAG{}_XCOMP {:.3f}'.format(calibration_instance[idx], factor * px[idx][0]))
    print('param set CAL_MAG{}_YCOMP {:.3f}'.format(calibration_instance[idx], factor * py[idx][0]))
    print('param set CAL_MAG{}_ZCOMP {:.3f}'.format(calibration_instance[idx], factor * pz[idx][0]))

#plot data

for idx in range(n_mag):
    fig = plt.figure(num=None, figsize=(25, 14), dpi=80, facecolor='w', edgecolor='k')
    fig.suptitle('Compensation Parameter Fit \n{} \nmag {} ID: {} (calibration instance {})'.format(log_name, idx, mag_id[idx], calibration_instance[idx]), fontsize=14, fontweight='bold')

    plt.subplot(1,3,1)
    plt.plot(power_resampled[idx], magX_body[idx], 'yo', power_resampled[idx], px[idx][0]*power_resampled[idx]+px[idx][1], '--k')
    plt.xlabel('current [kA]')
    plt.ylabel('mag X [G]')

    plt.subplot(1,3,2)
    plt.plot(power_resampled[idx], magY_body[idx], 'yo', power_resampled[idx], py[idx][0]*power_resampled[idx]+py[idx][1], '--k')
    plt.xlabel('current [kA]')
    plt.ylabel('mag Y [G]')

    plt.subplot(1,3,3)
    plt.plot(power_resampled[idx], magZ_body[idx], 'yo', power_resampled[idx], pz[idx][0]*power_resampled[idx]+pz[idx][1], '--k')
    plt.xlabel('current [kA]')
    plt.ylabel('mag Z [G]')

    # display results
    plt.figtext(0.24, 0.03, 'CAL_MAG{}_XCOMP: {:.3f} {}'.format(calibration_instance[idx],factor * px[idx][0],unit), horizontalalignment='center', fontsize=12, multialignment='left', bbox=dict(boxstyle="round", facecolor='#D8D8D8', ec="0.5", pad=0.5, alpha=1), fontweight='bold')
    plt.figtext(0.51, 0.03, 'CAL_MAG{}_YCOMP: {:.3f} {}'.format(calibration_instance[idx],factor * py[idx][0],unit), horizontalalignment='center', fontsize=12, multialignment='left', bbox=dict(boxstyle="round", facecolor='#D8D8D8', ec="0.5", pad=0.5, alpha=1), fontweight='bold')
    plt.figtext(0.79, 0.03, 'CAL_MAG{}_ZCOMP: {:.3f} {}'.format(calibration_instance[idx],factor * pz[idx][0],unit), horizontalalignment='center', fontsize=12, multialignment='left', bbox=dict(boxstyle="round", facecolor='#D8D8D8', ec="0.5", pad=0.5, alpha=1), fontweight='bold')


#compensation comparison plots
for idx in range(n_mag):
    fig = plt.figure(num=None, figsize=(25, 14), dpi=80, facecolor='w', edgecolor='k')
    fig.suptitle('Original Data vs. Compensation \n{}\nmag {} ID: {} (calibration instance {})'.format(log_name, idx, mag_id[idx], calibration_instance[idx]), fontsize=14, fontweight='bold')

    plt.subplot(3,1,1)
    original_x, = plt.plot(t_mag[idx], magX_body[idx], label='original')
    power_x, = plt.plot(t_mag[idx],magX_body[idx] - px[idx][0] * power_resampled[idx], label='compensated')
    plt.legend(handles=[original_x, power_x])
    plt.xlabel('Time [s]')
    plt.ylabel('Mag X corrected[G]')

    plt.subplot(3,1,2)
    original_y, = plt.plot(t_mag[idx], magY_body[idx], label='original')
    power_y, = plt.plot(t_mag[idx],magY_body[idx] - py[idx][0] * power_resampled[idx], label='compensated')
    plt.legend(handles=[original_y, power_y])
    plt.xlabel('Time [s]')
    plt.ylabel('Mag Y corrected[G]')

    plt.subplot(3,1,3)
    original_z, = plt.plot(t_mag[idx], magZ_body[idx], label='original')
    power_z, = plt.plot(t_mag[idx],magZ_body[idx] - pz[idx][0] * power_resampled[idx], label='compensated')
    plt.legend(handles=[original_z, power_z])
    plt.xlabel('Time [s]')
    plt.ylabel('Mag Z corrected[G]')

plt.show()

[(30, 'argparse.ArgumentParser', 'argparse.ArgumentParser', False, 'import argparse\n'), (45, 'pyulog.ULog', 'ULog', False, 'from pyulog import ULog\n'), (46, 'pyulog.px4.PX4ULog', 'PX4ULog', False, 'from pyulog.px4 import PX4ULog\n'), (272, 'matplotlib.pylab.show', 'plt.show', True, 'import matplotlib.pylab as plt\n'), (222, 'matplotlib.pylab.figure', 'plt.figure', True, 'import matplotlib.pylab as plt\n'), (225, 'matplotlib.pylab.subplot', 'plt.subplot', True, 'import matplotlib.pylab as plt\n'), (226, 'matplotlib.pylab.plot', 'plt.plot', True, 'import matplotlib.pylab as plt\n'), (227, 'matplotlib.pylab.xlabel', 'plt.xlabel', True, 'import matplotlib.pylab as plt\n'), (228, 'matplotlib.pylab.ylabel', 'plt.ylabel', True, 'import matplotlib.pylab as plt\n'), (230, 'matplotlib.pylab.subplot', 'plt.subplot', True, 'import matplotlib.pylab as plt\n'), (231, 'matplotlib.pylab.plot', 'plt.plot', True, 'import matplotlib.pylab as plt\n'), (232, 'matplotlib.pylab.xlabel', 'plt.xlabel', True, 'import matplotlib.pylab as plt\n'), (233, 'matplotlib.pylab.ylabel', 'plt.ylabel', True, 'import matplotlib.pylab as plt\n'), (235, 'matplotlib.pylab.subplot', 'plt.subplot', True, 'import matplotlib.pylab as plt\n'), (236, 'matplotlib.pylab.plot', 'plt.plot', True, 'import matplotlib.pylab as plt\n'), (237, 'matplotlib.pylab.xlabel', 'plt.xlabel', True, 'import matplotlib.pylab as plt\n'), (238, 'matplotlib.pylab.ylabel', 'plt.ylabel', True, 'import matplotlib.pylab as plt\n'), (248, 'matplotlib.pylab.figure', 'plt.figure', True, 'import matplotlib.pylab as plt\n'), (251, 'matplotlib.pylab.subplot', 'plt.subplot', True, 'import matplotlib.pylab as plt\n'), (252, 'matplotlib.pylab.plot', 'plt.plot', True, 'import matplotlib.pylab as plt\n'), (253, 'matplotlib.pylab.plot', 'plt.plot', True, 'import matplotlib.pylab as plt\n'), (254, 'matplotlib.pylab.legend', 'plt.legend', True, 'import matplotlib.pylab as plt\n'), (255, 'matplotlib.pylab.xlabel', 'plt.xlabel', True, 'import matplotlib.pylab as plt\n'), (256, 'matplotlib.pylab.ylabel', 'plt.ylabel', True, 'import matplotlib.pylab as plt\n'), (258, 'matplotlib.pylab.subplot', 'plt.subplot', True, 'import matplotlib.pylab as plt\n'), (259, 'matplotlib.pylab.plot', 'plt.plot', True, 'import matplotlib.pylab as plt\n'), (260, 'matplotlib.pylab.plot', 'plt.plot', True, 'import matplotlib.pylab as plt\n'), (261, 'matplotlib.pylab.legend', 'plt.legend', True, 'import matplotlib.pylab as plt\n'), (262, 'matplotlib.pylab.xlabel', 'plt.xlabel', True, 'import matplotlib.pylab as plt\n'), (263, 'matplotlib.pylab.ylabel', 'plt.ylabel', True, 'import matplotlib.pylab as plt\n'), (265, 'matplotlib.pylab.subplot', 'plt.subplot', True, 'import matplotlib.pylab as plt\n'), (266, 'matplotlib.pylab.plot', 'plt.plot', True, 'import matplotlib.pylab as plt\n'), (267, 'matplotlib.pylab.plot', 'plt.plot', True, 'import matplotlib.pylab as plt\n'), (268, 'matplotlib.pylab.legend', 'plt.legend', True, 'import matplotlib.pylab as plt\n'), (269, 'matplotlib.pylab.xlabel', 'plt.xlabel', True, 'import matplotlib.pylab as plt\n'), (270, 'matplotlib.pylab.ylabel', 'plt.ylabel', True, 'import matplotlib.pylab as plt\n'), (73, 'numpy.true_divide', 'np.true_divide', True, 'import numpy as np\n')]