import numpy as np
from PyQt5.QtWidgets import *
from PyQt5.QtGui import *
from PyQt5.QtCore import *

from ui import GUI



class Pens:
	axes = QPen(Qt.gray)

	ellipse = QPen(Qt.black)
	ellipse.setWidth(2)

	axis_linear = QPen(QColor(201, 141, 0))
	axis_linear.setWidth(2)
	axis_linear.setStyle(Qt.DashDotLine)

	axis_fast = QPen(QColor(51, 87, 123))
	axis_fast.setWidth(2)
	axis_fast.setStyle(Qt.DashDotLine)

	axis_slow = QPen(QColor(55, 123, 51))
	axis_slow.setWidth(2)
	axis_slow.setStyle(Qt.DashDotLine)

	alpha = QPen(Qt.red)
	theta = QPen(Qt.blue)

	radii = QPen(Qt.black)
	radii.setStyle(Qt.DashDotLine)



class ExpandingSpacer(QSpacerItem):
	def __init__(self):
		super().__init__(0, 0, QSizePolicy.Minimum, \
		                 QSizePolicy.Expanding)

class LayoutWrapper(QWidget):
	def __init__(self, layout):
		super().__init__()
		self.setLayout(layout)



class MeinLabel(QLabel):
	def __init__(self, *args):
		QLabel.__init__(self, *args)
		self.setFont(GUI.monospace)

class MeinGroßLabel(QLabel):
	def __init__(self, *args):
		QLabel.__init__(self, *args)
		self.setFont(GUI.bigfont)
		self.setAlignment(Qt.AlignHCenter)



class EllipseWidget(QWidget):
	def __init__(self, pol):
		QWidget.__init__(self)
		self.pol = pol
		self.state = None
		self.ellipse = None
		self.is_used = True

		self.setSizePolicy(QSizePolicy(QSizePolicy.Expanding, \
		                               QSizePolicy.Expanding))

	def paintEvent(self, event):
		P = QPainter(self)

		w, h = self.frameSize().width(), self.frameSize().height()
		cx, cy = w / 2, h / 2
		r = min(w, h) / 2

		# background
		if self.is_used:
			background = Qt.white
		else:
			background = QColor(0, 0, 0, 100)
		P.fillRect(QRect(0, 0, w, h), background)

		# coordinate axes
		P.setPen(Pens.axes)
		P.drawLine(cx - 20 * r, cy, cx + 20 * r, cy)
		P.drawLine(cx, cy - 20 * r, cx, cy + 20 * r)

		# polarizer axes
		P.translate(cx, cy)
		P.rotate(- (self.pol.angle + self.pol.delta) * 180 / np.pi)
		if self.pol.t2 == 0: # FIXME: half-assed
			P.setPen(Pens.axis_linear)
			P.drawLine(-20 * r, 0, 20 * r, 0)
		else:
			P.setPen(Pens.axis_fast)
			P.drawLine(-20 * r, 0, 20 * r, 0)
			P.drawText(0.92 * r, 0, "F")
			P.setPen(Pens.axis_slow)
			P.drawLine(0, -20 * r, 0, 20 * r)
			P.drawText(0, -0.92 * r, "S")

		if self.state is None:
			return

		if np.isnan(self.ellipse.a) or np.isnan(self.ellipse.b) \
		   or np.isnan(self.ellipse.alpha) or np.isnan(self.ellipse.theta):
			return

		P.resetTransform()
		r *= 0.88
		P.translate(cx, cy)
		P.scale(1, -1)

		# radii
		P.setPen(Pens.radii)
		csa = np.array([np.cos(self.ellipse.alpha), np.sin(self.ellipse.alpha)])
		xa = self.ellipse.a * r * csa
		P.drawLine(-xa[0], -xa[1], xa[0], xa[1])
		csb = np.array([np.cos(self.ellipse.alpha + np.pi / 2), \
		               np.sin(self.ellipse.alpha + np.pi / 2)])
		xb = self.ellipse.b * r * csb
		P.drawLine(-xb[0], -xb[1], xb[0], xb[1])

		def arc(x, y, r, t1, t2):
			P.drawArc(x - r / 2, y - r / 2, r, r, \
			          -180 * 16 / np.pi * t1, \
			          -180 * 16 / np.pi * t2)

		# azimuth
		P.setPen(Pens.alpha)
		P.drawLine(0, 0, xa[0], xa[1])
		arc(0, 0, self.ellipse.a * r, 0, self.ellipse.alpha)

		# ellipticity
		P.setPen(Pens.theta)
		P.drawLine(-xa[0], -xa[1], xb[0], xb[1])
		arc(-xa[0], -xa[1], r / 2, self.ellipse.alpha, abs(self.ellipse.theta))

		# the ellipse
		P.setPen(Pens.ellipse)
		path = QPainterPath()
		for i, t in enumerate(np.linspace(0, 2 * np.pi, 100)):
			x = r * np.real(np.exp(1j * t) * self.state)
			f = path.moveTo if t == 0 else path.lineTo # FML
			f(x[0], x[1]) # +Y should be upwards

			if abs(self.ellipse.theta) > 0.05 and i % 25 == 0:
				N = r * np.real(np.exp(1j * (t + 0.001)) * self.state) - x
				N /= np.linalg.norm(N)
				if np.any(np.isnan(N)):
					continue

				T = np.dot(np.array([[0, 1], [-1, 0]]), N)

				ax = x - N * 5 + T * 5
				P.drawLine(x[0], x[1], ax[0], ax[1])
				ax = x - N * 5 - T * 5
				P.drawLine(x[0], x[1], ax[0], ax[1])
		path.closeSubpath()
		P.drawPath(path)

		if abs(self.ellipse.theta) <= 0.05:
			x1 = xa - r / 20 * csb
			x2 = xa + r / 20 * csb
			P.drawLine(x1[0], x1[1], x2[0], x2[1])
			x1 = -xa - r / 20 * csb
			x2 = -xa + r / 20 * csb
			P.drawLine(x1[0], x1[1], x2[0], x2[1])



