import logging

logger = logging.getLogger()
logger.setLevel(logging.INFO)

import numpy as np

%matplotlib widget

from slam.utils.gates.custom_gates import ConversionGainGate
from slam.utils.polytopes.polytope_wrap import (
    monodromy_range_from_target,
    coverage_to_haar_expectation,
)
from slam.basis import CircuitTemplate, MixedOrderBasisCircuitTemplate
from slam.utils.visualize import (
    unitary_to_weyl,
    unitary_2dlist_weyl,
    coordinate_2dlist_weyl,
)
from slam.utils.gates.custom_gates import CustomCostGate
from qiskit.circuit.library import CXGate, SwapGate
from weylchamber import c1c2c3
from slam.utils.gates.snail_death_gate import SpeedLimitedGate
from tqdm import tqdm
import matplotlib.pyplot as plt

# # verifying that relative phase doesn't change 2Q gate location
# unitary = [ConversionGainGate(0, 0, p*0.5*np.pi, (1-p)*0.4*np.pi) for p in np.linspace(0,1,16)]
# print([u.cost() for u in unitary])
# unitary_to_weyl(*unitary);

# the first version of this file was exploring the conversion gain hamiltonian, let's also explore Heavy-Fluxonion, and Flux Qubit

from slam.utils.gates.duraton_scaling import cost_scaling
from slam.utils.gates.bare_candidates import (
    build_gates,
    get_group_name,
    plot_eharr,
    collect_data,
)
from slam.utils.gates.winner_selection import pick_winner

unitary_list, coordinate_list = build_gates()
coordinate_2dlist_weyl(*coordinate_list)

# # monkey patch, iterate through every gate and save its gc/gg mirror to base costs data set
# from slam.scripts.gate_exploration.bgatev2script import filename
# import h5py
# with h5py.File(filename, "a") as hf:
#     g = hf.require_group("bare_cost")
#     for gate, data in g.items():
#         gc, gg = data[0][2], data[0][3]
#         if gc == gg:
#             continue

#         # create mirror gate
#         mirror = ConversionGainGate(0,0, gg, gc)
#         #save a copy of the data with new gate
#         mirror_data = [mirror.params, list(data[1])] # scores stay the same
#         g.create_dataset(str(mirror), data=mirror_data)

# ## %%
# TODO put this in its own file
# """use 2 different 2Q gates and see minimum cost to build SWAP
# just see if a solution exists for 2Q gates only using 2 instances of each gate then stop (rather than using coverage set)"""
# from slam.utils.polytope_wrap import get_polytope_from_circuit, unitary_to_monodromy_coordinate
# from qiskit import QuantumCircuit
# from qiskit.circuit.library import SwapGate
# import itertools

# # interested in a hetero basis of 2 applications to build a SWAP target
# target_coords = unitary_to_monodromy_coordinate(SwapGate().to_matrix())

# #taking product over search space to build pairs of gates
# # a - might want to remove constraint of left side of x-axis only (I'm not sure if this makes a difference or not)
# # b - may need to increase density of gates to sample from
# unitary_pairs = list(itertools.product(unitary_list, repeat=2))

# #if find a solution, save it
# sol = []

# for gate1, gate2 in tqdm(unitary_pairs):

#     #basis can just be a circuit with 2Q gates
#     basis = QuantumCircuit(2)
#     basis.append(gate1, [0,1])
#     basis.append(gate2, [0,1])

#     circuit_polytope = get_polytope_from_circuit(basis)

#     if circuit_polytope.has_element(target_coords):
#         sol.append((gate1, gate2))
#         logging.info((gate1, gate2))

# print(sol)
# XXX

# template = MixedOrderBasisCircuitTemplate(base_gates=[CustomCostGate(CXGate(), "CX")], chatty_build=1)
# coverage_to_haar_expectation(template.coverage, chatty=1)

# set logger level to critical to suppress warnings
logger.setLevel(logging.CRITICAL)
# TODO I need to recompute everything, but to save time just do iswap family
filtered_unitary_list = list(filter(lambda x: x.params[2] == 0, unitary_list))
collect_data(filtered_unitary_list)

method = "linear"
duration_1q = 0.1
group_name = get_group_name(method, duration_1q)
cost_scaling(speed_method=method, duration_1q=duration_1q)
plot_eharr(group_name, metric=0)
pick_winner(group_name, metric=2)

method = "hardware"
duration_1q = 0.1
group_name = get_group_name(method, duration_1q)
cost_scaling(speed_method=method, duration_1q=duration_1q)
plot_eharr(group_name)
pick_winner(group_name)

method = "mid"
duration_1q = 0.1
group_name = get_group_name(method, duration_1q)
cost_scaling(speed_method=method, duration_1q=duration_1q)
# plot_eharr(group_name, metric=0)
pick_winner(group_name, metric=0)

method = "squared"
duration_1q = 0.1
group_name = get_group_name(method, duration_1q)
cost_scaling(speed_method=method, duration_1q=duration_1q)
plot_eharr(group_name, metric=0)
pick_winner(group_name, metric=0)