import numpy as np
import scipy.io.wavfile as wav
import matplotlib.pyplot as plt
import csv

from scipy import signal
from python_speech_features import base

# Approximate the values in the training and testing sets to their closest 
# n-bit floating-point representations.
def approximate_inputs(inputs, precision):
    
    # Determine the increment of the domain values from the desired precision.
    increment = np.power(2.0, -precision)
    
    # Initialize the domain of possible approximations.
    domain = np.linspace(0.0, 1.0-increment, 1.0/increment)
        
    # Find what indices of the domain are closest to the inputs.
    temp = np.round(inputs/increment)
    bound = np.power(2.0, precision) - 1.0
    temp[temp > bound] = bound
    indices = list(map(int, temp))

    # Return the approximated inputs.
    return domain[indices]

# Approximate the weights to their closest preset floating-point 
# representation.
def approximate_weights(weights):

    # Initialize the domain of possible approximations.
    positive = np.array([0.03125, 0.0625, 0.09375, 0.125, 0.15625, 0.25, 1.0])    
    negative = -positive[::-1]
    negative = np.append(negative, 0.0)
    domain = np.append(negative, positive)
    
    # Initialize the array to store the approximations.
    approximation = np.array([])
    
    # Find the array of closest approximations to the weights.
    for i in range(len(weights)):
        
        # Clamp the weight if it is greater than 1.0.        
        if weights[i] > 1.0:
            approximation = np.append(approximation, 1.0)
            continue
        
        # Otherwise, continue to find the closest approximation.
        error_initial = np.abs(domain[0] - weights[i])
        for j in range(1, len(domain)):
            error_next = np.abs(domain[j] - weights[i])
            if error_next > error_initial:
                approximation = np.append(approximation, domain[j-1])
                break
            else:
                if j == len(domain)-1:                
                    approximation = np.append(approximation, domain[j])
                    break
                error_initial = error_next
            
    # Return the approximated data.
    return approximation

# Resamples a signal from 44.1 kHz to another sampling rate.
def resample(path, new_rate):
    
    # Read the signal.
    (rate, sig) = wav.read(path + "raw/demo/demo_44100.wav")
    
    # Resample the signal with the desired rate.
    sig = signal.resample(sig, int(new_rate/rate*len(sig)))
    
    # Write the signal.
    wav.write(path + "raw/demo/demo.wav", new_rate, np.int16(sig))

# Generates the training and testing sets.
def generate_dataset(path, dictionary, people, dataset, cutoff, samples=7000):
    
    # Determine the type of dataset to generate.
    if "training" in dataset:
        filename = "training_set.csv"
    else:
        filename = "testing_set.csv"
      
    # Open a csv file.
    with open(path + filename, 'w') as file:
        writer = csv.writer(file, lineterminator="\n")
        
        # Iterate through all of the sample sets.
        for person in people:
        
            # Iterate through all words in the dictionary.
            for i in range(len(dictionary)):

                # Trim the raw signal.
                trim(path, dictionary[i], person, samples, cutoff) 
    
                # Read the signal.
                (rate, sig) = wav.read(path + "trimmed/" + person + "/" + \
                                       str(dictionary[i]) + ".wav")             
                
                # Compute the MFCCs of the signal and construct a feature 
                # vector.
                features = base.mfcc(sig, appendEnergy=False)
                features = (features.T)[1:13].T
                features = np.asarray(features).reshape(-1)
                
                # Normalize the feature vector.
                features = features/np.max(np.abs(features))/2 + 0.5
                
                # Approximate the values within the feature vector.
                features = approximate_inputs(features, 8)
                
                # Write the feature vector.
                writer.writerow([str(i)] + list(features))

# Plots the preprocessed signals. NOTE: Currently, this function only supports
# the comparison of a hard-coded number of sample sets. Change the arguments 
# to subplot to compare more.
def plot(path, dictionary, people, source):
    
    # Iterate through all words in the dictionary.
    for word in dictionary:     
        plt.figure()
        
        # Iterate through all of the sample sets.        
        for i in range(len(people)):
            
            # Read the signal.
            (rate, sig) = wav.read(path + source + "/" + people[i] + "/" + 
                                   str(word) + ".wav")       
            
            # Plot the signal.
            t = np.linspace(0, len(sig)/rate, len(sig))
            plt.subplot(3, 2, i+1)
            plt.plot(t, sig)
            plt.title("Name: " + str(people[i]) + " | Word: " + str(word))
            plt.xlabel("Time")     
            plt.ylabel("Amplitude")

# Trims the signals to isolate the speech of interest.
def trim(path, word, person, samples, cutoff):
           
    # Read the raw signal.
    (rate, sig) = wav.read(path + "raw/" + str(person) + "/" + \
                           str(word) + ".wav")

    # Scale the signal.
    sig_scaled = sig / np.max(np.abs(sig))
    
    # Set the bounds for trimming the signal.
    sig_scaled_abs = np.abs(sig_scaled)
    for n in range(0, len(sig)):
        if sig_scaled_abs[n] > cutoff:
            start = n
            end = start + samples
            break
        
    # Trim the signal.    
    sig_trimmed = np.array(sig[start:end], dtype=np.int16)
    while len(sig_trimmed) < samples:
        sig_trimmed = np.append(sig_trimmed, 0)
        
    # Write the trimmed signal.
    wav.write(path + "trimmed/" + person + "/" + str(word) + ".wav", rate, 
              sig_trimmed)
                  