import torch
from tqdm import tqdm
import torch

from .kv_cache_model import KVCacheModelLade, KVCacheModelSimpleWithGuess
from dualdec.cache_engine import CacheEngine


@torch.no_grad()
def dualdec(prefix : torch.Tensor, approx_model : torch.nn.Module, target_model : torch.nn.Module, ngram_cache : CacheEngine = None,
                         max_len : int = 512 , gamma : int = 4, window_size = 20, guess_set_size = 20, lookahead_level = 7, eos_token_id = 2, topk = 3) -> torch.Tensor:
    """
    Performs dual decoding with an approximate model and a target model to generate a sequence of tokens.

    Args:
        prefix (torch.Tensor): The initial sequence of tokens to start the generation from.
        approx_model (torch.nn.Module): The approximate model used for initial token generation. The model should support huggingface transformers model methods.
        target_model (torch.nn.Module): The target model used for refining the generated tokens. The model should support huggingface transformers model methods.
        ngram_cache (CacheEngine, optional): A cache engine for storing and retrieving n-gram predictions. Defaults to None, in which case a new cache engine is created.
        max_len (int, optional): The maximum length of the generated sequence. Defaults to 512.
        gamma (int, optional): The lookahead parameter for generation. Defaults to 4.
        window_size (int, optional): The window size used for n-gram generation. Defaults to 20. Currently, must be equal to guess_set_size.
        guess_set_size (int, optional): The size of the guess set for n-gram retrieving. Defaults to 20. Currently, must be equal to window_size.
        lookahead_level (int, optional): The level of lookahead decoding. Defaults to 7.
        eos_token_id (int, optional): The token id representing the end-of-sequence token. Defaults to 2. Should be given by tokenizer.eos_token_id.

    Returns:
        torch.Tensor: The generated sequence of tokens, including the initial prefix and any additional tokens generated by the function.
    """
    assert window_size == guess_set_size, "We only support same window_size and guess_set_size now. More combinations will be supported in the future."
    if ngram_cache == None:
        ngram_cache = CacheEngine(lookahead_level, guess_set_size)
    seq_len = prefix.shape[1]
    T = seq_len + max_len
    
    assert prefix.shape[0] == 1, "input batch size must be 1"

    assert approx_model.device == target_model.device
    
    device = target_model.device

    guess_size = lookahead_level - 1
    
    approx_model_cache = KVCacheModelLade(approx_model, window_size=window_size, guess_set_size=guess_set_size, lookahead_level=lookahead_level, topk=topk)
    target_model_cache = KVCacheModelSimpleWithGuess(target_model, lookahead_level=lookahead_level)

    # target_model_cache = KVCacheModelSimple(target_model)
    
    resample_count = 0
    target_sample_count = 0
    accepted_count = 0

    end_pos = None
    
    while prefix.shape[1] < T:
        # q = M_q[prefix + x_0, x_1, .., x_(gamma-2)]
        prefix_len = prefix.shape[1]


        x, out_len, guess = approx_model_cache.generate(prefix, ngram_cache, gamma)
        target_model_cache._forward_with_kvcache(x, guess)

        key_tok = int(x[:,-1])

        gen_len = out_len - prefix_len
        
        n = prefix_len + gen_len - 1

        # print(prefix)
        

        for i in range(gen_len):
            j = x[:, prefix_len + i]
            
            t_tok = target_model_cache._prob_history[:, prefix_len + i - 1, :].argmax(dim=-1, keepdim=True).to(j.device)
            if t_tok != j:
                remaining_target_tok = target_model_cache._prob_history[:, prefix_len + i: out_len, :].argmax(dim=-1) 
                remaining_approx_tok = x[:, prefix_len + i:]
                # from IPython import embed; embed()
                approx_model_cache.update_ngram_cache(remaining_approx_tok, remaining_target_tok) # TODO
                # reject
                n = prefix_len + i - 1
                break
            if j == eos_token_id:
                end_pos = prefix_len + i + 1
            accepted_count += 1
        
        # print(f"n : {n}, i : {i}, prefix_len + gamma - 1: {prefix_len + gamma - 1}")
        assert n >= prefix_len - 1, f"n {n}, prefix_len {prefix_len}"
        prefix = x[:, :n + 1]

        approx_model_cache.rollback(n+1)

        corr_ngram = [] # ngram corrected by target_model

        if n < prefix_len + gen_len - 1:
            t = target_model_cache._prob_history[:, n, :].argmax(dim=-1, keepdim=True)
            if t == eos_token_id:
                end_pos = n + 2

            first_tok = int(target_model_cache._prob_history[:, prefix_len + gen_len - 1, :].argmax(dim=-1))
            beg_pos = prefix_len + gen_len
            guess_num = len(guess) // guess_size
            for i in range(guess_num):
                real_ngram = tuple([first_tok] + target_model_cache._prob_history[0, beg_pos + i * guess_size : beg_pos + i * guess_size + guess_size - 1, :].argmax(dim=-1).tolist())
                corr_ngram.append(real_ngram)
            if len(corr_ngram) > 0:
                approx_model_cache.update_in_place(key_tok, corr_ngram)


            target_model_cache.rollback(n+1)
            prefix = torch.cat((prefix, t.to(prefix.device)), dim=1)
        else:
            # find the longest guess
            guess = [item for sublist in guess for item in sublist]
            guess_num = len(guess) // guess_size
            first_tok = int(target_model_cache._prob_history[:, n, :].argmax(dim=-1))
            beg_pos = prefix_len + gen_len
            candidate = [first_tok]
            longest_can_len = 1
            candidate_idx = -1
            tmp_end_pos = n + 2 if first_tok == eos_token_id else None
            loc_end_pos = None
            for i in range(guess_num):
                real_ngram = [first_tok] + target_model_cache._prob_history[0, beg_pos + i * guess_size : beg_pos + i * guess_size + guess_size, :].argmax(dim=-1).tolist()
                corr_ngram.append(tuple(real_ngram[:-1]))
                pred_ngram = guess[i * guess_size : (i + 1) * guess_size]
                ml = 0
                for j in range(guess_size):
                    ml = j
                    if real_ngram[j] == eos_token_id:
                        loc_end_pos = j
                    if real_ngram[j] != pred_ngram[j]:
                        break
                if ml + 1 > longest_can_len:
                    candidate = real_ngram[:ml + 1]
                    longest_can_len = ml + 1
                    candidate_idx = i
                    tmp_end_pos = loc_end_pos
            if tmp_end_pos is not None:
                end_pos = beg_pos + candidate_idx * guess_size + tmp_end_pos + 1
            candidate = torch.tensor([candidate], device=prefix.device)
            prefix = torch.cat((prefix, candidate), dim=1)
            if len(corr_ngram) > 0:
                approx_model_cache.update_in_place(key_tok, corr_ngram)
            if candidate_idx != -1:
                target_model_cache.confirm(n+1, beg_pos + candidate_idx * guess_size, candidate.shape[-1] - 1) # cur_len, start_pos, length
            else:
                target_model_cache.rollback(n+1)
        if end_pos is not None:
            break
        


    if end_pos is not None:
        prefix = prefix[:, :end_pos]
    return prefix[:, :T]