import networkx as nx
import numpy as np
import pandas as pd
from qiskit import QuantumRegister, ClassicalRegister, execute
from qiskit.providers.aer import AerSimulator
from functools import partial
from pathlib import Path
import copy
import pytest

from qiskit.quantum_info import Statevector

from QAOAKit import (
    opt_angles_for_graph,
    get_fixed_angles,
    get_graph_id,
    get_graph_from_id,
    angles_to_qaoa_format,
    beta_to_qaoa_format,
    gamma_to_qaoa_format,
    angles_to_qiskit_format,
    angles_to_qtensor_format,
    get_3_reg_dataset_table,
    get_3_reg_dataset_table_row,
    get_full_qaoa_dataset_table_row,
    get_full_qaoa_dataset_table,
    qaoa_maxcut_energy,
    angles_from_qiskit_format,
)
from QAOAKit.utils import (
    obj_from_statevector,
    maxcut_obj,
    isomorphic,
    load_weights_into_dataframe,
    load_weighted_results_into_dataframe,
    get_adjacency_matrix,
)
from QAOAKit.qaoa import get_maxcut_qaoa_circuit, get_maxcut_qaoa_qiskit_circuit
from QAOAKit.examples_utils import get_20_node_erdos_renyi_graphs

from qiskit_optimization import QuadraticProgram
from qiskit.algorithms.minimum_eigen_solvers.qaoa import QAOAAnsatz

test_utils_folder = Path(__file__).parent


def test_retrieval():
    for n in range(3, 10):
        G = nx.complete_graph(n)
        for p in range(1, 4):
            angles = opt_angles_for_graph(G, p)
            assert len(angles["beta"]) == p
            assert len(angles["gamma"]) == p
            assert isinstance(
                get_full_qaoa_dataset_table_row(G, p)[
                    ["C_opt", "graph_id", "C_{true opt}"]
                ],
                pd.Series,
            )


def test_fixed_angle_retrieval():
    for d in range(3, 5):
        for p in range(1, 4):
            angles = get_fixed_angles(d, p)
            assert type(angles) is dict
            assert len(angles["beta"]) == p
            assert len(angles["gamma"]) == p


def test_tables_consistency():
    p = 1
    G1 = nx.complete_graph(5)
    row = get_full_qaoa_dataset_table_row(G1, p)
    G2 = get_graph_from_id(row["graph_id"], 5)
    assert isomorphic(G1, G2)


def test_angle_conversion():
    full_qaoa_dataset_table = get_full_qaoa_dataset_table()
    for n_qubits in [3, 5]:
        for p in [2, 3]:
            df = full_qaoa_dataset_table.reset_index()
            df = df[(df["n"] == n_qubits) & (df["p_max"] == p)]
            for _, row in df.iterrows():
                obj = partial(maxcut_obj, w=get_adjacency_matrix(row["G"]))
                opt_cut = row["C_opt"]
                angles = angles_to_qaoa_format(
                    opt_angles_for_graph(row["G"], row["p_max"])
                )
                assert np.allclose(angles["beta"], beta_to_qaoa_format(row["beta"]))
                assert np.allclose(angles["gamma"], gamma_to_qaoa_format(row["gamma"]))
                qc = get_maxcut_qaoa_circuit(row["G"], angles["beta"], angles["gamma"])
                backend = AerSimulator(method="statevector")
                res = backend.run(qc).result()
                sv = res.get_statevector()
                obj_val = obj_from_statevector(sv, obj)
                assert np.isclose(opt_cut, obj_val)


def test_qaoa_maxcut_energy():
    full_qaoa_dataset_table = get_full_qaoa_dataset_table()
    for n_qubits in [3, 4]:
        p = 3
        df = full_qaoa_dataset_table.reset_index()
        df = df[(df["n"] == n_qubits) & (df["p_max"] == p)]
        for _, row in df.iterrows():
            assert np.isclose(
                row["C_opt"],
                qaoa_maxcut_energy(
                    row["G"],
                    beta_to_qaoa_format(row["beta"]),
                    gamma_to_qaoa_format(row["gamma"]),
                ),
            )


