# AUTOGENERATED! DO NOT EDIT! File to edit: ../../src/inference/infer_compilation.ipynb.

# %% auto 0
__all__ = ['split_U_to_tensor', 'get_new_unitary_indices', 'get_new_unitary_indices_batch', 'generate_comp_tensors',
           'check_correct_gates', 'check_correct_unitary_exact', 'check_correct_unitary_distance', 'get_gate_and_U_acc',
           'test_comp_acc', 'test_comp_acc_on_testset', 'test_comp_acc_on_rnd_samples', 'plot_hist_overview']

# %% ../../src/inference/infer_compilation.ipynb 2
from ..imports import *
from ..util import *
from .infer_misc import *
from .infer_gate_hist import get_tensor_gate_length
import genQC.platform.qcircuit_dataset_construction as data_con
from ..dataset.dataset_helper import check_duplicates_in_dataset, uniquify_tensor_dataset, shuffle_tensor_dataset
from ..platform.simulation.qcircuit_sim import instruction_name_to_qiskit_gate

from joblib import Parallel, delayed
import qiskit.quantum_info as qi

# %% ../../src/inference/infer_compilation.ipynb 4
def split_U_to_tensor(U: np.ndarray):
    U_r, U_i = torch.Tensor(np.real(U)), torch.Tensor(np.imag(U))
    U        = torch.stack([U_r, U_i], dim=0)
    return U

# %% ../../src/inference/infer_compilation.ipynb 5
def get_new_unitary_indices(Us, dataset, silent=False):
    if type(Us) == list: 
        Us = torch.stack([split_U_to_tensor(U) for U in Us]) #numpy to torch

    if not silent:
        print(f"- Checking {Us.shape[0]} unitaries for duplicates in dataset, {torch.unique(Us, dim=0).shape[0]} given unitaries are unique.")
    
    # need to check uniques only    
    Us_dataset = torch.unique(dataset.U, dim=0)

    # to vecs
    Us         = torch.reshape(Us        , [Us.shape[0]        , -1]).to(Us_dataset.device)  
    Us_dataset = torch.reshape(Us_dataset, [Us_dataset.shape[0], -1])  
    
    #---------------

    #check
    comp = ( Us_dataset.unsqueeze(dim=0) == Us.unsqueeze(dim=1) ) # gives [num of Us, num of dataset, ch]
    comp = torch.all(comp, dim=-1)                            # gives [num of Us, num of dataset]

    #reduce
    comp = torch.all(comp==False, dim=1) # gives indices that ARE NOT in datset
    # comp = torch.any(comp, dim=1)        # gives indices that ARE in datset

    #get indices
    comp = comp.nonzero().squeeze(dim=1)
    
    if not silent:
        print(f"- Checked {Us.shape[0]} given unitaries with dataset. Returned indices of {comp.shape[0]} not in dataset unitaries.")  
        
    return comp.cpu()

# %% ../../src/inference/infer_compilation.ipynb 6
def get_new_unitary_indices_batch(Us, dataset, auto_batch_size=32, silent=False, n_jobs=1):
    if type(Us) == list: 
        Us = torch.stack([split_U_to_tensor(U) for U in Us]) #numpy to torch

    if not silent:
        print(f"- Checking {Us.shape[0]} unitaries for duplicates in dataset, {torch.unique(Us, dim=0).shape[0]} given unitaries are unique.")

    #----------------------------------------
    samples     = Us.shape[0]
    num_batches = int(np.ceil(samples/auto_batch_size))

    Us_chunks = Us.chunk(num_batches)

    indices = []
    
    if n_jobs > 1:
        f = lambda Us_chunk: get_new_unitary_indices(Us_chunk, dataset, silent=True)
        indices = Parallel(n_jobs=n_jobs)(delayed(f)(Us_chunk) for Us_chunk in Us_chunks) 

    else:    
        for Us_chunk in Us_chunks:   
            comp = get_new_unitary_indices(Us_chunk, dataset, silent=True)
            indices.append(comp)
    
        indices = torch.cat(indices)

    if not silent:
        print(f"- Checked {samples} given unitaries with dataset. Returned indices of {indices.shape[0]} not in dataset unitaries.")  
    
    return indices

