"""Class to tokenize quantum circuits. Encode and decode quantum circuits into and from tensor representations."""

# AUTOGENERATED! DO NOT EDIT! File to edit: ../../../src/platform/tokenizer/circuits_tokenizer.ipynb.

# %% auto #0
__all__ = ['CircuitTokenizer']

# %% ../../../src/platform/tokenizer/circuits_tokenizer.ipynb #06272f6f-b4e3-4504-a90a-feebbf6ad821
from ...imports import *
from .base_tokenizer import BaseTokenizer, Vocabulary
from ..circuits_instructions import CircuitInstructions

# %% ../../../src/platform/tokenizer/circuits_tokenizer.ipynb #36032308-bd0e-4409-9db0-9d89fc258e5a
class CircuitTokenizer(BaseTokenizer):

    def __init__(self, vocabulary: Vocabulary, sign_labels: Optional[dict[str, int]] = None) -> None:   
        if 0 in vocabulary.values():        
            print(f"[WARNING]: The value 0 is reserved for background tokens, i.e. qubit time position which are not effected by gates.")
            print(f"[WARNING]: Automatically incrementing all vocabulary values by one ...")
            vocabulary = {k:v+1 for k,v in vocabulary.items()}
            assert 0 not in vocabulary.values()
        
        super().__init__(vocabulary)
        self.sign_labels = default(sign_labels, {"control_nodes": -1, "target_nodes": +1})
        
    def tokenize(self, instructions: CircuitInstructions) -> torch.Tensor | Tuple[torch.Tensor, torch.Tensor]:
        """Convert given instructions to a tensor. Identical to `CircuitTokenizer.encode`."""
        return self.encode(instructions=instructions)
    
    def encode(self, 
               instructions: CircuitInstructions, 
               max_gates: Optional[int] = None, 
               return_params_tensor: bool = True, 
               params_4pi_normalization: bool = True, 
               randomize_params: bool = False) -> torch.Tensor | Tuple[torch.Tensor, torch.Tensor]:
        """Convert given `CircuitInstructions` to a `torch.Tensor`."""

        assert len(instructions.tensor_shape) == 2
        num_of_qubits, time = instructions.tensor_shape
        max_gates           = default(max_gates, time)
        
        tensor = torch.zeros((num_of_qubits, max_gates), dtype=torch.int32) 
        params = []

        for t, instruction in zip(range(max_gates), instructions.data):  # this way we limit the number of gates even if there are more instructions
            
            if instruction.name not in self.vocabulary: raise Warning(f"`{instruction.name}` not in vocabulary.")
            
            params.append(instruction.params)
            
            gate_id = self.vocabulary[instruction.name]  
                             
            control_qubits, target_qubits = instruction.control_nodes, instruction.target_nodes
                           
            for bit in control_qubits:
                tensor[bit, t] = gate_id * self.sign_labels["control_nodes"]
            
            for bit in target_qubits:
                tensor[bit, t] = gate_id * self.sign_labels["target_nodes"]

        if return_params_tensor: 
            num_of_max_params = max([0] + [len(para) for para in params])
            params_tensor     = torch.zeros((num_of_max_params, max_gates), dtype=torch.float32)
            
            for t, para in enumerate(params):
                para = torch.tensor(para)

                if randomize_params:
                    para = 2.0*torch.rand_like(para) - 1.0   # rnd [-1, 1]
                
                elif params_4pi_normalization:        
                    para = para % (4.0*np.pi)              # limit to [0, 4pi]
                    para = (para-2.0*np.pi) / (2.0*np.pi)  # [0, 4pi] to [-1, +1] 

                params_tensor[:len(para), t] = para
                         
            return tensor, params_tensor       
        return tensor
        
    def decode(self, 
               tensor: torch.Tensor, 
               params_tensor: Optional[torch.Tensor] = None, 
               params_4pi_normalization: bool = True,
               ignore_errors: bool = False,
               place_error_placeholders: bool = False) -> CircuitInstructions:
        """Convert a given `torch.Tensor` to `CircuitInstructions`."""
     
        assert tensor.dim() == 2, f"{tensor.shape=}"
        num_of_qubits, time = tensor.shape
        
        instructions = CircuitInstructions(tensor_shape=tensor.shape)
        
        for t in range(time):         
            enc_time_slice = tensor[:, t] # contains all bits at time t   

            _gate_placed = False
            
            for gate_index, gate in self.vocabulary_inverse.items():   
            
                target_nodes  = (enc_time_slice == (self.sign_labels["target_nodes"]  * gate_index)).nonzero(as_tuple=True)[0]
                control_nodes = (enc_time_slice == (self.sign_labels["control_nodes"] * gate_index)).nonzero(as_tuple=True)[0]

                _gate_placed = False
                
                if target_nodes.nelement() > 0:                                   
                    params = []
                    if exists(params_tensor):
                        params = params_tensor[:, t]
                        if params_4pi_normalization:
                            params = (params+1.0) * 2.0*np.pi    # [-1, 1] to [0, 4pi]
                        params = params.tolist()

                    instructions.add_instruction(gate, control_nodes.tolist(), target_nodes.tolist(), params)
                    _gate_placed = True
                    
                    break  #break on first hit, per def only one gate allowed per t
              
                elif control_nodes.nelement() > 0: # no target but control means error
                    if not ignore_errors:                     
                        raise RuntimeError("target_nodes.nelement() <= 0 but control_nodes.nelement() > 0")

            if not _gate_placed and place_error_placeholders:
                # note we place a h gate with no qubits, so this is always an error
                instructions.add_instruction("h", [], [], [])
        
            #else # we are fine with tensors that have time steps with no action!
        
        return instructions

    @staticmethod
    def get_parametrized_tokens(vocabulary: Vocabulary) -> List[int]:
        parametrized_names     = "rx ry rz phase cp crx cry crz u u2 u3".split()
        non_parametrized_names = "x y z h cx cy cz ch ccx swap s sdg t tdg".split()
        
        parametrized_tokens = []
        for name, token in vocabulary.items():

            if name in parametrized_names:
                parametrized_tokens.append(token)
            elif name not in non_parametrized_names:
                raise NotImplementedError(f"Unknown gate {name}! Please add it to the known list.")

        return parametrized_tokens
