"""[Qiskit](https://github.com/Qiskit/qiskit) based quantum circuit backend."""

# AUTOGENERATED! DO NOT EDIT! File to edit: ../../../src/platform/backends/circuits_qiskit.ipynb.

# %% auto #0
__all__ = ['get_number_of_gate_params', 'instruction_name_to_qiskit_gate', 'get_target_control_qubits', 'CircuitsQiskitBackend']

# %% ../../../src/platform/backends/circuits_qiskit.ipynb #d8d25f86-82c1-482e-997e-cb08141fbcf5
from ...imports import *
from .base_backend import BaseBackend
from ..circuits_instructions import CircuitInstructions
from ..tokenizer.base_tokenizer import Vocabulary
from ...utils.config_loader import get_obj_from_str

import qiskit.circuit.library as ql
import qiskit.quantum_info as qi

from qiskit import QuantumCircuit, transpile
from qiskit.circuit.gate import Gate

# %% ../../../src/platform/backends/circuits_qiskit.ipynb #e0e6cce5-01ed-4842-bb24-323776cdae99
def get_number_of_gate_params(gate_cls: type[Gate]) -> int:
    # python: gives you the number of any arguments BEFORE *args, minus the ones that have a default, -1 for self parameter of classes
    return gate_cls.__init__.__code__.co_argcount - len(gate_cls.__init__.__defaults__) - 1

# %% ../../../src/platform/backends/circuits_qiskit.ipynb #6e7415a3-7e89-4efc-9a9e-80f9d8d637fe
def instruction_name_to_qiskit_gate(name: str) -> Gate:
    match name:
        case "swap": name = "Swap"
        case "cp":   name = "CPhase"
        case _:      name = name.upper()
        
    return get_obj_from_str(f"qiskit.circuit.library.standard_gates.{name}Gate")

# %% ../../../src/platform/backends/circuits_qiskit.ipynb #2486f09e-78e3-4dec-b9a3-21d22dfdcc4d
def get_target_control_qubits(qc: QuantumCircuit, gate: Gate) -> Tuple[List[int], List[int]]:
    """Get the target and control qubits of a Qiskit `Gate` of a `QuantumCircuit`."""
    
    acts_on_cnt = gate.operation.num_qubits
    acts_on     = [qc.find_bit(qubit).index for qubit in gate.qubits]      # order: (*control_qubits, *target_qubits)

    assert acts_on_cnt == len(acts_on), "error in: acts_on_cnt == len(acts_on)"

    num_ctrl_qubits = gate.operation.num_ctrl_qubits if hasattr(gate.operation, "num_ctrl_qubits") else 0 
    num_targ_qubits = acts_on_cnt - num_ctrl_qubits

    control_qubits, target_qubits = acts_on[:-num_targ_qubits], acts_on[-num_targ_qubits:]   
    return control_qubits, target_qubits  

