import networkx as nx
import numpy as np
import pytest

import pyquil.simulation.matrices as qmats
from pyquil import Program
from pyquil.api import QuantumComputer
from pyquil.device import NxDevice
from pyquil.experiment import ExperimentSetting, Experiment, zeros_state
from pyquil.gates import CNOT, H, I, MEASURE, PHASE, RX, RY, RZ, X
from pyquil.operator_estimation import measure_observables
from pyquil.paulis import sI, sX, sY, sZ
from pyquil.pyqvm import PyQVM
from pyquil.simulation._reference import ReferenceDensitySimulator, _is_valid_quantum_state
from pyquil.simulation.tools import lifted_gate_matrix
from pyquil.tests.utils import DummyCompiler


def test_qaoa_density():
    wf_true = [
        0.00167784 + 1.00210180e-05 * 1j,
        0.50000000 - 4.99997185e-01 * 1j,
        0.50000000 - 4.99997185e-01 * 1j,
        0.00167784 + 1.00210180e-05 * 1j,
    ]
    wf_true = np.reshape(np.array(wf_true), (4, 1))
    rho_true = np.dot(wf_true, np.conj(wf_true).T)
    prog = Program()
    prog.inst(
        [
            RY(np.pi / 2, 0),
            RX(np.pi, 0),
            RY(np.pi / 2, 1),
            RX(np.pi, 1),
            CNOT(0, 1),
            RX(-np.pi / 2, 1),
            RY(4.71572463191, 1),
            RX(np.pi / 2, 1),
            CNOT(0, 1),
            RX(-2 * 2.74973750579, 0),
            RX(-2 * 2.74973750579, 1),
        ]
    )

    qam = PyQVM(n_qubits=2, quantum_simulator_type=ReferenceDensitySimulator).execute(prog)
    rho = qam.wf_simulator.density
    np.testing.assert_allclose(rho_true, rho, atol=1e-8)


def test_larger_qaoa_density():
    prog = Program(
        H(0),
        H(1),
        H(2),
        H(3),
        X(0),
        PHASE(0.3928244130249029, 0),
        X(0),
        PHASE(0.3928244130249029, 0),
        CNOT(0, 1),
        RZ(0.78564882604980579, 1),
        CNOT(0, 1),
        X(0),
        PHASE(0.3928244130249029, 0),
        X(0),
        PHASE(0.3928244130249029, 0),
        CNOT(0, 3),
        RZ(0.78564882604980579, 3),
        CNOT(0, 3),
        X(0),
        PHASE(0.3928244130249029, 0),
        X(0),
        PHASE(0.3928244130249029, 0),
        CNOT(1, 2),
        RZ(0.78564882604980579, 2),
        CNOT(1, 2),
        X(0),
        PHASE(0.3928244130249029, 0),
        X(0),
        PHASE(0.3928244130249029, 0),
        CNOT(2, 3),
        RZ(0.78564882604980579, 3),
        CNOT(2, 3),
        H(0),
        RZ(-0.77868204192240842, 0),
        H(0),
        H(1),
        RZ(-0.77868204192240842, 1),
        H(1),
        H(2),
        RZ(-0.77868204192240842, 2),
        H(2),
        H(3),
        RZ(-0.77868204192240842, 3),
        H(3),
    )

    qam = PyQVM(n_qubits=4, quantum_simulator_type=ReferenceDensitySimulator).execute(prog)
    rho_test = qam.wf_simulator.density
    wf_true = np.array(
        [
            8.43771693e-05 - 0.1233845 * 1j,
            -1.24927731e-01 + 0.00329533 * 1j,
            -1.24927731e-01 + 0.00329533 * 1j,
            -2.50040954e-01 + 0.12661547 * 1j,
            -1.24927731e-01 + 0.00329533 * 1j,
            -4.99915497e-01 - 0.12363516 * 1j,
            -2.50040954e-01 + 0.12661547 * 1j,
            -1.24927731e-01 + 0.00329533 * 1j,
            -1.24927731e-01 + 0.00329533 * 1j,
            -2.50040954e-01 + 0.12661547 * 1j,
            -4.99915497e-01 - 0.12363516 * 1j,
            -1.24927731e-01 + 0.00329533 * 1j,
            -2.50040954e-01 + 0.12661547 * 1j,
            -1.24927731e-01 + 0.00329533 * 1j,
            -1.24927731e-01 + 0.00329533 * 1j,
            8.43771693e-05 - 0.1233845 * 1j,
        ]
    )

    wf_true = np.reshape(wf_true, (2 ** 4, 1))
    rho_true = np.dot(wf_true, np.conj(wf_true).T)
    np.testing.assert_allclose(rho_true, rho_test, atol=1e-8)


def _random_1q_density():
    state = np.random.random(2) + 1j * np.random.random()
    normalization = np.conj(state).T.dot(state)
    state /= np.sqrt(normalization)
    state = state.reshape((-1, 1))

    rho = state.dot(np.conj(state).T)
    assert np.isclose(np.trace(rho), 1.0)
    assert np.allclose(rho, np.conj(rho).T)
    return rho


def test_kraus_application_bitflip():
    p = 0.372
    qam = PyQVM(
        n_qubits=1,
        quantum_simulator_type=ReferenceDensitySimulator,
        post_gate_noise_probabilities={"bit_flip": p},
    )
    initial_density = _random_1q_density()
    qam.wf_simulator.density = initial_density
    qam.execute(Program(I(0)))
    final_density = (1 - p) * initial_density + p * qmats.X.dot(initial_density).dot(qmats.X)
    np.testing.assert_allclose(final_density, qam.wf_simulator.density)


def test_kraus_application_phaseflip():
    p = 0.372
    qam = PyQVM(
        n_qubits=1,
        quantum_simulator_type=ReferenceDensitySimulator,
        post_gate_noise_probabilities={"phase_flip": p},
    )
    initial_density = _random_1q_density()
    qam.wf_simulator.density = initial_density
    qam.execute(Program(I(0)))

    final_density = (1 - p) * initial_density + p * qmats.Z.dot(initial_density).dot(qmats.Z)
    np.testing.assert_allclose(final_density, qam.wf_simulator.density)


def test_kraus_application_bitphaseflip():
    p = 0.372
    qam = PyQVM(
        n_qubits=1,
        quantum_simulator_type=ReferenceDensitySimulator,
        post_gate_noise_probabilities={"bitphase_flip": p},
    )
    initial_density = _random_1q_density()