import numpy as np
from scipy import signal
import matplotlib.pyplot as plt
from pathlib import Path

import irutilities as irutil

# Load IR data
sessionName = "S1-M3969_npy"  # "S32-M441_npy"
sessionPath = Path('.').joinpath(sessionName)
posMic, posSrc, ir = irutil.loadIR(sessionPath)

# Sampling rate
samplerate = 48000

# Select IR data
srcIdx = 0
micIdx = 0
print("Source position (m): ", posSrc[srcIdx, :])
print("Mic position (m): ", posMic[micIdx, :])

Source position (m):  [2.  1.5 0. ]
Mic position (m):  [-0.5 -0.5 -0.2]

# Plot geometry
irutil.drawGeometry(posSrc, posMic)

# IR plots
ir_plt = ir[srcIdx, micIdx, :]
irutil.irPlots(ir_plt, samplerate)

RT60 (ms):  387.5

# Extract xy-plane at z=0
z = 0.0
posMic_z, ir_z = irutil.extract_plane(posMic, ir, z)
posMicX = np.unique(posMic_z[:, 0].round(4))
posMicY = np.unique(posMic_z[:, 1].round(4))
numXY = (posMicX.shape[0], posMicX.shape[0])
posMicXY, irXY, _ = irutil.sortIR(posMic_z, ir_z, numXY, posMicX, posMicY)
    
# Lowpass filter
maxFreq = 600
h = signal.firwin(numtaps=64, cutoff=maxFreq, fs=samplerate)
irXY_lp= signal.filtfilt(h, 1, irXY[srcIdx,:,:,:], axis=-1)

# Wave image
tIdx, _ = irutil.findPeak(irXY_lp, 0)
irutil.plotWave(posMicX, posMicY, irXY_lp, tIdx)
# irutil.plotWaveFronts(posMicX, irXY_lp, samplerate)

# Wave movie
irutil.movWave(sessionPath, posMicX, posMicY, irXY_lp, samplerate)