# AUTOGENERATED! DO NOT EDIT! File to edit: 01weylheisenberg.ipynb (unless otherwise specified).

__all__ = ['clock', 'shift', 'displace', 'weyl_heisenberg_indices', 'displacement_operators', 'weyl_heisenberg_states',
           'weyl_heisenberg_povm', 'to_weyl_heisenberg_basis', 'from_weyl_heisenberg_basis']

# Cell
import numpy as np
import qutip as qt

# Cell
def clock(d):
    r"""
    The clock operator $\hat{Z}$ for dimension $d$.
    """
    w = np.exp(2*np.pi*1j/d)
    return qt.Qobj(np.diag([w**i for i in range(d)]))

# Cell
def shift(d):
    r"""
    The shift operator $\hat{X}$ for dimension $d$.
    """
    return sum([qt.basis(d, i+1)*qt.basis(d, i).dag()\
                    if i != d-1 else qt.basis(d, 0)*qt.basis(d, i).dag()\
                        for i in range(d) for j in range(d)])/d

# Cell
def displace(d, a, b):
    r"""
    The displacement operator $\hat{D}_{a,b} = (-e^{\frac{i\pi}{d}})^{ab}\hat{X}^{b}\hat{Z}^{a}$ for dimension $d$.
    """
    Z, X = clock(d), shift(d)
    return (-np.exp(1j*np.pi/d))**(a*b)*X**b*Z**a

# Cell
def weyl_heisenberg_indices(d):
    r"""
    Returns a list with entries $(a, b)$ for $a, b \in [0, d)$.
    """
    return [(a,b) for b in range(d) for a in range(d)]

# Cell
def displacement_operators(d):
    r"""
    Returns a dictionary associating $(a, b)$ with $\hat{D}_{a,b}$ for $a, b \in [0, d)$.
    """
    return dict([((a,b), displace(d, a, b)) for a, b in weyl_heisenberg_indices(d)])

# Cell
def weyl_heisenberg_states(fiducial):
    r"""
    Applies the $d^2$ displacement operators to a fiducial state, which can be either
    a ket or a density matrix.
    """
    d = fiducial.shape[0]
    D = displacement_operators(d)
    return [D[(a,b)]*fiducial if fiducial.type == "ket" else\
            D[(a,b)]*fiducial*D[(a,b)].dag()\
                for a, b in weyl_heisenberg_indices(d)]

# Cell
def weyl_heisenberg_povm(fiducial):
    r"""
    Generates a Weyl-Heisenberg POVM by applying the $d^2$ displacement operators to a
    fiducial state and then, if the fiducial state is a ket $\mid \psi \rangle$, forming the projector $\mid \psi \rangle \langle \psi \mid$, and normalizing by $\frac{1}{d}$.

    Note that if the fiducial state is a density matrix, it may be the case that it is invariant under some displacement operators, in which case you'll run into problems!
    """
    return [(1/fiducial.shape[0])*(state*state.dag() if state.type=='ket' else state) for state in weyl_heisenberg_states(fiducial)]

# Cell
def to_weyl_heisenberg_basis(O, D=None):
    r"""
    Expands a $d \times d$ operator $O$ in the Weyl-Heisenberg basis with components:

    $$ O_{a,b} = \frac{1}{d} tr ( \hat{D}_{a,b}^{\dagger} \hat{O} ) $$

    Returns a dictionary associating $(a, b)$ with components.
    """
    d = O.shape[0]
    D = D if type(D) != type(None) else displacement_operators(d)
    return dict([(index, (D_.dag()*O).tr()/d) for index, D_ in D.items()])

# Cell
def from_weyl_heisenberg_basis(C, D=None):
    r"""
    Given a dictionary of Weyl-Heisenberg components, returns the operator $O$
    in the standard basis:

    $$\hat{O} = \sum_{a,b} O_{a,b}\hat{D}_{a,b}$$
    """
    d = int(np.sqrt(len(C)))
    D = D if type(D) != type(None) else displacement_operators(d)
    return sum([coeff*D[index] for index, coeff in C.items()])