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

# %% auto 0
__all__ = ['get_all_srvs', 'generate_srv_tensors', 'convert_tensors_to_srvs', 'get_srv_accuracy', 'true_sample_bin_dist',
           'test_srv_clr_distribution_bin_samples', 'test_srv_clr_distribution', 'test_guidance_dep',
           'test_srv_acc_vs_length', 'test_srv_acc_vs_maxLength', 'test_srv_length_distribution',
           'plot_srv_clr_distribution_hist', 'plot_srv_clr_distribution_bin_accuracy', 'plot_guidance_dep',
           'plot_srv_acc_vs_length', 'plot_srv_acc_vs_maxLength']

# %% ../../src/inference/infer_srv.ipynb 2
from ..imports import *
from .infer_misc import *
from .infer_gate_hist import get_circuit_gate_length
from ..platform.qcircuit_util import get_entanglement_bins
from ..platform.simulation.qcircuit_sim import schmidt_rank_vector
from .infer_compilation import generate_comp_tensors

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

# %% ../../src/inference/infer_srv.ipynb 4
def get_all_srvs(num_of_qubits):
    srvs = [x for x in itertools.product(*([[1,2]]*num_of_qubits))]
    srvs = np.array(srvs, dtype=int)[np.sum(srvs, axis=1)!=num_of_qubits+1].tolist()
    srvs = sorted(srvs, key=lambda x: sum(x))
    return srvs

# %% ../../src/inference/infer_srv.ipynb 5
def generate_srv_tensors(pipeline, prompt, samples, system_size, num_of_qubits, max_gates, g, no_bar=True, unique=False, auto_batch_size=512):
    if samples==0:
        out_tensor = torch.zeros((0, system_size, max_gates))
        return out_tensor
    
    #----------------------
    #prepare condtions
    
    prompt = str(prompt)
    c = pipeline.text_encoder.tokenize_and_push_to_device(prompt)

    #----------------------
    #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))
        
        latents = torch.randn((c_batch.shape[0], pipeline.model.clr_dim, system_size, max_gates))    
        out_tensor = pipeline(latents=latents, c=c_batch, g=g, no_bar=no_bar, enable_guidance=True)   
        out_tensor_list.append(out_tensor)
        
    out_tensor = torch.cat(out_tensor_list)
    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_srv_tensors) Generated {'unique_cnt ' if unique else ''}{out_tensor.shape[0]} tensors")

    return out_tensor

# %% ../../src/inference/infer_srv.ipynb 7
def convert_tensors_to_srvs(out_tensor, gate_pool, sort_srv=False, place_barrier=False, n_jobs=1): 
    qc_list, error_cnt = convert_tensors_to_circuits(out_tensor, gate_pool=gate_pool, place_barrier=place_barrier)
    
    srv_list = []
    
    #---------------------------------------------
    # This is a bottle-neck for more qubits, speed up with async
    
    if n_jobs > 1:
        assert sort_srv == False
        
        f = lambda qc: schmidt_rank_vector(qi.DensityMatrix(qc))
        # srv_list = Parallel(n_jobs=n_jobs, prefer="threads")(delayed(f)(qc) for qc in qc_list)  #prefer="threads"
        srv_list = Parallel(n_jobs=n_jobs)(delayed(f)(qc) for qc in qc_list) 
    
    else:  
        for qc in qc_list:            
            srv = schmidt_rank_vector(qi.DensityMatrix(qc))
            
            if sort_srv: srv = sorted(srv)      
            srv_list.append(srv)  
  
    return qc_list, error_cnt, srv_list

