"""Class to further tokenize tensor representations."""

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

# %% auto #0
__all__ = ['GatePairTokenizer', 'sort_config', 'get_topk_depth_unpacked']

# %% ../../../src/platform/tokenizer/tensor_tokenizer.ipynb #06272f6f-b4e3-4504-a90a-feebbf6ad821
from ...imports import *
from .base_tokenizer import BaseTokenizer, Vocabulary

# %% ../../../src/platform/tokenizer/tensor_tokenizer.ipynb #4a4b051b-eba3-4b3a-9bec-c5c44655bf26
class GatePairTokenizer(BaseTokenizer):

    def __init__(self, unique_class_values, zero_token, padding_token, device):
        super().__init__({})
        
        self.padding_token    = padding_token
        self.not_gates_tokens = torch.tensor([zero_token, padding_token]).to(device)

        self.current_tokens = torch.tensor(unique_class_values, device=device)
        self._current_depth = 0
        
        self.token_lookup = {}      #reduced forms, used for gadget extraction
        self.token_lookup_raw = {}  #the raw form, used for encoding
        
        self.token_depth  = {tok:0 for tok in self.current_tokens.cpu().tolist()}
        self.token_cnts   = {}
    
    def learn(self, tensors, max_depth, max_iters):
        # loop over get bets and then replace
        
        current_tensor = tensors
        self._current_depth = 0
        
        for i in tqdm(range(max_iters), total=max_iters):
        
            overlap_pairs          = self.extract_new_gate_overlap_pairs(current_tensor)
            overlap_pairs_std_form = self.standardize_overlap_pairs(overlap_pairs)

            top_pairs, topv = self.get_topk_pairs(overlap_pairs_std_form, k=1)
            top_pair = top_pairs[0]
       
            if top_pair.abs().sum() < 1:
                print("break: top_pair.abs().sum() < 1")
                break

            if topv < 2:
                print("break: no more pair with cnt > 1")
                break
            
            current_tensor = self.learn_step(current_tensor, top_pair, topv=topv)

            current_max_depth = max(self.token_depth.values())
            if current_max_depth > max_depth:
                print(f"break: max_depth {max_depth} reached")
                break
            
        print("break: max_iters reached")
        return current_tensor

    def to(self, device):
        self.not_gates_tokens = self.not_gates_tokens.to(device) 
        self.current_tokens   = self.current_tokens.to(device) 
        
        for k, v in self.token_lookup.items():
            self.token_lookup[k] = self.token_lookup[k].to(device) 
            
        for k, v in self.token_lookup_raw.items():
            self.token_lookup_raw[k] = self.token_lookup_raw[k].to(device) 
        
        return self

    def tokenize(self, tensors):
        """Identical to `GatePairTokenizer.encode`."""
        return self.encode(tensors=tensors)
    
    #---------------------------------------
    # Replace pairs with new tokens

    def learn_step(self, current_tensor, top_pair, new_tokens: Optional[torch.Tensor] = None, topv: Optional[torch.Tensor] = None):

        top_pair_reduced = top_pair[top_pair.abs().sum(-1)>0].unique_consecutive(dim=0)

        if not_exists(new_tokens):
            new_tokens = self.current_tokens.max() + 1 + torch.arange(top_pair_reduced.shape[0], device=current_tensor.device)
            self.current_tokens = torch.cat([self.current_tokens, new_tokens])

            key = tuple(new_tokens.cpu().tolist())
            self.token_lookup[key]     = top_pair_reduced  #top_pair[top_pair.sum(-1)>0]
            self.token_lookup_raw[key] = top_pair
            self.token_cnts[key]   = topv
            
            _current_depth = max(self.token_depth[k] for k in top_pair_reduced.flatten().cpu().tolist()) + 1
            
            for tok in new_tokens.cpu().tolist():
                self.token_depth[tok] = _current_depth

            if _current_depth > self._current_depth:
                self._current_depth = _current_depth
                print(f"New depth reached {self._current_depth}")
        else:
            assert top_pair_reduced.shape[0] == new_tokens.shape[0]
        
        # 1) Replace one all even pairs
        current_overlap_pairs = self.extract_current_gate_overlap_pairs(current_tensor, odd_pairs=False)
        current_tensor = self.replace_current_overlap_pairs(current_tensor, current_overlap_pairs, top_pair, top_pair_reduced, new_tokens, odd_pairs=False)
    
        # 2) Then Replace one all odd pairs
        current_overlap_pairs = self.extract_current_gate_overlap_pairs(current_tensor, odd_pairs=True)
        current_tensor = self.replace_current_overlap_pairs(current_tensor, current_overlap_pairs, top_pair, top_pair_reduced, new_tokens, odd_pairs=True)
        
        return current_tensor
    
    def extract_current_gate_overlap_pairs(self, current_tensor, odd_pairs: bool = True):
        # Extract overlap_pairs
        # ToDo optimize loops

        seq = current_tensor.shape[-1]
        seq_half = seq // 2 
        assert seq % 2 == 0

        overlap_pairs = []
        for current_tensor_i in current_tensor:
            _overlap_pairs = []

            if odd_pairs:
                for t in range(seq_half-1):
                    _overlap_pairs.append(current_tensor_i[:, 1+2*t:1+2*(t+1)])
            else:
                for t in range(seq_half):
                    _overlap_pairs.append(current_tensor_i[:, 2*t:2*(t+1)])
                
            overlap_pairs.append(torch.stack(_overlap_pairs))
            
        overlap_pairs = torch.stack(overlap_pairs)
        return overlap_pairs
        
    def replace_current_overlap_pairs(self, current_tensor, overlap_pairs, top_pair, top_pair_reduced, new_tokens, odd_pairs):

        overlap_pairs_std = self.standardize_overlap_pairs(overlap_pairs)
        is_top_overlap_pair = (overlap_pairs_std==top_pair).all(dim=(-1,-2), keepdim=False)

        new_tensor = torch.full_like(current_tensor, self.padding_token)
        
        for i in range(is_top_overlap_pair.shape[0]): #ToDo: this loop can be put in parallel! is batch dim
            t = 1 if odd_pairs else 0
            
            for j in range(is_top_overlap_pair.shape[1]):
        
                if is_top_overlap_pair[i, j]: #replace
                      
                    new_col = torch.zeros((current_tensor.shape[1]), dtype=overlap_pairs.dtype, device=overlap_pairs.device)
            
                    for new_token, top_pair_reduced_i in zip(new_tokens, top_pair_reduced):
                        ind = (overlap_pairs[i, j]==top_pair_reduced_i).all(-1)
                        new_col = torch.where(ind, new_token, new_col)
                    
                    new_col = new_col.unsqueeze(-1)
                    
                    tp1 = t + 1
                 
                else: #just copy old
                    new_col = overlap_pairs[i, j]
                    tp1 = t + 2
            
                new_tensor[i, :, t:tp1] = new_col
                t = tp1

            if odd_pairs:
                # copy first and last col
                new_tensor[..., 0]  = current_tensor[..., 0]
                new_tensor[..., -1] = current_tensor[..., -1]
            
        return new_tensor         
        
    #---------------------------------------
    # Find new pairs
    
    def extract_new_gate_overlap_pairs(self, current_tensor):
        #current_tensor = current_tensor.abs()

        isgate_token = 1 - torch.isin(current_tensor.abs(), self.not_gates_tokens.to(current_tensor.device)).int()

        # These are postions of the pairs (therefore shape-1) in which we have an overlap
        overlaps = isgate_token[..., :-1] + isgate_token[..., 1:]
        overlaps = (overlaps>1).int()

        # Number of overlaps two gates have!  we can say here only take 2 overlaps, or min 2, or min 1, eg..    ->  0 means parallel!!
        overlaps_cnt = torch.count_nonzero(overlaps, dim=1) 
        overlaps_ind = (overlaps_cnt>0)
    
        # Extract overlap_pairs
        # ToDo optimize loops
        
        overlap_pairs = []
        for current_tensor_i, overlaps_ind_i in zip(current_tensor, overlaps_ind):
            for t in range(current_tensor_i.shape[-1]-1):
                if overlaps_ind_i[t]:
                    overlap_pairs.append(current_tensor_i[:, t:t+2])

        overlap_pairs = torch.stack(overlap_pairs)
        return overlap_pairs

    def standardize_overlap_pairs(self, overlap_pairs):
        # Now we convert to std form, where the 1st gate gives the main order and the 2nd the secondory, this should remove all(?) redundant combinations!

        # 1) sort inner SECOND gate such that gate 2 is always on top
        inner_sorted_gate2, inner_sorted_gate2_indices =   torch.sort(overlap_pairs[..., 1], dim=-1, descending=True, stable=False)
        inner_sorted_gate1                             = torch.gather(overlap_pairs[..., 0], dim=-1, index=inner_sorted_gate2_indices)
        
        inner_overlap_pairs = torch.stack((inner_sorted_gate1, inner_sorted_gate2), dim=-1)
        
        # 2) sort outer FISRT gate such that gate 1 is always on top, NOTE WE NEED STABLE SORT TO CONSERVE INNER ORDER
        outer_sorted_gate1, outer_sorted_gate1_indices =   torch.sort(inner_overlap_pairs[..., 0], dim=-1, descending=True, stable=True)
        outer_sorted_gate2                             = torch.gather(inner_overlap_pairs[..., 1], dim=-1, index=outer_sorted_gate1_indices)
        
        overlap_pairs_std_form = torch.stack((outer_sorted_gate1, outer_sorted_gate2), dim=-1)

        return overlap_pairs_std_form.contiguous()

    def get_topk_pairs(self, overlap_pairs, k):
        # Now we can easily count the unique valid pairs! 
        pot_pairs, pot_pairs_cnts = overlap_pairs.unique(dim=0, return_counts=True)
        
        # Get topk best pairs
        topv, topi = torch.topk(pot_pairs_cnts, k)
        top_pairs  = pot_pairs[topi]
        
        return top_pairs, topv

    #---------------------------------------
    # Encoding

    def encode(self, tensors):
        # just replay all the pair replacements from learn, i.e. the vocab

        s = tensors.shape[1]
        current_tensor = tensors

        for new_tokens, top_pair in tqdm(self.token_lookup_raw.items()):
            top_pair = self.standardize_vocab_pair(top_pair, s, sort=True)
            new_tokens = torch.tensor(new_tokens, device=top_pair.device, dtype=top_pair.dtype)
            
            current_tensor = self.learn_step(current_tensor, top_pair, new_tokens=new_tokens)

        return current_tensor

    def standardize_vocab_pair(self, vocab_pair, s, sort: bool = True):
        
        if vocab_pair.shape[0]<2: # repeat for special gadgets which have full symetric sequential connection
            vocab_pair = vocab_pair.repeat(2, 1)
        
        vocab_pair = F.pad(vocab_pair, [0, 0, 0, s-vocab_pair.shape[0]]) # pad to full systemsize to have nice plotting

        if sort:
            vocab_pair = self.standardize_overlap_pairs(vocab_pair)
            
        return vocab_pair.contiguous()
        
    #---------------------------------------
    # Decoding

    def unpack_col(self, col):
        # col is [s, 1]
        s, _ = col.shape
    
        current_tokens = col.unique()
        current_tokens = current_tokens[current_tokens!=0]
        k = tuple(current_tokens.tolist())
        
        if k in self.token_lookup:
    
            # Unpack one col
            unpacked   = torch.zeros((s, 2), dtype=col.dtype, device=col.device)
            new_config = self.token_lookup[k]
            
            for current_token, new_config_i in zip(current_tokens, new_config):
                ind      = (col==current_token)
                unpacked = torch.where(ind, new_config_i, unpacked)
                 
            # Repeat unpacking for both new cols
            col1, col2 = unpacked.chunk(2, dim=-1)
    
            unpacked1 = self.unpack_col(col1)
            unpacked2 = self.unpack_col(col2)
    
            unpacked = torch.cat([unpacked1, unpacked2], dim=-1)
            return unpacked
        
        return col 

    def decode(self, tensor, cut_padding: bool = False):
        # split into cols we unpack, then recursively
        # tensor ... [s, t]
        assert tensor.dim() == 2
        
        cols     = tensor.chunk(tensor.shape[-1], dim=-1)
        unpacked = torch.cat([self.unpack_col(col) for col in cols], dim=-1)

        if cut_padding:
            # Cut from right as this was added padding in packing
            unpacked = unpacked[..., :tensor.shape[-1]]
        
        return unpacked

