# Copyright 2012 Free Software Foundation, Inc.
#
# This file is part of GNU Radio
#
# GNU Radio 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 3, or (at your option)
# any later version.
#
# GNU Radio 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 GNU Radio; see the file COPYING.  If not, write to
# the Free Software Foundation, Inc., 51 Franklin Street,
# Boston, MA 02110-1301, USA.
#

import sys
from PyQt4 import Qt, QtCore
from math import sin, cos, pi
import PyQt4.Qwt5 as Qwt
from PyQt4.Qwt5.anynumpy import *

class PzPlot(Qwt.QwtPlot):

    def __init__(self, *args):
        Qwt.QwtPlot.__init__(self, *args)


        self.ymax=0
        self.xmax=0
        self.ymin=0
        self.xmin=0
        self.setCanvasColor(Qt.Qt.darkCyan)

        grid = Qwt.QwtPlotGrid()
        grid.attach(self)
        grid.setMajPen(Qt.QPen(Qt.Qt.white, 0, Qt.Qt.DotLine))

        self.setAxisScale(Qwt.QwtPlot.xBottom, -3, 3)
        self.setAxisScale(Qwt.QwtPlot.yLeft, -2, 2)
        self.drawUnitcircle()



    def setCanvasColor(self, color):
        self.setCanvasBackground(color)
        self.replot()


    def drawUnitcircle(self):
        radius = 1.0
        steps = 1024

        angleStep = 2 * pi / steps
        x=[sin(a * angleStep) * radius for a in range(0, steps)]
        y=[cos(a * angleStep) * radius for a in range(0, steps)]

        curve = Qwt.QwtPlotCurve()
        curve.attach(self)
        curve.setPen(Qt.QPen(Qt.Qt.gray, 1, Qt.Qt.DotLine))
        curve.setData(x, y)


    def insertZeros(self, roots):
        self.removeallCurves()
        if len(roots):
            self.__insertZero(Qt.Qt.blue, roots.real,roots.imag)
            self.ymax =  max(roots.imag)
            self.ymin =  min(roots.imag)
            self.xmax =  max(roots.real)
            self.xmin =  min(roots.real)
            #To make the plot look good
            if self.xmax <= 1.3:
                self.xmax = 2
            if self.xmin >= -1.3:
                self.xmin = -2
            if self.ymax <= 1:
                self.ymax = 1.5
            if self.ymin >= -1:
                self.ymin = -1.5

            self.setAxisScale(Qwt.QwtPlot.xBottom, self.xmin, self.xmax)
            self.setAxisScale(Qwt.QwtPlot.yLeft, self.ymin, self.ymax)

    def insertPoles(self, roots):
            if len(roots):
                self.__insertPole(Qt.Qt.black, roots.real,roots.imag)
                ymax = max(roots.imag)
                ymax = max(ymax,self.ymax)
                ymin = min(roots.imag)
                ymin = min(ymin,self.ymin)
                xmax = max(roots.real)
                xmax = max(xmax,self.xmax)
                xmin = min(roots.real)
                xmin = min(xmin,self.xmin)

                #To make the plot look good
                if xmax <= 1.3:
                    xmax = 2
                else:
                    xmax = 1.2 * xmax

                if xmin >= -1.3:
                    xmin = -2
                else:
                    xmin = 1.2 * xmin

                if ymax <= 1:
                    ymax = 1.5
                else:
                    ymax = 1.2 * ymax

                if ymin >= -1:
                    ymin = -1.5
                else:
                    ymin = 1.2 * ymin

                self.setAxisScale(Qwt.QwtPlot.xBottom, xmin, xmax)
                self.setAxisScale(Qwt.QwtPlot.yLeft, ymin, ymax)
                self.drawUnitcircle()
                self.replot()

    def __insertZero(self, color, px, py):

        curve = Qwt.QwtPlotCurve()
        curve.attach(self)

        curve.setPen(Qt.QPen(Qt.Qt.white, 0, Qt.Qt.NoPen))
        curve.setSymbol(Qwt.QwtSymbol(Qwt.QwtSymbol.Ellipse,
                                      Qt.QBrush(Qt.Qt.gray),
                                      Qt.QPen(color),
                                      Qt.QSize(10, 10)))
        curve.setData(px, py)

    def __insertPole(self, color, px, py):
        curve = Qwt.QwtPlotCurve()
        curve.attach(self)

        curve.setPen(Qt.QPen(Qt.Qt.white, 0, Qt.Qt.NoPen))
        curve.setSymbol(Qwt.QwtSymbol(Qwt.QwtSymbol.XCross,
                                      Qt.QBrush(Qt.Qt.gray),
                                      Qt.QPen(color),
                                      Qt.QSize(7, 7)))
        curve.setData(px, py)

    def removeallCurves(self):
        for curve in self.itemList():
            if isinstance(curve, Qwt.QwtPlotCurve):
                curve.detach()
        self.replot()



