import numpy as _np
import scipy.sparse as _sp
from ._basis_utils import _shuffle_sites



####################################################
# set of helper functions to implement the partial # 
# trace of lattice density matrices. They do not   #
# have any checks and states are assumed to be     #
# in the non-symmetry reduced basis.               #
####################################################

def _lattice_partial_trace_pure(psi,sub_sys_A,L,sps,return_rdm="A"):
	"""
	This function computes the partial trace of a dense pure state psi over set of sites sub_sys_A and returns 
	reduced DM. Vectorisation available. 
	"""
	
	psi_v=_lattice_reshape_pure(psi,sub_sys_A,L,sps)

	if return_rdm == "A":
		return _np.squeeze(_np.einsum("...ij,...kj->...ik",psi_v,psi_v.conj())),None
	elif return_rdm == "B":
		return None,_np.squeeze(_np.einsum("...ji,...jk->...ik",psi_v.conj(),psi_v))
	elif return_rdm == "both":
		return _np.squeeze(_np.einsum("...ij,...kj->...ik",psi_v,psi_v.conj())),_np.squeeze(_np.einsum("...ji,...jk->...ik",psi_v.conj(),psi_v))

def _lattice_partial_trace_mixed(rho,sub_sys_A,L,sps,return_rdm="A"):
	"""
	This function computes the partial trace of a set of dense mixed states rho over set of sites sub_sys_A 
	and returns reduced DM. Vectorisation available. 
	"""
	rho_v=_lattice_reshape_mixed(rho,sub_sys_A,L,sps)
	if return_rdm == "A":
		return _np.einsum("...jlkl->...jk",rho_v),None
	elif return_rdm == "B":
		return None,_np.einsum("...ljlk->...jk",rho_v.conj())
	elif return_rdm == "both":
		return _np.einsum("...jlkl->...jk",rho_v),_np.einsum("...ljlk->...jk",rho_v.conj())

def _lattice_partial_trace_sparse_pure(psi,sub_sys_A,L,sps,return_rdm="A"):
	"""
	This function computes the partial trace of a sparse pure state psi over set of sites sub_sys_A and returns 
	reduced DM.
	"""
	psi=_lattice_reshape_sparse_pure(psi,sub_sys_A,L,sps)
	
	if return_rdm == "A":
		return psi.dot(psi.H),None
	elif return_rdm == "B":
		return None,psi.H.dot(psi)
	elif return_rdm == "both":
		return psi.dot(psi.H),psi.H.dot(psi)


def _lattice_reshape_pure(psi,sub_sys_A,L,sps):
	"""
	This function reshapes the dense pure state psi over the Hilbert space defined by sub_sys_A and its complement. 
	Vectorisation available. 
	"""
	extra_dims = psi.shape[:-1]
	n_dims = len(extra_dims)
	sub_sys_B = set(range(L))-set(sub_sys_A)

	sub_sys_A = tuple(sub_sys_A)
	sub_sys_B = tuple(sub_sys_B)

	L_A = len(sub_sys_A)
	L_B = len(sub_sys_B)

	Ns_A = (sps**L_A)
	Ns_B = (sps**L_B)
	T_tup = sub_sys_A+sub_sys_B
	psi_v = _shuffle_sites(sps,T_tup,psi)
	psi_v = psi_v.reshape(extra_dims+(Ns_A,Ns_B))

	return psi_v

'''
def _lattice_reshape_pure(psi,sub_sys_A,L,sps):
	"""
	This function reshapes the dense pure state psi over the Hilbert space defined by sub_sys_A and its complement. 
	Vectorisation available. 
	"""
	extra_dims = psi.shape[:-1]
	n_dims = len(extra_dims)
	sub_sys_B = set(range(L))-set(sub_sys_A)

	sub_sys_A = tuple(sub_sys_A)
	sub_sys_B = tuple(sub_sys_B)

	L_A = len(sub_sys_A)
	L_B = len(sub_sys_B)

	Ns_A = (sps**L_A)
	Ns_B = (sps**L_B)

	T_tup = sub_sys_A+sub_sys_B
	T_tup = tuple(range(n_dims)) + tuple(n_dims + s for s in T_tup)
	R_tup = extra_dims + tuple(sps for i in range(L))
	psi_v = psi.reshape(R_tup) # DM where index is given per site as rho_v[i_1,...,i_L,j_1,...j_L]
	psi_v = psi_v.transpose(T_tup) # take transpose to reshuffle indices
	psi_v = psi_v.reshape(extra_dims+(Ns_A,Ns_B))
	return psi_v
'''

def _lattice_reshape_mixed(rho,sub_sys_A,L,sps):
	"""
	This function reshapes the dense mixed state psi over the Hilbert space defined by sub_sys_A and its complement.
	Vectorisation available. 
	"""
	extra_dims = rho.shape[:-2]
	n_dims = len(extra_dims)
	sub_sys_B = set(range(L))-set(sub_sys_A)