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

import phys
from ui import GUI, Pens
from ui_widgets import *



class ElementEditorWindow(QMainWindow):
	windows = list()

	@staticmethod
	def open(row):
		win = ElementEditorWindow(row)
		win.show()
		ElementEditorWindow.windows.append(win)

	@staticmethod
	def close_all(pol=None):
		for win in ElementEditorWindow.windows:
			if pol is None or pol == win.pol:
				win.close()
				del(win)

	def __init__(self, row):
		super().__init__(flags=Qt.Dialog)
		self.row = row
		self.system = row.parent.system
		self.pol = row.pol

		self.setWindowTitle("Element editor")
		root = QVBoxLayout()
		root_fucking_random_container = QWidget()
		root_fucking_random_container.setLayout(root)
		self.setCentralWidget(root_fucking_random_container)

		# Name
		box = QHBoxLayout()
		root.addLayout(box)
		box.addWidget(QLabel("Name"))
		self.name = QLineEdit(self.pol.name)
		self.name.textChanged.connect(self.change_name)
		box.addWidget(self.name)

		# PHASE RETARDATION
		root.addWidget(MeinGroßLabel("Phase retardation"))

		self.phase_retardation = AngleSlider()
		self.phase_retardation.setValue(self.pol.phase_retardation * 180 / np.pi)
		self.phase_retardation.on_change = self.change_optics
		root.addLayout(self.phase_retardation)

		# TRANSMITTANCE
		root.addWidget(MeinGroßLabel("Transmittance"))

		box = QHBoxLayout()
		root.addLayout(box)

		box.addWidget(QLabel("T"))
		self.t1 = QLineEdit("%g" % self.pol.t1)
		self.t1.textChanged.connect(self.change_optics)
		box.addWidget(self.t1)
		self.t2 = QLineEdit("%g" % self.pol.t2)
		self.t2.textChanged.connect(self.change_optics)
		box.addWidget(self.t2)

		# ARRANGEMENT
		root.addWidget(MeinGroßLabel("Arrangement"))

		# Angle reference
		self.ref = QComboBox()
		self.populate_ref()
		self.ref.currentIndexChanged.connect(GUI.do_update)
		root.addWidget(self.ref)

		# Delta angle
		self.delta = AngleSlider("Delta")
		self.delta.on_change = self.change_delta
		root.addLayout(self.delta)

		root.addItem(ExpandingSpacer())

		# Done
		self.done = QPushButton("Done")
		self.done.pressed.connect(self.close)
		root.addWidget(self.done)

	def change_name(self):
		self.pol.name = self.name.text()
		self.row.optbox.name.setText(self.pol.name)

	def change_delta(self):
		try:
			self.pol.delta = self.delta.angle / 180 * np.pi
			GUI.do_update()
		except ValueError:
			pass

	def change_optics(self):
		try:
			self.pol.t1 = float(self.t1.text())
			self.pol.t2 = float(self.t2.text())
			self.pol.phase_retardation = \
				float(self.phase_retardation.angle) / 180 * np.pi
			GUI.do_update()
		except ValueError:
			pass

	def populate_ref(self):
		self.ref.addItem("Absolute", False)
		for i, pol in enumerate(self.system.elements):
			self.ref.addItem("Relative to %s" % pol.name)

		self.ref.setCurrentIndex(self.pol.ref)



class OptBox(QVBoxLayout):
	def __init__(self, row, pol):
		super().__init__()
		self.row = row

		# Name
		self.name = MeinGroßLabel(pol.name)
		self.addWidget(self.name)

		# Enable checkbox
		self.enable = QCheckBox("Enable")
		self.enable.setChecked(True)
		self.enable.stateChanged.connect(self.change_enable)
		self.addWidget(self.enable)

		# Angle
		self.angle = AngleSlider()
		self.angle.setValue(pol.angle * 180 / np.pi)
		self.angle.on_change = row.angle_change
		self.addLayout(self.angle)
		self.addItem(ExpandingSpacer())

	def change_enable(self):
		self.row.pol.enable = self.enable.isChecked()
		GUI.do_update()



