"""
pair_conformal.py
====================================
Module for calculating :math:`C_2` using linear algebra,
based on components generated by the method of recursive images
using conformal mapping to obtain :math:`V_{2\infty}`
"""
import scipy
import scipy.special
import numpy as np
import numpy.linalg
from infinite_conformal import fF
eps0= 8.854187817*10**-12
############### old code
'''
def fF(phi,k,accuracy_limit):
	"""
	Calculates the inverses of the incomplete elliptic integral of the first kind
	for a complex argument phi. This function uses Jacobis form (phi) rather than the trigonometric form.
	uses scipy.integrate.quad for intergration
	scipy.special.ellipkinc(phi, k*k) could have been used if phi was a real number

	Parameters
	----------
	phi : complex number
	k : real number
	accuracy_limit : limit for accuracy in quad

	Returns
	-------
	complex number
		value of incomplete integral
		"""

	outreal=scipy.integrate.quad(lambda x: fFargRe(k,phi,x), 0, 1,eps_r_of_layersel=accuracy_limit)
	outimag=scipy.integrate.quad(lambda x: fFargIm(k,phi,x), 0, 1,eps_r_of_layersel=accuracy_limit)
	return (outreal[0]+1j*outimag[0])*phi
def fFargRe(k,phi, x):
	"""Real part of the argument for the integral in fF()
		"""
	theta=phi*x
	return (1/np.sqrt(1-k*k*np.sin(theta)**2)).real
def fFargIm(k,phi, x):
	"""Imaginary part of the argument for the integral in fF()
		"""
	theta=phi*x
	return (1/np.sqrt(1-k*k*np.sin(theta)**2)).imag'''

def fdtDdz(a):
	"""Function for calculating :math:`\\frac{dt}{dz}` in order to calculate the electric field in x and y direction.
		Note that this function is disticnt from the one in infinite_conformal.py
		"""
	return 2/a
def fdwDdt(k,t):
	"""Function for calculating :math:`\\frac{dw}{dz}` in order to calculate the electric field in x and y direction.
		"""
	out = 1/np.sqrt(	(1-t**2)*(1-t**2*k**2)	)
	# because of the sqrt, there are two possible solutions, need to make sure we select the rigth one
	if out.imag>0:
		return out
	return -out
def fdwDdz(k,t,a):
	"""Function for calculating :math:`\\frac{dt}{dz}` in order to calculate the electric field in x and y direction.
		"""
	return fdwDdt(k,t)*fdtDdz(a)