# %% ../../src/inference/infer_compilation.ipynb 8
def generate_comp_tensors(pipeline, prompt, U, samples, system_size, num_of_qubits, max_gates, g, no_bar=True, unique=False, auto_batch_size=512):
    #----------------------
    #prepare condtions

    prompt = str(prompt)
    c = pipeline.text_encoder.tokenize_and_push_to_device(prompt)

    U = U.unsqueeze(0).to(pipeline.device)
    if system_size > num_of_qubits:
            n = 2**system_size
            pad = (0, n-U.shape[-1], 0, n-U.shape[-2]) 
            U   = F.pad(U, pad, "constant", 0)

    #----------------------
    #sample and post process to tensor encodings    
    
    batch_samples = [auto_batch_size] * int(np.floor(samples/auto_batch_size))
    if samples % auto_batch_size > 0: batch_samples.append(samples % auto_batch_size)
    if len(batch_samples) == 0: batch_samples.append(samples)

    out_tensor_list = []
    for batch_sample in batch_samples:     
        
        c_batch = c.repeat(batch_sample, *[1]*(c.dim()-1))
        U_batch = U.repeat(batch_sample, *[1]*(U.dim()-1))
        
        latents = torch.randn((c_batch.shape[0], pipeline.model.clr_dim, system_size, max_gates))    
        out_tensor = pipeline(latents=latents, c=c_batch, U=U_batch, g=g, no_bar=no_bar)   
        out_tensor_list.append(out_tensor)
        
    out_tensor = torch.cat(out_tensor_list)
    # out_tensor = pipeline(latents=latents, c=c, U=U, g=g, no_bar=no_bar)   

    out_tensor = pipeline.model.invert_clr(out_tensor)
    out_tensor = out_tensor[:, :num_of_qubits]
    
    if unique: out_tensor = torch.unique(out_tensor, dim=0)
    
    if not no_bar: print(f"[INFO]: (generate_comp_tensors) Generated {'unique_cnt ' if unique else ''}{out_tensor.shape[0]} tensors")

    return out_tensor

# %% ../../src/inference/infer_compilation.ipynb 10
def check_correct_gates(qc, num_of_qubits, gate_pool, max_gates, allowed_gate_clrs):
    tensor        = data_con.encode_circuit(qc, num_of_qubits, data_con.gate_pool_to_gate_classes(gate_pool), max_gates)
    gen_gate_clrs = torch.unique(tensor.abs()).tolist()     
    gate_corr     = set(gen_gate_clrs).issubset(set(allowed_gate_clrs))  # are gates correct?
    return gate_corr

# %% ../../src/inference/infer_compilation.ipynb 11
def check_correct_unitary_exact(qc, U):
    is_U = qi.Operator(qc).to_matrix()
    is_U = split_U_to_tensor(is_U)

    u_corr = torch.allclose(is_U, U)   # is U correct?
    return u_corr

# %% ../../src/inference/infer_compilation.ipynb 12
def check_correct_unitary_distance(qc, target_U, norms):
    is_U = qi.Operator(qc).to_matrix()
    is_U = torch.complex(torch.Tensor(np.real(is_U)), torch.Tensor(np.imag(is_U)))
    
    target_U = torch.complex(target_U[0], target_U[1])
    
    d = []
    for norm in norms:
        u_dist = norm.distance(is_U, target_U).item()
        d.append(u_dist)
    
    return d

# %% ../../src/inference/infer_compilation.ipynb 13
def get_gate_and_U_acc(out_tensor, allowed_gate_clrs, U, gate_pool, num_of_qubits, max_gates, norms=[], no_bar=True):

    if isinstance(gate_pool[0], str):
        gate_pool = [instruction_name_to_qiskit_gate(gate) for gate in gate_pool]
    
    #-------------------------
    #decode 
    qc_list, error_cnt = convert_tensors_to_circuits(out_tensor, gate_pool)
    if not no_bar: print(f"Error circuits: {error_cnt}")
    
    #-------------------------
    acc      = []  # combinded acc
    gate_acc = []  # only gates acc
    u_acc    = []  # only U acc
    u_norms  = []  # list of tuple(norms) for every qc
    
    comb_corr_qc = []
    gate_corr_qc = []
    u_corr_qc    = []
    
    #only check circuits that are non-error!
    for qc in qc_list:
        
        #---------------
        # check if in out_tensor only color that correspond to the condtion gate_pool     
        gate_corr = check_correct_gates(qc, num_of_qubits, gate_pool, max_gates, allowed_gate_clrs)
            
        #---------------
        # check unitary     
        u_corr = check_correct_unitary_exact(qc, U)           # true or false
        u_norm = check_correct_unitary_distance(qc, U, norms) # metrics values list

        #---------------           
        acc.append(gate_corr and u_corr)             
        gate_acc.append(gate_corr)   
        u_acc.append(u_corr)
        u_norms.append(u_norm)
        
        if gate_corr and u_corr: comb_corr_qc.append(qc)
        if gate_corr: gate_corr_qc.append(qc)
        if u_corr: u_corr_qc.append(qc)
        
    #average accuracy over sample
    acc      = np.mean(acc).item()
    gate_acc = np.mean(gate_acc).item()
    u_acc    = np.mean(u_acc).item()
    
    return acc, gate_acc, u_acc, np.array(u_norms), error_cnt, comb_corr_qc, gate_corr_qc, u_corr_qc, qc_list

