import numpy as np, scipy.optimize
import traceback
from PyQt5.QtWidgets import *
from PyQt5.QtGui import *
from PyQt5.QtCore import *

import phys, file

# global GUI-related shit
class GUI:
	table = None
	table_rows = list()
	monospace = QFont("monospace")
	bigfont = QFont("sans-serif", pointSize=10, weight=1000)

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)

from ui_widgets import *



class ElementEditorWindow(QMainWindow):
	windows = list()

	@staticmethod
	def open(pol):
		win = ElementEditorWindow(pol)
		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, pol):
		super().__init__(flags=Qt.Dialog)
		self.pol = 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)

		# OPTICS
		root.addWidget(MeinGroßLabel("Optics"))

		# Type combo
		self.type = QComboBox()
		self.type.addItem("Linear", "linear",)
		self.type.addItem("Quarterwave plate", "quarterwave")
		self.type.setCurrentIndex(0 if pol.type == "linear" else 1)
		self.type.currentIndexChanged.connect(self.change_type)
		root.addWidget(self.type)

		# 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_t)
		box.addWidget(self.t1)
		self.t2 = QLineEdit("%g" % self.pol.t2)
		self.t2.textChanged.connect(self.change_t)
		box.addWidget(self.t2)

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

		# Angle reference
		self.ref = QComboBox()
		self.populate_ref()
		self.ref.currentIndexChanged.connect(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.pol.row.optbox.name.setText(self.pol.name)

	def change_type(self):
		self.pol.set_type(self.type.currentData())
		update()

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

	def change_t(self):
		try:
			self.pol.t1 = float(self.t1.text())
			self.pol.t2 = float(self.t2.text())
			update()
		except ValueError:
			pass

	def populate_ref(self):
		self.ref.addItem("Absolute", False)
		for i, _ in enumerate(system.elements):
			if i >= self.pol.rownum - 1:
				continue
			self.ref.addItem("Relative to #%d" % (i + 1), i)

			if self.pol.ref is not False and i == self.pol.ref:
				self.ref.setCurrentIndex(self.pol.ref + 1)



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

		# Name
		self.name = MeinGroßLabel(pol.name)
		self.name.setContextMenuPolicy(Qt.CustomContextMenu)
		self.name.customContextMenuRequested.connect(lambda pos: OptBox.open_menu(self, pos))
		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 = lambda: TableRow.angle_change(row)
		self.addLayout(self.angle)
		self.addItem(ExpandingSpacer())

	def open_menu(self, pos):
		menu = QMenu()
		menu.addAction("Edit", lambda: ElementEditorWindow.open(self.row.pol))
		menu.addAction("Insert before", lambda: half_assed_element_creation(self.rownum - 1))
		menu.addAction("Insert after", lambda: half_assed_element_creation(self.rownum))
		menu.addAction("Delete", lambda: half_assed_element_deletion(self.rownum - 1))
		menu.exec_(self.name.mapToGlobal(pos))

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


class TableRow:
	def __init__(self, pol, grid, rownum):
		self.pol = pol
		pol.row = self
		pol.rownum = rownum
		self.grid = grid
		self.rownum = rownum
		self.angle = pol.angle
		self.delta = pol.delta
		self.ref = pol.ref

		self.optbox = OptBox(self, pol, rownum)
		grid.addLayout(self.optbox, rownum, 1)

		self.ellipse = EllipseWidget(pol)
		grid.addWidget(self.ellipse, rownum, 2)

		self.info = QVBoxLayout()
		self.info_angle = MeinLabel()
		self.info_angle.setAlignment(Qt.AlignLeft)
		self.info.addWidget(self.info_angle)
		self.info.addWidget(QLabel("Polarization state"))
		self.info_jones = MeinLabel()
		self.info_jones.setAlignment(Qt.AlignLeft)
		self.info.addWidget(self.info_jones)
		self.info.addItem(ExpandingSpacer())
		grid.addLayout(self.info, rownum, 3)

	@staticmethod
	def angle_change(row):
		row.angle = row.optbox.angle.angle / 180 * np.pi
		row.pol.angle = row.angle
		update()



def populate_table():
	GUI.input_intensity.setText("%g" % system.input_intensity)

	GUI.table = QGridLayout()
	GUI.table.setColumnStretch(1, 1)
	GUI.table.setColumnStretch(2, 2)
	GUI.table.setColumnStretch(3, 5)

	while GUI.opt_operand.count() > 0: # FIXME
		GUI.opt_operand.removeItem(0)

	GUI.table_rows = list()
	for i, pol in enumerate(system.elements):
		GUI.table_rows.append(TableRow(pol, GUI.table, i + 1))
		GUI.opt_operand.addItem("Element #%d" % (i + 1), i)

def update():
	for i, pol in enumerate(system.elements):
		row = GUI.table_rows[i]

		if pol.ref is not False:
			pol.angle += system.elements[pol.ref].angle

	system.recalculate()

	I = 1
	for i, pol in enumerate(system.elements):
		row = GUI.table_rows[i]

		# update all the diagrams
		row.ellipse.state = system.states[i]
		row.ellipse.ellipse = system.ellipses[i]
		row.ellipse.is_used = system.elements[i].enable
		row.ellipse.repaint()
		
		text = "%s at %f°" \
			% (pol.type, (pol.angle + pol.delta) * 180 / np.pi) 
		row.info_angle.setText(text)

		state = system.states[i]
		ellipse = system.ellipses[i]

		if state is None:
			I = 1
			row.info_jones.setText("Unpolarized")
		else:
			I = np.abs(state[0] * state[0].conjugate() \
				   + state[1] * state[1].conjugate())

		if state is None:
			continue
		
		A, B = state
		text = "I = %f\n" % I
		text += "%f ∠ %+f°\n%f ∠ %+f°\n" % \
		       (np.abs(A), 180 / np.pi * np.angle(A), \
			np.abs(B), 180 / np.pi * np.angle(B))
		text += "α = %f°\nφ = %+f°" % \
		       (180 / np.pi * ellipse.alpha, 180 / np.pi * ellipse.theta)
		text = text.replace("-", "- ").replace("+", "+ ")
		row.info_jones.setText(text)

	GUI.widok.intensity = I
	GUI.widok.repaint()

	#if GUI.auto_optimize.isChecked():
	#	optimize()


def half_assed_element_creation(index=None):
	pol = phys.Polarizer("linear")
	if index is None:
		system.elements.append(pol)
	else:
		system.elements.insert(index, pol)

	populate_table()
	GUI.table_frame = QFrame()
	GUI.table_frame.setLayout(GUI.table)
	GUI.scroll.setWidget(GUI.table_frame)

	ElementEditorWindow.open(pol)
	update()

def half_assed_clear():
	ElementEditorWindow.close_all()
	system.elements = list()
	populate_table()
	GUI.table_frame = QFrame()
	GUI.table_frame.setLayout(GUI.table)
	GUI.scroll.setWidget(GUI.table_frame)
	update()

def half_assed_element_deletion(index):
	ElementEditorWindow.close_all(pol=system.elements[pol])
	populate_table()
	GUI.table_frame = QFrame()
	GUI.table_frame.setLayout(GUI.table)
	GUI.scroll.setWidget(GUI.table_frame)
	update()


last_save_path = None
file_filter = "Wery Omportant Zejta (*.woz)"

def do_save(win, reuse_old):
	global system, last_save_path

	if reuse_old and last_save_path:
		path = last_save_path
	else:
		path, _ = QFileDialog.getSaveFileName(win, filter=file_filter)
		if path == "":
			return

	try:
		file.save_system(path, system)
	except:
		traceback.print_exc()

	last_save_path = path

def do_open(win):
	global system

	path, _ = QFileDialog.getOpenFileName(win, filter=file_filter)
	if path == "":
		return

	try:
		system = file.open_system(path)
	except:
		traceback.print_exc()

	populate_table()
	GUI.table_frame = QFrame()
	GUI.table_frame.setLayout(GUI.table)
	GUI.scroll.setWidget(GUI.table_frame)
	update()

def setup_menubar(win):
	menu = win.menuBar()

	# File
	open = QAction("&Open system", win)
	open.setShortcut("Ctrl+O")
	open.triggered.connect(lambda: do_open(win))
	save = QAction("&Save system", win)
	save.setShortcut("Ctrl+S")
	save.triggered.connect(lambda: do_save(win, True))
	save_as = QAction("&Save system as...", win)
	save_as.triggered.connect(lambda: do_save(win, False))
	close = QAction("Exit", win)
	close.setShortcut("Ctrl+Q")
	close.triggered.connect(exit)

	file = menu.addMenu("&File")
	file.addAction(open)
	file.addAction(save)
	file.addAction(save_as)
	file.addAction(close)

	# System
	add = QAction("&Add a new element", win)
	add.setShortcut("Ctrl+N")
	add.triggered.connect(lambda: half_assed_element_creation())
	clear = QAction("&Remove all elements", win)
	clear.triggered.connect(half_assed_clear)

	system = menu.addMenu("&System")
	system.addAction(add)
	system.addAction(clear)

	win.statusBar()

# FIXME: refactor
def optimize(which):
	if len(system.elements) == 0:
		return

	op_idx = GUI.opt_operand.currentData()
	op = system.elements[op_idx]

	def I(theta):
		op.angle = theta
		if op.ref is not False:
			op.angle += system.elements[op.ref].angle

		state = None
		for i, pol in enumerate(system.elements):
			if i >= len(system.ignore) or not system.ignore[i]:
				state = pol.mul(state)
		if state is None:
			return 1
		else:
			return float(np.abs(state[0] * state[0].conjugate() \
			                    + state[1] * state[1].conjugate()))

	if which == "min":
		f = I
	else:
		f = lambda theta: -I(theta)

	opt = scipy.optimize.minimize_scalar(f, bounds=[0, np.pi], method="bounded")
	GUI.table_rows[op_idx].optbox.angle.edit.setText("%g" % round(opt.x * 180 / np.pi, 3))



def change_input_intensity():
	try:
		system.input_intensity = float(GUI.input_intensity.text())
		update()
	except ValueError:
		pass

def setup(win, system_):
	global system
	system = system_

	# Needless to say, I'm getting real tired of this whole Layout/Widget
	# clusterfuck in Qt
	root = QVBoxLayout()
	root_fucking_random_container = QWidget()
	root_fucking_random_container.setLayout(root)
	win.setCentralWidget(root_fucking_random_container)

	# Top bar (input intensity)
	box = QHBoxLayout()
	root.addLayout(box)

	box.addWidget(QLabel("Input intensity"))
	GUI.input_intensity = QLineEdit("1")
	GUI.input_intensity.textChanged.connect(change_input_intensity)
	box.addWidget(GUI.input_intensity)

	hbox = QHBoxLayout()
	root.addLayout(hbox)

	rhs = QVBoxLayout()

	Widocques.image = QImage("jones.jpg")
	GUI.widok = Widocques()
	rhs.addWidget(GUI.widok)

	optbox = QHBoxLayout()
	rhs.addLayout(optbox)

	GUI.opt_operand = QComboBox()
	optbox.addWidget(GUI.opt_operand)
	button = QPushButton("Find a minimum")
	button.pressed.connect(lambda: optimize("min"))
	optbox.addWidget(button)
	button = QPushButton("Find a maximum")
	button.pressed.connect(lambda: optimize("max"))
	optbox.addWidget(button)

	GUI.scroll = QScrollArea()
	populate_table()
	GUI.table_frame = QFrame()
	GUI.table_frame.setLayout(GUI.table)
	GUI.scroll.setWidget(GUI.table_frame)

	hbox.addWidget(GUI.scroll)
	hbox.addLayout(rhs)

	setup_menubar(win)

	update()
	win.show()