import numpy as np                          # For convolution function
import matplotlib.pyplot as plt             # For plotting purposes
import scipy
from fontTools.misc.plistlib import end_true
# 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):
    N = 2**int(np.ceil(np.log2(len(x)+len(y))))
    x = np.pad(x, (0,N - len(x))) # Zero padding
    y = np.pad(y, (0,N - len(y))) # Zero padding

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

    H = (Y*np.conj(X))/(np.abs(X)**2+epsi)

    h = np.real(ifft(H)) # IFFT to find h[n]
    h = np.fft.fftshift(h)
    return h

def peak_ref(y):
  N, C = y.shape
  peaks = np.zeros(C, dtype=int)
  ref_ch = 4
  ref_sig = np.abs(y[:,ref_ch])
  ref_pk, _ = find_peaks(ref_sig, height= 0.5*np.max(ref_sig))

  if len(ref_pk) == 0:
      ref_peak = np.argmax(ref_sig)
  else:
      ref_peak = ref_pk[100]

  peaks[ref_ch] = ref_peak

  for ch in range(C):
      if ch == ref_ch:
          continue
      sig = np.abs(y[:,ch])

      start = max(0, ref_peak - 1500)
      end = min(N, ref_peak + 1500)

      local = sig[start:end]
      pk, _ = find_peaks(local, height= 0.9*np.max(sig))

      if len(pk) == 0:
          local_peak = np.argmax(local)
      else:
          local_peak = np.argmax(local)
      peaks[ch] = start + local_peak
  return peaks


def sig_comp(h1,Fs):
  center = len(h1) // 2
  h1_peak = np.argmax(np.abs(h1))

  sample_range = h1_peak - center
  time_dif = sample_range / Fs
  distance = time_dif * 34300 #cm
  return distance

def distance_calc(y1,y2,epsi,Fs, peak1, peak2):

  min_p = min(peak1, peak2)
  max_p = max(peak1, peak2)

  start = max(min_p - 800, 0)
  end = min(max_p + 800, len(y1))

  y1_crop = y1[start:end]
  y2_crop = y2[start:end]

  plt.figure(figsize=(10, 4))
  plt.plot(y1_crop, label='Channel 1 (cropped)')
  plt.plot(y2_crop, label='Channel 2 (cropped)')
  plt.title(f'Cropped signals for TDOA calculation\nPeaks: {peak1}, {peak2}')
  plt.xlabel('Samples (cropped)')
  plt.ylabel('Amplitude')
  plt.grid(True)
  plt.legend()
  plt.show()
  Lhat = 501
  h1 = ch3(y1_crop, y2_crop, Lhat, epsi)
  distance = sig_comp(h1, Fs)
  return distance

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, bounds = ([0,0,-1],[460,460,1]))
  return sol.x

def location(y,Fs):
  peaks = peak_ref(y)
  y1, y2, y3, y4, y5 = y.T
  p1, p2, p3, p4, p5 = peaks

  epsi = 0.0001
  d1 = distance_calc(y5,y1,epsi,Fs, p5, p1)
  d2 = distance_calc(y5,y2,epsi,Fs, p5, p2)
  d3 = distance_calc(y5,y3,epsi,Fs, p5, p3)
  d4 = distance_calc(y5,y4,epsi,Fs, p5,p4)

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

  ref_index = 4 #used to makes mic 5 the ref
  distances = np.array([d1,d2,d3,d4])
  print("Distances (d1, d2, d3, d4):", distances)
  print("Peak samples:   ", peaks)

  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")

    for i, file in enumerate(filenames):
        Fs, recording = wavfile.read(file)
        recording = recording.astype(np.float64) / np.max(np.abs(recording))
        print(f"\nRecording {i+1}: {file}")
        source_pos = location(recording, Fs)
        print("Estimated source position:", source_pos)

        num_channels = recording.shape[1]
        fig, axs = plt.subplots(num_channels, 1, figsize=(10, 2*num_channels), sharex=True)

        for ch in range(num_channels):
            axs[ch].plot(recording[:, ch])
            axs[ch].set_ylabel(f"Ch {ch+1}")
            axs[ch].grid(True)

        axs[-1].set_xlabel("Samples")
        plt.suptitle(f"Waveforms of all channels - Recording {i+1}")
        plt.tight_layout(rect=[0, 0.03, 1, 0.95])
        plt.show()