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

__all__ = ['random_ginibre', 'random_haar_povm', 'random_haar_effect']

# Cell
import numpy as np
import qutip as qt
import scipy as sc

from .povm import *

# Cell
def random_ginibre(n, m, real=False):
	return (np.random.randn(n, m) + 1j*np.random.randn(n,m))/np.sqrt(2) if not real else np.random.randn(n, m)/np.sqrt(2)

# Cell
def random_haar_povm(d, k=None, n=1, real=False):
    r"""
    Generates a Haar distributed random POVM for a Hilbert space of dimension $d$, with $k$ elements, and with "mixedness" $n$.

    $n$ must satisfy $d \leq kn$, and defaults to $n=1$, giving rank-1 POVM elements.

    $k$ defaults to $d^2$ if complex, $\frac{d(d+1)}{2}$ if real.
    """
    k = k if type(k) != type(None) else (d**2 if not real else int(d*(d+1)/2))
    povm = np.zeros((k, d, d), dtype=np.complex128) if not real else np.zeros((k, d, d))
    S = np.zeros(d, dtype=np.complex128) if not real else np.zeros(d)
    for i in range(k):
        Xi = random_ginibre(d, n, real=real)
        Wi = Xi @ Xi.conjugate().T
        povm[i, :, :] = Wi
        S = S + Wi
    S = sc.linalg.fractional_matrix_power(S, -1/2)
    for i in range(k):
        Wi = np.squeeze(povm[i, :, :])
        povm[i, :, :] = S @ Wi @ S
    return [qt.Qobj(e) for e in povm]

# Cell
def random_haar_effect(d, k=None, n=1, real=False):
    r"""
    Generates a Haar distributed random POVM effect of Hilbert space dimension $d$, as if it were part of a POVM of $k$ elements, with mixedness $n$.
    """
    k = k if type(k) != type(None) else (d**2 if not real else int(d*(d+1)/2))
    X = random_ginibre(d, n, real=real)
    W = X @ X.conjugate().T
    Y = random_ginibre(d, (k-1)*n, real=real)
    S = W + Y @ Y.conjugate().T
    S = sc.linalg.fractional_matrix_power(S, -1/2)
    return qt.Qobj(S @ W @ S.conjugate().T)