import sys
import struct
import os
import matplotlib.pyplot as plt
import numpy as np

#int to 2 byte Little Endian binary
def int_to_LE2(input_int: int) -> bytes:
	output_binary = struct.pack('<h', int(input_int))
	return output_binary

#2 byte Little Endian binary to integer
def LE2_to_int(input_binary: bytes) -> int:
	output_int = struct.unpack('<h', input_binary)
	return output_int[0]

#int to 4 byte Little Endian binary
def int_to_LE4(input_int: int) -> bytes:
	output_binary = struct.pack('<i', int(input_int))
	return output_binary

#4 byte Little Endian binary to integer
def LE4_to_int(input_binary: bytes) -> int:
	output_int = struct.unpack('<i', input_binary)
	return output_int[0]

#8 bit sample (signed byte) to integer
def sample8_to_int(input_binary: bytes) -> int:
	output_int = struct.unpack('<B', input_binary)
	return output_int[0]

#16 bit sample (unsigned short) to integer
def sample16_to_int(input_binary: bytes) -> int:
	output_int = struct.unpack('<h', input_binary)
	return output_int[0]

#32 bit sample (unsigned short) to integer
def sample32_to_int(input_binary: bytes) -> int:
	output_int = struct.unpack('<I', input_binary)
	return output_int[0]

#binary ascii to string
def ascii_to_string(input_binary: bytes) -> str:
	output_string = input_binary.decode('ascii')
	return output_string

