"""Functions to create a quantum circuit dataset."""

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

# %% auto #0
__all__ = ['CircuitConditionType', 'get_rnd_encoded_circuit', 'get_rnd_encoded_circuits', 'generate_circuit_dataset']

# %% ../../src/platform/circuits_generation.ipynb #c82c4d7b-e275-416e-b88d-3d527a9e3bff
from ..imports import *
from .backends.base_backend import BaseBackend
from .tokenizer.circuits_tokenizer import CircuitTokenizer, Vocabulary
from ..dataset.dataset_helper import get_unique_elements_indices
from ..utils.async_fn import MemoryMappedArray, Parallel, delayed

# %% ../../src/platform/circuits_generation.ipynb #297cefeb-ff48-4fdb-9df6-f02b43c518cb
class CircuitConditionType(enum.Enum):   
    SRV      = enum.auto() 
    UNITARY  = enum.auto() 

# %% ../../src/platform/circuits_generation.ipynb #c3c07f63-1d6e-4c6f-8079-f78eadccb172
def get_rnd_encoded_circuit(backend: BaseBackend, 
                            tokenizer: CircuitTokenizer,
                            condition: CircuitConditionType,
                            num_of_qubits: int, 
                            gate_pool: Optional[Sequence[str]],
                            min_gates: int, 
                            max_gates: int,  
                            rng: np.random.Generator, 
                            optimized: bool = True,
                            post_randomize_params: bool = True,
                            return_params: bool = True) -> Tuple[Any, torch.Tensor, ...]:   
    """Generate a random circuit with corresponding condition."""

    gate_pool = default(gate_pool, tokenizer.vocabulary)

    qc = backend.rnd_circuit(num_of_qubits, rng.integers(min_gates, max_gates+1), gate_pool, rng)  
    
    if optimized:          
        qc = backend.optimize_circuit(qc, gate_pool)   
        
    if post_randomize_params: 
        qc = backend.randomize_params(qc, rng)   
    
    match condition:
        case CircuitConditionType.SRV:     
            condition = torch.tensor(backend.schmidt_rank_vector(qc))
            
        case CircuitConditionType.UNITARY: 
            U         = backend.get_unitary(qc)   
            U_r, U_i  = torch.from_numpy(np.real(U)), torch.from_numpy(np.imag(U))
            condition = torch.stack([U_r, U_i], dim=0)
    
        case _: raise NotImplementedError(f"Not implemented given condition: {condition}")
    
    instructions = backend.backend_to_genqc(qc)
    enc_tuple    = tokenizer.encode(instructions, max_gates, return_params_tensor=return_params)  # qc_tensor, params_tensor

    return qc, condition, *enc_tuple

