# -*- coding: utf-8 -*-
"""
Created on Tue Feb  4 11:07:56 2020

@author: tartarotti_d-adm
"""

import numpy as np
import matplotlib.pyplot as plt

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

NAN = float('nan')

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):
        self.lines = {}
        self.yaxis = ((-2 * maxy, maxy), (-maxy, 2 * maxy))
        self.first = True
        self.fig = None
        
    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))
                plt.title("I/Q", fontsize=14)
                axleft.set_xlim(0, 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()
        self.fig.canvas.draw()
        self.fig.canvas.flush_events()