def test_qiskit_angle_conversion():
    full_qaoa_dataset_table = get_full_qaoa_dataset_table()
    for n_qubits in [3, 4]:
        p = 3
        df = full_qaoa_dataset_table.reset_index()
        df = df[(df["n"] == n_qubits) & (df["p_max"] == p)]
        for _, row in df.iterrows():
            angles = opt_angles_for_graph(row["G"], row["p_max"])
            G = row["G"]
            qc, C, offset = get_maxcut_qaoa_qiskit_circuit(
                G, p, angles_to_qiskit_format(angles)
            )
            backend = AerSimulator(method="statevector")
            sv = Statevector(backend.run(qc).result().get_statevector())
            obj_val = -(sv.expectation_value(C) + offset)
            opt_cut = row["C_opt"]
            assert np.isclose(opt_cut, obj_val)


def test_qiskit_qaoa_circuit():
    full_qaoa_dataset_table = get_full_qaoa_dataset_table()
    for n_qubits in [3, 4]:
        p = 1
        df = full_qaoa_dataset_table.reset_index()
        df = df[(df["n"] == n_qubits) & (df["p_max"] == p)]
        for _, row in df.iterrows():
            angles1 = opt_angles_for_graph(row["G"], row["p_max"])
            G = row["G"]
            qc1, C, offset = get_maxcut_qaoa_qiskit_circuit(
                G, p, angles_to_qiskit_format(angles1)
            )
            backend = AerSimulator(method="statevector")
            sv1 = Statevector(backend.run(qc1).result().get_statevector())
            angles2 = angles_to_qaoa_format(
                opt_angles_for_graph(row["G"], row["p_max"])
            )
            qc2 = get_maxcut_qaoa_circuit(row["G"], angles2["beta"], angles2["gamma"])
            sv2 = Statevector(backend.run(qc2).result().get_statevector())
            assert sv1.equiv(sv2)


def test_weighted_qaoa():
    G = nx.petersen_graph()
    np.random.seed(42)
    for (u, v) in G.edges():
        G[u][v]["weight"] = np.random.uniform(low=1, high=5)
    assert nx.is_weighted(G)
    backend = AerSimulator(method="statevector")
    p = 3
    angles = {
        "beta": np.random.uniform(low=0, high=np.pi / 2, size=3),
        "gamma": np.random.uniform(low=0, high=np.pi, size=3),
    }
    qc1, C, offset = get_maxcut_qaoa_qiskit_circuit(
        G, p, angles_to_qiskit_format(angles)
    )
    sv1 = Statevector(backend.run(qc1).result().get_statevector())
    angles2 = angles_to_qaoa_format(angles)
    qc2 = get_maxcut_qaoa_circuit(G, angles2["beta"], angles2["gamma"])
    sv2 = Statevector(backend.run(qc2).result().get_statevector())
    assert sv1.equiv(sv2)


def test_load_weighted_results():
    p = 2
    n = 5

    folder_path = Path(test_utils_folder, f"../data/weighted_angle_dat/p={p}/")
    df_weights = load_weights_into_dataframe(folder_path)
    df = load_weighted_results_into_dataframe(folder_path, p, n, df_weights)
    assert np.all(
        np.isclose(
            df.head(100).apply(
                lambda row: qaoa_maxcut_energy(
                    row["G"],
                    beta_to_qaoa_format(row["beta"]),
                    gamma_to_qaoa_format(row["gamma"]),
                ),
                axis=1,
            ),
            df.head(100)["C_opt"],
        )
    )


def test_qtensor_angle_conversion():
    qtensor = pytest.importorskip("qtensor")
    full_qaoa_dataset_table = get_full_qaoa_dataset_table()
    for n_qubits in [3, 4]:
        p = 3
        df = full_qaoa_dataset_table.reset_index()
        df = df[(df["n"] == n_qubits) & (df["p_max"] == p)]
        for _, row in df.iterrows():
            angles = angles_to_qtensor_format(
                opt_angles_for_graph(row["G"], row["p_max"])
            )
            G = row["G"]
            obj_val = qtensor.QAOA_energy(G, angles["gamma"], angles["beta"])
            opt_cut = row["C_opt"]
            assert np.isclose(opt_cut, obj_val)


