from __future__ import print_function, division, absolute_import

from ..basis import spin_basis_1d as _default_basis
from ..basis import isbasis as _isbasis

#from ._oputils import _get_matvec_function, matvec as _matvec

from ..tools.matvec import _matvec
from ..tools.matvec import _get_matvec_function

from ._make_hamiltonian import make_static
from ._make_hamiltonian import _check_almost_zero

from . import hamiltonian_core

# need linear algebra packages
import scipy.sparse.linalg as _sla
import scipy.linalg as _la
import scipy.sparse as _sp
import numpy as _np

import functools
from six import iteritems,itervalues,viewkeys

__all__=["quantum_operator","isquantum_operator"]

# function used to create Linearquantum_operator with fixed set of parameters.
def _quantum_operator_dot(op,pars,v):
	return op.dot(v,pars=pars,check=False)

class quantum_operator(object):
	"""Constructs parameter-dependent (hermitian and nonhermitian) operators.

		The `quantum_operator` class maps quantum operators to keys of a dictionary. When calling various methods
		of `quantum_operator`, it allows one to 'dynamically' specify the pre-factors of these operators.

		Examples
		---------

		It is often required to be able to handle a parameter-dependent Hamiltonian :math:`H(\\lambda)=H_1 + \\lambda H_2`, e.g.

		.. math::
			H_1=\sum_j J_{zz}S^z_jS^z_{j+1} + h_xS^x_j, \\qquad H_2=\\sum_j S^z_j

		The following code snippet shows how to use the `quantum_operator` class to vary the parameter :math:`\\lambda`
		without having to re-build the Hamiltonian every time.

		.. literalinclude:: ../../doc_examples/quantum_operator-example.py
			:linenos:
			:language: python
			:lines: 7-

	"""
	def __init__(self,input_dict,N=None,basis=None,shape=None,copy=True,check_symm=True,check_herm=True,check_pcon=True,matrix_formats={},dtype=_np.complex128,**basis_args):
		"""Intializes the `quantum_operator` object (parameter dependent quantum quantum_operators).

		Parameters
		-----------
		input_dict : dict
			The `values` of this dictionary contain quantum_operator lists, in the same format as the `static_list`
			argument of the `hamiltonian` class.

			The `keys` of this dictionary correspond to the parameter values, e.g. :math:`J_{zz},h_x`, and are
			used to specify the coupling strength during calls of the `quantum_operator` class methods.

			>>> # use "Jzz" and "hx" keys to specify the zz and x coupling strengths, respectively
			>>> input_dict = { "Jzz": [["zz",Jzz_bonds]], "hx" : [["x" ,hx_site ]] }

		N : int, optional
			Number of lattice sites for the `hamiltonian` object.
		dtype : 'type'
			Data type (e.g. numpy.float64) to construct the quantum_operator with.
		shape : tuple, optional
			Shape to create the `hamiltonian` object with. Default is `shape = None`.
		copy: bool, optional
			If set to `True`, this option creates a copy of the input array.
		check_symm : bool, optional
			Enable/Disable symmetry check on `static_list` and `dynamic_list`.
		check_herm : bool, optional
			Enable/Disable hermiticity check on `static_list` and `dynamic_list`.
		check_pcon : bool, optional
			Enable/Disable particle conservation check on `static_list` and `dynamic_list`.
		matrix_formats: dict, optional
			Dictionary of key,value pairs which, given a key associated with an operator in `input_dict`, the value of this key
			specifies the sparse matrix format {"csr","csc","dia","dense"}.
		kw_args : dict
			Optional additional arguments to pass to the `basis` class, if not already using a `basis` object
			to create the quantum_operator.

		"""
		self._is_dense = False
		self._ndim = 2
		self._basis = basis

		if not (dtype in hamiltonian_core.supported_dtypes):
			raise TypeError('hamiltonian does not support type: '+str(dtype))
		else:
			self._dtype=dtype

		opstr_dict = {}
		other_dict = {}
		self._quantum_operator = {}
		if isinstance(input_dict,dict):
			for key,op in iteritems(input_dict):
				if type(key) is not str:
					raise ValueError("keys to input_dict must be strings.")

				if type(op) not in [list,tuple]:
					raise ValueError("input_dict must contain values which are lists/tuples.")

				opstr_list = []
				other_list = []
				for ele in op:
					if hamiltonian_core._check_static(ele):
						opstr_list.append(ele)
					else:
						other_list.append(ele)

				if opstr_list:
					opstr_dict[key] = opstr_list
				if other_list:
					other_dict[key] = other_list
		else:
			raise ValueError("input_dict must be dictionary or another quantum_operator quantum_operators")

		if opstr_dict:
			# check if user input basis

			if basis is not None:
				if len(basis_args) > 0:
					wrong_keys = set(basis_args.keys())
					temp = ", ".join(["{}" for key in wrong_keys])
					raise ValueError(("unexpected optional argument(s): "+temp).format(*wrong_keys))

			# if not
			if basis is None:
				if N is None: # if L is missing
					raise Exception('if opstrs in use, argument N needed for basis class')

				if type(N) is not int: # if L is not int
					raise TypeError('argument N must be integer')

				basis=_default_basis(N,**basis_args)

			elif not _isbasis(basis):
				raise TypeError('expecting instance of basis class for argument: basis')

			static_opstr_list = []
			for key,opstr_list in iteritems(opstr_dict):
				static_opstr_list.extend(opstr_list)

			if check_herm:
				basis.check_hermitian(static_opstr_list, [])

			if check_symm:
				basis.check_symm(static_opstr_list,[])

			if check_pcon:
				basis.check_pcon(static_opstr_list,[])

			self._shape=(basis.Ns,basis.Ns)

			for key,opstr_list in iteritems(opstr_dict):
				O = make_static(basis,opstr_list,dtype)
				self._quantum_operator[key] = O

		if other_dict:
			if not hasattr(self,"_shape"):
				found = False
				if shape is None: # if no shape argument found, search to see if the inputs have shapes.
					for key,O_list in iteritems(other_dict):
						for O in O_list:
							try: # take the first shape found
								shape = O.shape
								found = True
								break
							except AttributeError:
								continue
				else:
					found = True

				if not found:
					raise ValueError('no dictionary entries have shape attribute.')
				if shape[0] != shape[1]:
					raise ValueError('quantum_operator must be square matrix')

				self._shape=shape

			for key,O_list in iteritems(other_dict):
				for i,O in enumerate(O_list):
					if _sp.issparse(O):
						self._mat_checks(O)
						if i == 0:
							self._quantum_operator[key] = O
						else:
							try:
								self._quantum_operator[key] += O
							except NotImplementedError:
								self._quantum_operator[key] = self._quantum_operator[key] + O

					elif O.__class__ is _np.ndarray:
						self._mat_checks(O)
						self._is_dense=True
						if i == 0:
							self._quantum_operator[key] = O
						else:
							try:
								self._quantum_operator[key] += O
							except NotImplementedError:
								self._quantum_operator[key] = self._quantum_operator[key] + O

					elif O.__class__ is _np.matrix:
						self._mat_checks(O)
						self._is_dense=True
						if i == 0:
							self._quantum_operator[key] = O
						else:
							try:
								self._quantum_operator[key] += O
							except NotImplementedError:
								self._quantum_operator[key] = self._quantum_operator[key] + O

					else:
						O = _np.asanyarray(O)
						self._mat_checks(O)
						if i == 0:
							self._quantum_operator[key] = O
						else:
							try:
								self._quantum_operator[key] += O
							except NotImplementedError:
								self._quantum_operator[key] = self._quantum_operator[key] + O

		else:
			if not hasattr(self,"_shape"):
				if shape is None:
					# check if user input basis
					basis=basis_args.get('basis')

					# if not
					if basis is None:
						if N is None: # if N is missing
							raise Exception("argument N or shape needed to create empty quantum_operator")

						if type(N) is not int: # if L is not int
							raise TypeError('argument N must be integer')

						basis=_default_basis(N,**basis_args)

					elif not _isbasis(basis):
						raise TypeError('expecting instance of basis class for argument: basis')

					shape = (basis.Ns,basis.Ns)

				else:
					basis=basis_args.get('basis')
					if not basis is None:
						raise ValueError("empty hamiltonian only accepts basis or shape, not both")

				if len(shape) != 2:
					raise ValueError('expecting ndim = 2')
				if shape[0] != shape[1]:
					raise ValueError('hamiltonian must be square matrix')

				self._shape=shape

		if basis is not None:
			self._basis = basis

		self._Ns = self._shape[0]

		keys = list(self._quantum_operator.keys())
		for key in keys:
			if _check_almost_zero(self._quantum_operator[key]):
				self._quantum_operator.pop(key)

		self.update_matrix_formats(matrix_formats)

	@property
	def get_operators(self,key):
		return self._quantum_operator[key]

	@property
	def shape(self):
		return self._shape

	@property
	def basis(self):
		""":obj:`basis`: basis used to build the `hamiltonian` object. Defaults to `None` if quantum_operator has
		no basis (i.e. was created externally and passed as a precalculated array).

		"""
		if self._basis is not None:
			return self._basis
		else:
			raise AttributeError("object has no attribute 'basis'")

	@property
	def ndim(self):
		"""int: number of dimensions, always equal to 2. """
		return self._ndim

	@property
	def Ns(self):
		"""int: number of states in the (symmetry-reduced) Hilbert space spanned by `basis`."""
		return self._Ns

	@property
	def get_shape(self):
		"""tuple: shape of the `quantum_operator` object, always equal to `(Ns,Ns)`."""
		return self._shape

	@property
	def is_dense(self):
		"""bool: `True` if the quantum_operator contains a dense matrix as a componnent of either
		the static or dynamic lists.

		"""
		return self._is_dense

	@property
	def dtype(self):
		"""type: data type of `quantum_operator` object."""
		return _np.dtype(self._dtype).name

	@property
	def T(self):
		""":obj:`quantum_operator`: transposes the operator matrix: :math:`H_{ij}\\mapsto H_{ji}`."""
		return self.transpose()

	@property
	def H(self):
		""":obj:`quantum_operator`: transposes and conjugates the operator matrix: :math:`H_{ij}\\mapsto H_{ji}^*`."""
		return self.getH()

	### state manipulation/observable routines

	def matvec(self,x):
		"""Matrix-vector multiplication.

		Performs the operation y=A*x where A is an MxN linear operator and x is a column vector or 1-d array.

		Notes
		-----
		This matvec wraps the user-specified matvec routine or overridden _matvec method to ensure that y has the correct shape and type.

		Parameters
		----------
		x : {matrix, ndarray}
			An array with shape (N,) or (N,1).

		Returns
		-------
		y : {matrix, ndarray}
			A matrix or ndarray with shape (M,) or (M,1) depending on the type and shape of the x argument.

		"""

		return self.dot(x)

	def rmatvec(self,x):
		"""Adjoint matrix-vector multiplication.

		Performs the operation y = A^H * x where A is an MxN linear operator and x is a column vector or 1-d array.

		Notes
		-----
		This rmatvec wraps the user-specified rmatvec routine or overridden _rmatvec method to ensure that y has the correct shape and type.

		Parameters
		----------
		x : {matrix, ndarray}
			An array with shape (M,) or (M,1).

		Returns
		-------
		y : {matrix, ndarray}
			A matrix or ndarray with shape (N,) or (N,1) depending on the type and shape of the x argument.

		"""
		return self.H.dot(x)

	def matmat(self,X):
		"""Matrix-matrix multiplication.

		Performs the operation y=A*X where A is an MxN linear operator and X dense N*K matrix or ndarray.

		Notes
		-----
		This matmat wraps any user-specified matmat routine or overridden _matmat method to ensure that y has the correct type.

		Parameters
		----------
		X : {matrix, ndarray}
			An array with shape (N,K).

		Returns
		-------
		Y : {matrix, ndarray}
			A matrix or ndarray with shape (M,K) depending on the type of the X argument.

		"""
		return self.dot(X)

	def dot(self,V,pars={},check=True,out=None,overwrite_out=True,a=1.0):
		"""Matrix-vector multiplication of `quantum_operator` quantum_operator for parameters `pars`, with state `V`.

		.. math::
			aH(\\lambda)|V\\rangle

		Notes
		-----

		Parameters
		-----------
		V : numpy.ndarray
			Vector (quantums tate) to multiply the `quantum_operator` quantum_operator with.
		pars : dict, optional
			Dictionary with same `keys` as `input_dict` and coupling strengths as `values`. Any missing `keys`
			are assumed to be set to unity.
		check : bool, optional
			Whether or not to do checks for shape compatibility.
		out : array_like, optional
			specify the output array for the the result. This is not supported if `V` is a sparse matrix.
		overwrite_out : bool, optional
			flag used to toggle between two different ways to treat `out`. If set to `True` all values in `out` will be overwritten with the result of the dot product.
			If `False` the result of the dot product will be added to the values of `out`.
		a : scalar, optional
			scalar to multiply the final product with: :math:`B = aHV`.

		Returns
		--------
		numpy.ndarray
			Vector corresponding to the `quantum_operator` quantum_operator applied on the state `V`.

		Examples
		---------
		>>> B = H.dot(A,pars=pars,check=True)

		corresponds to :math:`B = HA`.

		"""

		pars = self._check_scalar_pars(pars)

		if check:
			try:
				shape = V.shape
			except AttributeError:
				V =_np.asanyarray(V)
				shape = V.shape

			if shape[0] != self._shape[1]:
				raise ValueError("matrix dimension mismatch with shapes: {0} and {1}.".format(V.shape,self._shape))

			if V.ndim not in [1,2]:
				raise ValueError("Expecting  0< V.ndim < 3.")

		result_dtype = _np.result_type(V.dtype,self._dtype)

		if not (result_dtype in hamiltonian_core.supported_dtypes):
			raise TypeError('hamiltonian does not support type: '+str(dtype))

		if self.Ns <= 0:
			return _np.asarray([],dtype=result_dtype)

		if _sp.issparse(V):
			if out is not None:
				raise TypeError("'out' option does not apply for sparse inputs.")

			sparse_constuctor = getattr(_sp,V.get_format()+"_matrix")
			out = sparse_constuctor(V.shape,dtype=result_dtype)
			for key,J in pars.items():
				out = out + J*self._quantum_operator[key].dot(V)
			out = a*out

		else:
			if out is not None:
				try:
					if out.dtype != result_dtype:
						raise TypeError("'out' must be array with correct dtype and dimensions for output array.")
					if out.shape != V.shape:
						raise ValueError("'out' must be array with correct dtype and dimensions for output array.")
				except AttributeError:
					raise TypeError("'out' must be C-contiguous array with correct dtype and dimensions for output array.")

				if overwrite_out:
					out[...] = 0
			else:
				out = _np.zeros_like(V,dtype=result_dtype)

			eps = _np.finfo(self.dtype).eps
			V = _np.asarray(V,dtype=result_dtype)
			for key,J in pars.items():
				if _np.abs(J)>eps:
					self._matvec_functions[key](self._quantum_operator[key],V,overwrite_out=False,a=a*J,out=out)

		return out

	def rdot(self,V,pars={},check=False,out=None,overwrite_out=True,a=1.0):
		"""Vector-matrix multiplication of `quantum_operator` quantum_operator for parameters `pars`, with state `V`.

		.. math::
			a\\langle V]H(\\lambda)

		Parameters
		-----------
		V : numpy.ndarray
			Vector (quantums tate) to multiply the `quantum_operator` quantum_operator with.
		pars : dict, optional
			Dictionary with same `keys` as `input_dict` and coupling strengths as `values`. Any missing `keys`
			are assumed to be set to unity.
		check : bool, optional
			Whether or not to do checks for shape compatibility.
		out : array_like, optional
			specify the output array for the the result. This is not supported if `V` is a sparse matrix.
		overwrite_out : bool, optional
			flag used to toggle between two different ways to treat `out`. If set to `True` all values in `out` will be overwritten with the result.
			If `False` the result of the dot product will be added to the values of `out`.
		a : scalar, optional
			scalar to multiply the final product with: :math:`B = aVH`.

		Returns
		--------
		numpy.ndarray
			Vector corresponding to the `quantum_operator` quantum_operator applied on the state `V`.

		Examples
		---------
		>>> B = H.dot(A,pars=pars,check=True)

		corresponds to :math:`B = AH`.

		"""
		return self.transpose().dot(V.transpose(),pars=pars,check=check,out=out.T,overwrite_out=overwrite_out,a=a).transpose()

	def quant_fluct(self,V,pars={},check=True,enforce_pure=False):
		"""Calculates the quantum fluctuations (variance) of `hamiltonian` operator at time `time`, in state `V`.

		.. math::
			\\langle V|H^2(t=\\texttt{time})|V\\rangle - \\langle V|H(t=\\texttt{time})|V\\rangle^2

		Parameters
		-----------
		V : numpy.ndarray
			Depending on the shape, can be a single state or a collection of pure or mixed states
			[see `enforce_pure`].
		pars : dict, optional
			Dictionary with same `keys` as `input_dict` and coupling strengths as `values`. Any missing `keys`
			are assumed to be set to unity.

		enforce_pure : bool, optional
			Flag to enforce pure expectation value of `V` is a square matrix with multiple pure states
			in the columns.
		check : bool, optional

		Returns
		--------
		float
			Quantum fluctuations of `hamiltonian` operator in state `V`.

		Examples
		---------
		>>> H_fluct = H.quant_fluct(V,time=0,diagonal=False,check=True)

		corresponds to :math:`\\Delta H = \\sqrt{ \\langle V|H^2(t=\\texttt{time})|V\\rangle - \\langle V|H(t=\\texttt{time})|V\\rangle^2 }`.

		"""

		from .exp_op_core import isexp_op

		if self.Ns <= 0:
			return _np.asarray([])

		if hamiltonian_core.ishamiltonian(V):
			raise TypeError("Can't take expectation value of hamiltonian")

		if isexp_op(V):
			raise TypeError("Can't take expectation value of exp_op")

		# fluctuations =  expctH2 - expctH^2
		kwargs = dict(enforce_pure=enforce_pure)
		V_dot = self.dot(V,pars=pars,check=check)
		expt_value_sq = self._expt_value_core(V,V_dot,**kwargs)**2

		if V.shape[0] != V.shape[1] or enforce_pure:
			sq_expt_value = self._expt_value_core(V_dot,V_dot,**kwargs)
		else:
			V_dot = self.dot(V_dot,time=time,check=check)
			sq_expt_value = self._expt_value_core(V,V_dot,**kwargs)

		return sq_expt_value - expt_value_sq

	def expt_value(self,V,pars={},check=True,enforce_pure=False):
		"""Calculates expectation value of `hamiltonian` operator at time `time`, in state `V`.

		.. math::
			\\langle V|H(t=\\texttt{time})|V\\rangle

		Parameters
		-----------
		V : numpy.ndarray
			Depending on the shape, can be a single state or a collection of pure or mixed states
			[see `enforce_pure` argument of `basis.ent_entropy`].
		pars : dict, optional
			Dictionary with same `keys` as `input_dict` and coupling strengths as `values`. Any missing `keys`
			are assumed to be set to unity.
		enforce_pure : bool, optional
			Flag to enforce pure expectation value of `V` is a square matrix with multiple pure states
			in the columns.
		check : bool, optional

		Returns
		--------
		float
			Expectation value of `hamiltonian` operator in state `V`.

		Examples
		---------
		>>> H_expt = H.expt_value(V,time=0,diagonal=False,check=True)

		corresponds to :math:`H_{expt} = \\langle V|H(t=0)|V\\rangle`.

		"""
		from .exp_op_core import isexp_op

		if self.Ns <= 0:
			return _np.asarray([])

		if hamiltonian_core.ishamiltonian(V):
			raise TypeError("Can't take expectation value of hamiltonian")

		if isexp_op(V):
			raise TypeError("Can't take expectation value of exp_op")

		V_dot = self.dot(V,check=check,pars=pars)
		return self._expt_value_core(V,V_dot,enforce_pure=enforce_pure)

	def _expt_value_core(self,V_left,V_right,enforce_pure=False):
		if _sp.issparse(V_right):
			if V_left.shape[0] != V_left.shape[1] or enforce_pure: # pure states
				return _np.asscalar((V_left.H.dot(V_right)).toarray())
			else: # density matrix
				return V_right.diagonal().sum()
		else:
			V_right = _np.asarray(V_right).squeeze()
			if V_right.ndim == 1: # pure state
				return _np.vdot(V_left,V_right)
			elif V_left.shape[0] != V_left.shape[1] or enforce_pure: # multiple pure states
				return _np.einsum("ij,ij->j",V_left.conj(),V_right)
			else: # density matrix
				return V_right.trace()

	def matrix_ele(self,Vl,Vr,pars={},diagonal=False,check=True):
		"""Calculates matrix element of `quantum_operator` quantum_operator for parameters `pars` in states `Vl` and `Vr`.

		.. math::
			\\langle V_l|H(\\lambda)|V_r\\rangle

		Notes
		-----
		Taking the conjugate or transpose of the state `Vl` is done automatically.

		Parameters
		-----------
		Vl : numpy.ndarray
			Vector(s)/state(s) to multiple with on left side.
		Vl : numpy.ndarray
			Vector(s)/state(s) to multiple with on right side.
		pars : dict, optional
			Dictionary with same `keys` as `input_dict` and coupling strengths as `values`. Any missing `keys`
			are assumed to be set to unity.
		diagonal : bool, optional
			When set to `True`, returs only diagonal part of expectation value. Default is `diagonal = False`.
		check : bool,

		Returns
		--------
		float
			Matrix element of `quantum_operator` quantum_operator between the states `Vl` and `Vr`.

		Examples
		---------
		>>> H_lr = H.expt_value(Vl,Vr,pars=pars,diagonal=False,check=True)

		corresponds to :math:`H_{lr} = \\langle V_l|H(\\lambda=0)|V_r\\rangle`.

		"""
		Vr = self.dot(Vr,pars=pars,check=check)

		if check:
			try:
				shape = Vl.shape
			except AttributeError:
				Vl =_np.asanyarray(Vl)
				shape = Vl.shape

			if shape[0] != self._shape[1]:
				raise ValueError("matrix dimension mismatch with shapes: {0} and {1}.".format(V.shape,self._shape))

			if Vl.ndim > 2:
				raise ValueError("Expecting  0< V.ndim < 3.")

		if _sp.issparse(Vl):
			if diagonal:
				return Vl.H.dot(Vr).diagonal()
			else:
				return Vl.H.dot(Vr)
		else:
			if diagonal:
				return _np.einsum("ij,ij->j",Vl.conj(),Vr)
			else:
				return Vl.T.conj().dot(Vr)

	### Diagonalisation routines

	def eigsh(self,pars={},**eigsh_args):
		"""Computes SOME eigenvalues and eigenvectors of hermitian `quantum_operator` quantum_operator using SPARSE hermitian methods.

		This function method solves for eigenvalues and eigenvectors, but can only solve for a few of them accurately.
		It calls `scipy.sparse.linalg.eigsh <https://docs.scipy.org/doc/scipy/reference/generated/generated/scipy.sparse.linalg.eigsh.html>`_, which is a wrapper for ARPACK.

		Notes
		-----
		Assumes the quantum_operator is hermitian! If the flat `check_hermiticity = False` is used, we advise the user
		to reassure themselves of the hermiticity properties before use.

		Parameters
		-----------
		pars : dict, optional
			Dictionary with same `keys` as `input_dict` and coupling strengths as `values`. Any missing `keys`
			are assumed to be set to unity.
		eigsh_args :
			For all additional arguments see documentation of `scipy.sparse.linalg.eigsh <https://docs.scipy.org/doc/scipy/reference/generated/generated/scipy.sparse.linalg.eigsh.html>`_.

		Returns
		--------
		tuple
			Tuple containing the `(eigenvalues, eigenvectors)` of the `quantum_operator` quantum_operator.

		Examples
		---------
		>>> eigenvalues,eigenvectors = H.eigsh(pars=pars,**eigsh_args)

		"""
		if self.Ns == 0:
			return _np.array([]),_np.array([[]])

		return _sla.eigsh(self.tocsr(pars),**eigsh_args)

	def eigh(self,pars={},**eigh_args):
		"""Computes COMPLETE eigensystem of hermitian `quantum_operator` quantum_operator using DENSE hermitian methods.

		This function method solves for all eigenvalues and eigenvectors. It calls
		`numpy.linalg.eigh <https://docs.scipy.org/doc/numpy-1.10.1/reference/generated/numpy.linalg.eigh.html>`_,
		and uses wrapped LAPACK functions which are contained in the module py_lapack.

		Notes
		-----
		Assumes the quantum_operator is hermitian! If the flat `check_hermiticity = False` is used, we advise the user
		to reassure themselves of the hermiticity properties before use.

		Parameters
		-----------
		pars : dict, optional
			Dictionary with same `keys` as `input_dict` and coupling strengths as `values`. Any missing `keys`
			are assumed to be set to unity.
		eigh_args :
			For all additional arguments see documentation of `numpy.linalg.eigh <https://docs.scipy.org/doc/numpy-1.10.1/reference/generated/numpy.linalg.eigh.html>`_.

		Returns
		--------
		tuple
			Tuple containing the `(eigenvalues, eigenvectors)` of the `quantum_operator` quantum_operator.

		Examples
		---------
		>>> eigenvalues,eigenvectors = H.eigh(pars=pars,**eigh_args)

		"""
		eigh_args["overwrite_a"] = True

		if self.Ns <= 0:
			return _np.asarray([]),_np.asarray([[]])

		# fill dense array with hamiltonian
		H_dense = self.todense(pars=pars)
		# calculate eigh
		E,H_dense = _la.eigh(H_dense,**eigh_args)
		return E,H_dense

	def eigvalsh(self,pars={},**eigvalsh_args):
		"""Computes ALL eigenvalues of hermitian `quantum_operator` quantum_operator using DENSE hermitian methods.

		This function method solves for all eigenvalues. It calls
		`numpy.linalg.eigvalsh <https://docs.scipy.org/doc/numpy-1.10.1/reference/generated/numpy.linalg.eigvalsh.html#numpy.linalg.eigvalsh>`_,
		and uses wrapped LAPACK functions which are contained in the module py_lapack.

		Notes
		-----
		Assumes the quantum_operator is hermitian! If the flat `check_hermiticity = False` is used, we advise the user
		to reassure themselves of the hermiticity properties before use.

		Parameters
		-----------
		pars : dict, optional
			Dictionary with same `keys` as `input_dict` and coupling strengths as `values`. Any missing `keys`
			are assumed to be set to unity.
		eigvalsh_args :
			For all additional arguments see documentation of `numpy.linalg.eigvalsh <https://docs.scipy.org/doc/numpy-1.10.1/reference/generated/numpy.linalg.eigvalsh.html#numpy.linalg.eigvalsh>`_.

		Returns
		--------
		numpy.ndarray
			Eigenvalues of the `quantum_operator` quantum_operator.

		Examples
		---------
		>>> eigenvalues = H.eigvalsh(pars=pars,**eigvalsh_args)

		"""

		if self.Ns <= 0:
			return _np.asarray([])

		H_dense = self.todense(pars=pars)
		E = _np.linalg.eigvalsh(H_dense,**eigvalsh_args)
		#eigvalsh_args["overwrite_a"] = True
		#E = _la.eigvalsh(H_dense,**eigvalsh_args)
		return E

	### routines to change object type

	def tocsr(self,pars={}):
		"""Returns copy of a `quantum_operator` object for parameters `pars` as a `scipy.sparse.csr_matrix`.

		Casts the `quantum_operator` object as a
		`scipy.sparse.csr_matrix <https://docs.scipy.org/doc/scipy/reference/generated/scipy.sparse.csr_matrix.html>`_
		object.

		Parameters
		-----------
		pars : dict, optional
			Dictionary with same `keys` as `input_dict` and coupling strengths as `values`. Any missing `keys`
			are assumed to be set to unity.

		Returns
		--------
		:obj:`scipy.sparse.csr_matrix`

		Examples
		---------
		>>> H_csr=H.tocsr(pars=pars)

		"""
		pars = self._check_scalar_pars(pars)

		H = _sp.csr_matrix(self.get_shape,dtype=self._dtype)

		for key,J in pars.items():
			try:
				H += J*_sp.csr_matrix(self._quantum_operator[key])
			except:
				H = H + J*_sp.csr_matrix(self._quantum_operator[key])

		return H

	def tocsc(self,pars={}):
		"""Returns copy of a `quantum_operator` object for parameters `pars` as a `scipy.sparse.csc_matrix`.

		Casts the `quantum_operator` object as a
		`scipy.sparse.csc_matrix <https://docs.scipy.org/doc/scipy/reference/generated/scipy.sparse.csc_matrix.html>`_
		object.

		Parameters
		-----------
		pars : dict, optional
			Dictionary with same `keys` as `input_dict` and coupling strengths as `values`. Any missing `keys`
			are assumed to be set to unity.

		Returns
		--------
		:obj:`scipy.sparse.csc_matrix`

		Examples
		---------
		>>> H_csc=H.tocsc(pars=pars)

		"""
		pars = self._check_scalar_pars(pars)

		H = _sp.csc_matrix(self.get_shape,dtype=self._dtype)

		for key,J in pars.items():
			try:
				H += J*_sp.csc_matrix(self._quantum_operator[key])
			except:
				H = H + J*_sp.csc_matrix(self._quantum_operator[key])

		return H

	def todense(self,pars={},out=None):
		"""Returns copy of a `quantum_operator` object for parameters `pars` as a dense array.

		This function can overflow memory if not used carefully!

		Notes
		-----
		If the array dimension is too large, scipy may choose to cast the `quantum_operator` quantum_operator as a
		`numpy.matrix` instead of a `numpy.ndarray`. In such a case, one can use the `quantum_operator.toarray()`
		method.

		Parameters
		-----------
		pars : dict, optional
			Dictionary with same `keys` as `input_dict` and coupling strengths as `values`. Any missing `keys`
			are assumed to be set to unity.
		out : numpy.ndarray
			Array to fill in with the output.

		Returns
		--------
		obj
			Depending of size of array, can be either one of

			* `numpy.ndarray`.
			* `numpy.matrix`.

		Examples
		---------
		>>> H_dense=H.todense(pars=pars)

		"""
		pars = self._check_scalar_pars(pars)

		if out is None:
			out = _np.zeros(self._shape,dtype=self.dtype)
			out = _np.asmatrix(out)

		for key,J in pars.items():
			out += J * self._quantum_operator[key]

		return out

	def toarray(self,pars={},out=None):
		"""Returns copy of a `quantum_operator` object for parameters `pars` as a dense array.

		This function can overflow memory if not used carefully!

		Parameters
		-----------
		pars : dict, optional
			Dictionary with same `keys` as `input_dict` and coupling strengths as `values`. Any missing `keys`
			are assumed to be set to unity.
		out : numpy.ndarray
			Array to fill in with the output.

		Returns
		--------
		numpy.ndarray
			Dense array.

		Examples
		---------
		>>> H_dense=H.toarray(pars=pars)

		"""

		pars = self._check_scalar_pars(pars)

		if out is None:
			out = _np.zeros(self._shape,dtype=self.dtype)

		for key,J in pars.items():
			out += J * self._quantum_operator[key]

		return out

	def aslinearoperator(self,pars={}):
		"""Returns copy of a `quantum_operator` object for parameters `pars` as a `scipy.sparse.linalg.Linearquantum_operator`.

		Casts the `quantum_operator` object as a
		`scipy.sparse.linalg.Linearquantum_operator <https://docs.scipy.org/doc/scipy-0.16.0/reference/generated/scipy.sparse.linalg.Linearquantum_operator.html>`_
		object.

		Parameters
		-----------
		pars : dict, optional
			Dictionary with same `keys` as `input_dict` and coupling strengths as `values`. Any missing `keys`
			are assumed to be set to unity.

		Returns
		--------
		:obj:`scipy.sparse.linalg.Linearquantum_operator`

		Examples
		---------
		>>> H_aslinop=H.aslinearquantum_operator(pars=pars)

		"""
		pars = self._check_scalar_pars(pars)
		matvec = functools.partial(_quantum_operator_dot,self,pars)
		rmatvec = functools.partial(_quantum_operator_dot,self.H,pars)
		return _sla.LinearOperator(self.get_shape,matvec,rmatvec=rmatvec,matmat=matvec,dtype=self._dtype)

	def tohamiltonian(self,pars={}):
		"""Returns copy of a `quantum_operator` object for parameters `pars` as a `hamiltonian` object.

		Parameters
		-----------
		pars : dict, optional
			Dictionary with same `keys` as `input_dict` and coupling strengths as `values`. Any missing `keys`
			are assumed to be set to unity.

		Returns
		--------
		:obj:`hamiltonian`

		Examples
		---------
		>>> H_aslinop=H.tohamiltonian(pars=pars)

		"""
		pars = self._check_hamiltonian_pars(pars)

		static=[]
		dynamic=[]

		for key,J in pars.items():
			if type(J) is tuple and len(J) == 2:
				dynamic.append([self._quantum_operator[key],J[0],J[1]])
			else:
				if J == 1.0:
					static.append(self._quantum_operator[key])
				else:
					static.append(J*self._quantum_operator[key])

		return hamiltonian_core.hamiltonian(static,dynamic,dtype=self._dtype)

	def update_matrix_formats(self,matrix_formats):
		"""Change the internal structure of the matrices in-place.

		Parameters
		-----------
		matrix_formats: dict, optional
			Dictionary of key,value pairs which, given a key associated with an operator in `input_dict`, the value of this key
			specifies the sparse matrix format {"csr","csc","dia","dense"}.

		Examples
		---------
		Given `O` which has two operators defined by strings: 'Hx' for transverse field, and 'Hising' is the Ising part. The Ising part must be diagonal
		therefore it is useful to cast it to a DIA matrix format, while the transverse field is not diagonal so it is most efficient to use CSR matrix format.
		This can be accomplished by the following:

		>>> O.update_matrix_formats(dict(Hx="csr",Hising="dia"))

		"""
		if type(matrix_formats) is not dict:
			raise ValueError("matrix_formats must be a dictionary with the formats of the matrices being values and keys being the operator keys.")

		extra = set(matrix_formats.keys()) - set(self._quantum_operator.keys())
		if extra:
			raise ValueError("unexpected couplings: {}".format(extra))

		for key in self._quantum_operator.keys():
			if key in matrix_formats:
				fmt = matrix_formats[key]
				if fmt not in ["dia","csr","csc","dense"]:
					raise TypeError("sparse formats must be either 'csr','csc', 'dia' or 'dense'.")

				if fmt == "dense":
					O = self._quantum_operator[key]
					try:
						self._quantum_operator[key] = O.toarray()
					except AttributeError:
						self._quantum_operator[key] = _np.ascontiguousarray(O)
				else:
					sparse_constuctor = getattr(_sp,fmt+"_matrix")
					O = self._quantum_operator[key]
					if _sp.issparse(O):
						self._quantum_operator[key] = sparse_constuctor(O)

		self._update_matvecs()

	### algebra operations

	def transpose(self,copy=False):
		"""Transposes `quantum_operator` quantum_operator.

		Notes
		-----
		This function does NOT conjugate the quantum_operator.

		Returns
		--------
		:obj:`quantum_operator`
			:math:`H_{ij}\\mapsto H_{ji}`

		Examples
		---------

		>>> H_tran = H.transpose()

		"""

		new_dict = {key:[op.transpose()] for key,op in iteritems(self._quantum_operator)}
		return quantum_operator(new_dict,basis=self._basis,dtype=self._dtype,shape=self._shape,copy=copy)

	def conjugate(self):
		"""Conjugates `quantum_operator` quantum_operator.

		Notes
		-----
		This function does NOT transpose the quantum_operator.

		Returns
		--------
		:obj:`quantum_operator`
			:math:`H_{ij}\\mapsto H_{ij}^*`

		Examples
		---------

		>>> H_conj = H.conj()

		"""
		new_dict = {key:[op.conjugate()] for key,op in iteritems(self._quantum_operator)}
		return quantum_operator(new_dict,basis=self._basis,dtype=self._dtype,shape=self._shape,copy=False)

	def conj(self):
		"""Conjugates `quantum_operator` quantum_operator.

		Notes
		-----
		This function does NOT transpose the quantum_operator.

		Returns
		--------
		:obj:`quantum_operator`
			:math:`H_{ij}\\mapsto H_{ij}^*`

		Examples
		---------

		>>> H_conj = H.conj()

		"""
		return self.conjugate()

	def getH(self,copy=False):
		"""Calculates hermitian conjugate of `quantum_operator` quantum_operator.

		Parameters
		-----------
		copy : bool, optional
			Whether to return a deep copy of the original object. Default is `copy = False`.

		Returns
		--------
		:obj:`quantum_operator`
			:math:`H_{ij}\\mapsto H_{ij}^*`

		Examples
		---------

		>>> H_herm = H.getH()

		"""
		return self.conjugate().transpose(copy=copy)

	def copy(self):
		"""Returns a deep copy of `quantum_operator` object."""
		new_dict = {key:[op] for key,op in iteritems(self._quantum_operator)}
		return quantum_operator(new_dict,basis=self._basis,dtype=self._dtype,shape=self._shape,copy=True)

	def astype(self,dtype,copy=False,casting="unsafe"):
		""" Changes data type of `quantum_operator` object.

		Parameters
		-----------
		dtype : 'type'
			The data type (e.g. numpy.float64) to cast the Hamiltonian with.

		Returns
		`quantum_operator`
			quantum_operator with altered data type.

		Examples
		---------
		>>> H_cpx=H.astype(np.complex128)

		"""
		if dtype not in hamiltonian_core.supported_dtypes:
			raise ValueError("quantum_operator can only be cast to floating point types")

		new_dict = {key:[op.astype(dtype,copy=copy,casting=casting)] for key,op in iteritems(self._quantum_operator)}

		return quantum_operator(new_dict,basis=self._basis,dtype=dtype,shape=self._shape,copy=copy)

	### lin-alg operations

	def diagonal(self,pars={}):
		""" Returns diagonal of `quantum_operator` quantum_operator for parameters `pars`.

		Parameters
		-----------
		pars : dict, optional
			Dictionary with same `keys` as `input_dict` and coupling strengths as `values`. Any missing `keys`
			are assumed to be set to unity.

		Returns
		--------
		numpy.ndarray
			array containing the diagonal part of the operator :math:`diag_j = H_{jj}(\\lambda)`.

		Examples
		---------

		>>> H_diagonal = H.diagonal(pars=pars)

		"""
		pars = self._check_scalar_pars(pars)
		diag = _np.zeros(self.Ns,dtype=self._dtype)
		for key,value in iteritems(self._quantum_operator):
			diag += pars[key] * value.diagonal()
		return diag

	def trace(self,pars={}):
		""" Calculates trace of `quantum_operator` quantum_operator for parameters `pars`.

		Parameters
		-----------
		pars : dict, optional
			Dictionary with same `keys` as `input_dict` and coupling strengths as `values`. Any missing `keys`
			are assumed to be set to unity.

		Returns
		--------
		float
			Trace of quantum_operator :math:`\\sum_{j=1}^{Ns} H_{jj}(\\lambda)`.

		Examples
		---------

		>>> H_tr = H.trace(pars=pars)

		"""
		pars = self._check_scalar_pars(pars)
		tr = 0.0
		for key,value in iteritems(self._quantum_operator):
			try:
				tr += pars[key] * value.trace()
			except AttributeError:
				tr += pars[key] * value.diagonal().sum()
		return tr

	def __str__(self):
		s = ""
		for key,op in iteritems(self._quantum_operator):
			s = s + ("{}:\n{}\n".format(key,op))

		return s

	def __repr__(self):
		return "<{} x {} quspin.operator.quantum_operator with {} operator(s)>".format(self.shape[0],self.shape[1],len(self._quantum_operator))

	def __call__(self,**pars):
		pars = self._check_scalar_pars(pars)
		if self.is_dense:
			return self.todense(pars)
		else:
			return self.tocsr(pars)

	def __neg__(self):
		return self.__imul__(-1)

	def __iadd__(self,other):
		self._is_dense = self._is_dense or other._is_dense
		if isinstance(other,quantum_operator):
			for key,value in iteritems(other._quantum_operator):
				if key in self._quantum_operator:
					self._quantum_operator[key] = self._quantum_operator[key] + value
				else:
					self._quantum_operator[key] = value

				if _check_almost_zero(self._quantum_operator[key]):
					self._quantum_operator.pop(key)

			self._update_matvecs()
			return self
		elif other == 0:
			return self
		else:
			return NotImplemented

	def __add__(self,other):
		result_dtype = _np.result_type(self._dtype, other.dtype)
		new = self.astype(result_dtype,copy=True)
		new += other
		return new

	def __isub__(self,other):
		self._is_dense = self._is_dense or other._is_dense
		if isinstance(other,quantum_operator):
			for key,value in iteritems(other._quantum_operator):
				if key in self._quantum_operator:
					self._quantum_operator[key] = self._quantum_operator[key] - value
				else:
					self._quantum_operator[key] = -value

				if _check_almost_zero(self._quantum_operator[key]):
					self._quantum_operator.pop(key)

			self._update_matvecs()
			return self
		elif other == 0:
			return self
		else:
			return NotImplemented

	def __sub__(self,other):
		result_dtype = _np.result_type(self._dtype, other.dtype)
		new = self.astype(result_dtype,copy=True)
		new -= other
		return new

	def __imul__(self,other):
		if isinstance(other,quantum_operator):
			return NotImplemented
		elif not _np.isscalar(other):
			return NotImplemented
		else:
			for op in itervalues(self._quantum_operator):
				op *= other

			self._update_matvecs()
			return self

	def __mul__(self,other):
		result_dtype = _np.result_type(self._dtype, other.dtype)
		new = self.astype(result_dtype,copy=True)
		new *= other
		return new

	def __idiv__(self,other):
		if isinstance(other,quantum_operator):
			return NotImplemented
		elif not _np.isscalar(other):
			return NotImplemented
		else:
			for op in itervalues(self._quantum_operator):
				op /= other
			self._update_matvecs()
			return self

	def __div__(self,other):
		result_dtype = _np.result_type(self._dtype, other.dtype)
		new = self.astype(result_dtype,copy=True)
		new /= other
		return new

	def _check_hamiltonian_pars(self,pars):

		if not isinstance(pars,dict):
			raise ValueError("expecing dictionary for parameters.")

		extra = set(pars.keys()) - set(self._quantum_operator.keys())
		if extra:
			raise ValueError("unexpected couplings: {}".format(extra))

		missing = set(self._quantum_operator.keys()) - set(pars.keys())
		for key in missing:
			pars[key] = _np.array(1,dtype=_np.int32)

		for key,J in pars.items():
			if type(J) is tuple:
				if len(J) != 2:
					raise ValueError("expecting parameters to be either scalar or tuple of function and arguements of function.")
			else:
				J = _np.array(J)
				if J.ndim > 0:
					raise ValueError("expecting parameters to be either scalar or tuple of function and arguements of function.")

		return pars

	def _check_scalar_pars(self,pars):

		if not isinstance(pars,dict):
			raise ValueError("expecing dictionary for parameters.")

		extra = set(pars.keys()) - set(self._quantum_operator.keys())
		if extra:
			raise ValueError("unexpected couplings: {}".format(extra))

		missing =  set(self._quantum_operator.keys()) - set(pars.keys())
		for key in missing:
			pars[key] = 1.0

		return pars

	# checks
	def _mat_checks(self,other,casting="same_kind"):
		try:
			if other.shape != self._shape: # only accepts square matricies
				raise ValueError('shapes do not match')
			if not _np.can_cast(other.dtype,self._dtype,casting=casting):
				raise ValueError('cannot cast types')
		except AttributeError:
			if other._shape != self._shape: # only accepts square matricies
				raise ValueError('shapes do not match')
			if not _np.can_cast(other.dtype,self._dtype,casting=casting):
				raise ValueError('cannot cast types')

	def _update_matvecs(self):
		self._matvec_functions = {}

		for key in self._quantum_operator.keys():
			self._matvec_functions[key] = _get_matvec_function(self._quantum_operator[key])

