# AUTOGENERATED! DO NOT EDIT! File to edit: ../../src/dataset/qc_dataset.ipynb.

# %% auto 0
__all__ = ['Qc_Config_Dataset_config', 'Qc_Config_Dataset']

# %% ../../src/dataset/qc_dataset.ipynb 3
from ..imports import *
from .config_dataset import Config_Dataset, Config_Dataset_config
from ..config_loader import *
from .dataset_helper import *
from ..platform.qcircuit_dataset_construction import decode_circuit
from ..platform.simulation.qcircuit_sim import schmidt_rank_vector, instruction_name_to_qiskit_gate
import qiskit.quantum_info as qi

# %% ../../src/dataset/qc_dataset.ipynb 4
@dataclass
class Qc_Config_Dataset_config(Config_Dataset_config):
    optimized: bool
    dataset_to_gpu: bool
    random_samples: int  
    num_of_qubits: int  
    min_gates: int 
    max_gates: int 
    gate_pool: list[str] 

# %% ../../src/dataset/qc_dataset.ipynb 5
class Qc_Config_Dataset(Config_Dataset):
    """Dataset for quantum circuits, access `gate_pool` directly and all other paras with `.params_config`"""
    
    req_params     = [f.name for f in dataclasses.fields(Qc_Config_Dataset_config)]
    add_balance_fn = None
    
    def __init__(self, device: torch.device=torch.device("cpu"), **parameters):
        super().__init__(device, **parameters)            
        self.gate_pool = parameters["gate_pool"] #[get_obj_from_str(gate) for gate in parameters["gate_pool"]]           
            
    @property
    def params_config(self):
        params_config = super().params_config            
        #params_config["gate_pool"] = [class_to_str(gate) for gate in params_config["gate_pool"]]
        params_config = Qc_Config_Dataset_config(**params_config)
        return params_config   
    
    #----------------------------
       
    def x_y_preprocess(self, balance_max=None, shuffle=False, max_samples=None):
        #params_config = self.params_config
        #if params_config.dataset_to_gpu: self.to("cuda")
         
        z_proc = []
        for k,v in self.store_dict.items(): 
            if k != "x" and k != "y":
                z_proc.append(getattr(self, k))
                              
        x_proc, y_proc = self.x, self.y
        
        #---------------------
        if shuffle:
            x_proc, y_proc, *z_proc = shuffle_tensor_dataset(x_proc, y_proc, *z_proc)
                
        if exists(max_samples):
            x_proc = x_proc[:max_samples]
            y_proc = y_proc[:max_samples]
            z_proc = (iz[:max_samples] for iz in z_proc)   
        
        #---------------------
        t = self.store_dict["y"]
        if exists(balance_max): 
            if t == "tensor" or t == "numpy": x_proc, y_proc, *z_proc = balance_tensor_dataset(x_proc, y_proc, *z_proc, make_unique=True, 
                                                                                               samples=balance_max, add_balance_fn=self.add_balance_fn) 
            else:                             print(f"[WARNING]: Unsupported y type: `{t}`. Not balancing dataset!")
        else: print(f"[INFO]: Not balancing dataset!  {balance_max=}")
          
        #---------------------
        if shuffle:
            x_proc, y_proc, *z_proc = shuffle_tensor_dataset(x_proc, y_proc, *z_proc)
            
        return x_proc, y_proc, *z_proc
    
    def valid_split(self, x, y, *z, p_valid=0.1):
        splits = max(int(x.shape[0] * p_valid), 1)
        x, x_valid =    x[splits:].clone(), x[:splits].clone() 
        
        t = self.store_dict["y"]
        if   t == "tensor" : y, y_valid =    y[splits:].clone(), y[:splits].clone()      
        elif t == "numpy":   y, y_valid =    y[splits:]        , y[:splits]      
            
        else: raise NotImplementedError("Not implemented")
        
        try:
            z       = list(iz[splits:].clone() for iz in z)
            z_valid = list(iz[:splits].clone() for iz in z)     
        except:
            z       = list(iz[splits:] for iz in z)
            z_valid = list(iz[:splits] for iz in z)   
        
        return x, x_valid, y, y_valid, (z, z_valid)
    
    def get_dataloaders(self, batch_size, p_valid=0.1, balance_max=None, max_samples=None, y_on_cpu=False):
 
        excepts = []
        if y_on_cpu: excepts.append("y")
        if self.params_config.dataset_to_gpu: self.to("cuda", excepts=excepts)
               
        x_proc, y_proc, *z_proc              = self.x_y_preprocess(balance_max=balance_max, max_samples=max_samples)       
        x, x_valid, y, y_valid, (z, z_valid) = self.valid_split(x_proc, y_proc, *z_proc, p_valid=p_valid)
               
        ds       = TensorDataset(x, y, *z)
        ds_valid = TensorDataset(x_valid, y_valid, *z_valid)
        
        if self.params_config.dataset_to_gpu: 
            train_loader = DataLoader(dataset=ds      , batch_size=batch_size, shuffle=True)
            valid_loader = DataLoader(dataset=ds_valid, batch_size=batch_size, shuffle=True)

        else:              
            train_loader = DataLoader(dataset=ds      , batch_size=batch_size, shuffle=True, pin_memory=True, num_workers=12)
            valid_loader = DataLoader(dataset=ds_valid, batch_size=batch_size, shuffle=True, pin_memory=True, num_workers=12)

        self.dataloaders = DataLoaders(train_loader, valid_loader)        
        return self.dataloaders

    #----------------------------
  
    def plot_example(self):
        params_config = self.params_config
        enc_tensor = self.x[0]
        
        while enc_tensor.dim()>2: enc_tensor=enc_tensor[0]

        params = None
        if hasattr(self, "params"): params=self.params[0]

        if isinstance(self.gate_pool[0], str):
            gate_pool = [instruction_name_to_qiskit_gate(gate) for gate in self.gate_pool]
        else:
            gate_pool = self.gate_pool
        
        qc = decode_circuit(enc_tensor, gate_pool, params_tensor=params)
 
        t = self.store_dict["y"]
        if   t == "tensor"     : label = self.y[0].cpu().tolist()
        elif t == "tensor_list": 
            print("Not implemented")
            return
        else : 
            label = self.y[0]#.tolist()
            while len(label.shape)>0: label=label[0]
        
        print(f"Label: ``{label}``    SRV is: {schmidt_rank_vector(qi.DensityMatrix(qc))}")
        display(qc.draw("mpl", plot_barriers=False))

    def plot_distribution(self):
        if hasattr(self, "dataloaders"): x, y, *z = self.dataloaders.train.dataset.tensors  
        else:                            x, y     = self.x, self.y
            
        t = self.store_dict["y"]
        if   t == "tensor"     : data={"svr":[iy for iy in y.cpu().tolist()]}
        elif t == "numpy": data={"svr":[iy for iy in y.tolist()]}
        else:   #  list tensor_list        
            print("Not implemented")
            return
                        
        print("Train dataset (x, y):", x.shape, y.shape)
        print("Train uniques  x    :", torch.unique(x, dim=0).shape)              
              
        #real data distribution     
        df = pd.DataFrame(data)   
        cnts = df['svr'].value_counts(normalize=True)
        for n,v in zip(cnts.index, cnts.values): print(f"{n}: {v*100:.1f}%")    
        ax = df['svr'].value_counts().plot(kind='bar')