def test_from_qiskit_conversion():
    angles = {
        "beta": np.random.uniform(low=0, high=np.pi / 2, size=3),
        "gamma": np.random.uniform(low=0, high=np.pi, size=3),
    }
    angles2 = angles_from_qiskit_format(angles_to_qiskit_format(copy.deepcopy(angles)))
    assert np.all(np.isclose(angles["beta"], angles2["beta"]))
    assert np.all(np.isclose(angles["gamma"], angles2["gamma"]))


def test_3_reg_table():
    df = get_3_reg_dataset_table().reset_index()
    assert len(df) >= 1000

    # test that the angles are correct
    for _, row in df[(df["p_max"] == 1) | (df["p_max"] == 2)].head(50).iterrows():
        G = row["G"]
        angles = angles_to_qaoa_format({"beta": row["beta"], "gamma": row["gamma"]})
        assert np.isclose(
            qaoa_maxcut_energy(G, angles["beta"], angles["gamma"]), row["C_opt"]
        )

    # test that the optima match the full qaoa table
    for _, row in df[(df["p_max"] == 1) | (df["p_max"] == 2)].head(50).iterrows():
        G = row["G"]
        angles = angles_to_qaoa_format({"beta": row["beta"], "gamma": row["gamma"]})
        assert np.isclose(
            qaoa_maxcut_energy(G, angles["beta"], angles["gamma"]), row["C_opt"]
        )


def test_fixed_angles_3_reg():
    df = get_3_reg_dataset_table().sample(n=10).reset_index()

    for _, row in df.iterrows():
        angles = angles_to_qaoa_format(get_fixed_angles(3, row["p_max"]))
        assert np.isclose(
            qaoa_maxcut_energy(row["G"], angles["beta"], angles["gamma"]),
            row["C_fixed"],
            rtol=1e-4,
        )


def test_get_opt_angles_large_3_reg():
    for s in range(10):
        G = nx.random_regular_graph(3, 10, seed=s)
        for p in [1, 2]:
            angles = angles_to_qaoa_format(opt_angles_for_graph(G, p))
            row = get_3_reg_dataset_table_row(G, p)
            assert np.isclose(
                row["C_opt"], qaoa_maxcut_energy(G, angles["beta"], angles["gamma"])
            )


def test_example_in_README():
    # build graph
    G = nx.star_graph(5)
    # grab optimal angles
    p = 3
    angles = angles_to_qaoa_format(opt_angles_for_graph(G, p))
    # build circuit
    qc = get_maxcut_qaoa_circuit(G, angles["beta"], angles["gamma"])
    qc.measure_all()
    # run circuit
    backend = AerSimulator()
    assert isinstance(backend.run(qc).result().get_counts(), dict)


def test_no_save_state():
    n = 6
    qc = get_maxcut_qaoa_circuit(nx.star_graph(n - 1), [0.1], [0.1], save_state=False)
    creg = ClassicalRegister(n)
    creg_ancilla = ClassicalRegister(1)
    qreg_ancilla = QuantumRegister(1)
    qc.add_register(creg)
    qc.add_register(creg_ancilla)
    qc.add_register(qreg_ancilla)
    qc.h(qreg_ancilla)
    for i in range(n):
        qc.cx(qreg_ancilla, i)
    qc.h(qreg_ancilla)

    qc.measure(range(6), creg)
    qc.measure(qreg_ancilla, creg_ancilla)

    counts = (
        execute(qc, AerSimulator(), basis_gates=["u1", "u2", "u3", "cx"])
        .result()
        .get_counts()
    )
    assert isinstance(counts, dict)


def test_examples_erdos_renyi():
    df = get_20_node_erdos_renyi_graphs()
    assert isinstance(df, pd.DataFrame)
    assert len(df) == 30