def isquantum_operator(obj):
	"""Checks if instance is object of `quantum_operator` class.

	Parameters
	-----------
	obj :
		Arbitraty python object.

	Returns
	--------
	bool
		Can be either of the following:

		* `True`: `obj` is an instance of `quantum_operator` class.
		* `False`: `obj` is NOT an instance of `quantum_operator` class.

	"""
	return isinstance(obj,quantum_operator)


#!/usr/bin/python
# -*- coding: utf-8 -*-

"""Tools for reading/writing of files and performing operations."""

import csv
import functools
import multiprocessing
import glob
import os
import sys

import numpy as np

from pyrvt.peak_calculators import get_peak_calculator, get_region
from pyrvt import motions

# Try to load the modules required for reading/writing files
try:
    import xlrd
except ImportError:
    xlrd = None

try:
    import xlwt
except ImportError:
    xlwt = None

try:
    import openpyxl
except ImportError:
    openpyxl = None

PARAMETER_NAMES = [
    ('magnitude', 'Magnitude'),
    ('distance', 'Distance (km)'),
    ('vs30', 'Vs30 (m/s)'),
    ('kappa', 'Site Atten., Kappa0 (sec)'),  # Site Atten., κ₀
    ('duration', 'Duration (sec)'),
    ('region', 'Region'),
]

