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

from ui import *

def R(theta):
	return np.array([[np.cos(theta), np.sin(theta)],
			 [-np.sin(theta), np.cos(theta)]])


class Ellipse:
	def __init__(self, state):
		# FIXME: a less brute-force way of doing this
		if state is None:
			self.alpha = np.nan
			self.theta = np.nan
			self.e = np.nan
			self.a = np.nan
			self.b = np.nan
			return

		def x(theta):
			return np.real(np.exp(1j * theta) * state)

		def r(theta):
			return np.linalg.norm(x(theta))

		def angle(x):
			a = np.arctan2(x[1], x[0])
			if a < 0:
				a += 2 * np.pi
			if a > np.pi:
				a -= np.pi
			return a

		opt = scipy.optimize.minimize_scalar(r, bounds=[0, np.pi], \
		                                     method="bounded")
		self.b = r(opt.x)
		opt = scipy.optimize.minimize_scalar(lambda x: -r(x), \
		      bounds=[0, np.pi], method="bounded")
		self.a = r(opt.x)

		self.alpha = angle(x(opt.x))

		self.e = self.b / self.a
		self.theta = np.arctan(self.e)
		if self.alpha > angle(x(opt.x + 0.001)):
			self.theta *= -1

class Polarizer:
	def __init__(self, type, delta=0):
		self.name = "New element" # FIXME
		self.type = type
		self.angle = 0
		self.delta = delta
		self.ref = False
		self.t1 = 1
		self.t2 = 0
		self.enable = True
		self.set_type(type)

	def set_type(self, type):
		if type == "linear":
			self.M = np.array([[1, 0], [0, 0]])
		elif type == "quarterwave":
			self.M = np.exp(-1j / 4 * np.pi) * \
			         np.array([[1, 0], [0, 1j]])
		else:
			raise ValueError("bad Polarizer type: %s" % type)
		self.type = type

	def mul(self, state):
		# unpolarized light
		if state is None:
			if self.type == "linear":
				return np.dot(R(-self.angle - self.delta), \
				              np.array([[1], [0]])) * np.sqrt(self.t1)
			else:
				return None

		if type == "linear":
			A = np.sqrt(np.array([[self.t1, 0], [0, self.t2]]))
		else:
			A = np.sqrt(np.array([[self.t1, 0], [0, self.t1]]))

		M = np.matmul(R(-self.angle - self.delta), \
		    np.matmul(np.matmul(self.M, A), R(self.angle + self.delta)))
		return np.dot(M, state)

class System:
	def __init__(self):
		self.elements = list()
		self.input_intensity = 1

	def recalculate(system):
		system.states = [None] * len(system.elements)
		system.ellipses = list()

		state = None
		for i, pol in enumerate(system.elements):
			if pol.enable:
				new_state = pol.mul(state)
				if state is None and new_state is not None:
					state = new_state * np.sqrt(system.input_intensity)
				else:
					state = new_state
			system.states[i] = state
			system.ellipses.append(Ellipse(state))