class AngleSlider(QHBoxLayout):
	def __init__(self, label=None):
		super().__init__()
		self.angle = 0

		if label:
			self.addWidget(QLabel(label))

		self.slider = QSlider(Qt.Horizontal)
		self.slider.setMinimumWidth(100)
		self.slider.setMinimum(0)
		self.slider.setMaximum(180)
		self.slider.setTickPosition(QSlider.TicksBelow)
		self.slider.setTickInterval(45)
		self.slider.valueChanged.connect(\
			lambda: AngleSlider.change(self, "bar"))
		self.addWidget(self.slider)

		self.edit = QLineEdit()
		self.edit.setText("0")
		self.edit.setFixedWidth(45)
		self.edit.textChanged.connect(\
			lambda: AngleSlider.change(self, "edit"))
		self.addWidget(self.edit)
		self.addWidget(QLabel("°"))

	def setValue(self, angle):
		self.angle = angle
		self.slider.blockSignals(True)
		self.slider.setValue(angle)
		self.slider.blockSignals(False)
		self.edit.blockSignals(True)
		self.edit.setText("%g" % angle)
		self.edit.blockSignals(False)

	@staticmethod
	def change(self, source):
		if source == "edit":
			try:
				angle = float(self.edit.text())
			except ValueError:
				return
			self.slider.blockSignals(True)
			self.slider.setValue(angle)
			self.slider.blockSignals(False)
		elif source == "bar":
			angle = self.slider.value()
			self.edit.blockSignals(True)
			self.edit.setText("%g" % angle)
			self.edit.blockSignals(False)
		self.angle = angle
		self.on_change()



class PoincareWidget(QWidget):
	def __init__(self,):
		QWidget.__init__(self)
		self.setSizePolicy(QSizePolicy(QSizePolicy.Expanding, \
		                               QSizePolicy.Expanding))

	def paintEvent(self, event):
		pass