"""
Functions for TF analysis.

@author: mje
@email: mads [] cnru.dk
"""

# import mne
from mne.time_frequency import (psd_multitaper, tfr_multitaper, tfr_morlet,
                                cwt_morlet)
from mne.viz import iter_topography
import matplotlib.pyplot as plt
import numpy as np


def calc_psd_epochs(epochs, plot=False):
    """Calculate PSD for epoch.

    Parameters
    ----------
    epochs : list of epochs
    plot : bool
        To show plot of the psds. It will be average for each condition that is shown.

    Returns
    -------
    psds_vol : numpy array
        The psds for the voluntary condition.
    psds_invol : numpy array
        The psds for the involuntary condition.
    """
    tmin, tmax = -0.5, 0.5
    fmin, fmax = 2, 90
    # n_fft = 2048  # the FFT size (n_fft). Ideally a power of 2
    psds_vol, freqs = psd_multitaper(epochs["voluntary"],
                                     tmin=tmin, tmax=tmax,
                                     fmin=fmin, fmax=fmax)
    psds_inv, freqs = psd_multitaper(epochs["involuntary"],
                                     tmin=tmin, tmax=tmax,
                                     fmin=fmin, fmax=fmax)

    psds_vol = 20 * np.log10(psds_vol)  # scale to dB
    psds_inv = 20 * np.log10(psds_inv)  # scale to dB

    if plot:
        def my_callback(ax, ch_idx):
            """Executed once you click on one of the channels in the plot."""
            ax.plot(freqs, psds_vol_plot[ch_idx], color='red',
                    label="voluntary")
            ax.plot(freqs, psds_inv_plot[ch_idx], color='blue',
                    label="involuntary")
            ax.set_xlabel = 'Frequency (Hz)'
            ax.set_ylabel = 'Power (dB)'
            ax.legend()

        psds_vol_plot = psds_vol.copy().mean(axis=0)
        psds_inv_plot = psds_inv.copy().mean(axis=0)

        for ax, idx in iter_topography(epochs.info,
                                       fig_facecolor='k',
                                       axis_facecolor='k',
                                       axis_spinecolor='k',
                                       on_pick=my_callback):
            ax.plot(psds_vol_plot[idx], color='red', label="voluntary")
            ax.plot(psds_inv_plot[idx], color='blue', label="involuntary")
        plt.legend()
        plt.gcf().suptitle('Power spectral densities')
        plt.show()

    return psds_vol, psds_inv, freqs


def multitaper_analysis(epochs):
    """

    Parameters
    ----------
    epochs : list of epochs

    Returns
    -------
    result : numpy array
        The result of the multitaper analysis.

    """
    frequencies = np.arange(6., 90., 2.)
    n_cycles = frequencies / 2.
    time_bandwidth = 4  # Same time-smoothing as (1), 7 tapers.
    power, plv = tfr_multitaper(epochs, freqs=frequencies, n_cycles=n_cycles,
                                time_bandwidth=time_bandwidth, return_itc=True)

    return power, plv


def morlet_analysis(epochs):
    """

    Parameters
    ----------
    epochs : list of epochs

    Returns
    -------
    result : numpy array
        The result of the multitaper analysis.

    """
    frequencies = np.arange(6., 90., 2.)
    n_cycles = frequencies / 2.
    power, plv = tfr_morlet(epochs, freqs=frequencies, n_cycles=4,
                            return_itc=True,
                            verbose=True)

    return power, plv


def single_trial_tf(epochs):
    """
    Parameters:
    ----------
    epochs

    Returns
    -------

    """

    results = []
    frequencies = np.arange(6., 90., 2.)
    n_cycles = 4.
    for j in range(len(epochs)):
        tfr = cwt_morlet(epochs.get_data()[j],
                          sfreq=epochs.info["sfreq"],
                          freqs=frequencies,
                          use_fft=True,
                          n_cycles=n_cycles,
                          zero_mean=False)
        results.append(tfr)
    return results


@@ -67,18 +67,16 @@
 dk%0A%22%22%22%0A%0A
-# 
 import m
@@ -2953,32 +2953,44 @@
 _analysis(epochs
+, n_cycles=4
 ):%0A    %22%22%22%0A%0A    
@@ -3173,33 +3173,33 @@
 = np.arange(6., 
-9
+3
 0., 2.)%0A    n_cy
@@ -3185,32 +3185,34 @@
 6., 30., 2.)%0A   
+ #
  n_cycles = freq
@@ -3220,24 +3220,25 @@
 encies / 2.%0A
+%0A
     power, p
@@ -3289,17 +3289,24 @@
 _cycles=
-4
+n_cycles
 ,%0A      
@@ -3444,24 +3444,26 @@
 s):%0A    %22%22%22%0A
+%0A%0A
     Paramete
@@ -3464,17 +3464,16 @@
 rameters
-:
 %0A    ---
@@ -3486,24 +3486,88 @@
 -%0A    epochs
+ : Epochs object%0A        The epochs to calculate TF analysis on.
 %0A%0A    Return
@@ -3580,16 +3580,41 @@
 -------%0A
+    results : numpy array
 %0A    %22%22%22
@@ -3656,38 +3656,38 @@
 = np.arange(6., 
-9
+3
 0., 
-2
+1
 .)%0A    n_cycles 
@@ -3688,17 +3688,17 @@
 ycles = 
-4
+5
 .%0A    fo
@@ -3796,17 +3796,16 @@
         
- 
 sfreq=ep
@@ -3820,25 +3820,24 @@
 o%5B%22sfreq%22%5D,%0A
- 
             
@@ -3893,17 +3893,16 @@
         
- 
 use_fft=
@@ -3903,25 +3903,24 @@
 e_fft=True,%0A
- 
             
@@ -3943,33 +3943,32 @@
 ycles=n_cycles,%0A
- 
                 

Update for minor things