def read_events(fname, response_type):
    """Read data from the file an Excel work book.

    Parameters
    ----------
    fname : str
        Filename of the input file.
    response_type : str
        Type of response. Valid options are: 'psa' for psuedo-spectral
        acceleration, or 'fa' for Fourier amplitude.

    Returns
    -------
    ext : str
        Extension of input file
    reference : :class:`numpy.ndarray`
        Reference of the response. This is either period (sec) for
        response_type 'psa' or frequency (Hz) for response_type 'fa'
    events : List[dict]
        List of events read from the file. See ``Note`` in
        :func:`.calc_compatible_spectra` for more information on structure of
        the dictionaries.

    """
    assert response_type in ['psa', 'fa']

    ext = os.path.splitext(fname)[1].lower()
    # Load the file depending on the format
    if ext == '.csv':

        def parse(s):
            try:
                return float(s)
            except ValueError:
                return s

        with open(fname) as fp:
            reader = csv.reader(fp)
            rows = [[parse(r) for r in row] for row in reader]
    elif ext == '.xls':
        if xlrd is None:
            raise RuntimeError('xlrd is required to open an xls file')
        wb = xlrd.open_workbook(fname)
        ws = wb.sheet_by_index(0)
        rows = [ws.row_values(i) for i in range(ws.nrows)]
    elif ext == '.xlsx':
        if openpyxl is None:
            raise RuntimeError('openpyxl is required to open an xlsx file')
        wb = openpyxl.load_workbook(fname, read_only=True)
        ws = wb.worksheets[0]
        rows = [[r.value for r in row] for row in ws.rows]
        # Close the file so that it may be deleted if needed. This is only
        # important so that in the test cases the temporary .xlsx file can be
        # deleted.
        wb._archive.close()
    else:
        raise NotImplementedError

    parameters = {
        key: rows[i][1:]
        for i, (key, label) in enumerate(PARAMETER_NAMES)
    }

    event_row = len(parameters) + 1
    event_count = len(rows[0]) - 1

    reference = np.array([row[0] for row in rows[event_row:]])

    events = []
    for i in range(event_count):
        resps = np.array([row[i + 1] for row in rows[event_row:]])
        # Extract the appropriate attributes
        e = {k: v[i] for k, v in parameters.items()}
        e[response_type] = resps

        if 'region' in e:
            e['region'] = get_region(e['region'])

        events.append(e)

    return ext, reference, events

