import numpy as np


def softmax(x):
    """Compute softmax values for each sets of scores in x."""
    e_x = np.exp(x - np.max(x))
    return e_x / e_x.sum()


def reduce_mul(l):
    out = 1.0
    for x in l:
        out *= x
    return out


def decode_step(step, encoder_logits):
    words_prob = encoder_logits[step]
    words_prob = softmax(words_prob)
    ouput_step = [(idx, prob) for idx, prob in enumerate(words_prob)]
    ouput_step = sorted(ouput_step, key=lambda x: x[1], reverse=True)
    return ouput_step


def check_exceed(seq, mask, max_count):
    count = 0
    for i, word in enumerate(seq):
        sign_index = word[0]
        if sign_index > 0 and mask[i]:
            count += 1
    if count > max_count:
        return False
    else:
        return True


def beam_search_step(step, encoder_logits, top_seqs, mask, k, max_count):
    all_seqs = []
    for seq in top_seqs:
        seq_score = reduce_mul([_score for _, _score in seq])
        # get current step using encoder_context & seq
        current_step = decode_step(step, encoder_logits)
        for i, word in enumerate(current_step):
            if i >= k:
                break
            word_index, word_score = word
            score = seq_score * word_score
            rs_seq = seq + [word]
            all_seqs.append((rs_seq, score))
    all_seqs = sorted(all_seqs, key=lambda seq: seq[1], reverse=True)
    # Expression constraint
    filtered_seqs = [seq for seq, _ in all_seqs if check_exceed(seq, mask, max_count)]
    # topk_seqs = [seq for seq, _ in all_seqs[:k]]
    topk_seqs = [seq for seq in filtered_seqs[:k]]
    return topk_seqs


def beam_search(encoder_logits, mask, beam_size, max_count):
    max_len = sum(mask)
    # START
    top_seqs = [[(0, 1.0)]]
    # loop
    for i in range(1, max_len + 1):
        top_seqs = beam_search_step(i, encoder_logits, top_seqs, mask, beam_size, max_count)

    number_indices_list, sign_indices_list, scores_list = [], [], []
    for seq in top_seqs:
        number_indices, sign_indices = [], []
        for i, word in enumerate(seq):
            sign_index, score = word
            if sign_index > 0 and mask[i]:
                number_indices.append(i)
                sign_indices.append(sign_index)
        if number_indices == [] and sign_indices == []:
            continue
        number_indices_list.append(number_indices)
        sign_indices_list.append(sign_indices)
        seq_score = reduce_mul([_score for _, _score in seq])
        scores_list.append(seq_score)
    if scores_list != []:
        scores_list = softmax(np.array(scores_list))
    return number_indices_list, sign_indices_list, scores_list.tolist()

