from enum import Enum
from typing import Sequence, Tuple, overload

from ..common import WordAlignmentMatrix

class Aligner:
    def get_src_word_index(self, word: str) -> int: ...
    def get_trg_word_index(self, word: str) -> int: ...
    @overload
    def get_best_alignment(self, src_sentence: str, trg_sentence: str) -> Tuple[float, WordAlignmentMatrix]: ...
    @overload
    def get_best_alignment(
        self, src_sentence: Sequence[int], trg_sentence: Sequence[int]
    ) -> Tuple[float, WordAlignmentMatrix]: ...
    @overload
    def get_best_alignment(
        self, src_sentence: Sequence[str], trg_sentence: Sequence[str]
    ) -> Tuple[float, WordAlignmentMatrix]: ...
    @overload
    def get_best_alignments(
        self, src_sentences: Sequence[Sequence[int]], trg_sentences: Sequence[Sequence[int]]
    ) -> Sequence[Tuple[float, WordAlignmentMatrix]]: ...
    @overload
    def get_best_alignments(
        self, src_sentences: Sequence[Sequence[str]], trg_sentences: Sequence[Sequence[str]]
    ) -> Sequence[Tuple[float, WordAlignmentMatrix]]: ...

class SymmetrizationHeuristic(Enum):
    NONE = ...
    UNION = ...
    INTERSECTION = ...
    OCH = ...
    GROW = ...
    GROW_DIAG = ...
    GROW_DIAG_FINAL = ...
    GROW_DIAG_FINAL_AND = ...

class SymmetrizedAligner(Aligner):
    def __init__(self, direct_aligner: Aligner, inverse_aligner: Aligner) -> None: ...
    @property
    def heuristic(self) -> SymmetrizationHeuristic: ...
    @heuristic.setter
    def heuristic(self, value: SymmetrizationHeuristic) -> None: ...

class AlignmentModel(Aligner):
    @property
    def num_sentence_pairs(self) -> int: ...
    @property
    def src_vocab_size(self) -> int: ...
    @property
    def trg_vocab_size(self) -> int: ...
    @property
    def variational_bayes(self) -> bool: ...
    @variational_bayes.setter
    def variational_bayes(self, value: bool) -> None: ...
    def add_sentence_pair(
        self, src_sentence: Sequence[str], trg_sentence: Sequence[str], count: float = 1
    ) -> Tuple[int, int]: ...
    def get_sentence_pair(self, n: int) -> Tuple[Sequence[str], Sequence[str], float]: ...
    def add_src_word(self, word: str) -> int: ...
    def get_src_word(self, word_index: int) -> str: ...
    def src_word_exists(self, word: str) -> bool: ...
    def add_trg_word(self, word: str) -> int: ...
    def get_trg_word(self, word_index: int) -> str: ...
    def trg_word_exists(self, word: str) -> bool: ...
    def start_training(self) -> None: ...
    def train(self) -> None: ...
    def end_training(self) -> None: ...
    def get_sentence_length_log_prob(self, src_length: int, trg_length: int) -> float: ...
    def get_sentence_length_prob(self, src_length: int, trg_length: int) -> float: ...
    def get_translation_log_prob(self, src_word_index: int, trg_word_index: int) -> float: ...
    def get_translation_prob(self, src_word_index: int, trg_word_index: int) -> float: ...
    def get_translations(self, s: int, threshold: float = 0) -> Sequence[Tuple[int, float]]: ...
    def load(self, prefix_filename: str) -> None: ...
    def print(self, prefix_filename: str) -> None: ...
    def clear(self) -> None: ...

class IncrAlignmentModel(AlignmentModel):
    def start_incr_training(self, sentence_pair_range: Tuple[int, int]) -> None: ...
    def incr_train(self, sentence_pair_range: Tuple[int, int]) -> None: ...
    def end_incr_training(self) -> None: ...

class Ibm1AlignmentModel(AlignmentModel):
    def __init__(self) -> None: ...

class IncrIbm1AlignmentModel(Ibm1AlignmentModel, IncrAlignmentModel):
    def __init__(self) -> None: ...

class Ibm2AlignmentModel(Ibm1AlignmentModel, AlignmentModel):
    @overload
    def __init__(self) -> None: ...
    @overload
    def __init__(self, model: Ibm1AlignmentModel) -> None: ...
    def get_alignment_log_prob(self, j: int, src_length: int, trg_length: int, i: int) -> float: ...
    def get_alignment_prob(self, j: int, src_length: int, trg_length: int, i: int) -> float: ...

class IncrIbm2AlignmentModel(Ibm2AlignmentModel, IncrAlignmentModel):
    def __init__(self) -> None: ...

class HmmAlignmentModel(Ibm1AlignmentModel, AlignmentModel):
    @overload
    def __init__(self) -> None: ...
    @overload
    def __init__(self, model: Ibm1AlignmentModel) -> None: ...
    def get_alignment_log_prob(self, prev_i: int, src_length: int, i: int) -> float: ...
    def get_alignment_prob(self, prev_i: int, src_length: int, i: int) -> float: ...

class IncrHmmAlignmentModel(HmmAlignmentModel, IncrAlignmentModel):
    def __init__(self) -> None: ...

class FastAlignModel(IncrAlignmentModel):
    def __init__(self) -> None: ...
    def get_alignment_log_prob(self, j: int, src_length: int, trg_length: int, i: int) -> float: ...
    def get_alignment_prob(self, j: int, src_length: int, trg_length: int, i: int) -> float: ...

class Ibm3AlignmentModel(Ibm2AlignmentModel, Ibm1AlignmentModel, AlignmentModel):
    @overload
    def __init__(self) -> None: ...
    @overload
    def __init__(self, model: HmmAlignmentModel) -> None: ...
    @overload
    def __init__(self, model: Ibm2AlignmentModel) -> None: ...
    def get_distortion_log_prob(self, i: int, src_length: int, trg_length: int, j: int) -> float: ...
    def get_distortion_prob(self, i: int, src_length: int, trg_length: int, j: int) -> float: ...
    def get_fertility_log_prob(self, src_word_index: int, fertility: int) -> float: ...
    def get_fertility_prob(self, src_word_index: int, fertility: int) -> float: ...

class Ibm4AlignmentModel(Ibm3AlignmentModel, Ibm2AlignmentModel, Ibm1AlignmentModel, AlignmentModel):
    @overload
    def __init__(self) -> None: ...
    @overload
    def __init__(self, model: Ibm3AlignmentModel) -> None: ...
    def get_head_distortion_log_prob(
        self, src_word_class: int, trg_word_class: int, trg_length: int, dj: int
    ) -> float: ...
    def get_head_distortion_prob(self, src_word_class: int, trg_word_class: int, trg_length: int, dj: int) -> float: ...
    def get_nonhead_distortion_log_prob(self, trg_word_class: int, trg_length: int, dj: int) -> float: ...
    def get_nonhead_distortion_prob(self, trg_word_class: int, trg_length: int, dj: int) -> float: ...

__all__ = [
    "Aligner",
    "AlignmentModel",
    "FastAlignModel",
    "HmmAlignmentModel",
    "Ibm1AlignmentModel",
    "Ibm2AlignmentModel",
    "Ibm3AlignmentModel",
    "Ibm4AlignmentModel",
    "IncrAlignmentModel",
    "IncrHmmAlignmentModel",
    "IncrIbm1AlignmentModel",
    "IncrIbm2AlignmentModel",
    "SymmetrizationHeuristic",
    "SymmetrizedAligner",
]