class CanvasPicker(Qt.QObject):
    curveChanged = QtCore.pyqtSignal(tuple)
    mouseposChanged = QtCore.pyqtSignal(tuple)
    def __init__(self, plot):
        Qt.QObject.__init__(self, plot)
        self.__selectedCurve = None
        self.__selectedPoint = -1
        self.__selectedcPoint = -1
        self.__addedZero = -1
        self.__addedcZero = -1
        self.changeConjugate = False
        self.enableZeroadd= False
        self.enablePoleadd= False
        self.enablepzDelete= False
        self.iir = False
        self.__plot = plot

        canvas = plot.canvas()
        canvas.installEventFilter(self)

        # We want the focus, but no focus rect.
        # The selected point will be highlighted instead.
        canvas.setFocusPolicy(Qt.Qt.StrongFocus)
        canvas.setCursor(Qt.Qt.PointingHandCursor)
        canvas.setFocusIndicator(Qwt.QwtPlotCanvas.ItemFocusIndicator)
        canvas.setFocus()

        canvas.setWhatsThis(
            'All points can be moved using the left mouse button '
            'or with these keys:\n\n'
            '- Up: Select next curve\n'
            '- Down: Select previous curve\n'
            '- Left, "-": Select next point\n'
            '- Right, "+": Select previous point\n'
            '- 7, 8, 9, 4, 6, 1, 2, 3: Move selected point'
            )

        self.__shiftCurveCursor(True)


    def event(self, event):
        if event.type() == Qt.QEvent.User:
            self.__showCursor(True)
            return True
        try:
            return Qt.QObject.event(event)
        except TypeError:
            return False



    def set_conjugate(self):
        self.changeConjugate = not(self.changeConjugate)

    def set_iir(self,val=True):
        self.iir = val

    def add_zero(self):
        self.enableZeroadd = not(self.enableZeroadd)

    def add_pole(self):
        #Adding pole support only for iir
        if self.iir:
            self.enablePoleadd = not(self.enablePoleadd)

    def delete_pz(self):
        self.enablepzDelete = not(self.enablepzDelete)

    def eventFilter(self, object, event):

        if event.type() == Qt.QEvent.FocusIn:
            self.__showCursor(True)
        if event.type() == Qt.QEvent.FocusOut:
            self.__showCursor(False)

        if event.type() == Qt.QEvent.Paint:
            Qt.QApplication.postEvent(
                self, Qt.QEvent(Qt.QEvent.User))
        elif event.type() == Qt.QEvent.MouseButtonPress:
            if self.enableZeroadd or self.enablePoleadd:
                self.__drawAddedzero_pole(True, event.pos())
            else:
                self.__select(event.pos())
            return True
        elif event.type() == Qt.QEvent.MouseMove:
            curve = self.__selectedCurve
            if curve:
                tp=(self.__plot.invTransform(curve.xAxis(), event.pos().x()),
                    self.__plot.invTransform(curve.xAxis(), event.pos().y()))
                self.mouseposChanged.emit(tp)
            self.__move(event.pos())
            return True

        if event.type() == Qt.QEvent.KeyPress:
            delta = 5
            key = event.key()
            if key == Qt.Qt.Key_Up:
                self.__shiftCurveCursor(True)
                return True
            elif key == Qt.Qt.Key_Down:
                self.__shiftCurveCursor(False)
                return True
            elif key == Qt.Qt.Key_Right or key == Qt.Qt.Key_Plus:
                if self.__selectedCurve:
                    self.__shiftPointCursor(True)
                else:
                    self.__shiftCurveCursor(True)
                return True
            elif key == Qt.Qt.Key_Left or key == Qt.Qt.Key_Minus:
                if self.__selectedCurve:
                    self.__shiftPointCursor(False)
                else:
                    self.__shiftCurveCursor(True)
                return True

            if key == Qt.Qt.Key_1:
                self.__moveBy(-delta, delta)
            elif key == Qt.Qt.Key_2:
                self.__moveBy(0, delta)
            elif key == Qt.Qt.Key_3:
                self.__moveBy(delta, delta)
            elif key == Qt.Qt.Key_4:
                self.__moveBy(-delta, 0)
            elif key == Qt.Qt.Key_6:
                self.__moveBy(delta, 0)
            elif key == Qt.Qt.Key_7:
                self.__moveBy(-delta, -delta)
            elif key == Qt.Qt.Key_8:
                self.__moveBy(0, -delta)
            elif key == Qt.Qt.Key_9:
                self.__moveBy(delta, -delta)

        return Qwt.QwtPlot.eventFilter(self, object, event)


    def __select(self, pos):
        found, distance, point = None, 1e100, -1

        for curve in self.__plot.itemList():
            if isinstance(curve, Qwt.QwtPlotCurve):
                if curve.symbol().style() != Qwt.QwtSymbol.NoSymbol:
                    i, d = curve.closestPoint(pos)
                    if d < distance:
                        found = curve
                        point = i
                        distance = d

        self.__showCursor(False)
        self.__selectedCurve = None
        self.__selectedPoint = -1
        self.__selectedcPoint = -1

        if found and distance < 10:
            self.__selectedCurve = found
            self.__selectedPoint = point
            #search for conjugate point if enabled
            if self.changeConjugate:
                j=self.__searchConjugate(found.x(point),found.y(point))
                self.__selectedcPoint = j
            #delete zero/pole if enabled
            if self.enablepzDelete:
                self.__deleteZero()
            self.__showCursor(True)
        else:
            self.mouseposChanged.emit((None,None))

    def __deleteZero(self):
        curve = self.__selectedCurve
        xData = zeros(curve.dataSize(), Float)
        yData = zeros(curve.dataSize(), Float)

        for i in range(curve.dataSize()):
            xData[i] = curve.x(i)
            yData[i] = curve.y(i)

        if(self.__selectedPoint != -1):
            xData=delete(xData, self.__selectedPoint)
            yData=delete(yData, self.__selectedPoint)
        #one less to accomodate previous delete
        if(self.__selectedcPoint != -1):
            xData=delete(xData, self.__selectedcPoint-1)
            yData=delete(yData, self.__selectedcPoint-1)

        curve.setData(xData, yData)
        self.__plot.replot()
        px=[]
        py=[]
        for c in self.__plot.itemList():
            if isinstance(c, Qwt.QwtPlotCurve):
                px.append([c.x(i) for i in range(c.dataSize())])
                py.append([c.y(i) for i in range(c.dataSize())])
        tp=(vectorize(complex)(px[0],py[0]),vectorize(complex)(px[1],py[1]))
        self.curveChanged.emit(tp)

    def __moveBy(self, dx, dy):
        if dx == 0 and dy == 0:
            return

        curve = self.__selectedCurve
        if not curve:
            return

        x = self.__plot.transform(
            curve.xAxis(), curve.x(self.__selectedPoint)) + dx
        y = self.__plot.transform(
            curve.yAxis(), curve.y(self.__selectedPoint)) + dy
        self.__move(Qt.QPoint(x, y))


    def __move(self, pos):
        curve = self.__selectedCurve
        if not curve:
            return

        xData = zeros(curve.dataSize(), Float)
        yData = zeros(curve.dataSize(), Float)

        #poles made unmovable when FIR design is active
        move_enable=True
        if not self.iir:
            if self.__selectedCurve.symbol().style() == Qwt.QwtSymbol.XCross:
                move_enable=False
        if move_enable:
            for i in range(curve.dataSize()):
                if i == self.__selectedPoint:
                    xData[i] = self.__plot.invTransform(curve.xAxis(), pos.x())
                    yData[i] = self.__plot.invTransform(curve.yAxis(), pos.y())
                    if(self.__selectedcPoint != -1):
                        xData[self.__selectedcPoint] = xData[i]
                        yData[self.__selectedcPoint] = -yData[i]
                elif i != self.__selectedcPoint:
                    xData[i] = curve.x(i)
                    yData[i] = curve.y(i)
            curve.setData(xData, yData)
            self.__plot.replot()
            px=[]
            py=[]
            for c in self.__plot.itemList():
                if isinstance(c, Qwt.QwtPlotCurve):
                    px.append([c.x(i) for i in range(c.dataSize())])
                    py.append([c.y(i) for i in range(c.dataSize())])
            tp=(vectorize(complex)(px[0],py[0]),vectorize(complex)(px[1],py[1]))
            self.curveChanged.emit(tp)
            self.__showCursor(True)

    def __searchConjugate(self, x, y):
        curve = self.__selectedCurve
        for i in range(curve.dataSize()):
            if (round(curve.x(i),8) == round(x,8) and round(curve.y(i),8) == -round(y,8)):
                if (y !=0):
                    return i
        return -1

    def __drawAddedzero_pole(self, showIt, pos):
        editcurve=''
        for curve in self.__plot.itemList():
            if isinstance(curve, Qwt.QwtPlotCurve):
                if self.enableZeroadd:
                    if curve.symbol().style() == Qwt.QwtSymbol.Ellipse:
                        editcurve=curve
                if self.enablePoleadd:
                    if curve.symbol().style() == Qwt.QwtSymbol.XCross:
                        editcurve=curve
        if not editcurve:
            return

        if self.changeConjugate:
            extrapoints=2
        else:
            extrapoints=1

        xData = zeros(editcurve.dataSize()+extrapoints, Float)
        yData = zeros(editcurve.dataSize()+extrapoints, Float)

        for i in range(editcurve.dataSize()):
            xData[i] = editcurve.x(i)
            yData[i] = editcurve.y(i)
        xData[i+1] = self.__plot.invTransform(editcurve.xAxis(), pos.x())
        yData[i+1] = self.__plot.invTransform(editcurve.yAxis(), pos.y())

        if self.changeConjugate:
            xData[i+2] = xData[i+1]
            yData[i+2] = -yData[i+1]
            self.__addedcZero=i+2

        editcurve.setData(xData, yData)

        self.__addedZero=i+1
        symbol = Qwt.QwtSymbol(editcurve.symbol())
        newSymbol = Qwt.QwtSymbol(symbol)
        newSymbol.setPen(Qt.QPen(Qt.Qt.red))
        doReplot = self.__plot.autoReplot()

        self.__plot.setAutoReplot(False)
        editcurve.setSymbol(newSymbol)

        editcurve.draw(self.__addedZero, self.__addedZero)
        if self.changeConjugate:
            editcurve.draw(self.__addedcZero, self.__addedcZero)

        editcurve.setSymbol(symbol)
        self.__plot.setAutoReplot(doReplot)
        px=[]
        py=[]
        for c in self.__plot.itemList():
            if isinstance(c, Qwt.QwtPlotCurve):
                px.append([c.x(i) for i in range(c.dataSize())])
                py.append([c.y(i) for i in range(c.dataSize())])
        tp=(vectorize(complex)(px[0],py[0]),vectorize(complex)(px[1],py[1]))
        self.curveChanged.emit(tp)
        self.__showCursor(True)

    def __showCursor(self, showIt):
        curve = self.__selectedCurve
        if not curve:
            return

        # Use copy constructors to defeat the reference semantics.
        symbol = Qwt.QwtSymbol(curve.symbol())
        newSymbol = Qwt.QwtSymbol(symbol)
        if showIt:
            newSymbol.setBrush(symbol.brush().color().dark(150))

        doReplot = self.__plot.autoReplot()

        self.__plot.setAutoReplot(False)
        curve.setSymbol(newSymbol)

        curve.draw(self.__selectedPoint, self.__selectedPoint)
        if (self.__selectedcPoint != -1):
            curve.draw(self.__selectedcPoint, self.__selectedcPoint)

        curve.setSymbol(symbol)
        self.__plot.setAutoReplot(doReplot)

    def __shiftCurveCursor(self, up):
        curves = [curve for curve in self.__plot.itemList()
                  if isinstance(curve, Qwt.QwtPlotCurve)]

        if not curves:
            return

        if self.__selectedCurve in curves:
            index = curves.index(self.__selectedCurve)
            if up:
                index += 1
            else:
                index -= 1
            # keep index within [0, len(curves))
            index += len(curves)
            index %= len(curves)
        else:
            index = 0

        self.__showCursor(False)
        self.__selectedPoint = 0
        self.__selectedCurve = curves[index]
        self.__showCursor(True)


    def __shiftPointCursor(self, up):
        curve = self.__selectedCurve
        if not curve:
            return

        if up:
            index = self.__selectedPoint + 1
        else:
            index = self.__selectedPoint - 1
        # keep index within [0, curve.dataSize())
        index += curve.dataSize()
        index %= curve.dataSize()
        if index != self.__selectedPoint:
            self.__showCursor(False)
            self.__selectedPoint = index
            self.__showCursor(True)