def write_events(fname, reference, reference_label, response_type,
                 response_label, events):
    """Write the events to a file.

    Parameters
    ----------
    fname : str
        Save the events to this file. The directory is created if needed.
    reference : array_like
        Periods of the oscillator response shared across all events.
    reference_label : str
        Label of the reference (e.g., 'Frequency (Hz)').
    response_type : str
        Type of response. Valid options: `psa` for pseudo-spectral
        accleration, or `fa` for Fourier amplitude.
    response_label : str
        Label of the response type (e.g., 'Fourier Ampl. (g/sec)')
    events : List[dict]
        Events to write to file. See ``Note`` in
        :func:`.compute_compatible_spectra` for more information.

    Raises
    ------
    NotImplementedError:
        If extension is not supported

    """
    # Create the rows of output
    rows = []
    # Output the parameters
    for key, label in PARAMETER_NAMES:
        rows.append([label] + [e[key] for e in events])

    rows.append([reference_label] + len(events) * [response_label])

    # Output the response spectra
    for i in range(len(reference)):
        rows.append([reference[i]] + [e[response_type][i] for e in events])

    # Create the directory
    dirname = os.path.dirname(fname)
    if not os.path.exists(dirname):
        os.makedirs(dirname)

    # Write the file
    ext = os.path.splitext(fname)[1].lower()

    if ext == '.csv':
        if sys.version_info < (3, 1):
            fp = open(fname, 'wt')
        else:
            fp = open(fname, 'wt', newline='')
        writer = csv.writer(fp)
        writer.writerows(rows)
        fp.close()
    elif ext == '.xls':
        if xlwt is None:
            raise RuntimeError('xlwt is required to open an xls file')
        wb = xlwt.Workbook()
        ws = wb.add_sheet('Sheet 1')
        for i, row in enumerate(rows):
            for j, cell in enumerate(row):
                ws.write(i, j, cell)
        wb.save(fname)
    elif ext == '.xlsx':
        if openpyxl is None:
            raise RuntimeError('openpyxl is required to open an xlsx file')
        wb = openpyxl.Workbook()
        ws = wb.worksheets[0]
        for row in rows:
            ws.append(row)
        wb.save(fname)
    else:
        raise NotImplementedError