# %% ../../src/platform/circuits_generation.ipynb #d8271eeb-3c07-4a94-8cc1-3a44b9cb37d3
def get_rnd_encoded_circuits(backend: BaseBackend, 
                             tokenizer: CircuitTokenizer,
                             condition: CircuitConditionType,                                  
                             samples: int, 
                             num_of_qubits: int, 
                             min_gates: int, 
                             max_gates: int,                       
                             min_sub_gate_pool_cnt: int = 1,
                             max_sub_gate_pool_cnt: Optional[int] = None,
                             fixed_sub_gate_pool: Optional[Sequence[str]] = None,
                             max_num_params: Optional[int] = None,
                             filter_unique: bool = True,
                             optimized: bool = True,                                                         
                             post_randomize_params: bool = True,
                             return_params: bool = True,
                             silent: bool = False) -> Tuple[torch.Tensor, ...]:
    """
    Generate ´samples´ number of random circuits with corresponding condition. 
    Creates prompts for conditioning.
    """

    if condition not in [CircuitConditionType.SRV, CircuitConditionType.UNITARY]:
        raise NotImplementedError(f"Not implemented {condition}")
    
    sub_gate_pool = fixed_sub_gate_pool
    gate_pool     = list(tokenizer.vocabulary)

    rng     = np.random.default_rng() 
    n       = len(gate_pool) + 1
    c_range = np.arange(n-1)

    if exists(max_sub_gate_pool_cnt):
        max_sub_gate_pool_cnt = max(min_sub_gate_pool_cnt, min(max_sub_gate_pool_cnt+1, n))
    else:                             
        max_sub_gate_pool_cnt = n

    #------------------
    # Generate single circuits sequentially
    
    x = []
    y = []    
    
    if condition is CircuitConditionType.UNITARY: 
        u = []
    
    if return_params: 
        p = []   # Note: params is of different size -> keep list

    for i in tqdm(range(samples), disable=silent):

        if not exists(fixed_sub_gate_pool):    
            sub_gate_pool_cnt = rng.integers(min_sub_gate_pool_cnt, max_sub_gate_pool_cnt)        
            sub_gate_pool_ind = rng.choice(c_range, size=sub_gate_pool_cnt, replace=False)   

            #NOTE: with this we have always the same ordering of the prompt gates!!
            sub_gate_pool_ind = np.sort(sub_gate_pool_ind)  
            
            sub_gate_pool     = [gate_pool[ind] for ind in sub_gate_pool_ind]

        val = get_rnd_encoded_circuit(backend=backend,
                                      tokenizer=tokenizer,
                                      condition=condition,
                                      num_of_qubits=num_of_qubits,
                                      gate_pool=sub_gate_pool, 
                                      min_gates=min_gates, 
                                      max_gates=max_gates,
                                      rng=rng,
                                      optimized=optimized,
                                      post_randomize_params=post_randomize_params,
                                      return_params=return_params)
        
        if return_params:
            _, cond, qc_tensor, params_tensor = val    
            p.append(params_tensor)
        else:
            _, cond, qc_tensor = val

        x.append(qc_tensor)
        
        match condition:
            case CircuitConditionType.SRV:      
                label = f"Generate SRV: {cond.tolist()}"     
                
            case CircuitConditionType.UNITARY:  
                label = f"Compile using: {[str(gate) for gate in sub_gate_pool]}"           
                u.append(cond)
                
            case _: raise NotImplementedError(f"Not implemented given condition: {condition}")
      
        y.append(label)
        
    #------------------
    # Make tensors unique and combine tensors and arrays

    x = torch.stack(x, dim=0)
    y = np.array(y)

    if condition is CircuitConditionType.UNITARY: 
        u = torch.stack(u, dim=0) 
    
    if filter_unique:
        tensor_unique, tensor_indices = get_unique_elements_indices(x)
        
        x = x[tensor_indices]
        y = y[tensor_indices]         
        if return_params:                             p = [p[i] for i in tensor_indices.tolist()]    
        if condition is CircuitConditionType.UNITARY: u = u[tensor_indices]

        if not silent: 
            print(f"[INFO]: Generated unique circuits: {tensor_unique.shape[0]}.")
    
    if not exists(max_num_params):
        p_max_para  = max(pi.shape[0] for pi in p)
        p_min_value = min(pi.min() if pi.numel()>0 else 0 for pi in p)
        p_max_value = max(pi.max() if pi.numel()>0 else 0 for pi in p)
        if not silent: print(f"[INFO]: No max_num_params provided, infered {p_max_para=}, {p_min_value=} and {p_max_value=}.")
    else:
        if not silent: print(f"[INFO]: Using provided {max_num_params=}.")
        p_max_para = max_num_params
        
    p_t = torch.zeros((len(p), p_max_para, max_gates))
    for i,pi in enumerate(p):
        p_t[i, :pi.shape[0], :pi.shape[1]] = pi
    p = p_t

    if return_params: 
        return x, y, u, p
    return x, y, u

