"""Sampling functions for model inference."""

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

# %% auto #0
__all__ = ['get_batch_samples', 'batched_sampling', 'prepare_prompts', 'generate_tensors', 'generate_compilation_tensors',
           'decode_tensors_to_backend']

# %% ../../src/inference/sampling.ipynb #54fea3fd-4a6f-43cc-a2ff-b2cb35d2d626
from ..imports import *
from ..utils.async_fn import run_parallel_jobs
from ..platform.simulation import Simulator 
from ..platform.tokenizer.base_tokenizer import BaseTokenizer
from ..pipeline.pipeline import Pipeline

# %% ../../src/inference/sampling.ipynb #ff3a4d77-0335-4ffd-9190-900b0efe24dd
def get_batch_samples(samples: int, auto_batch_size: int = 512) -> list[int]:
    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)
                                                     
    assert sum(batch_samples) == samples
    return batch_samples

# %% ../../src/inference/sampling.ipynb #8a47a3b2-84a4-44b0-846b-bba6f9bd878b
def batched_sampling(pipeline: Pipeline,
                     cond_kwargs: dict[str, torch.Tensor], 
                     samples: int, 
                     system_size: int, 
                     num_of_qubits: int, 
                     max_gates: int, 
                     g: float = 1.0, 
                     init_latents: Optional[torch.Tensor] = None,
                     no_bar: bool = True, 
                     unique: bool = False, 
                     auto_batch_size: int = 512, 
                     enable_params: bool = True, 
                     reduce_spatial: bool = True,
                     return_predicted_x0: bool = False):
    
    """ e.g. cond_kwargs.keys = {"c", "micro_cond", "negative_c", "U"} """

    assert "c" in cond_kwargs
    
    c_in = cond_kwargs["c"].shape[0]
    if c_in == 1:
        # Same conditions for all samples
        for cond in cond_kwargs.values():
            assert cond.shape[0] == 1

        cond_kwargs = {kw : val.repeat(auto_batch_size, *[1]*(val.dim()-1)) 
                       for kw, val in cond_kwargs.items()}
    
    else:
        # Different conditions for all samples
        for cond in cond_kwargs.values():
            assert cond.shape[0] == samples
        
    cond_kwargs = {kw:val.to(pipeline.device) 
                   for kw, val in cond_kwargs.items()}    

    #----------------------------------------
    if exists(init_latents):
        assert init_latents.shape[0] == samples
        init_latents = init_latents.to(pipeline.device)
    
    #----------------------------------------
    
    # Sample and post process to tensor encodings    
    batch_samples = get_batch_samples(samples=samples, auto_batch_size=auto_batch_size)

    #----------------------------------------
    
    off = 0
    out_tensor_list   = []
    predicted_x0_list = []
    
    for batch_sample in batch_samples:  
        #------------
        if c_in == 1:
            # Same conditions for all samples
            _cond_kwargs = {kw:val[:batch_sample] 
                            for kw, val in cond_kwargs.items()}       
        else:
            # Different conditions for all samples
            _cond_kwargs = {kw:val[off:off+batch_sample] 
                            for kw, val in cond_kwargs.items()}
            
        #------------
        if exists(init_latents):
            latents = init_latents[off:off+batch_sample] 
            
        else:
            if pipeline.embedder.channel_last:
                latents = torch.randn((batch_sample, system_size, max_gates, pipeline.model.params_config.clr_dim))    
            else:
                latents = torch.randn((batch_sample, pipeline.model.params_config.clr_dim, system_size, max_gates))    

        off += batch_sample

        #------------
        out_tensor = pipeline.denoising(latents=latents, 
                                        g=g, 
                                        no_bar=no_bar, 
                                        # enable_guidance=True, 
                                        return_predicted_x0=return_predicted_x0,
                                        **_cond_kwargs)   

        if return_predicted_x0:
            out_tensor, predicted_x0 = out_tensor

        out_tensor_list.append(out_tensor)

        if return_predicted_x0:
            # predicted_x0 ... [timesteps, *out_tensor.shape]
            predicted_x0_list.append(predicted_x0)

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

    out_tensor_raw = torch.cat(out_tensor_list).to(pipeline.device)

    if return_predicted_x0:
        predicted_x0_raw = torch.cat(predicted_x0_list, dim=1).to(pipeline.device)

    if enable_params: out_tensor, params = pipeline.embedder.invert(out_tensor_raw, reduce_spatial=reduce_spatial)
    else:             out_tensor         = pipeline.embedder.invert(out_tensor_raw)
              
    #----------------------------------------
    
    out_tensor = out_tensor[:, :num_of_qubits] 
    
    if unique: 
        if enable_params: 
            raise NotImplementedError("We have unique and enable_params enabled, how should we handle that?")
        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")

    if enable_params:       
        if return_predicted_x0:
            return out_tensor, params, predicted_x0_raw
        return out_tensor, params
        
    elif return_predicted_x0:
        return out_tensor, predicted_x0_raw
        
    return out_tensor

