# Path: openfermioncirq/experiments/hfvqe/third_party/higham.py
# @np.vectorize
# def heaviside(x, bias=0) -> int:
#     """
#     Heaviside function Theta(x - bias)
# 
#     returns 1 if x >= bias else 0
# 
#     :param x: floating point number as input to heavisde
#     :param bias: shift on the heaviside function
#     :return: 1 or 0 int
#     """
#     indicator = 1 if x >= bias else 0
#     return indicator
# 
# def higham_polynomial(eigenvalues, shift):
#     """
#     Calculate the higham_polynomial
# 
#     :param eigenvalues:  vector of eigenvalues
#     :param shift: where to put the bias for the heaviside function
#     """
#     heaviside_indicator = np.asarray(heaviside(eigenvalues, bias=shift))
#     return heaviside_indicator.T.dot(eigenvalues - shift)
# 
# def higham_root(eigenvalues, target_trace, epsilon=1.0E-15):
#     """
#     Find the root of f(sigma) = sum_{j}Theta(l_{i} - sigma)(l_{i} - sigma) = T
# 
#     :param eigenvalues: ordered list of eigenvalues from least to greatest
#     :param target_trace: trace to maintain on new matrix
#     :param epsilon: precision on bisection linesearch
#     """
#     if target_trace < 0.0:
#         raise ValueError("Target trace needs to be a non-negative number")
# 
#     # when we want the trace to be zero
#     if np.isclose(target_trace, 0.0):
#         return eigenvalues[-1]
# 
#     # find top sigma
#     sigma = eigenvalues[-1]
#     while higham_polynomial(eigenvalues, sigma) < target_trace:
#         sigma -= eigenvalues[-1]
# 
#     sigma_low = sigma
#     sigma_high = eigenvalues[-1]
# 
#     while sigma_high - sigma_low >= epsilon:
#         midpoint = sigma_high - (sigma_high - sigma_low) / 2.0
#         if higham_polynomial(eigenvalues, midpoint) < target_trace:
#             sigma_high = midpoint
#         else:
#             sigma_low = midpoint
# 
#     return sigma_high
# 
# def map_to_tensor(mat):
#     if mat.ndim != 2:
#         raise TypeError("I only map matrices to rank-4 tensors with symmetric support")
#     dim = int(np.sqrt(mat.shape[0]))
#     tensor_form = np.zeros((dim, dim, dim, dim))
#     for p, q, r, s in product(range(dim), repeat=4):
#         tensor_form[p, q, r, s] = mat[p * dim + q, r * dim + s]
#     return tensor_form
# 
# def map_to_matrix(mat):
#     if mat.ndim != 4:
#         raise TypeError("I only map rank-4 tensors to matices with symmetric support")
#     dim = mat.shape[0]
#     matform = np.zeros((dim**2, dim**2))
#     for p, q, r, s in product(range(dim), repeat=4):
#         assert np.isclose(mat[p, q, r, s].imag, 0.0)
#         matform[p * dim + q, r * dim + s] = mat[p, q, r, s].real
#     return matform
# 
# def fixed_trace_positive_projection(bmat, target_trace):
#     """
#     Perform the positive projection with fixed trace
# 
#     :param bmat:  Symmetric matrix to perform positive projection on
#     :param target_trace:  What the trace should be
#     :return: new matrix that has the target trace and is positive semidefinite
#     """
#     bmat = np.asarray(bmat)
#     map_to_four_tensor = False
#     if bmat.ndim == 4:
#         bmat = map_to_matrix(bmat)
#         map_to_four_tensor = True
# 
#     # symmeterize bmat
#     if np.allclose(bmat - bmat.conj().T, np.zeros_like(bmat)):
#         bmat = 0.5 * (bmat + bmat.conj().T)
# 
#     w, v = np.linalg.eigh(bmat)
#     if np.all(w >= -1.0*float(1.0E-14)) and np.isclose(np.sum(w), target_trace):
#         purified_matrix = bmat
#     else:
#         sigma = higham_root(w, target_trace)
#         shifted_eigs = np.multiply(heaviside(w - sigma), (w - sigma))
#         purified_matrix = np.zeros_like(bmat)
#         for i in range(w.shape[0]):
#             purified_matrix += shifted_eigs[i] * v[:, [i]].dot(v[:, [i]].conj().T)
# 
#     if map_to_four_tensor:
#         purified_matrix = map_to_tensor(purified_matrix)
# 
#     return purified_matrix


from itertools import product
from openfermioncirq.experiments.hfvqe.third_party.higham import (heaviside,
                                                                  higham_polynomial,
                                                                  higham_root,
                                                                  map_to_tensor,
                                                                  map_to_matrix,
                                                                  fixed_trace_positive_projection)
import numpy as np
import pytest
    assert np.isclose(higham_polynomial(w, shift=sigma), target_trace)

    with pytest.raises(ValueError):
        higham_root(w, target_trace=-1)

    tw = higham_root(w, target_trace=0)
    assert np.isclose(tw, w[-1])


def test_matrix_2_tensor():
    dim = 10
    np.random.seed(42)
    mat = np.random.random((dim**2, dim**2))
    mat = 0.5 * (mat + mat.T)
    tensor = map_to_tensor(mat)
    for p, q, r, s in product(range(dim), repeat=4):
        assert np.isclose(tensor[p, q, r, s], mat[p * dim + q, r * dim + s])

    test_mat = map_to_matrix(tensor)
    assert np.allclose(test_mat, mat)

    with pytest.raises(TypeError):
        map_to_tensor(np.zeros((4, 4, 4, 4)))

    with pytest.raises(TypeError):
        map_to_matrix(np.zeros((4, 4)))


def test_reconstruction():
    dim = 20