def _calc_fa(target_freqs, damping, method, event):
    """Calculate the fourier amplitudes for an event.

    Note that this is intended as a helper function to be called by
    multiprocessing.Pool.
    """
    event_keys = ['magnitude', 'distance', 'region']
    event_kwds = {key: event[key] for key in event_keys}
    crm = motions.CompatibleRvtMotion(
        target_freqs,
        event['psa'],
        duration=event['duration'],
        osc_damping=damping,
        event_kwds=event_kwds,
        peak_calculator=get_peak_calculator(method, event_kwds))
    psa_calc = crm.calc_osc_accels(target_freqs, damping)
    return crm, psa_calc

def calc_compatible_spectra(method, periods, events, damping=0.05):
    """Compute the response spectrum compatible motions.

    Parameters
    ----------
    method : str
        RVT peak factor method, see
        :func:`~.peak_calculators.get_peak_calculator`.
    periods : array_like
        Periods of the oscillator response shared across all events.
    events : List[dict]
        All events to consider. See ``Note``.
    damping : float, optional
        Fractional damping of the oscillator (decimal). Default value of 0.05
        for a damping ratio of 5%.

    Returns
    -------
    :class:`numpy.ndarray`
        Frequency of the computed Fourier amplitude spectra.

    Note
    ----
    Each event dictionary should have the following keys:

    - **psa** : :class:`numpy.ndarray` -- pseudo-spectral accelerations. This
      is the target for the :class:`~.motions.CompatibleRvtMotion`.
    - **duration** : float, optional -- duration of the ground motion
    - **magnitude** : float, optional -- earthquake magnitude
    - **distance** : float, optional -- earthquake distance (km)
    - **region** : str, optional -- earthquake source region, see
      :func:`~.peak_calculators.get_region` If no duration is provided one
      is estimated from the magnitude, distance, and region.

    The `events` dictionary is modified by this function and adds the
    following keys:

    - **duration** : float -- duration of the ground motion if one was not
      specified
    - **fa** : :class:`numpy.ndarray` -- Fourier amplitude spectra in units of
      g/sec
    - **psa_calc** : :class:`numpy.ndarray` -- Pseudo-spectral acceleration
      calculated from `fa`. This will differ slightly from `psa_target`.

    """
    target_freqs = 1. / periods
    with multiprocessing.Pool() as pool:
        results = pool.map(
            functools.partial(_calc_fa, target_freqs, damping, method), events)

    # Copy values back into the dictionary
    for event, (crm, psa_calc) in zip(events, results):
        if not event['duration']:
            event['duration'] = crm.duration
        event['fa'] = crm.fourier_amps
        event['psa_calc'] = psa_calc

    # Return the frequency from one of the computed motions.
    freqs = results[0][0].freqs
    return freqs