# %% ../../src/inference/sampling.ipynb #fe9a3240-a3e6-4901-870d-677a3ea27376
def prepare_prompts(pipeline: Pipeline, 
                    prompt: str | Sequence[str], 
                    negative_prompt: Optional[str | Sequence[str]] = None):
    
    # Prepare conditions
    c = pipeline.text_encoder.tokenize_and_push_to_device(prompt)

    if exists(negative_prompt):
        negative_c = pipeline.text_encoder.tokenize_and_push_to_device(negative_prompt)
        assert negative_c.shape[0] == 1
    else:
        negative_c = None

    return c, negative_c

# %% ../../src/inference/sampling.ipynb #ece2684d-7a76-47f4-8791-73d71867863a
def generate_tensors(pipeline: Pipeline, 
                     prompt: str | Sequence[str], 
                     samples: int,
                     system_size: int, 
                     num_of_qubits: int, 
                     max_gates: int, 
                     g: float = 1.0, 
                     init_latents: Optional[torch.Tensor] = None,
                     no_bar: bool = True, 
                     unique: bool = False, 
                     auto_batch_size: int = 512, 
                     enable_params: bool = False,
                     reduce_spatial: bool = True,
                     return_predicted_x0: bool = False,
                     negative_prompt: Optional[str | Sequence[str]] = None,
                     micro_cond: Optional[torch.Tensor] = None) -> torch.Tensor:

    if exists(micro_cond):
        raise NotImplementedError()

    # Prepare conditions
    c, negative_c = prepare_prompts(pipeline, prompt, negative_prompt)
    
    cond_kwargs = {"c":c}
    if exists(negative_c): cond_kwargs["negative_c"] = negative_c
    if exists(micro_cond): cond_kwargs["micro_cond"] = micro_cond

    # Perform sampling
    out = batched_sampling(pipeline=pipeline, 
                           cond_kwargs=cond_kwargs, 
                           samples=samples, 
                           system_size=system_size,
                           num_of_qubits=num_of_qubits,
                           max_gates=max_gates,
                           g=g,  
                           init_latents=init_latents,
                           no_bar=no_bar,
                           unique=unique,
                           auto_batch_size=auto_batch_size,
                           enable_params=enable_params,
                           reduce_spatial=reduce_spatial,
                           return_predicted_x0=return_predicted_x0)
    return out