# %% ../../src/platform/circuits_generation.ipynb #9abc01cb-136e-4d87-bb91-beae08e0bb33
def generate_circuit_dataset(backend: BaseBackend, 
                             tokenizer: CircuitTokenizer,
                             condition: CircuitConditionType,                                  
                             total_samples: int,    
                             num_of_qubits: int, 
                             min_gates: int, 
                             max_gates: int,  
                             batch_samples: int = 128,  
                             n_jobs: int = 1,
                             unitary_dtype: torch.dtype = torch.float16,
                             min_sub_gate_pool_cnt: int = 1,
                             max_sub_gate_pool_cnt: Optional[int] = None,
                             fixed_sub_gate_pool: Optional[Sequence[str]] = None,
                             max_num_params: Optional[int] = None,
                             filter_unique: bool = True,
                             optimized: bool = True,                                                         
                             post_randomize_params: bool = True,
                             return_params: bool = True) -> Tuple[torch.Tensor, ...]:
    """
    Generates ´samples´ number of random circuits with corresponding condition. 
    Supports large scale dataset with large unitaries. Uses memory mapping and parallelization.

    - ´unitary_dtype´ only relevant for ´condition=CircuitConditionType.UNITARY´   
    """

    if condition not in [CircuitConditionType.UNITARY]:
        raise NotImplementedError(f"Not implemented {condition=}")

    if not return_params:
        raise NotImplementedError(f"Not implemented {return_params=}")

    total_samples = int(total_samples)
    batch_samples = min(int(batch_samples), total_samples)
    njobs         = max(min(n_jobs, total_samples//batch_samples), 1)

    #------------------
    # Check data sizes

    gen_data = functools.partial(get_rnd_encoded_circuits, 
                                 backend=backend, 
                                 tokenizer=tokenizer,
                                 condition=condition,
                                 samples=batch_samples,        
                                 num_of_qubits=num_of_qubits, 
                                 min_gates=min_gates, 
                                 max_gates=max_gates,
                                 min_sub_gate_pool_cnt=min_sub_gate_pool_cnt,
                                 max_sub_gate_pool_cnt=max_sub_gate_pool_cnt,
                                 fixed_sub_gate_pool=fixed_sub_gate_pool,
                                 max_num_params=max_num_params,
                                 filter_unique=filter_unique,
                                 optimized=optimized,
                                 post_randomize_params=post_randomize_params,
                                 return_params=return_params,
                                 silent=True)
    
    x, y, u, p = gen_data()
    x_global = torch.zeros((total_samples, *x.shape[1:]), dtype=x.dtype) 
    y_global =     np.empty(total_samples, dtype=y.dtype)  
    u_global = torch.zeros((total_samples, 2, u.shape[-2], u.shape[-1]), dtype=unitary_dtype) 
    p_global = torch.zeros((total_samples, *p.shape[1:]), dtype=p.dtype) 

    #------------------
    # Run memory mapped parallel generation

    def _f(idx, x_map, y_map, u_map, p_map):
        x, y, u, p = gen_data()
     
        off  = x.shape[0]
        idx *= batch_samples
    
        x_map[idx:idx+off] = x
        y_map[idx:idx+off] = y
        u_map[idx:idx+off] = u
        p_map[idx:idx+off] = p
    

    def _scope():
        x_map = MemoryMappedArray(x_global)
        y_map = MemoryMappedArray(y_global, type="numpy")
        u_map = MemoryMappedArray(u_global)
        p_map = MemoryMappedArray(p_global)
        
        with Parallel(n_jobs=n_jobs) as parallel:
            loop_set = range(int(np.floor(total_samples/batch_samples)))
            _        = parallel(delayed(_f)(idx, x_map.obj_memmap, y_map.obj_memmap, u_map.obj_memmap, p_map.obj_memmap) for idx in loop_set)

        return x_map.get_obj(), y_map.get_obj(), u_map.get_obj(), p_map.get_obj()
        
    (x_global, x_file), (y_global, y_file), (u_global, u_file), (p_global, p_file) = _scope()
        
    MemoryMappedArray.clean([x_file, y_file, u_file, p_file])

    #------------------
    # Collect results and remove the holes
 
    x_global_nonzero = torch.logical_not((x_global==0).all(-1).all(-1)).nonzero(as_tuple=True)
    print(f"[INFO]: Generated {x_global_nonzero[0].shape[0]} valid circuits.")

    # complex indexing makes copy not view
    x_global = x_global[x_global_nonzero]#.contiguous().clone()
    u_global = u_global[x_global_nonzero]#.contiguous().clone()
    p_global = p_global[x_global_nonzero]#.contiguous().clone()
    y_global = y_global[x_global_nonzero] #np.ascontiguousarray(y_global[x_global_nonzero])

    if filter_unique:
        tensor_unique, tensor_indices = get_unique_elements_indices(x_global)
        
        x_global = x_global[tensor_indices]
        y_global = y_global[tensor_indices]
        u_global = u_global[tensor_indices]
        p_global = p_global[tensor_indices]

        print(f"[INFO]: After filtering unique circuits: {x_global.shape[0]}.")
    
    return x_global, y_global, u_global, p_global
