# *****************************************************************************
#
# This program is free software; you can redistribute it and/or modify it under
# the terms of the GNU General Public License as published by the Free Software
# Foundation; either version 2 of the License, or (at your option) any later
# version.
#
# This program is distributed in the hope that it will be useful, but WITHOUT
# ANY WARRANTY; without even the implied warranty of MERCHANTABILITY or FITNESS
# FOR A PARTICULAR PURPOSE.  See the GNU General Public License for more
# details.
#
# You should have received a copy of the GNU General Public License along with
# this program; if not, write to the Free Software Foundation, Inc.,
# 59 Temple Place, Suite 330, Boston, MA  02111-1307  USA
#
# Module authors:
#   Damaris Tartarotti Maimone
#   Markus Zolliker <markus.zolliker@psi.ch>
#
# *****************************************************************************
"""support for ultrasound plot clients"""

import numpy as np
import matplotlib.pyplot as plt
# disable the behaviour of raising the window to the front each time it is updated
plt.rcParams["figure.raise_window"] = False


NAN = float('nan')


class Pause:
    """allows to leave the plot loop when the window is closed

    Usage:
        pause = Pause(fig)

        # do initial plots
        plt.show()
        while pause(0.5):
            # do plot updates
            plt.draw()
    """
    def __init__(self, fig):
        fig.canvas.mpl_connect('close_event', self.on_close)
        self.running = True

    def on_close(self, event):
        self.running = False

    def __call__(self, interval):
        try:
            plt.pause(interval)
        except Exception:
            pass
        return self.running


def rect(x1, x2, y1, y2):
    return np.array([[x1,x2,x2,x1,x1],[y1,y1,y2,y2,y1]])


def rects(intervals, y12):
    result = [rect(*intervals[0], *y12)]
    for x12 in intervals[1:]:
        result.append([[NAN],[NAN]])
        result.append(rect(*x12, *y12))
    return np.concatenate(result, axis=1)    


class Plot:
    def __init__(self, maxy, maxx=None):
        self.lines = {}
        self.yaxis = ((-2 * maxy, maxy), (-maxy, 2 * maxy))
        self.maxx = maxx
        self.first = True
        self.fig = None

    def pause(self, interval):
        """will be overridden when figure is created"""
        return False

    def set_line(self, iax, name, data, fmt, **kwds):
        """
        plot or update a line

        when called with self.first = True: plot the line
        when called with self.first = False: update the line

        iax: 0: left, 1: right yaxis
        name: the name of the line. used also as label for legend, if not starting with underscore
        data: data[0] and data[1] are used for x/y data respectively
        fmt, other keywords: forwarded to <axis>.plot
        """
        # ax: 0: left, 1: right
        if self.first:
            if name.startswith('_'):
                label = '_nolegend_'
            else:
                label = name
            self.lines[name], = self.ax[iax].plot(data[0], data[1], fmt, label=label, **kwds)
        else:
            self.lines[name].set_data(data[0:2])

    def close(self):
        if self.fig:
            plt.close(self.fig)
            self.fig = None
            self.first = True
        
    def plot(self, curves, rois=None, average=None):
        boxes = rects(rois[1:], self.yaxis[0])
        pbox = rect(*rois[0], *self.yaxis[1])
        rbox = rect(*rois[1], *self.yaxis[0])

        pshift = self.yaxis[0][1] * 0.5
        pulse = curves[3] - pshift
        # normalized to 0.8 * pshift:
        #pulse = (curves[3] / np.max(curves[3]))* pshift * 0.8 - pshift

        try:
            if self.first:
                plt.ion()
                self.fig, axleft = plt.subplots(figsize=(15,7))
                self.pause = Pause(self.fig)
                plt.title("I/Q", fontsize=14)
                axleft.set_xlim(0, self.maxx or curves[0][-1])
                self.ax = [axleft, axleft.twinx()]
                self.ax[0].axhline(y=0, color='#cccccc')  # show x-axis line
                self.ax[1].axhline(y=0, color='#cccccc')
                self.ax[0].set_ylim(*self.yaxis[0])
                self.ax[1].set_ylim(*self.yaxis[1])

            self.set_line(0, "I", curves, 'b-')  # using curves [0] and [1]
            self.set_line(0, "_Iaverage", average, 'b.')

            self.set_line(0, "Ampl", (curves[0],np.sqrt(curves[1]**2+curves[2]**2)), '#808080')

            self.set_line(1, "Q", (curves[0], curves[2]), 'g-')
            self.set_line(1, "_Qaverage", (average[0], average[2]), 'g.')

            self.set_line(0, "pulse", (curves[0], pulse), 'c-')

            self.set_line(0, "roi's", boxes, 'm-')
            self.set_line(1, "pulse reg", pbox, 'k-')
            self.set_line(0, "ctrl reg", rbox, 'r-')

            if self.first:
                self.fig.legend(fontsize=12)
                plt.tight_layout()
        finally:
            self.first = False

        plt.draw()
        # TODO: do not know why this is needed:
        self.fig.canvas.draw()
        self.fig.canvas.flush_events()