# %% ../../../src/platform/tokenizer/tensor_tokenizer.ipynb #ce69148d-37df-47e3-9c63-5c9149d938ac
def sort_config(vocab_config):
    """Sort a vocab_config for nicer plotting."""
    
    t = vocab_config.shape[-1]
    all_inds = set(range(t))

    # Sort one ind, gather the rest
    for i in reversed(range(t)):
        gather_inds = all_inds - {i}

        sorted_gates = [None] * t
        
        sorted_gate_i, sorted_gate_i_indices = torch.sort(vocab_config[..., i], dim=-1, descending=True, stable=True)
        sorted_gates[i] = sorted_gate_i
        
        for gather_ind in gather_inds:
            sorted_gates[gather_ind] = torch.gather(vocab_config[..., gather_ind], dim=-1, index=sorted_gate_i_indices)

        vocab_config = torch.stack(sorted_gates, dim=-1)
    
    return vocab_config

# %% ../../../src/platform/tokenizer/tensor_tokenizer.ipynb #2aca2f79-49b4-43d4-8c8e-6d5759f77db0
def get_topk_depth_unpacked(gate_pair_tokenizer, s, use_raw=False, standardize=True):
    """Useful for plotting."""
    
    # Sort into depths
    unpacked_vocab_configs_depths      = {}
    unpacked_vocab_configs_cnts_depths = {}

    if use_raw:
        iters = zip(gate_pair_tokenizer.token_lookup_raw.items(), gate_pair_tokenizer.token_cnts.values())
    else:
        iters = zip(gate_pair_tokenizer.token_lookup.items(), gate_pair_tokenizer.token_cnts.values())
    
    for (vocab_tokens, vocab_config), vocab_config_cnts in tqdm(iters, total=len(gate_pair_tokenizer.token_cnts)):
    
        tok = vocab_tokens[0]
        token_depth = gate_pair_tokenizer.token_depth[tok]

        if standardize:
            vocab_config = gate_pair_tokenizer.standardize_vocab_pair(vocab_config, s, sort=False)
        unpacked_vocab_config = gate_pair_tokenizer.decode(vocab_config)
    
        #--------
        unpacked_vocab_config = sort_config(unpacked_vocab_config)
     
        if token_depth not in unpacked_vocab_configs_depths:
            unpacked_vocab_configs_depths[token_depth]      = []
            unpacked_vocab_configs_cnts_depths[token_depth] = []
            
        unpacked_vocab_configs_depths[token_depth].append(unpacked_vocab_config)
        unpacked_vocab_configs_cnts_depths[token_depth].append(vocab_config_cnts)

    return unpacked_vocab_configs_depths, unpacked_vocab_configs_cnts_depths