class single_recursive_images:
	"""
		A class that houses the potential, capacitance and electric fields from the method of recursive images.
		The values obtained are unphysical, as the potential at the electrodes is not constant.

		Parameters
		----------
		eta : float
			cover fraction of the electrodes
		interface_of_electrodes : int
			interface for the electrodes
		thickness_of_layers : list of floats
			thicknesses of the layers, this list will be 2 shorter than eps_x_of_layers and eps_y_of_layers, as the outermost layers have no defined thickness, but are infinite
		eps_x_of_layers : list of floats
			in-plane dielectric constant of the layers
		eps_y_of_layers : list of floats
			out-of-plane dielectric constant of the layers
		max_reflections : int, optional
			maximum number of reflections to considder, defaults to 8
		accuracy_limit : float, optional
			reflections with  less than accuracy_limit are ignored, defaults to 10**-15

		"""
	_V_dicts=dict()#_V_dicts[k][y][x]=V_{I,inf}(x,y)
	def __init__(self,a,b,interface_of_electrodes,thickness_of_layers,eps_x_of_layers,eps_y_of_layers,max_reflections=8,accuracy_limit=10**-7):
		self.b=b
		self.a=a
		self.eta=b/(a+b)
		self.interface_of_electrodes=interface_of_electrodes
		self.thickness_of_layers=thickness_of_layers
		# calculate the y-coordinate of each interface
		self.y_of_interfaces=[0]
		for T in thickness_of_layers:
			self.y_of_interfaces.append(self.y_of_interfaces[-1]+T)
		self.eps_x_of_layers=eps_x_of_layers
		self.eps_y_of_layers=eps_y_of_layers
		self.number_of_layers=len(self.eps_y_of_layers)
		# calculate the reflection coefficients for each interface
		self.r_pos_dir=[] # reflection coefficients for potential going in positive direction
		self.r_neg_dir=[] # reflection coefficients for potential going in negative direction
		self.t_pos_dir=[] # transmission coefficients for potential going in positive direction
		self.t_neg_dir=[] # transmission coefficients for potential going in negative direction
		for i in range(self.number_of_layers-1):
			eps1=np.sqrt(self.eps_y_of_layers[i]*self.eps_x_of_layers[i])
			eps2=np.sqrt(self.eps_y_of_layers[i+1]*self.eps_x_of_layers[i+1])
			self.r_pos_dir.append((eps1-eps2)/(eps1+eps2))
			self.t_pos_dir.append(self.r_pos_dir[-1]+1)
			self.r_neg_dir.append(-self.r_pos_dir[-1])
			self.t_neg_dir.append(self.r_neg_dir[-1]+1)
		# calculate eps_r of all layers
		self.eps_r_of_layers=[]
		for i in range(self.number_of_layers):
			if self.eps_y_of_layers[i]>0:
				self.eps_r_of_layers.append((self.eps_x_of_layers[i]/self.eps_y_of_layers[i])**0.5)
			else:
				self.eps_r_of_layers.append(1)
		# calculate k for the electrodes
		# k is the argument for the eliptical integrals
		self.max_reflections=max_reflections
		self.accuracy_limit=accuracy_limit
		self.k=(1-self.eta)/(1+self.eta) #external

		self.Kk= scipy.special.ellipk(float(self.k**2))
		self.Kpk= scipy.special.ellipk(1-float(self.k**2))
		if not self.k in self._V_dicts:
			self._V_dicts[self.k]=dict()
		self.Vdict=self._V_dicts[self.k] #Vdict[y][x]=V_{I,inf}(x,y)
		self.tree=[]
	def get_tree(self):
		if len(self.tree)==0:
			self.make_tree()
		return self.tree
	def make_tree(self):
		"""
		Function for building the tree\n
		The  tree consists of a list of lists where
		the main list iterates over the layers and
		the sublists contain a series of cases defined as:\n
		[dist prior, direction, amplitude]\n
		dist prior: the distance traveled to get to the current layer\n
		direction: the direction of the potential throught the layer (reflection switches the direction)\n
		amplitude: multiplication of all the reflection and transmission coefficients so far\n
		The cases are calculated iteratively, staring with the initial projected fields
		(with positive direction in the layer above the electrodes, and negative direction below them)\n
			"""

		''' tree will contain the full tree, and the reflections are added iteratively
			temp_tree_1 and temp_tree_2 are used to keep track of what fields will be
			adressed in the next iteration'''
		tree=[]
		temp_tree_2=[]
		for layer in range(self.number_of_layers):
			tree.append([])
			temp_tree_2.append([])
		#[dist prior,direction,amplitude,displacementx]
		''' add initial fields'''
		temp_tree_2[self.interface_of_electrodes].append([0,-1,1])
		tree[self.interface_of_electrodes].append([0,-1,1])
		temp_tree_2[self.interface_of_electrodes+1].append([0,1,1])
		tree[self.interface_of_electrodes+1].append([0,1,1])
		for step in range(self.max_reflections):
			temp_tree_1=temp_tree_2
			temp_tree_2=[]
			for layer in range(self.number_of_layers):
				temp_tree_2.append([])
			''' at this point
			temp_tree_1 contains the fields that should generate additional reflections
			temp_tree_2 is empty, and cases for the next iteration are added here'''
			for layer in range(self.number_of_layers-2): #layer-2 because the outermost layers have infinite thickness and cannot generate reflections
				for case in temp_tree_1[layer+1]:
					dist_prior=case[0]
					dist_add=self.thickness_of_layers[layer]*self.eps_r_of_layers[layer+1]
					direction=case[1]
					amplitude=case[2]
					if abs(amplitude*np.exp(-np.pi*(dist_prior+dist_add)))>self.accuracy_limit:
						if direction==1:
							temp_tree_2[layer+1].append([dist_prior+dist_add,-direction,amplitude*self.r_pos_dir[layer+1]])
							temp_tree_2[layer+2].append([dist_prior+dist_add,direction,amplitude*self.t_pos_dir[layer+1]])
							tree[layer+1].append(temp_tree_2[layer+1][-1]) # add to actual tree
							tree[layer+2].append(temp_tree_2[layer+2][-1]) # add to actual tree
						if direction==-1:
							temp_tree_2[layer+1].append([dist_prior+dist_add,-direction,amplitude*self.r_neg_dir[layer]])
							temp_tree_2[layer].append([dist_prior+dist_add,direction,amplitude*self.t_neg_dir[layer]])
							tree[layer+1].append(temp_tree_2[layer+1][-1]) # add to actual tree
							tree[layer].append(temp_tree_2[layer][-1]) # add to actual tree
		self.tree=tree
	def get_C(self):
		"""
		Returns
		-------
		float
			capacitace
			"""
		eps_m_below=self.eps_y_of_layers[self.interface_of_electrodes]*self.eps_r_of_layers[self.interface_of_electrodes]
		eps_m_above=self.eps_y_of_layers[self.interface_of_electrodes+1]*self.eps_r_of_layers[self.interface_of_electrodes+1]
		return (eps_m_below+eps_m_above)*self.Kpk/self.Kk*eps0/2
	def get_C_int_Ex(self):
		"""
		Function for calculating the capacitance by integrating :math:`\\varepsilon_xE_x` at :math:`x=0`.\n
		Used for testing, as it should give same output as get_C().\n
		For all practical applications get_C() should be used instead\n

		Returns
		-------
		float
			capacitace
			"""
		if self.eps_x_of_layers[0]>0:
			G,error=scipy.integrate.quad(lambda y: self.get_Ex(0,y), -80, 0)
			C = G*self.eps_x_of_layers[0]*eps0
		else:
			C = 0
		for i in range(self.number_of_layers-2):
			if self.eps_x_of_layers[i+1]>0:
				G,error=scipy.integrate.quad(lambda y: self.get_Ex(0,y), self.y_of_interfaces[i], self.y_of_interfaces[i+1])
				C+= G*self.eps_x_of_layers[i+1]*eps0
		if self.eps_x_of_layers[-1]>0:
			G,error=scipy.integrate.quad(lambda y: self.get_Ex(0,y), self.y_of_interfaces[-1], 80)
			C+= G*self.eps_x_of_layers[-1]*eps0
		return -C

	def get_V_Ex_Ey(self,x,y,get_V=1,get_Ex=1,get_Ey=1): # accepts 'x' as a list, but 'y' must be single value
		"""
		Function for calculating the the potential and electric fields at coordinates (x,y)

		Parameters
		----------
		x : float or list of floats
			x-coordiate(s)
		y : float
			y-coordiate
		get_V : bool, optional
			V is only calculated if this flag is set to True, default: True
		get_Ex : bool, optional
			Ex is only calculated if this flag is set to True, default: True
		get_Ey : bool, optional
			Ey is only calculated if this flag is set to True, default: True

		Returns
		-------
		list of float for V, Ex, Ey
			"""
		tree=self.get_tree()
		#print('donetree',time.time()-t0)
		x=np.array(x)
		x=np.atleast_1d(x)
		V=np.zeros(x.size)
		Ex=np.zeros(x.size)
		Ey=np.zeros(x.size)
		layer=0

		while layer<len(self.y_of_interfaces) and self.y_of_interfaces[layer]<=y:
			layer+=1
		if self.eps_y_of_layers[layer]==0:
			return 0,0,0
		#nn=[]
		for case in tree[layer]:
			#dist_prior=#[dist prior,direction]
			dist_prior=case[0]
			direction=case[1]
			amplitude=case[2]
			if amplitude==0: continue
			if direction==1:
				#interfaces[layer] is the interface below the layer
				Y=dist_prior+self.eps_r_of_layers[layer]*(y-self.y_of_interfaces[layer-1])
			else: #direction==-1
				Y=dist_prior-self.eps_r_of_layers[layer]*(y-self.y_of_interfaces[layer])
			for i,XX in enumerate(x):
				if get_Ex or get_Ey:
					z=abs(XX)+1j*Y # posistion on the z-plane
					t=z*2/self.a # posistion on the t-plane
					dwDdz=fdwDdz(self.k,t,self.a) # :math:`\\frac{dw}{dz}`
					if get_Ex:
						Ex[i]+=1/self.Kk*dwDdz.real*amplitude
					if get_Ey:
						Ey[i]+=1/self.Kk*dwDdz.imag*direction*amplitude*self.eps_r_of_layers[layer]
				if get_V:
					Yp=abs(Y)
					if not Yp in self.Vdict:
						self.Vdict[Yp]={}
					if not XX in self.Vdict[Yp]:
						z=XX+1j*Yp# posistion on the z-plane
						t=z*2/self.a # posistion on the t-plane
						F=fF(np.arcsin(t),self.k,self.accuracy_limit/amplitude) # heavy lifting for transformation to w-plane
						self.Vdict[Yp][XX]=((1/self.Kk)*F.real)
					V[i]+=self.Vdict[Yp][XX]*amplitude
			####################################
		return V/2, Ex/2, Ey/2
	def get_V(self,x,y):
		"""
		calls get_V_Ex_Ey() to calculate V
		"""
		V,Ex,Ey=self.get_V_Ex_Ey(x,y,1,0,0)
		return V
	def get_Ex(self,x,y):
		"""
		calls get_V_Ex_Ey() to calculate Ex
		"""
		V,Ex,Ey=self.get_V_Ex_Ey(x,y,0,1,0)
		return Ex
	def get_Ey(self,x,y):
		"""
		calls get_V_Ex_Ey() to calculate Ey
		"""
		V,Ex,Ey=self.get_V_Ex_Ey(x,y,0,0,1)
		return Ey