class SystemTableRow:
	def __init__(self, parent, pol, row_number):
		self.parent = parent
		self.pol = pol
		self.row_number = row_number

		self.angle = pol.angle
		self.delta = pol.delta
		self.ref = pol.ref

		# Option box (column 1)
		self.optbox = OptBox(self, pol)
		parent.setCellWidget(row_number, 0, LayoutWrapper(self.optbox))

		# Ellipse diagram (column 2)
		self.ellipse = EllipseWidget(pol)
		parent.setCellWidget(row_number, 1, self.ellipse)

		# Jones vector (column 3)
		box = QVBoxLayout()
		parent.setCellWidget(row_number, 2, LayoutWrapper(box))

		grid = QGridLayout()
		grid.setColumnStretch(2, 10)
		box.addLayout(grid)
		box.addItem(ExpandingSpacer())

		for i, label in enumerate(["|Ex|", "arg Ex", "|Ey|", "arg Ey", "δ"]):
			grid.addWidget(QLabel(label), 1 + i, 1)

		self.jones = list()
		for i in range(5):
			label = MeinLabel("-")
			grid.addWidget(label, 1 + i, 2)
			self.jones.append(label)

		# Stokes parameters (column 4)
		box = QVBoxLayout()
		parent.setCellWidget(row_number, 3, LayoutWrapper(box))

		grid = QGridLayout()
		grid.setColumnStretch(2, 10)
		box.addLayout(grid)
		box.addItem(ExpandingSpacer())

		for i, label in enumerate(["I", "Q", "U", "V", "α", "θ"]):
			grid.addWidget(QLabel(label), 1 + i, 1)

		self.stokes = list()
		for i in range(6):
			label = MeinLabel("-")
			grid.addWidget(label, 1 + i, 2)
			self.stokes.append(label)

	def angle_change(self):
		self.angle = self.optbox.angle.angle / 180 * np.pi
		self.pol.angle = self.angle
		GUI.do_update()



class SystemTable(QTableWidget):
	def __init__(self):
		super().__init__()
		self.rows = list()

		self.setColumnCount(4) # this won't change

		self.hh = self.horizontalHeader()
		self.hh.setSectionResizeMode(0, QHeaderView.ResizeToContents)
		self.hh.resizeSection(1, 150)
		self.hh.resizeSection(2, 150)
		self.hh.setStretchLastSection(True)
		self.setHorizontalHeaderLabels([
			"Settings", "Diagram", "State", "Stokes"
		])

		self.vh = self.verticalHeader()
		self.vh.setSectionResizeMode(QHeaderView.ResizeToContents)

		# Smooth scrolling
		self.setHorizontalScrollMode(QAbstractItemView.ScrollPerPixel)
		self.setVerticalScrollMode(QAbstractItemView.ScrollPerPixel)

		# Select only entire rows
		self.setSelectionBehavior(QAbstractItemView.SelectRows)
	
		self.setContextMenuPolicy(Qt.CustomContextMenu)
		self.customContextMenuRequested.connect(self.context_menu)

	def context_menu(self, pos):
		indices = list()
		for lo, hi in [(r.topRow(), r.bottomRow()) \
		               for r in self.selectedRanges()]:
			indices.extend(list(range(lo, hi + 1)))

		menu = QMenu()

		if len(indices) == 0:
			if len(self.rows) > 0:
				menu.addAction("Add a new element at the top", \
				               lambda: self.insert_row(None, True))
				menu.addAction("Add a new element at the bottom", \
				               lambda: self.insert_row(None, False))
			else:
				menu.addAction("Add a new element", lambda: self.insert_row(None, False))
		elif len(indices) == 1:
			i = indices[0]
			menu.addAction("Edit %s" % self.rows[i].pol.name, \
			               lambda: ElementEditorWindow.open(self.rows[i]))
			menu.addSeparator()
			menu.addAction("Add a new element before", lambda: self.insert_row(self.rows[i], True))
			menu.addAction("Add a new element after", lambda: self.insert_row(self.rows[i], False))
			deltext = "Delete %s" % self.rows[i].pol.name
		else:
			deltext = "Delete %d elements" % len(indices)

		if len(indices) >= 1:
			menu.addAction(deltext, lambda: GUI.table.delete_rows(indices))

		menu.addSeparator()

		menu.exec_(self.mapToGlobal(pos))

	def clear(self):
		ElementEditorWindow.close_all()
		self.system.elements = list()
		self.rows = list()
		self.clearContents()
		self.setRowCount(0)
		GUI.do_update()

	def renumber_rows(self):
		for i, row in enumerate(self.rows):
			row.row_number = i

	def insert_row(self, wrt, before, pol=None):
		if wrt is not None:
			i = wrt.row_number + 0 if before else 1
		else:
			i = 0 if before else len(self.rows)

		if pol is None:
			pol = phys.Polarizer()
			editor = True
		else:
			editor = False

		self.system.elements.insert(i, pol)
		self.insertRow(i)
		row = SystemTableRow(self, pol, i)
		self.rows.insert(i, row)
		self.renumber_rows()
		GUI.do_update()

		if editor:
			ElementEditorWindow.open(row)

	def delete_rows(self, indices):
		for i in sorted(indices):
			ElementEditorWindow.close_all(pol=self.rows[i].pol)

		offs = 0
		for i in sorted(indices):
			self.removeRow(i - offs)
			del(self.rows[i - offs])
			del(self.system.elements[i - offs])
			offs += 1
	
		self.renumber_rows()
		GUI.do_update()

	def populate(self, system):
		self.system = system

		self.rows = list()
		self.clearContents()
		self.setRowCount(len(system.elements))

		for i, pol in enumerate(system.elements):
			self.rows.append(SystemTableRow(self, pol, i))