import numpy as np
from scipy.signal import butter, lfilter
import pyramid
import video
import numpy as np
from scipy.fft import fft, fftfreq

import numpy as np
from scipy import fftpack, signal
from numpy.fft import fft, fftfreq

# Temporal bandpass filter with Fast-Fourier Transform
def fft_filter(video, freq_min, freq_max, fps):
    fft = fftpack.fft(video, axis=0)
    frequencies = fftpack.fftfreq(video.shape[0], d=1.0 / fps)
    bound_low = (np.abs(frequencies - freq_min)).argmin()
    bound_high = (np.abs(frequencies - freq_max)).argmin()

    # Zero out frequencies outside the desired range
    fft[:bound_low] = 0
    fft[bound_high:] = 0

    # Apply inverse FFT to get the filtered video
    filtered_video = np.abs(fftpack.ifft(fft, axis=0))
    return filtered_video, fft, frequencies

def find_breath_rate(fft, freqs, freq_min, freq_max):
    fft_maximums = []

    # Iterate through the frequencies and accumulate FFT amplitude
    for i in range(fft.shape[0]):
        if freq_min <= freqs[i] <= freq_max:
            fft_maximums.append(np.abs(fft[i]).max())
        else:
            fft_maximums.append(0)

    peaks, _ = signal.find_peaks(fft_maximums)
    
    if len(peaks) == 0:
        return 0  # No peaks found

    max_peak = peaks[np.argmax([fft_maximums[peak] for peak in peaks])]
    return freqs[max_peak] * 60  # Convert frequency (Hz) to breaths per minute (bpm)


def butter_bandpass(lowcut, highcut, fs, order=1):
    """
    Calculates the Butterworth bandpass filter coefficients.
    
    :param lowcut: Low frequency cutoff (e.g., 0.1 Hz for breath detection).
    :param highcut: High frequency cutoff (e.g., 0.5 Hz for breath detection).
    :param fs: Video frame rate (sampling frequency).
    :param order: Filter order (default is 1).
    :return: Numerator (b) and denominator (a) polynomials of the IIR filter.
    """
    low = lowcut / (0.5 * fs)  # Normalize the frequencies by Nyquist frequency
    high = highcut / (0.5 * fs)
    b, a = butter(order, [low, high], btype='band')
    return b, a

def apply_butter(laplace_video_list, levels, alpha, cutoff=20, low=0.1, high=0.5, fps=30, width=512, height=512, linearAttenuation=True):
    """
    Applies the Butterworth filter to the video sequence, magnifies the filtered video sequence,
    and attenuates spatial frequencies for breath detection.
    
    :param laplace_video_list: Laplace video pyramid.
    :param levels: Pyramid levels.
    :param alpha: Magnification factor.
    :param cutoff: Spatial frequencies cutoff factor.
    :param low: Temporal low frequency cutoff (related to the breath rate).
    :param high: Temporal high frequency cutoff.
    :param fps: Video frame rate.
    :param width: Video frame width.
    :param height: Video frame height.
    :param linearAttenuation: Whether to apply linear attenuation.
    :return: List of filtered video frames.
    """
    print('Applying Butterworth filter...')
    filtered_video_list = []
    b, a = butter_bandpass(low, high, fps, order=1)

    # Spatial wavelength (lambda)
    lambda1 = (width ** 2 + height ** 2) ** 0.5
    delta = cutoff / 8 / (1 + alpha)

    for i in range(levels):  # Iterate through pyramid levels
        current_alpha = lambda1 / (8 * delta) - 1  # Alpha calculation
        current_alpha /= 2

        # Apply the Butterworth filter to the temporal image sequence
        filtered = lfilter(b, a, laplace_video_list[i], axis=0)

        # Ignore the lowest and highest pyramid levels
        if i == levels - 1 or i == 0:
            filtered *= 0

        # Spatial frequencies attenuation
        if current_alpha > alpha:
            filtered *= alpha
        else:
            filtered *= current_alpha if linearAttenuation else 0

        filtered_video_list.append(filtered)
        lambda1 /= 2  # Decrease lambda for the next level

    return filtered_video_list

def start(video_frames, alpha=50, cutoff=16, low=0.1, high=0.5, linearAttenuation=True, chromAttenuation=0.4, fps=30, width=512, height=512, time_window = 5):
    """
    Performs motion magnification on the video frames by applying Butterworth bandpass filter and saves the output video.
    This can be used for detecting breathing patterns.
    
    :param video_frames: Numpy array containing video frames (shape: [num_frames, height, width, channels])
    :param alpha: Magnification factor (e.g., 50 for subtle motion like breathing).
    :param cutoff: Spatial frequencies cutoff factor.
    :param low: Temporal low frequency cutoff (e.g., 0.1 Hz for breaths).
    :param high: Temporal high frequency cutoff (e.g., 0.5 Hz for breaths).
    :param fps: Frames per second of the video.
    :param width: Width of the video frames.
    :param height: Height of the video frames.
    :return: Filtered video frames and estimated breathing rate.
    """

    # Convert RGB to YIQ for luminance-based processing
    yiq_video = video.rgb2yiq(video_frames)

    # Determine pyramid levels for the video
    levels = video.calculate_pyramid_levels(width, height)

    # Build Laplacian pyramid for each video frame
    lap_video_list = pyramid.laplacian_video_pyramid(yiq_video, levels)

    # Apply Butterworth filter for breath frequencies (low=0.1 Hz, high=0.5 Hz)
    filtered_video_list = apply_butter(lap_video_list, levels, alpha, cutoff, low, high, fps, width, height, linearAttenuation)

    # Reconstruct the magnified video from the filtered pyramid
    final_video = pyramid.reconstruct(filtered_video_list, levels)
    filtered_video, fft_data, frequencies = fft_filter(final_video.mean(axis=(1, 2, 3)), low, high, fps)
    bpm = find_breath_rate(fft_data, frequencies, low, high)

    final_video += yiq_video
    final_rgb = video.yiq2rgb(final_video)

    final_rgb = np.clip(final_rgb, 0, 255).astype(np.uint8)

    return final_rgb, bpm




