"""
Matplotlib plotting functions.

Convenience wrappers around common matplotlib pyplot plot calls. This module
simplifies the creation of filter frequency response curves and general
signal plots over time.
"""

from __future__ import print_function, division

import matplotlib.pyplot as plt
import numpy as np
import sensormotion.signal

from scipy.signal import freqz


def plot_filter_response(frequency, sample_rate, filter_type, filter_order=2,
                         show_grid=True, fig_size=(10, 5)):
    """Plot filter frequency response.

    Generate a plot showing the frequency response curve of a filter with the
    specified parameters.

    Parameters
    ----------
    frequency : int or tuple of ints
        The cutoff frequency for the filter. If `filter_type` is set as
        'bandpass' then this needs to be a tuple of integers representing
        the lower and upper bound frequencies. For example, for a bandpass
        filter with range of 2Hz and 10Hz, you would pass in the tuple (2, 10).
        For filter types with a single cutoff frequency then a single integer
        should be used.
    sample_rate : float
        Sampling rate of the signal.
    filter_type : {'lowpass', 'highpass', 'bandpass', 'bandstop'}
        Type of filter to build.
    filter_order: int, optional
        Order of the filter.
    show_grid : bool, optional
        Toggle to show grid lines on the plot.
    fig_size : tuple, optional
        Tuple containing the width and height of the resulting figure.
    """

    b, a = sensormotion.signal.build_filter(frequency, sample_rate,
                                            filter_type, filter_order)

    # Plot the frequency response
    w, h = freqz(b, a, worN=8000)
    f, axarr = plt.subplots(figsize=fig_size)
    axarr.plot(0.5*sample_rate*w/np.pi, np.abs(h), 'b')
    axarr.set_xlim(0, 0.5*sample_rate)
    axarr.set_xlabel('Frequency (Hz)')
    axarr.grid(show_grid)

    # Add lines and markers at the cutoff frequency/frequencies
    if filter_type == 'bandpass':
        for i in range(len(frequency)):
            axarr.axvline(frequency[i], color='k')
            axarr.plot(frequency[i], 0.5*np.sqrt(2), 'ko')
    else:
        axarr.axvline(frequency, color='k')
        axarr.plot(frequency, 0.5*np.sqrt(2), 'ko')

    plt.suptitle("Filter Frequency Response", size=16)
    plt.show()


def plot_signal(time, signal, title='', xlab='', ylab='',
                line_width=1, alpha=1, color='k',
                subplots=False, show_grid=True, fig_size=(10, 5)):
    """Plot signals over time.

    Convenience wrapper around pyplot to quickly create plots of signals
    over time. This is useful if you want to avoid remembering the syntax
    for matplotlib.

    Parameters
    ----------
    time : ndarray
        Time vector of the signal (X-axis)
    signal : ndarray or list of dicts
        If plotting a single line, you can pass in a single array.

        If you want to plot multiple lines (either on the same plot, or as
        subplots) you should pass in a list of dictionaries. Each dictionary
        represents a separate line and contains options for that line. The
        dictionary can hold the follow keys: 'data' (required), 'label',
        'color', 'alpha', 'line_width'.

        For example, a 2 line plot can be created like this:

        >>> plot.plot_signal(time_array, [{'data': line1_data,
                                           'label': 'line 1',
                                           'color': 'b'},
                                          {'data': line2_data,
                                           'label': 'second line',
                                           'line_Width': 4}]
                            )
    title : str, optional
        Title of the plot.
    xlab : str, optional
        Label for the x axis.
    ylab : str, optional
        Label for the y axis.
    line_width : int or float, optional
        Width of the plot lines.
    alpha : int or float, optional
        Alpha of the plot lines.
    color : str, optional
        Colour of the plot lines. Receives a string following the matplotlib
        colors API: http://matplotlib.org/api/colors_api.html
    subplots : bool, optional
        If multiple lines are passed in, should they be displayed on the same
        plot (False)? Or should each be displayed in separate subplots (True)?
    show_grid : bool, optional
        Toggle to display grid lines on the plot.
    fig_size : tuple, optional
        Set the figure size of the resulting plot.
    """

    if type(signal) == list:  # Multiple lines to be plotted
        if subplots:
            f, axarr = plt.subplots(len(signal), 1, figsize=fig_size)
        else:
            f, axarr = plt.subplots(figsize=fig_size)

        for i, line in enumerate(signal):  # Iterate through each plot line
            cur_data = line['data']

            # Get options for the current line
            try:
                cur_label = line['label']
            except KeyError:
                print('Warning: Label missing for signal')
                cur_label = ''
            try:
                cur_color = line['color']
            except KeyError:
                cur_color = color
            try:
                cur_alpha = line['alpha']
            except KeyError:
                cur_alpha = alpha
            try:
                cur_linewidth = line['line_width']
            except KeyError:
                cur_linewidth = line_width

            if subplots:  # Show lines in separate plots, in the same figure
                axarr[i].plot(time, cur_data, label=cur_label,
                              linewidth=cur_linewidth,
                              alpha=cur_alpha, color=cur_color)

                axarr[i].set_xlim(min(time), max(time))
                axarr[i].set_xlabel(xlab)
                axarr[i].set_ylabel(ylab)
                axarr[i].grid(show_grid)
                axarr[i].legend()
                f.subplots_adjust(hspace=0.5)
            else:  # Show all lines on the same plot
                axarr.plot(time, cur_data, label=cur_label,
                           linewidth=cur_linewidth,
                           alpha=cur_alpha, color=cur_color)

                axarr.set_xlim(min(time), max(time))
                axarr.set_xlabel(xlab)
                axarr.set_ylabel(ylab)
                axarr.grid(show_grid)
                axarr.legend()
    else:  # Single line plot
        f, axarr = plt.subplots(figsize=fig_size)
        axarr.plot(time, signal, linewidth=line_width,
                   alpha=alpha, color=color)
        axarr.set_xlim(min(time), max(time))
        axarr.set_xlabel(xlab)
        axarr.set_ylabel(ylab)
        axarr.grid(show_grid)

    plt.suptitle(title, size=16)
    plt.show()