def operation_psa2fa(src,
                     dst,
                     damping,
                     method='LP99',
                     fixed_spacing=True,
                     verbose=True):
    """Compute the accel. response spectrum from a Fourier amplitude spectrum.

    Parameters
    ----------
    src : str
        Source for the pseudo-spectral accelerations (PSA). This can be a
        filename or pattern used in :func:`glob.glob`.
    dst : str
        Destination directory for the output PSA. The directory is created if
        it does not exist.
    damping : float
        Fractional damping of the oscillator ( decimal).
    method : str
        RVT peak factor method, see
        :func:`~.peak_calculators.get_peak_calculator`.
    fixed_spacing : bool, optional
        If `True`, then the periods are interpolated to 301 points equally
        space in log-space from 0.01 to 10.
    verbose : bool, optional
        Print status of calculation.

    """
    for filename_src in glob.iglob(src):
        if verbose:
            print('Processing:', filename_src)
        ext, periods, events = read_events(filename_src, 'psa')

        if fixed_spacing:
            # Interpolate the periods to a smaller range
            _periods = np.logspace(-2, 1, 301)

            for e in events:
                e['psa'] = np.exp(
                    np.interp(_periods, periods, np.log(e['psa'])))

            periods = _periods

        # Compute the FA from the PSA
        freqs = calc_compatible_spectra(
            method, periods, events, damping=damping)

        if not os.path.exists(dst):
            os.makedirs(dst)

        basename = os.path.basename(filename_src)
        pathname_dst = os.path.join(dst, basename.rsplit('_', 1)[0])

        write_events(pathname_dst + '_sa' + ext, periods, 'Period (s)',
                     'psa_calc', 'Sa (g)', events)
        write_events(pathname_dst + '_fa' + ext, freqs, 'Frequency (Hz)', 'fa',
                     'FA (g-s)', events)

