# Path: nmt_chainer/translation/beam_search.py
# class BeamSearchParams:
#     """
#     Record holding parameters controling the beam search
#         beam_width: number of partial translation kept in each beam
#         beam_pruning_margin: maximum score difference accepted within a beam
#         force_finish: force the generation of EOS if we did not find a translation after nb_steps steps
#         use_unfinished_translation_if_none_found: will ureturn unfinished translation if we did not find a translation after nb_steps steps
#     """
#     beam_width:int = 20
#     beam_pruning_margin:Optional[float] = None
#     beam_score_length_normalization:str = "none" # none simple google
#     beam_score_length_normalization_strength:float = 0.2
#     beam_score_coverage_penalty:str = "none" # none google
#     beam_score_coverage_penalty_strength: float=0.2
#     force_finish:bool = False
#     use_unfinished_translation_if_none_found:bool = False
#     always_consider_eos_and_placeholders:bool = False


from nmt_chainer.translation.beam_search import BeamSearchParams
from chainer import cuda, Function, gradient_check, Variable, optimizers, serializers
from chainer import Link, Chain, ChainList
import numpy as np
import chainer
import chainer.functions as F
import chainer.links as L
import nmt_chainer.models.encoders
import nmt_chainer.models.encoder_decoder
import nmt_chainer.models.attention
import nmt_chainer.utilities.utils as utils
import nmt_chainer.models.decoder_cells as decoder_cells
import logging
        import nmt_chainer.translation.evaluation as evaluation
        import nmt_chainer.translation.beam_search as beam_search

        loss, attn = encdec(src_batch, tgt_batch, src_mask)

        total_loss_naive = 0
        total_length = 0
        for i in range(len(trg_data)):
            raw_s_raw = trg_data[i]
            raw_s = [raw_s_raw[0], raw_s_raw[1] + [eos_idx]]
            input_seq = [Variable(np.array([v], dtype=np.int32)) for v in raw_s[0]]
            tgt_seq = [Variable(np.array([v], dtype=np.int32)) for v in raw_s[1]]
            loss_naive, attn_naive = encdec.naive_call(input_seq, tgt_seq, None)
            total_loss_naive += float(loss_naive.data) * len(raw_s[1])
            total_length += len(raw_s[1])

        assert abs(total_loss_naive / total_length - float(loss.data)) < 1e-6


class TestBeamSearch:
    def test_1(self):
        Vi, Ei, Hi, Vo, Eo, Ho, Ha, Hl = 29, 37, 13, 53, 7, 12, 19, 33
        encdec = nmt_chainer.models.encoder_decoder.EncoderDecoder(
            Vi, Ei, Hi, Vo, Eo, Ho, Ha, Hl)
        eos_idx = Vo - 1
        src_data = [[2, 3, 3, 4, 4, 5], [1, 3, 8, 9, 2]]
#         best1_gen = evaluation.beam_search_translate(encdec, eos_idx, src_data, beam_width = 10, nb_steps = 15, gpu = None, beam_opt = False,
#                           need_attention = False)
#         best2_gen = evaluation.beam_search_translate(encdec, eos_idx, src_data, beam_width = 10, nb_steps = 15, gpu = None, beam_opt = True,
#                           need_attention = False)

        # TODO: not much point to this test now that beam_opt distinction is removed

# Path: nmt_chainer/translation/beam_search.py
# class BeamSearchParams:
#     """
#     Record holding parameters controling the beam search
#         beam_width: number of partial translation kept in each beam
#         beam_pruning_margin: maximum score difference accepted within a beam
#         force_finish: force the generation of EOS if we did not find a translation after nb_steps steps
#         use_unfinished_translation_if_none_found: will ureturn unfinished translation if we did not find a translation after nb_steps steps
#     """
#     beam_width:int = 20
#     beam_pruning_margin:Optional[float] = None
#     beam_score_length_normalization:str = "none" # none simple google
#     beam_score_length_normalization_strength:float = 0.2
#     beam_score_coverage_penalty:str = "none" # none google
#     beam_score_coverage_penalty_strength: float=0.2
#     force_finish:bool = False
#     use_unfinished_translation_if_none_found:bool = False
#     always_consider_eos_and_placeholders:bool = False


# Path: tests/suite1/models_test.py
from nmt_chainer.translation.beam_search import BeamSearchParams
from chainer import cuda, Function, gradient_check, Variable, optimizers, serializers
from chainer import Link, Chain, ChainList
import numpy as np
import chainer
import chainer.functions as F
import chainer.links as L
import nmt_chainer.models.encoders
import nmt_chainer.models.encoder_decoder
import nmt_chainer.models.attention
import nmt_chainer.utilities.utils as utils
import nmt_chainer.models.decoder_cells as decoder_cells
import logging
        import nmt_chainer.translation.evaluation as evaluation
        import nmt_chainer.translation.beam_search as beam_search


        loss, attn = encdec(src_batch, tgt_batch, src_mask)

        total_loss_naive = 0
        total_length = 0
        for i in range(len(trg_data)):
            raw_s_raw = trg_data[i]
            raw_s = [raw_s_raw[0], raw_s_raw[1] + [eos_idx]]
            input_seq = [Variable(np.array([v], dtype=np.int32)) for v in raw_s[0]]
            tgt_seq = [Variable(np.array([v], dtype=np.int32)) for v in raw_s[1]]
            loss_naive, attn_naive = encdec.naive_call(input_seq, tgt_seq, None)
            total_loss_naive += float(loss_naive.data) * len(raw_s[1])
            total_length += len(raw_s[1])

        assert abs(total_loss_naive / total_length - float(loss.data)) < 1e-6


class TestBeamSearch:
    def test_1(self):
        Vi, Ei, Hi, Vo, Eo, Ho, Ha, Hl = 29, 37, 13, 53, 7, 12, 19, 33
        encdec = nmt_chainer.models.encoder_decoder.EncoderDecoder(
            Vi, Ei, Hi, Vo, Eo, Ho, Ha, Hl)
        eos_idx = Vo - 1
        src_data = [[2, 3, 3, 4, 4, 5], [1, 3, 8, 9, 2]]
#         best1_gen = evaluation.beam_search_translate(encdec, eos_idx, src_data, beam_width = 10, nb_steps = 15, gpu = None, beam_opt = False,
#                           need_attention = False)
#         best2_gen = evaluation.beam_search_translate(encdec, eos_idx, src_data, beam_width = 10, nb_steps = 15, gpu = None, beam_opt = True,
#                           need_attention = False)

        # TODO: not much point to this test now that beam_opt distinction is removed

        best1_gen = evaluation.beam_search_translate(encdec, eos_idx, src_data, beam_search.BeamSearchParams(beam_width=10), nb_steps=15, gpu=None,