# %% ../../../src/platform/backends/circuits_qiskit.ipynb #31d639ab-a606-4d19-a998-60eca46dd0a2
class CircuitsQiskitBackend(BaseBackend):

    BASIC_BACKEND_TYPE = type[QuantumCircuit]
    
    def backend_to_genqc(self, qc: QuantumCircuit, ignore_barriers: bool = True) -> CircuitInstructions:
        """Convert a given Qiskit `QuantumCircuit` to genQC `CircuitInstructions`."""

        if ignore_barriers:
            gates = []
            for gate in qc.data:
                if gate.operation.name != "barrier": 
                    gates.append(gate)
        else:
            gates = qc.data
        
        instructions = CircuitInstructions(tensor_shape=torch.Size([qc.num_qubits, len(gates)]))
        
        for gate in gates:          
            control_qubits, target_qubits = get_target_control_qubits(qc, gate)      

            #Correction to qiskit v1.3.1
            #`cp` is symmetric but qiskit uses
            #`cp(theta, control_qubit, target_qubit)` target and control
            #We make it only target connections like `swap` gates
            if gate.operation.name == "cp":
                target_qubits.extend(control_qubits)
                control_qubits = []
            
            instructions.add_instruction(gate.operation.name, control_qubits, target_qubits, gate.operation.params)

        return instructions

    def genqc_to_backend(self, 
                         instructions: CircuitInstructions, 
                         place_barriers: bool = True,
                         ignore_errors: bool = False,
                         place_error_placeholders: bool = False) -> QuantumCircuit:  
        """Convert given genQC `CircuitInstructions` to a Qiskit `QuantumCircuit`."""
        
        gate_classes = {name:instruction_name_to_qiskit_gate(name) for name in instructions.instruction_names_set}        
        qc           = QuantumCircuit(instructions.num_qubits)  
        
        for instruction in instructions.data:
            gate_cls         = gate_classes[instruction.name]      
            num_of_paramters = get_number_of_gate_params(gate_cls)   

            control_qubits, target_qubits = instruction.control_nodes, instruction.target_nodes
            params                        = instruction.params[:num_of_paramters]
            
            try:
                qc.append(gate_cls(*params), [*control_qubits, *target_qubits], [])        
            except Exception as err:
                if ignore_errors: continue
                elif place_error_placeholders:
                    qc.append(ql.UnitaryGate(np.eye(2**instructions.num_qubits), label="Err"), range(instructions.num_qubits))
                    # qc.append(ql.UnitaryGate(np.eye(2), label="Err"), [0])
                    continue
                raise err
            
            if place_barriers: qc.barrier()
    
        return qc 

    def get_unitary(self, qc: QuantumCircuit, remove_global_phase: bool = True) -> np.ndarray:
        """Return the unitary matrix of a `QuantumCircuit`."""
        U = qi.Operator(qc).to_matrix().astype(np.complex128)
        if remove_global_phase:
            U *= np.exp(-1j * qc.global_phase) 
        return U

    def schmidt_rank_vector(self, qc: Optional[QuantumCircuit] = None, densityMatrix: Optional[qi.DensityMatrix] = None) -> List[int]:   
        """Return the SRV of a `qi.DensityMatrix`."""

        if not exists(densityMatrix):
            densityMatrix = qi.DensityMatrix(qc)
        
        systems_cnt = len(densityMatrix.dims())   
        total_trace = set(range(systems_cnt))    
        rank_vector = []
        
        for i in range(systems_cnt): 
            trace = list(total_trace - {i})
            red_densityMatrix = qi.partial_trace(densityMatrix, trace)        
            # r = np.count_nonzero(np.linalg.eigvals(red_densityMatrix) > 1e-14) # was slower during testing   
            r = np.linalg.matrix_rank(red_densityMatrix, hermitian=True).item()       
            rank_vector.append(r)
        
        return rank_vector

    def optimize_circuit(self, 
                         qc: QuantumCircuit, 
                         vocabulary: Vocabulary, 
                         optimization_level: int = 1, 
                         silent: bool = True) -> QuantumCircuit:
        """Use `qiskit.compiler.transpile` to optimize a circuit."""
        
        if optimization_level == 0: 
            return qc
                
        while optimization_level > 0:
            try:
                qc_opt = transpile(qc, optimization_level=optimization_level, basis_gates=vocabulary)
                return qc_opt
                
            except Exception as er: 
                if not silent: print(er)
                pass 
    
            optimization_level -= 1
    
        return qc

    def rnd_circuit(self, 
                    num_of_qubits: int, 
                    num_of_gates:int, 
                    gate_pool: Union[Sequence[Gate], Sequence[str]], 
                    rng: np.random.Generator) -> QuantumCircuit:
        """Create a random `QuantumCircuit`."""

        qc = QuantumCircuit(num_of_qubits)    
        gate_indices = rng.choice(len(gate_pool), num_of_gates)

        gate_pool = list(gate_pool)
        if isinstance(gate_pool[0], str):
            gate_pool = [instruction_name_to_qiskit_gate(gate) for gate in gate_pool]  
        
        for gate_index in gate_indices:
            gate_cls         = gate_pool[gate_index]
            num_of_paramters = get_number_of_gate_params(gate_cls)         
            params           = rng.uniform(low=0, high=4.0*np.pi, size=num_of_paramters) if num_of_paramters > 0 else [] 
            
            gate = gate_cls(*params)      
            act_qubits = rng.choice(num_of_qubits, gate.num_qubits, replace=False)   # order: (*act_qubits)=(*control_qubits, *target_qubits)   
            qc.append(gate, [*act_qubits], [])
        
        return qc

    def randomize_params(self, qc: QuantumCircuit, rng: np.random.Generator) -> QuantumCircuit:
        """Randomize all parameters of a `QuantumCircuit`. This creates a new `QuantumCircuit` and therefore deletes global phase."""
        
        qc_new = QuantumCircuit(qc.num_qubits)
        
        for gate in qc.data:
            gate_cls                       = instruction_name_to_qiskit_gate(gate.operation.name)
            control_qubits, target_qubits  = get_target_control_qubits(qc, gate)
            params                         = rng.uniform(low=0, high=4.0*np.pi, size=len(gate.operation.params))
            
            qc_new.append(gate_cls(*params), [*control_qubits, *target_qubits], [])           
            
        return qc_new

    def draw(self, qc: QuantumCircuit, **kwargs) -> None:
        """Draw the given `QuantumCircuit` using Qiskit.""" 
        return qc.draw("mpl", **kwargs)
        # plt.show()