def _calc_psa(osc_freqs, damping, method, freqs, event):
    """Calculate the response spectra for an event.

    Note that this is intended as a helper function to be called by
    multiprocessing.Pool.
    """
    m = motions.RvtMotion(
        freqs=freqs,
        fourier_amps=event['fa'],
        duration=event['duration'],
        peak_calculator=get_peak_calculator(method,
                                            dict(
                                                region=event['region'],
                                                mag=event['magnitude'],
                                                dist=event['distance'])))
    psa = m.calc_osc_accels(osc_freqs, damping)
    return psa

def operation_fa2psa(src,
                     dst,
                     damping,
                     method='LP99',
                     fixed_spacing=True,
                     verbose=True):
    """Compute the Fourier amplitude spectrum from a accel. response spectrum.

    Parameters
    ----------
    src : str
        Source for the Fourier amplitudes. This can be a filename or pattern
        used in :func:`glob.glob`.
    dst : str
        Destination directory for the output PSA. The directory is created if
        it does not exist.
    damping : float
        Fractional damping of the oscillator (decimal).
    method : str
        RVT peak factor method, see
        :func:`~.peak_calculators.get_peak_calculator`.
    fixed_spacing : bool, optional
        If `True`, then the periods are interpolated to 301 points equally
        space in log-space from 0.01 to 10.

    """
    if fixed_spacing:
        periods = np.logspace(-2, 1, 301)
        osc_freqs = 1. / periods

    for filename_src in glob.iglob(src):
        if verbose:
            print('Processing:', filename_src)
        ext, freqs, events = read_events(filename_src, 'fa')

        if not fixed_spacing:
            osc_freqs = freqs
            periods = 1. / osc_freqs

        with multiprocessing.Pool() as pool:
            psas = pool.map(
                functools.partial(_calc_psa, osc_freqs, damping, method,
                                  freqs), events)

        for event, psa in zip(events, psas):
            event['psa'] = psa

        if not os.path.exists(dst):
            os.makedirs(dst)

        basename = os.path.basename(filename_src)

        pathname_dst = os.path.join(dst, basename.rsplit('_', 1)[0])

        write_events(pathname_dst + '_sa' + ext, periods, 'Period (s)', 'psa',
                     'PSA (g)', events)


# -*- coding: utf-8 -*-
import os
import numpy as np
import pandas as pd
import pickle
import pystan
import bokeh.plotting
import tqdm
from .stats import compute_statistics