class multiple_recursive_images:
	"""
	A class generates and houses instances of the class single_recursive_images and uses a linear combination of these
	to approximate physcal values for the capacitance, potential, and electric fields of an electrode pair.

	Parameters
	----------
	etas : list of float
		cover fraction of the electrode(s) must be a list of the same length as there are interfaces, i.e. 1 shorter than eps_x_of_layers
		each value must be 0=<eta=<1. 0 indicates that no electrode exists at the corresponding interface
	thickness_of_layers : list of floats
		thicknesses of the layers, this list will be 2 shorter than eps_x_of_layers and eps_y_of_layers, as the outermost layers have no defined thickness, but are infinite
	eps_x_of_layers : list of floats
		in-plane dielectric constant of the layers
	eps_y_of_layers : list of floats
		out-of-plane dielectric constant of the layers
	LAcomp: int or list of ints, optional
		the number of components used for each set of electrodes. If list, the list must be the same length as etas, and have a number larger than one for each nonzero value in etas.
		If not list the same value vill be used for all electrodes. Defaults to 8
	max_reflections : int, optional
		maximum number of reflections to considder, defaults to 8
	voltages : list of floats or None, optional
		If a list of floats is provided, the voltage at each set of electrodes will be set to +/- the corresponding value in the list. the list must have the same length as etas.
		If None, the potential at each set of electrodes with a cover fraction less than 1 is set to +/-0.5, for a total difference of 1V. The potential at electrodes with a cover fraction of 1 is set to 0, as this is assumed to be grounded.
		Defaults to None. NB! the calcualted capacitance assumed a potential of +/-0.5 at the electrodes, and if a different value is used, calculated capacitance will reflect the charge at the electrodes rather than the capacitance!
	periods : list of floats or none, optional
		list of the same length as there are interfaces that declares the period for electrodes on that interface.
		If None, all electrodes have a period of one. Electrodes extend from period*(1-eta)/2 to period*(1+eta/2) in this class.
		Typical use is when approximating an infinite grounded electrode. The infinite electrode is then set to have a period > 1.
		Defaults to None
	accuracy_limit : float, optional
		reflections with amplitude less than accuracy_limit are ignored, defaults to 10**-15

	"""
	def __init__(self,etas,thickness_of_layers,eps_x_of_layers,eps_y_of_layers,LAcomp=8,max_reflections=8, voltages = None,periods=None,accuracy_limit=10**-10):
		# eta = vector of floats, t = vector of floats, eps_x_of_layers = vector of floats, eps_y_of_layers = vector of floats, LAcomp = int, max_reflections = int,
		# eps_x_of_layers and eps_y_of_layers refer to materials, the length must therefore be at least 2
		# eta refers to interfaces, the length must therefore be at least 1
		# t refers to the thicknes of layers of finite thickness. This vector may have 0 elements.
		# LAcomp must be an int, or a vector of same length as eta
		# voltages must be 'None' or a vector of same length as eta, and is uset to set voltage ratios if multiple sets of electrodes are used
		self.accuracy_limit=accuracy_limit
		self.etas=etas
		if periods==None:
			self.periods=np.ones(len(self.etas))
		else:
			self.periods=periods

		self.thickness_of_layers=thickness_of_layers
		self.eps_x_of_layers=eps_x_of_layers
		self.eps_y_of_layers=eps_y_of_layers
		self.LAcomp=LAcomp
		self.max_reflections=max_reflections
		self.voltages=voltages
		if self.voltages == None: # create voltages vector, all electrode sets should have V=0.5, unless they are continious, then they should have 0. If no electrodes are present at the interface, then the voltage is also set to 0
			self.voltages = []
			for eta in self.etas:
				if eta == 0:
					self.voltages.append(0)
				elif eta == 1:
					self.voltages.append(0)
				else:
					self.voltages.append(0.5)
		# declare the step size for for eta in each set of electrodes
		self.electrodesteps=[]
		if isinstance(LAcomp, (int,)):
			for eta in self.etas:
				self.electrodesteps.append(eta/self.LAcomp)
		else:
			for eta, LA in zip(self.etas, self.LAcomp):
				if LA>0:
					self.electrodesteps.append(eta/LA)
				else:
					self.electrodesteps.append(0)
		# calculate the relevant locations, as and bs for the linear algebra aproach
		# for each case, the a and b is changed so that the outermost part of the electrode remains in place
		self.xpoints=[] #<- list of list of points of interest
		self.xpointVs=[]
		self.caseAs=[] #<- list of lists of a of individual cases
		self.caseBs=[] #<- list of lists of b of individual cases
		for eta, step, V,period in zip(self.etas, self.electrodesteps,self.voltages,self.periods):
			self.xpoints.append([])
			self.caseAs.append([])
			self.caseBs.append([])
			l=step*0.5
			while l<eta and l<1-step: # the second condition prevents the last from beeing added if the electrodes are continious.
				self.xpoints[-1].append((-l+0.5+eta/2)*period)
				self.caseBs[-1].append((l+step*0.5)*period)
				self.caseAs[-1].append((1+eta-2*l-step)*period)
				self.xpointVs.append(V)
				l+=step
		self.C=-1
	def get_C(self):
		"""
		Function for calculate the total capacitance using linear algebra.
		To do this it needs to first build the different single_recursive_images objects
		Then use them to calculate the potential at self.xpoints, and use linear algebra to find the weights.
		Finally, the capacitance is calculated by sum(weights*single_recursive_images.get_C)

		Returns
		-------
		float
			capacitace
			"""
		if self.C==-1:
			self.Cs=[]
			self.single=[]
			self.Vs=[]
			# electrodes may exist on multiple interfaces, so iterate
			for interface, _ in enumerate(self.caseAs):
				for  a,b in zip(self.caseAs[interface],self.caseBs[interface]):
					self.single.append(single_recursive_images(a,b,interface,self.thickness_of_layers,self.eps_x_of_layers,self.eps_y_of_layers,self.max_reflections,self.accuracy_limit))
					self.Vs.append(np.array([]))
					# calculate the potential for the last case at all designated locations
					# iterate over designated locations: first iterate over layers
					y=0 # need y to keep track of location as we iterate over layers
					for interfaceOfXpoint, _ in enumerate(self.caseAs):
						if len(self.xpoints[interfaceOfXpoint])>0: # only calculate if there are any points
							# utilize the fact that get_V suppoprts a list of x coordinates as unput to
							#    calculate all points on this interface by one function call
							V=self.single[-1].get_V(self.xpoints[interfaceOfXpoint],y)
							# concatenate the potentials at this interface with previous potensials at other interfaces
							self.Vs[-1]=np.concatenate((self.Vs[-1], V))
						if interfaceOfXpoint<len(self.thickness_of_layers):
							# need y to keep track of location as we iterate over layers
							y+=self.thickness_of_layers[interfaceOfXpoint]
					# get the corresponding capacitance associtaed with these potentials
					self.Cs.append(self.single[-1].get_C())
			# solve the linear algebra expression for X: 	AX=B
			A=np.array(self.Vs).transpose()
			B=np.array(self.xpointVs)
			self.weights=numpy.linalg.solve(A,B)
			# calculate the capacitance
			self.C=np.dot(np.array(self.Cs),np.array(self.weights))
		return self.C
	def get_C_int_Ex(self):
		"""
		Function for calculating the capacitance by integrating :math:`\\varepsilon_xE_x` at :math:`x=0`.\n
		Used for testing, as it should give same output as get_C().\n
		For all practical applications get_C() should be used instead\n

		Returns
		-------
		float
			capacitace
			"""
		if self.C==-1:
			self.get_C()
		Cs=[]
		for case in self.single:
			Cs.append(case.get_C_int_Ex())
		C=np.dot(np.array(Cs),np.array(self.weights))
		return C
	def get_V_Ex_Ey(self,x,y,get_V=1,get_Ex=1,get_Ey=1):
		"""
		Function for calculating the the potential and electric fields at coordinates (x,y)

		Parameters
		----------
		x : float or list of floats
			x-coordiate(s)
		y : float
			y-coordiate
		get_V : bool, optional
			V is only calculated if this flag is set to True, default: True
		get_Ex : bool, optional
			Ex is only calculated if this flag is set to True, default: True
		get_Ey : bool, optional
			Ey is only calculated if this flag is set to True, default: True

		Returns
		-------
		list of float for V, Ex, Ey
			"""
		if self.C==-1:
			self.get_C()
		V=0
		Ex=0
		Ey=0
		for case, component in zip(self.single, self.weights):
			Vi,Exi,Eyi=case.get_V_Ex_Ey(x,y,get_V,get_Ex,get_Ey)
			V+=Vi*component
			#print(Vi, component)
			Ex+=Exi*component
			Ey+=Eyi*component
		return(V,Ex,Ey)
	def get_V(self,x,y):
		"""
		calls get_V_Ex_Ey() to calculate V
		"""
		V,Ex,Ey=self.get_V_Ex_Ey(x,y,1,0,0)
		return V
	def get_Ex(self,x,y):
		"""
		calls get_V_Ex_Ey() to calculate Ex
		"""
		V,Ex,Ey=self.get_V_Ex_Ey(x,y,0,1,0)
		return Ex
	def get_Ey(self,x,y):
		"""
		calls get_V_Ex_Ey() to calculate Ey
		"""
		V,Ex,Ey=self.get_V_Ex_Ey(x,y,0,0,1)
		return Ey