# %% ../../src/inference/sampling.ipynb #3e25a4f6-c9fc-404c-b415-b7b68be998bd
def generate_compilation_tensors(pipeline: Pipeline, 
                                 prompt: str | Sequence[str], 
                                 U: torch.Tensor, 
                                 samples: int, 
                                 system_size: int, 
                                 num_of_qubits: int, 
                                 max_gates: int, 
                                 g: float = 1.0, 
                                 tensor_prod_pad: bool = True,
                                 init_latents: Optional[torch.Tensor] = None,
                                 no_bar: bool = True, 
                                 unique: bool = False, 
                                 auto_batch_size: int = 512, 
                                 enable_params: bool = True, 
                                 reduce_spatial: bool = True,
                                 return_predicted_x0: bool = False,
                                 negative_prompt: Optional[str | Sequence[str]] = None,
                                 negative_u: Optional[torch.Tensor] = None,
                                 micro_cond: Optional[torch.Tensor] = None) -> torch.Tensor:
    """
    Samples tensor encodings from the DM for the given sample parameters.

    What kind of unitary padding we have depends on what we used for model training, so it depends on the concrete model weights.
    """

    if torch.is_complex(U):
        U = torch.stack([U.real, U.imag], dim=-3)
    
    if exists(micro_cond):
        raise NotImplementedError()

    # Prepare conditions
    c, negative_c = prepare_prompts(pipeline, prompt, negative_prompt)

    cond_kwargs = {"c":c}
    if exists(negative_c): cond_kwargs["negative_c"] = negative_c
    if exists(micro_cond): cond_kwargs["micro_cond"] = micro_cond

    def tensor_pad(U):
        # Prepare unitary condition
        assert U.dim() in [3, 4]
        if U.dim() == 3: 
            # [2, N, N] to [1, 2, N, N]
            U = U.unsqueeze(0)
        
        if system_size > num_of_qubits:
            N = 2**system_size
     
            if tensor_prod_pad:
                # Pad with identity tensor product, assume Big Endian 
                
                U_pad = torch.zeros((U.shape[0], 2, N, N), device=U.device, dtype=U.dtype)
    
                U_side = U.shape[-1]
                for jj in range(N//U_side):  
                    _slice = slice(U_side * jj, U_side * (jj+1))
                    U_pad[..., _slice, _slice] = U         
    
                U = U_pad
            
            else:
                # zero pad
                pad = (0, N-U.shape[-1], 0, N-U.shape[-2]) 
                U   = F.pad(U, pad, "constant", 0)
        return U
    
    cond_kwargs["U"] = tensor_pad(U)
    if exists(negative_u): 
        cond_kwargs["negative_u"] = tensor_pad(negative_u)
        
    # Perform sampling
    out = batched_sampling(pipeline=pipeline, 
                           cond_kwargs=cond_kwargs, 
                           samples=samples, 
                           system_size=system_size,
                           num_of_qubits=num_of_qubits,
                           max_gates=max_gates,
                           g=g,  
                           init_latents=init_latents,
                           no_bar=no_bar,
                           unique=unique,
                           auto_batch_size=auto_batch_size,
                           enable_params=enable_params,
                           reduce_spatial=reduce_spatial,
                           return_predicted_x0=return_predicted_x0)
    return out

# %% ../../src/inference/sampling.ipynb #fc14c810-6793-4f64-a421-7cc902ec38b8
def decode_tensors_to_backend(simulator: Simulator, 
                              tokenizer: BaseTokenizer, 
                              tensors: torch.Tensor, 
                              params: Optional[torch.Tensor] = None, 
                              silent: bool = True,
                              n_jobs: int = 1,
                              filter_errs: bool = True,
                              return_tensors: bool = False,
                    ) -> tuple[Sequence[any], int] | tuple[Sequence[any], int, torch.Tensor]:
    tensors = tensors.cpu()

    if exists(params):
        params  = params.cpu()
        iter_pack = zip(tensors, params)
        _decode   = lambda x, p: tokenizer.decode(x, p)
        
    else:
        iter_pack = zip(tensors, )
        _decode   = lambda x: tokenizer.decode(x)
    
    def _f(iter_vars):
        try:
            instructions = _decode(*iter_vars)
            backend_obj  = simulator.backend.genqc_to_backend(instructions, place_barriers=False)
            return backend_obj
        except Exception as err:
            if silent: return None
            raise err
        
    pot_qcs = run_parallel_jobs(_f, iter_pack, n_jobs)

    if filter_errs:
        backend_obj_list = [pot_qc for pot_qc in pot_qcs if exists(pot_qc)]
        err_cnt          = sum(1 for pot_qc in pot_qcs if not_exists(pot_qc))
        assert len(backend_obj_list) + err_cnt == len(pot_qcs)

        if return_tensors:
            tensors = tensors[torch.tensor([exists(pot_qc) for pot_qc in pot_qcs])]
        
    else:
        backend_obj_list = pot_qcs
        err_cnt = None

    if return_tensors:
        return backend_obj_list, err_cnt, tensors
    return backend_obj_list, err_cnt
