from qiskit.transpiler.coupling import CouplingMap
from mirror_gates.pass_managers import Mirage, QiskitLevel3
from transpile_benchy.metrics.gate_counts import (
    DepthMetric,
    TotalMetric,
    TotalSwaps,
)
from qiskit.transpiler import PassManager
import numpy as np
from mirror_gates.logging import transpile_benchy_logger
from mirror_gates.sqiswap_decomposer import SiSwapDecomposePass
from qiskit.transpiler.passmanager import PassManager
from qiskit import Aer, transpile
from qiskit.quantum_info import state_fidelity
from qiskit_aer import AerSimulator

# Import from Qiskit Aer noise module
from qiskit_aer.noise import (
    NoiseModel,
    QuantumError,
    ReadoutError,
    pauli_error,
    depolarizing_error,
    thermal_relaxation_error,
    RelaxationNoisePass,
)
from qiskit.circuit import Delay
from qiskit.transpiler.passes import PadDelay, ASAPSchedule
from qiskit import QuantumCircuit

from transpile_benchy.metrics.gate_counts import (
    DepthMetric,
    TotalMetric,
    TotalSwaps,
)
from qiskit.circuit.library import iSwapGate
from qiskit.transpiler import CouplingMap
from mirror_gates.pass_managers import Mirage, QiskitLevel3
from mirror_gates.utilities import SubsMetric
from mirror_gates.logging import transpile_benchy_logger

N = 8
coupling_map = CouplingMap.from_line(8)

from transpile_benchy.library import CircuitLibrary

library = CircuitLibrary.from_txt("../../circuits/small_circuits.txt")

# new metric for simulating fidelity
from transpile_benchy.metrics.abc_metrics import MetricInterface
from qiskit.transpiler.basepasses import AnalysisPass
from mirror_gates.noisy_fidelity import get_noisy_fidelity
from qiskit.converters import dag_to_circuit
from mirror_gates.utilities import DoNothing

# TODO: we could move gate times and T1,T2 to be parameters
# currently, are just hardcoded at the top of the file


class FidelityAnalysis(AnalysisPass):
    def __init__(self, coupling_map, sqrt_iswap_basis):
        super().__init__()
        self.coupling_map = coupling_map
        self.sqrt_iswap_basis = sqrt_iswap_basis

    def run(self, dag):
        qc = dag_to_circuit(dag)
        fid, duration, _, expected_fid = get_noisy_fidelity(
            qc, self.coupling_map, sqrt_iswap_basis=self.sqrt_iswap_basis
        )
        self.property_set["noisy_fidelity"] = fid
        self.property_set["scheduled_duration"] = duration
        self.property_set["expected_fidelity"] = expected_fid
        return dag


class FidelityMetric(MetricInterface):
    """Calculate the noisy simulated fidelity."""

    def __init__(self, coupling_map):
        """Initialize the metric."""
        super().__init__(name="noisy_fidelity", pretty_name="Noisy Fidelity")
        self.use_geometric_mean = True
        self.coupling_map = coupling_map

    def _construct_pass(self, basis_gate):
        """Return the pass associated with this metric."""
        sqrt_iswap_basis = basis_gate == iSwapGate().power(1 / 2)
        return FidelityAnalysis(self.coupling_map, sqrt_iswap_basis)


class DurationMetric(MetricInterface):
    """Calculate the noisy simulated fidelity."""

    def __init__(self):
        """Initialize the metric."""
        super().__init__(name="scheduled_duration", pretty_name="Circuit Duration (ns)")
        self.use_geometric_mean = True

    def _construct_pass(self):
        """Return the pass associated with this metric."""
        return DoNothing()


class ExpectedFidelityMetric(MetricInterface):
    """Calculate the noisy simulated fidelity."""

    def __init__(self):
        """Initialize the metric."""
        super().__init__(name="expected_fidelity", pretty_name="Expected Fidelity")
        self.use_geometric_mean = True

    def _construct_pass(self):
        """Return the pass associated with this metric."""
        return DoNothing()

# XXX set consolidate to False
# this is allowed only because my pass manager will preserve consolidation
# see post_stage, I call fastconsolidate manually

# NOTE: use TotalSwaps to verify baselines have > 0 swaps
# otherwise, there is no room for improvement.
# we can include these if we want to show our methods will still work
# but somewhat trivial since we just append VF2Layout
metrics = [
    DepthMetric(consolidate=False),
    TotalMetric(consolidate=False),
    TotalSwaps(consolidate=False),
    SubsMetric(),
    FidelityMetric(coupling_map),
    DurationMetric(),
    ExpectedFidelityMetric(),
]

transpilers = [
    # QiskitLevel3(coupling_map, cx_basis=True),
    # Mirage(
    #     coupling_map,
    #     logger=transpile_benchy_logger,
    #     cx_basis=True,
    # ),
    QiskitLevel3(coupling_map, cx_basis=False),
    Mirage(
        coupling_map,
        logger=transpile_benchy_logger,
        cx_basis=False,
    ),
]

from transpile_benchy.benchmark import Benchmark

benchmark = Benchmark(
    transpilers=transpilers,
    circuit_library=library,
    metrics=metrics,
    logger=transpile_benchy_logger,
    num_runs=1,
)

benchmark.run()

print(benchmark)

benchmark.summary_statistics(transpilers[0], transpilers[1])

from transpile_benchy.render import plot_benchmark

plot_benchmark(
    benchmark,
    save=1,
    legend_show=1,
    filename="twolocal",
    color_override=[0, 3],
)