class StanModel(object):
    r"""
    Custom StanModel class for crafting and sampling from Stan
    models.
    """

    def __init__(self, file, data_dict=None, samples=None, force_compile=False):
        """
        Parameters
        ----------
        model: str
            Relative path to saved Stan model code. To deter bad habits,
            this class does not accept a string as the model code. Save
            your Stan models.
        data_dict: dictionary
            Dictonary of all data block parameters for the model.
        force_compile: bool
            If True, model will be forced to compile. If False,
            a precompiled file will be loaded if present.
        """
        if ".pkl" in file:
            s = _load(file)
            self.model = s[0]
            self.samples = s[1]
        else:
            self.model = loadStanModel(file, force=force_compile)
            self.data = data_dict
            self.samples = samples
            self.df = None

    def sample(self, data_dict=None, iter=2000, chains=4, return_df=True, **kwargs):
        """
        Samples the assembled model given the supplied data dictionary
        and returns output as a dataframe.
        """
        if data_dict == None:
            data_dict = self.data
        self.chains = chains
        self.iter = iter
        self.samples = self.model.sampling(
            data_dict, chains=chains, iter=iter, **kwargs
        )
        if return_df:
            self.df = self.samples.to_dataframe(diagnostics=True)
            return [self.samples, self.df]
        else:
            return self.samples

    # Pickling objects
    def dump(fname):
        """Saves StanFit4Model object and sampling summary as a pickled dictionary."""
        with open(f"{fname.split('.')[0]}.pkl", "wb") as _file:
            pickle.dump({"model": self.model, "fit": self.samples}, _file, protocol=-1)

    def _load(fname):
        with open(file, "rb") as _file:
            fit_dict = pickle.load(_file)
        self.model = fit_dict[0]
        self.samples = fit_dict[1]
        return [self.model, self.samples]

    def summarize_parameters(self, parnames=[], mass_frac=0.95):
        """
        Summarizes all or a subset of parameters from a Stan model.

        Parameters
        ----------
        parnames: list
            List of desired parnames. If left empty, all parameters
            are summarized and returned.
        mass_frac: float [0, 1]
            The probability mass fraction for the HPD. Default is
            the 95% credible region.

        Returns
        -------
        summary_df: pandas DataFrame
            Dataframe of summarized parameters. The columns are as
            follows:
                parameter = name of parameter in Stan model
                dimension = index (dimension) of the parameter
                mean = mean of samples
                median = median of samples
                mode = parameter value when the log posterior is maximized
                hpd_min = minimum bound of the highest probability density
                    defined by the mass fraction.
                hpd_max = upper bound of the highest probability density
                    defined by the mass fraction
        """
        # Extract the sampling information and find the mode
        samples = self.samples
        fit = samples.extract()
        mode_ind = np.argmax(fit["lp__"])

        # Get a list of all parameters defined in the model and assign a dimension
        pars = samples.model_pars

        # Convert the dimensions for each parameter to integers.
        _dims = []
        for d in samples.par_dims:
            if len(d) == 0:
                _dims.append(1)
            else:
                _dims.append(int(d[0]))

        par_dims = {p: v for p, v in zip(pars, _dims)}
        if len(parnames) != 0:
            pars = parnames
            desired_pars = {k: v for k, v in par_dims.items() if k in parnames}
            par_dims = desired_pars

        # Iterate through each parameter and compute the aggregate properties.
        df = pd.DataFrame(
            [],
            columns=[
                "parameter",
                "dimension",
                "mean" "mode",
                "median",
                "hpd_min",
                "hpd_max",
                "mass_fraction",
            ],
        )
        for par, dim in par_dims.items():
            par_samples = fit[par]
            if dim == 1:
                par_samples = par_samples[:, np.newaxis]
            for j in range(dim):
                # Compute the summary statistics
                par_mode = par_samples[:, j][mode_ind]
                par_mean = np.mean(par_samples[:, j])
                par_median = np.median(par_samples[:, j])
                hpd_min, hpd_max = compute_hpd(par_samples[:, j], mass_frac=mass_frac)

                # Assemble a dictionary to append to the data frame
                par_dict = {
                    "parameter": par,
                    "dimension": j + 1,
                    "mean": par_mean,
                    "mode": par_mode,
                    "median": par_median,
                    "hpd_min": hpd_min,
                    "hpd_max": hpd_max,
                    "mass_fraction": mass_frac,
                }
                df = df.append(par_dict, ignore_index=True)
        df["dimension"] = df["dimension"].astype(int)
        return df

def loadStanModel(fname, force=False):
    """Loads a precompiled Stan model. If no compiled model is found, one will be saved."""
    # Identify the model name and directory structure
    rel, sm_dir = fname.split("/stan/")
    sm_name = sm_dir.split(".stan")[0]
    pkl_name = f"{rel}/stan/{sm_name}.pkl"
    # Check if the model is precompiled
    if (os.path.exists(pkl_name) == True) and (force != True):
        print("Found precompiled model. Loading...")
        model = pickle.load(open(pkl_name, "rb"))
        print("finished!")
    else:
        print("Precompiled model not found. Compiling model...")
        _path = rel + "/stan/"
        model = pystan.StanModel(fname, include_paths=_path)
        print("finished!")
        with open(pkl_name, "wb") as f:
            pickle.dump(model, f)
    return model

def infer_empirical_bohr(
    data,
    model,
    groupby=["mutant", "repressors", "operator", "IPTGuM"],
    verbose=True,
    force_compile=False,
    **kwargs,
):
    """
    Infers the empirical bohr parameter (and relevant correction) for a collection of
    fold-change measurements

    Parameters
    ----------
    data: pandas DataFrame object
        The data from which the empirical bohr will be determined. This should have at least
        a fold-change column and a grouping parameter.
    model: str
        Path to Stan model to load. Model will be compiled if `force_compile`==True.
    groupby: list, optional
        List of identifiers by which to group the supplied data. Default groups by
        'mutant', 'repressors', 'operator', and 'IPTGuM'
    verbose: bool
        If true, the progress will be printed to screen as a bar.
    force_compile: bool
        If True, the stan model will be recompiled.
    **kwargs: keyword arguments
        kwargs to be passed to the sampler.

    Returns
    -------
    statistics: pandas DataFrame
        Dataframe of statistics for relevant parameters.
    """

    # Load the stan model and compile if needed.
    model = StanModel(model, force_compile=force_compile)

    # Make a storage list for the individual statistics
    fc_stats = []

    # Make a quiet or loud iterator.
    if verbose:
        iter = tqdm.tqdm(data.groupby(groupby))
    else:
        iter = data.groupby(groupby)

    # Iter through each grouping and infer
    for g, d in iter:
        # Define parameters of the reference state
        ref = d["ref_bohr"].unique()[0]
        fc_ref = (1 + np.exp(-ref)) ** -1

        # Assemble the data dictionary and sample the posterior
        data_dict = {"N": len(d), "foldchange": d["fold_change"]}
        fit, samples = model.sample(data_dict, **kwargs)

        # Identify the extrema
        extrema = (samples["fc_mu"] < samples["fc_sigma"]).astype(int) + (
            1 - samples["fc_mu"] < samples["fc_sigma"]
        ).astype(int)

        # Compute the empirical bohr parameter and the delta bohr
        samples["empirical_bohr"] = -np.log((samples["fc_mu"]) ** -1 - 1)
        samples["delta_bohr"] = ref - samples["empirical_bohr"]

        # Compute the delta F error of the reference, given the sigma
        delta_F_ref_upper = np.nan_to_num(
            ref + np.log((fc_ref + samples["fc_sigma"]) ** -1 - 1)
        )
        delta_F_ref_lower = np.nan_to_num(
            ref + np.log((fc_ref - samples["fc_sigma"]) ** -1 - 1)
        )
        samples["correction"] = (delta_F_ref_upper + delta_F_ref_lower) * extrema
        samples["delta_bohr_corrected"] = samples["delta_bohr"] - samples["correction"]

        _dbohr_stats = compute_statistics(
            samples,
            varnames=[
                "delta_bohr",
                "empirical_bohr",
                "fc_mu",
                "fc_sigma",
                "delta_bohr_corrected",
                "correction",
            ],
            logprob_name="lp__",
        )
        _dbohr_stats["mutant"] = g[0]
        _dbohr_stats["repressors"] = g[1]
        _dbohr_stats["operator"] = g[2]
        _dbohr_stats["IPTGuM"] = g[3]
        _dbohr_stats["class"] = d["class"].unique()[0]
        fc_stats.append(_dbohr_stats)

    return pd.concat(fc_stats)