import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
%matplotlib inline
import math as m
from scipy.signal import argrelextrema as argex
plt.style.use('ggplot')

data_dir = '../data/'
trial_data = np.loadtxt(data_dir+'20171010_cavendish_trial.txt', delimiter='\t')

plt.plot(trial_data[:,0], trial_data[:,1])
plt.title("Trial Data from Cavendish Balance")
plt.ylabel("Anglular Positon (mrads)")
plt.xlabel("Time (s)")

"""
Explanation: Analysis of Preliminary Trail Data Pulled from Cavendish Balance
End of explanation
"""


plt.plot(trial_data[3900:,0], trial_data[3900:,1])
plt.title("Trial Data from Cavendish Balance, Adjusted")
plt.ylabel("Anglular Positon (mrads)")
plt.xlabel("Time (s)")

np.savetxt(data_dir+'20171010_cavendish_trial_useable.txt', trial_data[3200:,], delimiter=',')

x_data = trial_data[3900:,0]
y_data = trial_data[3900:,1]

#I want to try to extract the peaks of the corresponding sine waves to show the exponential decay and use that 
#to fit my curve to.

angles = []
time = []

for i in range(1, len(y_data)-1):
    if y_data[i] >= np.average(y_data):
        if (y_data[i] > y_data[i-1]) and (y_data[i] > y_data[i+1]):
            angles.append(float(y_data[i]))
            time.append(float(x_data[i]))
            
inds = argex(np.array(angles), np.greater)
times = np.array(time)[inds]
turning_points = np.array(angles)[inds]
plt.plot(times, turning_points, 'g--')
plt.plot(x_data, y_data, 'b-')

"""
Explanation: The weird behavior at the beginning occured when we were making an alteration to the experimental setup itself (doing one of our many adjustments to try and zero out our $\theta_e$. We can ignore this as it is not indicative of our data and look at where it moves into a damped harmonic oscialltion.
End of explanation
"""


from scipy.optimize import curve_fit

def decay(t, a, b, w, phi, theta_0):
    return a*np.exp(-b*t)*np.cos(w*t + phi) + theta_0

popt, pcov = curve_fit(decay, x_data, y_data , p0 = (-32, 1.3e-3, 3e-2, -6e-1, 0))

popt

plt.plot(times, turning_points, 'b--', label = 'Decay from Data')
plt.plot(x_data, y_data, 'r-', linewidth = 5, label = "Raw Data")
plt.plot(x_data, decay(x_data, *popt), 'g-', label = 'Fit of a*np.exp(-b*t)*np.cos(w*t + phi) + theta_0')
plt.title("Free Oscillation Data, No Large Masses")
plt.ylabel("Angle, Not Calibrated (mrad)")
plt.xlabel("Time (s)")
plt.legend()

round(popt[1],5)

"""
Explanation: Trying out the scipy.optimzie library to fit this to a decaying sinuisodal curve.
End of explanation
"""


