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

__all__ = ['povm_gram', 'povm_phi', 'dm_probs', 'probs_dm', 'conditional_probs', 'quantum_inner_product', 'tensor_povm',
           'upgrade_povm', 'apply_dims', 'implement_povm', 'discriminator_povm']

# Cell
import numpy as np
import qutip as qt
from functools import reduce

from .weyl_heisenberg import *
from .sics import *
from .krauss import *
from .utils import *

# Cell
def povm_gram(E):
    r"""
    The Gram matrix is the matrix of inner products. Given POVM elements $\{\hat{E}\}$, we form the matrix:

    $$ G_{i,j} = tr(\frac{\hat{E}_{i}}{tr \hat{E}_{i}} \frac{\hat{E}_{j}}{tr \hat{E}_{j}})$$
    """
    return np.array([[((a/a.tr())*(b/b.tr())).tr() for b in E] for a in E]).real

# Cell
def povm_phi(E):
    r"""
    Given POVM elements $\{\hat{E}\}$, we first form the matrix:

    $$ \hat{\Phi}^{-1}_{i,j} = tr(\hat{E}_{i} \frac{\hat{E}_{j}}{tr \hat{E}_{j}}) $$

    And then take its inverse to construct $\hat{\Phi}$, the magical quantum transition matrix.
    """
    phi_inv = np.array([[(a*b/b.tr()).tr() for b in E] for a in E]).real
    return np.linalg.inv(phi_inv)

# Cell
def dm_probs(rho, E):
    r"""
    Given a density matrix $\rho$, expands it in the basis provided by POVM elements $\{\hat{E}\}$, giving a probability vector $\vec{p}$.

    $$p_{i} = tr(\hat{E}_{i}\rho) $$
    """
    return np.array([(e*rho).tr() for e in E]).real

# Cell
def probs_dm(p, E, phi=None):
    r"""
    Given a probability vector $\vec{p}$ and a POVM $\{\hat{E}\}$, recovers the density matrix $\rho$. If it's not provided, we first construct $\hat{\Phi}$, the magical quantum coherence matrix, and then form the vector of quasiprobabilities $\vec{q} = \hat{\Phi}\vec{p}$. Then:

    $$ \rho = \sum_{i} q_{i}\frac{\hat{E_{i}}}{tr E_{i}}$$
    """
    phi = phi if type(phi) != type(None) else povm_phi(E)
    return sum([c*E[i]/E[i].tr() for i, c in enumerate(phi @ p)])

# Cell
def conditional_probs(A, B):
    r"""
    Given two POVM's $\{\hat{A}\}$ and $\{\hat{B}\}$ (or PVM's), constructs the matrix of conditional probabilities $r(j|i)$ for outcome $A_{j}$ given outcome $B_{i}$:

    $$\hat{R}_{j,i} = tr(\hat{A}_{j}\frac{\hat{B}_{i}}{tr \hat{B}_{i}})$$
    """
    return np.array([[((A[j]*B[i]/B[i].tr())).tr() for i in range(len(B))] for j in range(len(A))]).real

# Cell
def quantum_inner_product(r, s, phi=None):
    r"""
    The quantum inner product expressed in terms of probability vectors. If $\hat{\Phi}$ is not provided,
    we use the SIC-POVM $\hat{\Phi}$ of the appropriate dimensionality.

    $$ tr(\sigma\rho) = d \vec{s} \hat{\Phi} \vec{r}$$

    Where $\vec{r}$ is the probability vector for $\rho$ and $\vec{s}$ is the probability vector for $\sigma$ with
    respect to the same POVM.
    """
    d = int(np.sqrt(len(p)))
    phi = phi if type(phi) != type(None) else sic_gram(d)
    return d*(s @ phi @ p).real


# Cell
def tensor_povm(*povms):
    r"""
    Forms the tensor product of a list of POVM's, which is itself a POVM.
    """
    def tensor_two(A, B):
        return [qt.tensor(b, a) for b in B for a in A]
    return reduce(lambda A, B: tensor_two(A, B), povms)

# Cell
def upgrade_povm(E, i, dims):
    r"""
    Upgrades a POVM to act on the $i^{th}$ subspace of a tensor product space whose subspaces are given by a list `dims`. If `dims` is an
    integer, we assume it refers to the number of subspaces all of the same dimensionality as the POVM.
    """
    if type(dims) == int:
        dims = [E[0].shape[0]]*dims
    return [qt.tensor(*[e if i == j else qt.identity(d) for j, d in enumerate(dims)]) for e in E]

# Cell
def apply_dims(E, dims):
    r"""
    Helper function which sets the tensor dimensions of each POVM element to `dims`.
    """
    for e in E:
        e.dims = [dims, dims]
    return E

# Cell
def implement_povm(E):
    r"""
    Returns a unitary operator $\hat{U}$ implementing a given POVM on $H_{d} \otimes H_{n}$, where $d$ is the dimensionality
    of the original system and $n$ is the dimensionality of the auxilliary system and which is the same as the number of POVM elements.
    """
    n, d = len(E), E[0].shape[0]
    aux_projectors = [qt.tensor(qt.identity(d), qt.basis(n, i)*qt.basis(n, i).dag()) for i in range(n)]
    V = sum([qt.tensor(E[i].sqrtm(), qt.basis(n, i)) for i in range(n)])
    povm_elements = [V.dag()*aux_projectors[i]*V for i in range(n)]
    assert np.all([np.allclose(E[i], povm_elements[i]) for i in range(n)])
    Q, R = np.linalg.qr(V, mode="complete")
    for i in range(d):
        Q.T[[i,n*i]] = Q.T[[n*i,i]]
    U = qt.Qobj(-Q)
    U.dims = [[d, n],[d, n]]
    #assert np.allclose(V, U*qt.tensor(qt.identity(d), qt.basis(n, 0)))
    return U

# Cell
def discriminator_povm(a, b):
    r"""
    Returns a non informationally complete POVM which has the special property
    of distinguishing between two arbitrary states $\mid a \rangle$ and $\mid b\rangle$, which are not necessarily orthogonal (which is impossible with a standard PVM).

    It has three elements:

    $$ \hat{F}_{a} = \frac{1}{1+\mid\langle a \mid b \rangle\mid}(\hat{I} - \mid b \rangle \langle b \mid) $$
    $$ \hat{F}_{b} = \frac{1}{1+\mid\langle a \mid b \rangle\mid}(\hat{I} - \mid a \rangle \langle a \mid) $$
    $$ \hat{F}_{?} = \hat{I} - \hat{F}_{a} - \hat{F}_{b} $$

    The first tests for "not B", the second tests for "not A", and the third outcome represents an inconclusive result.
    """
    d = a.shape[0]
    p = abs(a.overlap(b))
    Fa = (1/(1+p))*(qt.identity(d) - b*b.dag())
    Fb = (1/(1+p))*(qt.identity(d) - a*a.dag())
    Fq = qt.identity(d) - Fa - Fb
    return [Fa, Fb, Fq]