import numpy as np
from scipy.signal import argrelextrema


def straight_line(x, p1, p2):
    """
    Define a straight line between two points P1(x_min, y_min) and P2(x_max, y_max). The line will
    have the shape y = mx + c, being m the slope and c the offset.
    x :
    p1: (x_1, y_1)
    p2: (x_2, y_2)
    ----
    return: y
    """
    # unfold the x and y values of each point
    x_1, y_1 = p1
    x_2, y_2 = p2

    # define the slope
    m = (y_2 - y_1)/(x_2 - x_1)

    # get the offset
    c = y_1 - m*x_1

    # get the value of y
    y = m*np.array(x) + c

    return y


def compute_aperiodic_beta(psd_single_trial, frequencies):
    """
    Computes the 1/f component of the power spectrum in the beta band. It takes the single trial
    psd and computes the mean of the log of the psd in the beta band. Then, it computes the log of
    the frequencies and finds the local minima. It returns the line that goes from the first to the
    last local minima, somehow the 1/f assuming it is linear in the log-log space.
    :param psd_single_trial: array of the shape (n_trials, n_channels, n_freqs) with the psd of each
    site.
    :param frequencies: array of frequencies of the shape (n_freqs,)
    ----
    :return: aperiodic_all_channels: array of the same shape (n_channels, n_freqs) with the 1/f component
    """
    # compute the log of the frequencies
    log_freq = np.log10(frequencies)

    # compute the log of the psd
    log_psd_single_trial = np.log10(psd_single_trial)

    # calculate the mean of all the log single-trial-psd inside the site
    mean_log_psd = log_psd_single_trial.mean(axis=0)  # (n_channels, n_freq)

    # aperiodic of all channels
    aperiodic_all_channels = []
    for channel in range(mean_log_psd.shape[0]):
        # find local minima
        local_mins = argrelextrema(mean_log_psd[channel], np.less)[0]

        # check that there are local minima
        if len(local_mins) != 0:
            # get the two coordinates of the y
            y_start = mean_log_psd[channel, local_mins[0]]
            y_end = mean_log_psd[channel, local_mins[-1]]
            x_start = frequencies[local_mins[0]]
            x_end = frequencies[local_mins[-1]]

        else:
            print(f'no local mins in channel {channel}')
            # if there are no local minimum, take as the first minimum the minimum w freq < 12Hz
            idxs_freqs_x_start = np.where(frequencies < 13)
            idx_x_start = np.argmin(mean_log_psd[channel, idxs_freqs_x_start])
            x_start = frequencies[idx_x_start]

            # set the x_end the same as x_start - it will be corrected later
            x_end = x_start
            # get the y_start
            y_start = mean_log_psd[channel, idx_x_start]
            # get the y_end
            y_end = y_start

        # use the minimum value between 35-50Hz as the min if there is no local minimum or if the
        # minimum is in between the beta band - because there are two peaks.
        if np.logical_or(x_start == x_end, x_end < 30):
            idx_freqs_x_end = np.where(np.logical_and(frequencies > 35, frequencies < 51))[0]
            idx_x_end = idx_freqs_x_end[np.argmin(mean_log_psd[channel, idx_freqs_x_end])]
            y_end = mean_log_psd[channel, idx_x_end]
            x_end = frequencies[idx_x_end]
            print(f'correcting x_end - ch: {channel}')

        # if the first minimum is after 12Hz, then set the minimum at the first frequency value
        # after 10Hz.
        if x_start > 12:
            x_start = frequencies[np.where(frequencies > 10)[0][0]]
            y_start = mean_log_psd[channel, np.where(frequencies == x_start)[0]][0]
            print(f'correcting x_start - ch: {channel}')

        # compute the aperiodic line
        aperiodic = straight_line(log_freq, (np.log10(x_start), y_start), (np.log10(x_end), y_end))
        aperiodic_all_channels.append(aperiodic)

    return np.array(aperiodic_all_channels)