import feather
from allan_variance_utils import *


file_path = './imu0.feather'

# Load csv into np array
df = feather.read_dataframe(file_path)  # also, you can read in csv file with pandas
freq = 400  # define your data's sample rate

# Transform gyro unit from [rad/s] to [deg/s]
df.iloc[:, 1:4] = df.iloc[:, 1:4] * 180 / PI
gyro_x = df.iloc[:, 1].to_numpy()

# Pass data into function
# Note that the function only accept one axis;
tau_x, allan_var_x = allan_variance(data=gyro_x, f=freq, max_clusters=200)
allan_dev_x = np.sqrt(allan_var_x)
plot_result(tau_x, allan_dev_x)

# Another version to compute Allan variance
allan_vars_x = allan_variance2(data=gyro_x, f=freq)
periods_x, allan_var_x = allan_vars_x[:, 0], allan_vars_x[:, 1]
allan_dev_x = np.sqrt(allan_var_x)
plot_result2(periods_x, allan_dev_x)

log_tau = np.log10(tau)
log_allan_dev = np.log10(allan_dev)
dlog_allan_dev = np.diff(log_allan_dev) / np.diff(log_tau)

slope_arw = -0.5
argmin_abs_i = np.argmin(np.abs(dlog_allan_dev - slope_arw))  # index that are most closed to -0.5

# Find the y-intercept of the line
intercept = log_allan_dev[argmin_abs_i] - slope_arw * log_tau[argmin_abs_i]

def _linear_func(x, a, b):
    return a * x + b

def _fit_intercept(x, y, a, tau):
    log_x, log_y = np.log(x), np.log(y)
    # a in range [m, m+0.001]; b in range [-Inf, Inf]
    coefs, _ = curve_fit(_linear_func, log_x, log_y, bounds=([a, -np.inf], [a + 0.001, np.inf]))
    poly = np.poly1d(coefs)
    print('Fitting polynomial equation:', np.poly1d(poly))
    y_fit = lambda x: np.exp(poly(np.log(x)))
    return y_fit(tau), y_fit

# White noise(velocity/angle random walk)
bp_wn = np.where(periods == 10)[0][0]  # white noise break point for short.
wn, fit_func_wn = _fit_intercept(periods[0:bp_wn], allan_dev[0:bp_wn], -0.5, 1.0)