#!/usr/bin/env python

import numpy, time
from control.matlab import *
import matplotlib.pyplot as plt

def simulate_heater(total_time,timestep,Kp,Ki,Kd):
	num_steps = numpy.ceil(total_time/timestep)
	timestamps = linspace(0,total_time,num=num_steps)

	open_loop = tf([1], [1000, 50])

	controller = tf([Kd, Kp, Ki], [1, 0])
	system = open_loop*controller

	#bode(system)

	time_series, timestamps = step(system,timestamps)

	return time_series,timestamps

#PID values

Kp = 0
Ki = 0
Kd = 0

max_power = 0
max_power_limit = 0.8
increment = 100

total_time = 20
timestep = .1


while max_power<=max_power_limit:

	#increase Kp
	Kp = Kp + increment

	#get total_time s of data
	time_series, timestamps = simulate_heater(total_time,timestep,Kp,Ki,Kd)

	#for debugging
	plt.plot(timestamps, time_series)
	plt.xlabel('Time (s)')
	plt.ylabel('Magnitude')
	plt.title('Time Data')
	plt.show()

	data_fft = numpy.fft.fft(time_series, norm = "ortho")

	#convert to just magnitude
	data_fft = numpy.absolute(data_fft)

	freqs = numpy.fft.fftfreq(len(time_series), d=timestep)

	#get rid of DC component (need to modify this part to include envelope around 0)
	zero_index = numpy.where(freqs==0)
	data_fft[zero_index] = 0

	#for debugging
	plt.plot(freqs, data_fft)
	plt.xlabel('Frequency (Hz)')
	plt.ylabel('Magnitude')
	plt.title('Frequency Data')
	plt.show()

	#find frequency w/ maximum

	max_index = numpy.argmax(data_fft)
	oscillation_freq = freqs[max_index]
	Tu = 1/oscillation_freq
	max_mag = data_fft[max_index]

	print ('Max Power: ' + str(max_mag))
	print ('Oscillation Frequency: ' + str(oscillation_freq))

#calculate Ziegler Nichols PID depending on control type
#for now just do this one type
Kp = .33*Kp
Ki = Tu/2
Kd = Tu/3

print ('Kp: ' + str(Kp))
print ('Ki: ' + str(Ki))
print ('Kd: ' + str(Kd))