# %% ../../src/inference/infer_compilation.ipynb 15
def test_comp_acc(pipeline, samples, system_size, gate_pool, num_of_qubits, max_gates, g, str_cond_to_gate_indices: callable, Us, ys, train_dataset=None, norms=[]):
    
    if exists(train_dataset):
        not_dups_ind = get_new_unitary_indices_batch(Us, train_dataset)   
        Us           = [Us[i] for i in not_dups_ind]
        ys           = [ys[i] for i in not_dups_ind]
        
    #--------------------
    acc_s         = []
    gate_acc_s    = []
    u_acc_s       = []
    u_norms_s     = []
    uniques_cnt_s = []
    error_cnt_s   = []
    
    num_found_distinct_circuits_s = []
    
    for U,y in tqdm(zip(Us,ys), total=len(Us)):
        
        allowed_gate_clrs = str_cond_to_gate_indices(y)         
        if isinstance(U, np.ndarray):
            U = split_U_to_tensor(U)
    
        out_tensor = generate_comp_tensors(pipeline, y, U, samples, system_size, num_of_qubits, max_gates, g, unique=False)
        outs       = get_gate_and_U_acc(out_tensor, allowed_gate_clrs, U, gate_pool, num_of_qubits, max_gates, norms) 
        
        acc, gate_acc, u_acc, u_norms, error_cnt, comb_corr_qc, gate_corr_qc, u_corr_qc, qc_list = outs

        if len(qc_list) > 0:
            uniques_cnt = torch.stack([data_con.encode_circuit(qc, num_of_qubits, data_con.gate_pool_to_gate_classes(gate_pool), max_gates) for qc in qc_list]).unique(dim=0).shape[0] #how many uniques in sample (not counting erroro circuits)
            # uniques_cnt = out_tensor.shape[0] - error_cnt #was with unique acc definition
        else:
            uniques_cnt = 0

        if len(comb_corr_qc) > 0:
            num_found_distinct_circuits = torch.stack([data_con.encode_circuit(qc, num_of_qubits, data_con.gate_pool_to_gate_classes(gate_pool), max_gates) for qc in comb_corr_qc]).unique(dim=0).shape[0] #how many distinct exact solutions we have
        else:
            num_found_distinct_circuits = 0
    
        #--------------------
        acc_s.append(acc)
        gate_acc_s.append(gate_acc)
        u_acc_s.append(u_acc)
        u_norms_s.append(u_norms)
        uniques_cnt_s.append(uniques_cnt)
        error_cnt_s.append(error_cnt)
        num_found_distinct_circuits_s.append(num_found_distinct_circuits)
        
    solved_tasks = np.count_nonzero(num_found_distinct_circuits_s)
    print(f"Solved {solved_tasks} correctly (at least one qc) that is {100*solved_tasks/len(num_found_distinct_circuits_s):0.2f}%")
        
    return acc_s, gate_acc_s, u_acc_s, u_norms_s, uniques_cnt_s, error_cnt_s, num_found_distinct_circuits_s           

# %% ../../src/inference/infer_compilation.ipynb 16
def test_comp_acc_on_testset(pipeline, samples, num_of_U, system_size, gate_pool, num_of_qubits, max_gates, g, str_cond_to_gate_indices: callable, 
                             prompt_mod: callable, test_dataset, train_dataset=None, norms=[], fix_y=None):
    '''returns: acc_s, gate_acc_s, u_acc_s, uniques_cnt_s, error_cnt_s, num_found_circuits_s, task_qc_len_s'''
    
    if hasattr(test_dataset, "z"): # mixed dataset has padding but a z record! 
        Us, ys, zs    = uniquify_tensor_dataset(test_dataset.U, test_dataset.y, test_dataset.z) 
        Us, ys, zs    = shuffle_tensor_dataset(Us, ys, zs)
        Us, ys, zs    = Us[:num_of_U], ys[:num_of_U], zs[:num_of_U]
        task_qc_len_s = zs[:, 1]
    
    else:                          # not mixed dataset has no padding
        Us, ys, xs    = uniquify_tensor_dataset(test_dataset.U, test_dataset.y, test_dataset.x) 
        Us, ys, xs    = shuffle_tensor_dataset(Us, ys, xs)
        Us, ys, xs    = Us[:num_of_U], ys[:num_of_U], xs[:num_of_U]
        task_qc_len_s = get_tensor_gate_length(xs) 

    if exists(fix_y): ys  = [fix_y for y in ys]
    else:             ys  = [prompt_mod(y) for y in ys]
        
    
    print(f"Picked {Us.shape[0]} unitaries from test set")
    print(f"Sample task: {ys[0]}")
    print(Us[0])
    print(xs[0])
    
    out = test_comp_acc(pipeline, samples, system_size, gate_pool, num_of_qubits, max_gates, g, str_cond_to_gate_indices, Us.cpu(), ys, train_dataset, norms)  
    return *out, task_qc_len_s

