import numpy as np

import matplotlib.pyplot as plt
import fft

class fft_meas(object):
	"""docstring for fft_meas"""
	# 61000-4-7 says 200ms
	# 10-cycles for 50Hz
	# 12-cycles for 60Hz
	def __init__(self, fs=50e3, tn=0.2):

		self.fs = fs
		self.ts = 1/fs
		self.tn = tn
		self.N = int(int(fs*tn))
		self.data = np.zeros(self.N)
		self.idx = -1
		self.time = 0

		self.freq = 0 
		self.a = 0
		self.b = 0
		self.c = 0
		self.Y_C = 0

		# 61000-4-7: 3.2.3
		# Upto the 40 Harmonic per 3.3.1 Note 2
		# Harmonic order
		self.Y_H = np.zeros(40)
		# 61000-4-7: 3.2.3
		# Upto the 40 Harmonic per 3.3.2 Note 2
		# Grouped Harmonic order
		self.Y_g = np.zeros(40)
		# 61000-4-7: 3.2.2 
		self.h = np.zeros(40)

		self.THD = 0
		self.THDG = 0
		

		self.phi = 0
		self.base_freq = 0


	def step(self, y, base_freq, unit):

		self.idx += 1

		self.data[self.idx] = y

		if self.idx == self.N-1:

			self.freq , self.a, self.b , self.c, self.Y_C, self.phi = fft.fft(x=self.data, fs=self.fs)
			
			self.calc_Y_H(base_freq)
			self.calc_THD()

			self.calc_Y_g(base_freq)
			self.calc_THDG()



			self.idx = -1


	def calc_Y_H(self, base_freq):
		# 61000-4-7: 3.2.3
			if np.abs(base_freq-50) < np.abs(base_freq-60):
				self.base_freq = 50
			else:
				self.base_freq = 60			

			if self.base_freq == 50:
				for k in range(len(self.Y_H)):
					self.Y_H[k] = self.Y_C[int(10*k)]
					self.h[k] = self.freq[int(10*k)]
			elif self.base_freq == 60:
				for k in range(len(self.Y_H)):
					self.Y_H[k] = self.Y_C[int(12*k)]
					self.h[k] = self.freq[int(12*k)]

	def calc_THD(self):
		# 61000-4-7: 3.3.1
		self.THD = 0
		for k in range(2, len(self.Y_H)):
			self.THD += np.power(self.Y_H[k]/self.Y_H[1],2)
			self.THD =  np.sqrt(self.THD)

	def calc_Y_g(self, base_freq):
		# 61000-4-7: 3.2.4
			if np.abs(base_freq-50) < np.abs(base_freq-60):
				self.base_freq = 50
			else:
				self.base_freq = 60			

			if self.base_freq == 50:
				for k in range(len(self.Y_g)):
					self.Y_g[k] = self.Y_C[int(10*k)-1] + self.Y_C[int(10*k)] + self.Y_C[int(10*k)+1]
					self.h[k] = self.freq[int(10*k)]
			elif self.base_freq == 60:
				for k in range(len(self.Y_g)):
					self.Y_g[k] = self.Y_C[int(12*k)-1] + self.Y_C[int(12*k)] + self.Y_C[int(12*k)+1]
					self.h[k] = self.freq[int(12*k)]

	def calc_THDG(self):
		# 61000-4-7: 3.3.1
		self.THDG = 0
		for k in range(2, len(self.Y_H)):
			self.THDG += np.power(self.Y_g[k]/self.Y_g[1],2)
			self.THDG =  np.sqrt(self.THDG)
		print(self.THDG)
		
	def plot(self):
		#fs, a, b, c, YC, phi = ftt.fft(Ph1, sample_freq)

		plt.plot(self.freq, self.c)
		plt.xlabel("x")
		plt.ylabel("y")
		plt.title("Simple plot")
		plt.show()
		
	def plot_Y_H(self):

		plt.stem(self.h, self.Y_H, use_line_collection=True)
		plt.xlabel("Harmonic order h")
		plt.ylabel("RMS value Y_H,h")
		plt.title("Harmonic spectrum (IEC 61000-4-7: 3.2.3)")
		plt.grid(True)
		plt.show()