# %% ../../src/inference/infer_srv.ipynb 9
def get_srv_accuracy(srv_list, target_srv):
    if not isinstance(srv_list  , (torch.Tensor, torch.IntTensor, torch.FloatTensor, torch.LongTensor)):   srv_list = torch.tensor(srv_list)
    if not isinstance(target_srv, (torch.Tensor, torch.IntTensor, torch.FloatTensor, torch.LongTensor)): target_srv = torch.tensor(target_srv, device=srv_list.device)
    
    srv_uniques, srv_uniques_cnt = torch.unique(srv_list, dim=0, return_counts=True)
    
    if srv_uniques.numel() == 0: return 0

    comp  = torch.all(target_srv==srv_uniques, dim=1)
    index = comp.nonzero().squeeze() 
        
    if index.dim() == 0: correct_srv_percentage = srv_uniques_cnt[index]/srv_uniques_cnt.sum()           
    else:                correct_srv_percentage = 0 
 
    return correct_srv_percentage

# %% ../../src/inference/infer_srv.ipynb 11
def true_sample_bin_dist(samples_per_bin, bin_size):
    true_samples     = [max(samples_per_bin//bin_size, 1) for i in range(bin_size)]

    if samples_per_bin-sum(true_samples) > 0:
        true_samples[0] += (samples_per_bin-sum(true_samples))
    
    # assert sum(true_samples)==samples_per_bin
    # assert len(true_samples)==bin_size

    # print(f"{true_samples=}")
    
    return true_samples

# %% ../../src/inference/infer_srv.ipynb 12
def test_srv_clr_distribution_bin_samples(pipeline, samples_per_bin, system_size, num_of_qubits, max_gates, g, gate_pool, silent=False, device="cpu", U=None, prompt_mod: callable=lambda c: c,
                                         only_diag=False, n_jobs=1):
    dist_srvs = get_all_srvs(num_of_qubits)  
    cond_srvs = dist_srvs

    values = torch.zeros((len(cond_srvs), len(dist_srvs)), device=device)

    #---------------------

    ent_bins, ent_labels = get_entanglement_bins(num_of_qubits)

    i = 0
    
    for ent_bin in tqdm(ent_bins, total=len(ent_bins)):

        true_samples = true_sample_bin_dist(samples_per_bin, len(ent_bin))
        
        for ind,srv in tqdm(enumerate(ent_bin), total=len(ent_bin)): 
            if exists(U): out_tensor = generate_comp_tensors(pipeline, prompt_mod(srv), U, true_samples[ind], system_size, num_of_qubits, max_gates, g=g, unique=False)
            else:         out_tensor = generate_srv_tensors( pipeline, prompt_mod(srv),    true_samples[ind], system_size, num_of_qubits, max_gates, g=g, unique=False)

            qc_list, error_cnt, svr_list = convert_tensors_to_srvs(out_tensor, gate_pool, n_jobs=n_jobs)
 
            if only_diag:
                values[i, i] = get_srv_accuracy(svr_list, srv)
            else:
                for j, dist_srv in enumerate(dist_srvs):      
                    values[i, j] = get_srv_accuracy(svr_list, dist_srv)
    
            if not silent:
                print(f"{cond_srv}:      unique_cnt {out_tensor.unique(dim=0).shape[0]}      error_cnt {error_cnt}      acc {values[i, i]:.2f}")

            i += 1
            
    return values

# %% ../../src/inference/infer_srv.ipynb 13
def test_srv_clr_distribution(pipeline, samples_per_srv, system_size, num_of_qubits, max_gates, g, gate_pool, silent=False, device="cpu", U=None, prompt_mod: callable=lambda c: c,
                             dist_srvs=None, cond_srvs=None, only_diag=False, n_jobs=1):
    if not exists(dist_srvs):
        dist_srvs = get_all_srvs(num_of_qubits)  
        
    if not exists(cond_srvs):
        cond_srvs = dist_srvs

    values = torch.zeros((len(cond_srvs), len(dist_srvs)), device=device)

    #---------------------
      
    for i, cond_srv in tqdm(enumerate(cond_srvs), total=len(cond_srvs)):    

        if exists(U): out_tensor = generate_comp_tensors(pipeline, prompt_mod(cond_srv), U, samples_per_srv, system_size, num_of_qubits, max_gates, g=g, unique=False)
        else:         out_tensor = generate_srv_tensors( pipeline, prompt_mod(cond_srv),    samples_per_srv, system_size, num_of_qubits, max_gates, g=g, unique=False)
        
        qc_list, error_cnt, svr_list = convert_tensors_to_srvs(out_tensor, gate_pool, n_jobs=n_jobs)
                 
        if only_diag:
            values[i, i] = get_srv_accuracy(svr_list, srv)
        else:
            for j, dist_srv in enumerate(dist_srvs):      
                values[i, j] = get_srv_accuracy(svr_list, dist_srv)

        if not silent:
            print(f"{cond_srv}:      unique_cnt {out_tensor.unique(dim=0).shape[0]}      error_cnt {error_cnt}      acc {values[i, i]:.2f}")

    return values

# %% ../../src/inference/infer_srv.ipynb 14
def test_guidance_dep(pipeline, srvs, samples, system_size, num_of_qubits, max_gates, gs, gate_pool, prompt_mod: callable=lambda c: c, U=None, n_jobs=1):        
    guidance_dep_out = []

    for srv in srvs:
        unique_percentage_list      = []
        error_cnt_percentage_list   = []
        correct_srv_percentage_list = []
        
        for g in tqdm(gs):   
            if exists(U): out_tensor = generate_comp_tensors(pipeline, prompt_mod(srv), U, samples, system_size, num_of_qubits, max_gates, g=g, unique=False)
            else:         out_tensor = generate_srv_tensors( pipeline, prompt_mod(srv),    samples, system_size, num_of_qubits, max_gates, g=g, unique=False)
                        
            #---------------------------------
            #calculate the copy percentage, dataset and sample? 
            
            unique_percentage  = out_tensor.unique(dim=0).shape[0]/out_tensor.shape[0]             
            unique_percentage_list.append(unique_percentage)
                            
            #---------------------------------
            #decode tensors, get srv
            
            qc_list, error_cnt, svr_list = convert_tensors_to_srvs(out_tensor, gate_pool, n_jobs=n_jobs)
            error_cnt_percentage_list.append(error_cnt/out_tensor.shape[0])
            
            #---------------------------------
            #record the correct number       
            
            correct_srv_percentage = get_srv_accuracy(svr_list, srv)   
            correct_srv_percentage_list.append(correct_srv_percentage)

        guidance_dep_out.append((unique_percentage_list, error_cnt_percentage_list, correct_srv_percentage_list))
    
    return guidance_dep_out

# %% ../../src/inference/infer_srv.ipynb 15
def test_srv_acc_vs_length(pipeline, samples, system_size, num_of_qubits, max_gates, g, gate_pool, prompt_mod: callable=lambda c: c, U=None, n_jobs=1):
    ent_bins, ent_labels = get_entanglement_bins(num_of_qubits)
    
    ent_ls   = []
    ent_accs = []
    ent_cnts = []
    
    for ent_bin in tqdm(ent_bins, total=len(ent_bins)):
        ls_acc = dict() #keep track over bins
        ls_cnt = dict()

        true_samples = true_sample_bin_dist(samples_per_bin, len(ent_bin))
        
        for ind,srv in enumerate(ent_bin): 
            if exists(U): out_tensor = generate_comp_tensors(pipeline, prompt_mod(srv), U, true_samples[ind], system_size, num_of_qubits, max_gates, g=g, unique=False)
            else:         out_tensor = generate_srv_tensors( pipeline, prompt_mod(srv),    true_samples[ind], system_size, num_of_qubits, max_gates, g=g, unique=False)
                                   
            qc_list, error_cnt, svr_list = convert_tensors_to_srvs(out_tensor, gate_pool, n_jobs=n_jobs)     

            lengths = get_circuit_gate_length(qc_list) #work in qc space to check only non errors

            if lengths.numel() < 1: continue
            
            for l in lengths.unique(): #range(lengths.min(), lengths.max()):
                indices = (lengths==l).nonzero().squeeze() 
                               
                if indices.numel() > 0:                    
                    srvs = torch.tensor(svr_list)[indices] 
                    if indices.dim() == 0: srvs = srvs.unsqueeze(0)
                                            
                    acc = get_srv_accuracy(srvs, srv)
                    
                    #----------    
                    t = ls_acc.pop(l, [])
                    t.append(acc)
                    ls_acc[l] = t
                    
                    t = ls_cnt.pop(l, 0)
                    t += srvs.shape[0]
                    ls_cnt[l] = t
                           
        ls   = sorted(ls_acc)                   # sorted keys (l)
        accs = [np.mean(ls_acc[l]) for l in ls] # average acc per l
        cnts = [np.sum(ls_cnt[l]) for l in ls]
        
        ent_ls.append(ls)
        ent_accs.append(accs)
        ent_cnts.append(cnts)

    return ent_ls, ent_accs, ent_cnts, ent_labels

# %% ../../src/inference/infer_srv.ipynb 16
def test_srv_acc_vs_maxLength(pipeline, samples_per_bin, system_size, num_of_qubits, max_gates_list, g, gate_pool, prompt_mod: callable=lambda c: c, U=None, n_jobs=1):
    ent_bins, ent_labels = get_entanglement_bins(num_of_qubits)
    
    ent_accs = []
    for ent_bin in tqdm(ent_bins, total=len(ent_bins)):
        
        true_samples = true_sample_bin_dist(samples_per_bin, len(ent_bin))
        
        bin_accs = []
        for max_gates in max_gates_list:

            accs = []
            for ind,srv in enumerate(ent_bin): 
                if exists(U): out_tensor = generate_comp_tensors(pipeline, prompt_mod(srv), U, true_samples[ind], system_size, num_of_qubits, max_gates, g=g, unique=False)
                else:         out_tensor = generate_srv_tensors( pipeline, prompt_mod(srv),    true_samples[ind], system_size, num_of_qubits, max_gates, g=g, unique=False)
                                   
                qc_list, error_cnt, svr_list = convert_tensors_to_srvs(out_tensor, gate_pool, n_jobs=n_jobs)     

                acc = get_srv_accuracy(svr_list, srv)
                
                accs.append(acc)       
            bin_accs.append(np.mean(accs))     
        ent_accs.append(bin_accs)
        
    return ent_accs, ent_labels

# %% ../../src/inference/infer_srv.ipynb 17
def test_srv_length_distribution(pipeline, samples_per_bin, system_size, num_of_qubits, max_gates, g, gate_pool, silent=False, U=None, prompt_mod: callable=lambda c: c, n_jobs=1):
    ent_bins, ent_labels = get_entanglement_bins(num_of_qubits)

    ls = []
    
    for ent_bin in tqdm(ent_bins, total=len(ent_bins)):

        true_samples = true_sample_bin_dist(samples_per_bin, len(ent_bin))
        
        bin_ls = []
        
        for ind,srv in tqdm(enumerate(ent_bin), total=len(ent_bin)): 
            if exists(U): out_tensor = generate_comp_tensors(pipeline, prompt_mod(srv), U, true_samples[ind], system_size, num_of_qubits, max_gates, g=g, unique=False)
            else:         out_tensor = generate_srv_tensors( pipeline, prompt_mod(srv),    true_samples[ind], system_size, num_of_qubits, max_gates, g=g, unique=False)

            qc_list, error_cnt, svr_list = convert_tensors_to_srvs(out_tensor, gate_pool, n_jobs=n_jobs)

            qc_ls = get_circuit_gate_length(qc_list) #tensor [qcs]
            bin_ls.append(qc_ls)

        ls.append(torch.cat(bin_ls))

    return ls #[ent_bins, num_of_non_err_samples]

# %% ../../src/inference/infer_srv.ipynb 19
def plot_srv_clr_distribution_hist(values, samples, num_of_qubits, save=False, dist_srvs=None, cond_srvs=None):
    if not exists(dist_srvs):
        dist_srvs = get_all_srvs(num_of_qubits)  
        
    if not exists(cond_srvs):
        cond_srvs = dist_srvs

    n         = len(dist_srvs)
    values    = values.cpu()
    
    fig = plt.figure(figsize=(12,12))#, constrained_layout=True) 
    plt.title(f"Generated samples per condition: {samples}")
    plt.ylabel(r"Condition")
    plt.xlabel(r"Generated distribution")

    #--------------------------------------------
    if num_of_qubits < 6 or 0:
        plt.yticks(range(len(cond_srvs)), [str(b) for b in cond_srvs])
        plt.xticks(range(n), [str(b) for b in dist_srvs], rotation=90 if n>3 else 0)
    else:
        plt.yticks([])
        plt.xticks([])

    #--------------------------------------------
    plt.imshow(values, vmin=0, vmax=1)
    # plt.imshow(values.cpu(), norm="log")
    plt.colorbar()

    #--------------------------------------------
    #print acc
    x_shift = 1*40 if num_of_qubits==5 else 0
    if num_of_qubits < 4:
        for i in range(n): 
            plt.text(x_shift+i, i, f"{values[i, i]:0.2f}", color='black', ha='center', va='center', fontsize="large")

    #--------------------------------------------
    #draw rects
    off = 0.5
    for i in range(2, num_of_qubits):
        w = scipy.special.comb(num_of_qubits, i, exact=True)
        plt.gca().add_patch(plt.Rectangle((off, off), w, w, ls="-", ec="white", fc="none")) #, transform=plt.gca().transAxes))
        off += w

    #--------------------------------------------
    #print average acc for rects
    off = 0
    for i in [0]+list(range(2, num_of_qubits+1)):
        w = scipy.special.comb(num_of_qubits, i, exact=True)   
        d1 = off
        d2 = d1 + w
        mean_acc = values[d1:d2, d1:d2].diag().mean()
        plt.text(off+2*w/3, off+w/7, f"{mean_acc:0.2f}", color='red', ha='center', va='center', fontsize="x-large")   
        off += w

    #--------------------------------------------
    if save:
        plt.savefig('plot_srv_clr_distribution_hist.svg', bbox_inches='tight')
      
    plt.show()

# %% ../../src/inference/infer_srv.ipynb 20
def plot_srv_clr_distribution_bin_accuracy(values, samples, num_of_qubits, save=False, plot_percentages=False, trainSet_srv=None):
    values = values.cpu().diag()
    ent_bins, ent_labels = get_entanglement_bins(num_of_qubits)
    
    n = sum(len(srvs) for srvs in ent_bins)
    x = np.arange(n)  # the label locations
    width = 0.8

    #------------------------
    fig = plt.figure(figsize=(6.6, 4), constrained_layout=True) 
    # plt.title(f"Generated samples per condition: {samples}", fontsize=14)
    plt.ylabel(r"Accuracy", fontsize=25)
    plt.yticks(fontsize=14)
    plt.xticks([])
 
    i = 0
    for j,(label, srvs) in enumerate(zip(ent_labels, ent_bins)):   
        label = f"{sum(srvs[0])-num_of_qubits}"
        incre = len(srvs)
        rects = plt.bar(x[i:i+incre], values[i:i+incre], width, label=label)
        i += incre
        if plot_percentages: plt.gca().bar_label(rects, padding=3, fmt="%0.2f")

    ncols = len(ent_labels)//2+1 if len(ent_labels) > 5 else len(ent_labels)
    leg1 = plt.legend(loc="lower center", fontsize=14, ncols=ncols, title="# of entangled qubits:", title_fontsize=14,bbox_to_anchor=(0.5, 1.01))
    ax = fig.add_artist(leg1)
    
    if exists(trainSet_srv):
        if trainSet_srv.shape[-1]==num_of_qubits:
            srvs = []
            for s in ent_bins: srvs.extend(s)
            
            dataset_percentages = [get_srv_accuracy(trainSet_srv, srv).cpu() for srv in srvs]
            xmin = x - width*0.55
            xmax = x + width*0.55            
            ag = plt.hlines(dataset_percentages, xmin, xmax, label="Random sampling" , color="black", linestyle="-", linewidths=2.3)
        
            plt.legend(handles=[ag], fontsize=14, frameon=False)

    ymin, ymax = plt.ylim()    
    plt.ylim(ymin, ymax+0.04)

    if save:
        plt.savefig(f"plot_srv_clr_distribution_bin_accuracy.svg", bbox_inches='tight', transparent=True)
        
    plt.show()

# %% ../../src/inference/infer_srv.ipynb 21
def plot_guidance_dep(srvs, gs, guidance_dep_out, samples, save=False):    
    assert len(srvs) == len(guidance_dep_out)

    n = len(srvs)
    fig, axs = plt.subplots(1, n, figsize=(12, 5), squeeze=False, constrained_layout=True)   
    fig.suptitle(fr"Generated {samples} samples per $g$ and SRV")
    
    for i,srv in enumerate(srvs):  
        unique_percentage_list, error_cnt_percentage_list, correct_srv_percentage_list = guidance_dep_out[i]
                     
        #---------------------------------
        #plot now gs vs the numbers   
            
        plt.sca(axs[0, i])
        plt.xlabel(r"Guidance scale $g$")     
        plt.title(f"SRV = {srv}")
        plt.plot(gs, unique_percentage_list     , label="Unique tensors percentage")
        plt.plot(gs, error_cnt_percentage_list  , label="Error circuits percentage")
        plt.plot(gs, correct_srv_percentage_list, label="Correct SRV percentage")
        
        if i == (n-1): plt.legend() 

    if save:
        plt.savefig("plot_guidance_dep.svg", bbox_inches='tight', transparent=True)
    
    plt.show()

# %% ../../src/inference/infer_srv.ipynb 22
def plot_srv_acc_vs_length(ent_ls, ent_accs, ent_cnts, ent_labels, samples, plot_dist=True, save=False):
    fig, axs = plt.subplots(2 if plot_dist else 1, 1, figsize=(12, 7), squeeze=False, constrained_layout=True) 

    #-------------------
    plt.sca(axs[0,0])
    plt.title(f"Generated samples per entanglement: {samples}")
    plt.ylabel("Accuracy")
    plt.xlabel("Gate number")
    for i,ent_label in enumerate(ent_labels): 
        plt.plot(ent_ls[i], ent_accs[i], label=f"{ent_label}")
    plt.legend()

    #-------------------
    if plot_dist:
        plt.sca(axs[1,0])
        plt.title(f"Used samples per l to calculate accuracy, should match gate distribution")
        plt.ylabel("Used samples")
        plt.xlabel("Gate number")
        for i,ent_label in enumerate(ent_labels): 
            plt.plot(ent_ls[i], ent_cnts[i], label=f"{ent_label}")
        plt.legend()

    #-------------------
    if save:
        plt.savefig('plot_srv_acc_vs_length.svg', bbox_inches='tight')
         
    plt.show()

# %% ../../src/inference/infer_srv.ipynb 23
def plot_srv_acc_vs_maxLength(ent_accs, ent_labels, max_gates_list, samples, plot_dist=True, save=False):
    fig = plt.figure(figsize=(12, 4), constrained_layout=True) 
    
    plt.title(f"Generated samples per maxGates per entanglement: {samples}")
    plt.ylabel("Accuracy")
    plt.xlabel("Max number of gates / tensor size")
    plt.xticks(max_gates_list)
    
    for ent_acc,ent_label in zip(ent_accs, ent_labels): 
        plt.plot(max_gates_list, ent_acc, label=f"{ent_label}")
        
    plt.legend()

    if save:
        plt.savefig('plot_srv_acc_vs_length.svg', bbox_inches='tight')
         
    plt.show()