# %% ../../src/inference/infer_compilation.ipynb 17
def test_comp_acc_on_rnd_samples(pipeline, samples, num_of_U, system_size, gate_pool, num_of_qubits, max_gates, g, str_cond_to_gate_indices: callable,
                                 prompt_mod: callable, # takes a single prompt and returns it modified
                                 rnd_min_gates, rnd_max_gates,
                                 train_dataset=None, norms=[], fix_y=None):
    '''returns: acc_s, gate_acc_s, u_acc_s, uniques_cnt_s, error_cnt_s, num_found_circuits_s, task_qc_len_s'''
    
    enc_t, ys, Us = data_con.gen_compilation_rndGates_dataset(num_of_U, num_of_qubits, rnd_min_gates, rnd_max_gates, gate_pool)
    task_qc_len_s = get_tensor_gate_length(enc_t) #should give a complexity meassure, longer circuits have a more complex unitary to compile?            

    if exists(fix_y): ys  = [fix_y for y in ys]
    else:             ys  = [prompt_mod(y) for y in ys]
    
    print(f"Sample task: {ys[0]}")
    print(split_U_to_tensor(Us[0]))
    print(enc_t[0])
        
    out = test_comp_acc(pipeline, samples, system_size, gate_pool, num_of_qubits, max_gates, g, str_cond_to_gate_indices, Us, ys, train_dataset, norms)
    return *out, task_qc_len_s

# %% ../../src/inference/infer_compilation.ipynb 19
def plot_hist_overview(out_tuple, num_of_samples_per_U, rnd_min_gates, rnd_max_gates, max_gates, num_of_qubits):
    acc_s, gate_acc_s, u_acc_s, u_norms_s, uniques_cnt_s, error_cnt_s, num_found_circuits_s, task_qc_len_s = out_tuple          
    
    if not exists(rnd_min_gates): rnd_min_gates = ""
    if not exists(rnd_max_gates): rnd_max_gates = ""
    
    fig, axs = plt.subplots(2, 3, figsize=(13, 6.4), squeeze=False, constrained_layout=True)  
    fig.suptitle(f"Histogram of compilation accuracies       (Unitary cnt={len(acc_s)}, samples_per_U={num_of_samples_per_U}  {rnd_min_gates=}  {rnd_max_gates=}  {max_gates=}  qubits={num_of_qubits})")

    n       = 20
    density = False
    bins    = np.linspace(0,1, n+1)

    #-----------------
    plt.sca(axs[0, 0])
    plt.title("Combined accuracy")
    plt.xlabel(r"Accuracy")
    plt.ylabel(r"Bin population" if density==False else "Accuracy distribution")
    plt.hist(acc_s, density=density, bins=n*4)

    #-----------------
    plt.sca(axs[0, 1])
    plt.title("Unitary accuracy")
    plt.xlabel(r"Accuracy")
    plt.hist(u_acc_s, density=density, bins=bins)

    #-----------------
    plt.sca(axs[0, 2])
    plt.title("Gate accuracy")
    plt.xlabel(r"Accuracy")
    plt.hist(gate_acc_s, density=density, bins=bins)

    #-----------------
    plt.sca(axs[1, 0])
    plt.title("Generated unique circuits")
    plt.ylabel(r"Bin population" if density==False else "Number distribution")
    plt.xlabel(r"Number of unique circuits")
    plt.hist(uniques_cnt_s, density=density, bins=n)
    
    #-----------------
    plt.sca(axs[1, 1])
    plt.title("Generated error circuits")
    plt.xlabel(r"Number of error circuits")
    plt.hist(error_cnt_s, density=density, bins=n)
    
    #-----------------
    plt.sca(axs[1, 2])
    plt.title("Absolute number of distinct correct circuits")
    plt.xlabel(r"Number of found circuits")
    plt.hist(num_found_circuits_s, density=density, bins=n*4)
    
    #-----------------
    plt.show()