class WaveObject(object):
	HEADER_SIZE = 44

	NP_DATA_TYPES = {
	8 : np.dtype('<u1'),
	16 : np.dtype('<i2'),
	32 : np.dtype('<i4')
	}

	#Just did this to make it a collapsible
	def init_header(self):
		#name [data, descriptor, type]
		self.header = {
		'f_type' : [b'RIFF', 'File type', 'ascii'],
		'f_size' : [b'', 'File size', 'LE4'],
		'f_fmt' : [b'WAVE', 'File format', 'ascii'],
		's_chunk_1_id' : [b'fmt ', 'Sub chunk 1 ID', 'ascii'],
		's_chunk_1_size' : [b'\x10\x00\x00\x00', 'Sub chunk 1 size', 'LE4'],
		'audio_fmt' : [b'\x01\x00', 'Audio format', 'LE2'],
		'n_channels' : [b'', 'Number of channels', 'LE2'],
		'sample_rate' : [b'', 'Sample rate', 'LE4'],
		'byte_rate' : [b'', 'Byte rate', 'LE4'],
		'block_align' : [b'', 'Block align', 'LE2'],
		'bit_depth' : [b'', 'Bit depth', 'LE2'],
		's_chunk_2_id' : [b'data', 'Sub chunk 2 ID', 'ascii'],
		'data_size' : [b'', 'Data size', 'LE4']}

	#Opens specified file if it exists, creates it otherwise
	#Defaults to 'newfile.wav'
	def __init__(self, file_name: str = 'newfile.wav'):
		self.init_header()
		if os.path.isfile(file_name):
			self.open_wav(file_name)
		else:
			self.create_wav(file_name)

	#Opens a wav file, checks it, then rebuilds the header
	def open_wav(self, file_name: str):
		with open(file_name, mode='rb') as f:
			header_data = f.read(self.HEADER_SIZE)
			raw_data = f.read()

		self.build_header(header_data)
		#check that file is a WAV file
		bit_depth = LE2_to_int(self.header['bit_depth'][0])
		n_channels = LE2_to_int(self.header['n_channels'][0])

		if (self.header['f_type'][0] != b'RIFF') or (self.header['f_fmt'][0] != b'WAVE'):
			print("File is not a valid WAV file")
			sys.exit(1)

		if (bit_depth % 8 != 0) or (bit_depth > 32):
			print("Bit depth is invalid (" + self.header['bit_depth'][0] + ")")
			sys.exit(1)

		dt = self.NP_DATA_TYPES[bit_depth]
		raw_data = np.fromstring(raw_data, dt)
		raw_data.shape = (-1, n_channels)
		self.data = raw_data
		self.rebuild_header()

	#creates a new wav file from specified file name, defaults values to normal values and then builds the header
	def create_wav(self, file_name: str, n_channels: int = 2, sample_rate: int = 44100, bit_depth: int = 16):
		if n_channels < 0:
			print("Number of channels is invalid (" + n_channels + ")")
			sys.exit(1)
		if sample_rate < 0:
			print("Sample rate is invalid (" + sample_rate + ")")
			sys.exit(1)
		if (bit_depth % 8 != 0) or (bit_depth > 32):
			print("Bit depth is invalid (" + bit_depth + ")")
			sys.exit(1)

		self.header['n_channels'][0] = int_to_LE2(n_channels)
		self.header['sample_rate'][0] = int_to_LE4(sample_rate)
		self.header['bit_depth'][0] = int_to_LE2(bit_depth)

		dt = self.NP_DATA_TYPES[bit_depth]
		raw_data = np.array([], dt)
		raw_data.shape = (-1, n_channels)
		self.rebuild_header()
		print("Created new file: " + file_name)

	#Takes WAV header binay as input, writes to self.header
	def build_header(self, header_data: bytes):
		index = 0
		for item,entry in self.header.items():
			if (entry[2] == 'ascii') or (entry[2] == 'LE4'):
				self.header[item][0] = header_data[index:index+4]
				index += 4
			elif entry[2] == 'LE2':
				self.header[item][0] = header_data[index:index+2]
				index += 2

	#Prints header info from self.header
	def print_header(self, raw_binary: int = 0):
		for item,entry in self.header.items():
			if raw_binary:
				print(entry[1] + ": " + str(entry[0]))
			elif entry[2] == 'ascii':
				print(entry[1] + ": " + ascii_to_string(entry[0]))
			elif entry[2] == 'LE2':
				print(entry[1] + ": " + str(LE2_to_int(entry[0])))
			elif entry[2] == 'LE4':
				print(entry[1] + ": " + str(LE4_to_int(entry[0])))
		print('Total samples: ' + str(len(self.data)))

	#Rebuilds the header information when there is a change
	#Does it in the order byte rate, block align, data size
	#Additionally updates meta-info of the header that is not written to the file
	def rebuild_header(self):
		sample_rate = LE4_to_int(self.header['sample_rate'][0])
		bit_depth = LE2_to_int(self.header['bit_depth'][0])
		num_channels = LE2_to_int(self.header['n_channels'][0])
		try:
			total_samples = len(self.data) * 2
		except None:
			total_samples = 0

		#should be sample rate * bit depth * channels / 8 (LE4)
		byte_rate = sample_rate * bit_depth * num_channels / 8
		self.header['byte_rate'][0] = int_to_LE4(byte_rate)

		#should be bit depth * channels / 8 (LE2)
		block_align = bit_depth * num_channels / 8
		self.header['block_align'][0] = int_to_LE2(block_align)

		#should be toal samples * bit depth / 8 (LE4)
		data_size = total_samples * bit_depth / 8
		self.header['data_size'][0] = int_to_LE4(data_size)

		#size of data + header size - 8 because first 8 bytes don't count for some reason (LE4)
		file_size = data_size + self.HEADER_SIZE - 8
		self.header['f_size'][0] = int_to_LE4(file_size)

	#Plots the data, takes a start and end time as an input, defaults to entire file
	def plot_waveform(self, start: int = 0, end: int = -1):
		sample_rate = LE4_to_int(self.header['sample_rate'][0])
		n_channels = LE2_to_int(self.header['n_channels'][0])
		bit_depth = LE2_to_int(self.header['bit_depth'][0])

		start_sample = start * sample_rate

		if end == -1:
			end_sample = len(self.data)
		else:
			end_sample = end * sample_rate

		samples = self.data[start_sample:end_sample]
		x = np.arange(start_sample, end_sample)

		plt.figure('Waveforms')

		for c in range(0, n_channels):
			line_label = 'Channel ' + str(c)
			plt.plot(x, samples[:,c], label = line_label)
		plt.legend()
		plt.xlabel('Sample')
		plt.ylabel('Amplitude')
		plt.show()

	#merge audio streams and prevent clipping
	#def merge_audio(self, stream_a: array, stream_b: array):
	#	bit_depth = LE2_to_int(self.header['bit_depth'][0])


	#Amplifies each value of the array by a certain percentage
	def amplify(self, magnitude: float = 1.25):
		bit_depth = LE2_to_int(self.header['bit_depth'][0])
		dt = self.NP_DATA_TYPES[bit_depth]
		samples = np.array(self.data, float)
		samples = samples * magnitude
		self.data = np.array(samples, dt)

	#Adds reverb to the data, takes strength which is how much the reverb decays and time which is in seconds
	def add_reverb(self, strength: float = 0.25, delay: float = 1.5):
		sample_rate = LE4_to_int(self.header['sample_rate'][0])
		n_channels = LE2_to_int(self.header['n_channels'][0])
		bit_depth = LE2_to_int(self.header['bit_depth'][0])
		dt = self.NP_DATA_TYPES[bit_depth]

		#get length of reverb effect in samples
		reverb_length = int(sample_rate * delay)


		#create arrays
		samples = np.array(self.data, float)
		reverb = np.array(samples, float)
		buffer = np.zeros([reverb_length, n_channels], float)
		#apply reverb, shift downwards
		reverb = reverb * strength
		reverb = np.insert(reverb, 0, buffer, 0)
		#add buffer to end of samples
		samples = np.append(samples, buffer, 0)
		#apply reverb to original data, prevent clipping
		data = 2 * (samples + reverb) - (samples * reverb / (2 ^ int(bit_depth / 2)))
		#convert sata back to proper data type and apply
		self.data = np.array(data, dt)
		self.rebuild_header()

	#Takes a new sample rate as an input then writes it to the header
	def change_sample_rate(self, new_sample_rate: int):
		if new_sample_rate <= 0:
			print("Sample rate is invalid (" + new_sample_rate + ") value not being changed")
		else:
			self.header['sample_rate'][0] = int_to_LE4(new_sample_rate)
			self.rebuild_header()

	#Takes new bit depth as an input then writes it to the header
	def change_bit_depth(self, new_bit_depth: int):
		if (new_bit_depth % 8 != 0) or (new_bit_depth > 32):
			print("Bit depth is invalid (" + new_bit_depth + ") value not being changed")
		else:
			self.header['bit_depth'][0] = int_to_LE2(new_bit_depth)
			self.rebuild_header()

	#Takes channel number as an input then writes it to the header
	def change_num_channels(self, new_num_channels: int):
		if new_num_channels <= 0:
			print("Channel number is invalid (" + new_num_channels + ") value not being changed")
		else:
			self.header['n_channels'][0] = int_to_LE2(new_num_channels)
			self.rebuild_header()

	#Takes a file name and start/end as args, defaults to first 100 samples
	def dump_samples(self, file_name: str, start: int = 0, end: int = 100):
		if len(self.data) < end:
			print("Specified end (" + str(end) + ") is greater than number of samples")
			exit(1)
		samples = self.data[start:end]
		with open(file_name, mode='w+') as f:
			for s in samples:
				f.write(str(s) + "\n")

	#Takes the file name as an input and writes all data to that file
	def write_file(self, file_name: str):
		with open(file_name, mode='bw+') as f:
			for item,entry in self.header.items():
				f.write(entry[0])
			f.write(self.data.tostring())
		print("Wrote to file " + file_name)