# Path: openfermioncirq/experiments/hfvqe/third_party/higham.py
# @np.vectorize
# def heaviside(x, bias=0) -> int:
#     """
#     Heaviside function Theta(x - bias)
# 
#     returns 1 if x >= bias else 0
# 
#     :param x: floating point number as input to heavisde
#     :param bias: shift on the heaviside function
#     :return: 1 or 0 int
#     """
#     indicator = 1 if x >= bias else 0
#     return indicator
# 
# def higham_polynomial(eigenvalues, shift):
#     """
#     Calculate the higham_polynomial
# 
#     :param eigenvalues:  vector of eigenvalues
#     :param shift: where to put the bias for the heaviside function
#     """
#     heaviside_indicator = np.asarray(heaviside(eigenvalues, bias=shift))
#     return heaviside_indicator.T.dot(eigenvalues - shift)
# 
# def higham_root(eigenvalues, target_trace, epsilon=1.0E-15):
#     """
#     Find the root of f(sigma) = sum_{j}Theta(l_{i} - sigma)(l_{i} - sigma) = T
# 
#     :param eigenvalues: ordered list of eigenvalues from least to greatest
#     :param target_trace: trace to maintain on new matrix
#     :param epsilon: precision on bisection linesearch
#     """
#     if target_trace < 0.0:
#         raise ValueError("Target trace needs to be a non-negative number")
# 
#     # when we want the trace to be zero
#     if np.isclose(target_trace, 0.0):
#         return eigenvalues[-1]
# 
#     # find top sigma
#     sigma = eigenvalues[-1]
#     while higham_polynomial(eigenvalues, sigma) < target_trace:
#         sigma -= eigenvalues[-1]
# 
#     sigma_low = sigma
#     sigma_high = eigenvalues[-1]
# 
#     while sigma_high - sigma_low >= epsilon:
#         midpoint = sigma_high - (sigma_high - sigma_low) / 2.0
#         if higham_polynomial(eigenvalues, midpoint) < target_trace:
#             sigma_high = midpoint
#         else:
#             sigma_low = midpoint
# 
#     return sigma_high
# 
# def map_to_tensor(mat):
#     if mat.ndim != 2:
#         raise TypeError("I only map matrices to rank-4 tensors with symmetric support")
#     dim = int(np.sqrt(mat.shape[0]))
#     tensor_form = np.zeros((dim, dim, dim, dim))
#     for p, q, r, s in product(range(dim), repeat=4):
#         tensor_form[p, q, r, s] = mat[p * dim + q, r * dim + s]
#     return tensor_form
# 
# def map_to_matrix(mat):
#     if mat.ndim != 4:
#         raise TypeError("I only map rank-4 tensors to matices with symmetric support")
#     dim = mat.shape[0]
#     matform = np.zeros((dim**2, dim**2))
#     for p, q, r, s in product(range(dim), repeat=4):
#         assert np.isclose(mat[p, q, r, s].imag, 0.0)
#         matform[p * dim + q, r * dim + s] = mat[p, q, r, s].real
#     return matform
# 
# def fixed_trace_positive_projection(bmat, target_trace):
#     """
#     Perform the positive projection with fixed trace
# 
#     :param bmat:  Symmetric matrix to perform positive projection on
#     :param target_trace:  What the trace should be
#     :return: new matrix that has the target trace and is positive semidefinite
#     """
#     bmat = np.asarray(bmat)
#     map_to_four_tensor = False
#     if bmat.ndim == 4:
#         bmat = map_to_matrix(bmat)
#         map_to_four_tensor = True
# 
#     # symmeterize bmat
#     if np.allclose(bmat - bmat.conj().T, np.zeros_like(bmat)):
#         bmat = 0.5 * (bmat + bmat.conj().T)
# 
#     w, v = np.linalg.eigh(bmat)
#     if np.all(w >= -1.0*float(1.0E-14)) and np.isclose(np.sum(w), target_trace):
#         purified_matrix = bmat
#     else:
#         sigma = higham_root(w, target_trace)
#         shifted_eigs = np.multiply(heaviside(w - sigma), (w - sigma))
#         purified_matrix = np.zeros_like(bmat)
#         for i in range(w.shape[0]):
#             purified_matrix += shifted_eigs[i] * v[:, [i]].dot(v[:, [i]].conj().T)
# 
#     if map_to_four_tensor:
#         purified_matrix = map_to_tensor(purified_matrix)
# 
#     return purified_matrix


# Path: openfermioncirq/experiments/hfvqe/third_party/higham_test.py
from itertools import product
from openfermioncirq.experiments.hfvqe.third_party.higham import (heaviside,
                                                                  higham_polynomial,
                                                                  higham_root,
                                                                  map_to_tensor,
                                                                  map_to_matrix,
                                                                  fixed_trace_positive_projection)
import numpy as np
import pytest

    assert np.isclose(higham_polynomial(w, shift=sigma), target_trace)

    with pytest.raises(ValueError):
        higham_root(w, target_trace=-1)

    tw = higham_root(w, target_trace=0)
    assert np.isclose(tw, w[-1])


def test_matrix_2_tensor():
    dim = 10
    np.random.seed(42)
    mat = np.random.random((dim**2, dim**2))
    mat = 0.5 * (mat + mat.T)
    tensor = map_to_tensor(mat)
    for p, q, r, s in product(range(dim), repeat=4):
        assert np.isclose(tensor[p, q, r, s], mat[p * dim + q, r * dim + s])

    test_mat = map_to_matrix(tensor)
    assert np.allclose(test_mat, mat)

    with pytest.raises(TypeError):
        map_to_tensor(np.zeros((4, 4, 4, 4)))

    with pytest.raises(TypeError):
        map_to_matrix(np.zeros((4, 4)))


def test_reconstruction():
    dim = 20

    np.random.seed(42)