import numpy as np                          # For convolution function
import matplotlib.pyplot as plt             # For plotting purposes
import scipy
# import samplerate

from scipy import signal
from scipy.optimize import least_squares
from scipy.io import wavfile                # For reading .wav files
from scipy.signal import find_peaks         # For filter function
from scipy.fft import fft, ifft             # For fft and ifft

# def wavaudioread(filename, fs):
#     fs_wav, y_wav = wavfile.read(filename)
#     y = samplerate.resample(y_wav, fs / fs_wav, "sinc_best")
#
#     return y

# ch3 function
def ch3(x,y,Lhat,epsi):

    # your code here (read the above text carefully for all the steps):
    N = max(len(x), len(y)) + 50
    x = np.concatenate((x, np.zeros(N - len(x)))) # Zero padding
    y = np.concatenate((y, np.zeros(N - len(y))))

    X = fft(x) # FFTs to find X[k] and Y[k]
    Y = fft(y)

    H = Y / X # Computation of H[k]

    threshold = epsi * max(np.absolute(X))
    ii = np.absolute(X) < threshold

    H[ii] = 0

    h = np.real(ifft(H)) # IFFT to find h[n]
    h[np.abs(h) < 1e-12] = 0
    # Truncation to length Lhat (optional and actually not recommended before you inspect the entire h)
    h = h[0:Lhat]

    return h

def signal_Crop(y1):
   scale = 300
   pk = np.argmax(y1)
   start = int(max(pk - scale,0))
   end = int(pk + scale)
   return start, end

def sig_comp(h1,Fs):
  h1_peak, _ = find_peaks(h1, height = np.max(h1)*0.99)

  peak_1 = h1_peak[0]

  sample_range = np.array([peak_1])
  time_dif = sample_range / Fs
  distance = abs(time_dif) * 34300

  return distance, peak_1

def distance_calc(y1,y2,epsi,Fs):
  start1, end1 = signal_Crop(y1)
  start2, end2 = signal_Crop(y2)

  start_true= min(start1, start2)
  end_true = max(end1, end2)

  y1_crop = y1[start_true:end_true]
  y2_crop = y2[start_true:end_true]

  Lhat = len(y1_crop)

  h1 = ch3(y1_crop, y2_crop, Lhat, epsi)
  distance, peak_1 = sig_comp(h1, Fs)
  print(distance)
  return distance[0]

def location_calc(mic_coords, ref_index, distances, start_point = [230,230,0]):
  ref = mic_coords[ref_index]
  other_indices = [i for i in range(mic_coords.shape[0]) if i != ref_index]

  def residuals(S):
    S = np.array([S[0],S[1],0])
    res = []
    for i, idx in enumerate(other_indices):
      mic = mic_coords[idx]
      diff = np.linalg.norm(S-mic) - np.linalg.norm(S-ref) - distances[i]
      res.append(diff)
    return res

  sol = least_squares(residuals, start_point)
  return sol.x

def location(y,Fs):
  y1 = y[0]
  y2 = y[1]
  y3 = y[2]
  y4 = y[3]
  y5 = y[4]
  epsi = 0.01
  d1 = distance_calc(y5,y1,epsi,Fs)
  d2 = distance_calc(y5,y2,epsi,Fs)
  d3 = distance_calc(y5,y3,epsi,Fs)
  d4 = distance_calc(y5,y4,epsi,Fs)

  mic_coords = np.array([
   [0,0,25], # mic 1
   [0,460,25], # mic 2
   [460,460,25], # mic 3
   [460,0,25], # mic 4
   [0,230,55] # mic 5
  ])

  ref_index = 4 #used to makes mic 5 the ref
  distances = np.array([d1,d2,d3,d4])

  source_pos = location_calc(mic_coords, ref_index, distances)
  return source_pos

  ref_index = 4 #used to makes mic 5 the ref
  distances = np.array([d1,d2,d3,d4])

  source_pos = location_calc(mic_coords, ref_index, distances)
  return source_pos

if __name__ == "__main__":
    # Coordinates of the recordings
    record_x = [64, 82, 109, 143, 150, 178, 232]
    record_y = [40, 399, 76, 296, 185, 439, 275]

    # List to store filenames
    filenames = []

    # Generate filenames based on coordinates
    for i in range(len(record_x)):
        real_x = record_x[i]
        real_y = record_y[i]
        filenames.append(f"../files/Student Recordings/record_x{real_x}_y{real_y}.wav")

    # Load the first recording
    Fs, recording = wavfile.read(filenames[0])
    print(location(recording, Fs))
    print(Fs)