import numpy as np                          # For convolution function
import matplotlib.pyplot as plt             # For plotting for tests

from scipy.io import wavfile                # For reading .wav files for testing
from scipy.signal import find_peaks         # For filter function
from scipy.fft import fft, ifft             # For fft and ifft
from scipy.optimize import least_squares    # For estimating KITT's location

def recording_crop_normalize(recordings, ref_mic):
    # Finding the last peak in the recording of the chosen reference microphone
    ref_sig = recordings[:,ref_mic]
    ref_peaks = find_peaks(ref_sig, height= 0.5*np.max(ref_sig))
    ref_peak = ref_peaks[-1]

    # Cropping all recordings to show only the peaks around the reference peak
    start = ref_peak - 1500
    end = ref_peak + 1500
    recordings_cropped = recordings[start:end]

    # Normalizing all recordings after they are cropped
    amplitude, sample, mic = recordings_cropped.shape
    recordings_cropped_normalized = np.zeros((amplitude, sample, mic))
    for i in range(mic):
        recordings_cropped_normalized[:, i] = recordings_cropped[:, i]/max(recordings_cropped[:, i])
    return recordings_cropped_normalized

def ch3(x, y, epsilon):
    # Find both x (recording) and y (reference recording) in the frequency domain
    padded_length = max(len(x), len(y))
    X = fft(x, padded_length-len(x))
    Y = fft(y, padded_length-len(y))

    # Perform the deconvolution in the frequency domain
    H = (Y*np.conj(X))/(np.abs(X)**2+epsilon)

    # Find the channel estimation the time domain and centres it
    channel_estimate = np.real(ifft(H))
    channel_estimate = np.fft.fftshift(channel_estimate)
    return channel_estimate

def distance_calc(channel_estimate, sample_frequency):
    # Finding the location of the peak in the channel estimate relative to the reference peak
    center = len(channel_estimate)//2
    peak = np.argmax(np.abs(channel_estimate))
    sample_range = peak - center

    # Calculating the Time Difference of Arrival (TDOA) from found peak location and using it together with the speed of sound to calculate the distance
    time_dif = sample_range / sample_frequency
    distance = time_dif * 34300  # cm
    return distance

def location_estimation(mic_coords, ref_mic, distances, start_point = [230,230,0]):
    # Using the location of the reference microphone as the refence point
    ref_point = mic_coords[ref_mic]
    other_indices = [i for i in range(mic_coords.shape[0]) if i != ref_mic]

    # Generating the residuals function that is to be minimized
    def residuals_function(point):
        point = np.array([point[0],point[1],0])
        residuals = []
        for i, idx in enumerate(other_indices):
          mic = mic_coords[idx]
          residual = (np.linalg.norm(point-mic) - np.linalg.norm(point-ref_point)) - distances[i]
          residuals.append(residual)
        return residuals

    # Uses the least squares method to minimize the difference between the estimated location and the location calculated from the microphone recordings.
    location = least_squares(residuals_function, start_point, bounds = ([0,0,-1],[460,460,1]))
    return location.x