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 jones_to_stokes(E):
	I = np.abs(E[0] ** 2) + np.abs(E[1] ** 2)
	Q = np.abs(E[0] ** 2) - np.abs(E[1] ** 2)
	U = 2 * np.real(E[0] * np.conjugate(E[1]))
	V = 2 * np.imag(E[0] * np.conjugate(E[1]))
	return np.array([I, Q, U, V])


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


class Ellipse:
	def __init__(self, state):
		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)

		V = jones_to_stokes(state)

		self.alpha = np.arctan2(V[2], V[1]) / 2
		if self.alpha < 0:
			self.alpha += np.pi

		R = np.sqrt(V[1] ** 2 + V[2] ** 2 + V[3] ** 2) 
		self.theta = np.arcsin(V[3] / R) / 2

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

	def matrix(self):
		A = np.sqrt(np.array([[self.t1, 0], [0, self.t2]])) # FIXME: half-assed again
		M =  np.array([[1, 0], [0, np.exp(1j * self.phase_retardation),]])
		return np.matmul(R(-self.angle - self.delta),
			np.matmul(np.matmul(M, A), R(self.angle + self.delta)))

	def mul(self, state):
		# unpolarized light
		if state is None:
			if self.t2 == 0: # FIXME: this is half-assed
				return np.dot(R(-self.angle - self.delta), \
				              np.array([[1], [0]])) * np.sqrt(self.t1)
			else:
				return None
		return np.dot(self.matrix(), state)

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

	def recalculate(system):
		system.states = [None] * len(system.elements)
		system.Vs = [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))