"""Expert for character-level string transduction.

Given an input string x, a target string t, alignment index i, and a partial
prediction y, it returns optimal cost-to-go for all valid edit actions.

Also includes ActionVocabulary class for compatibility with maxwell dictionary.
Class stores valid edit actions for given dataset.
"""

import abc
import argparse
import dataclasses
from typing import (
    Any,
    Dict,
    Iterable,
    Iterator,
    List,
    Sequence,
    Set,
    Tuple,
)

import numpy
from maxwell import actions, sed

from .. import data, defaults, special


class Error(Exception):
    pass


class ActionError(Error):
    pass


class AlignerError(Error):
    pass


class ActionVocabulary:
    """Manages encoding of action vocabulary for transducer training."""

    # TODO: Port more of the logic to the dataset class.
    i2w: Dict[int, actions.Edit]
    w2i: Dict[actions.Edit, int]
    start_vocab_idx: int
    target_characters: Set[Any]

    def __init__(self, index: data.indexes.Index):
        self.target_characters = set()
        self.i2w = [
            actions.Start(),
            actions.End(),
            actions.ConditionalDel(),
            actions.ConditionalCopy(),
        ]
        self.start_vocab_idx = len(self.i2w)
        self.w2i = {w: i for i, w in enumerate(self.i2w)}
        # Uses index from dataset to create action vocabulary.
        self.encode_actions([index(t) for t in index.target_vocabulary])
        # Sets unknown character decoding.
        self.encode_actions([special.UNK_IDX])
        # Adds source characters if index has tied embeddings.
        if index.tie_embeddings:
            self.encode_actions([index(s) for s in index.source_vocabulary])

    def encode(self, symb: actions.Edit) -> int:
        """Returns index referencing symbol in encoding table.

        If the symbol is not encoded, it is added to the encoding table.

        Args:
            symb (actions.Edit): edit action to add to encoding table.

        Returns:
            int: index of symbol in encoding table.
        """
        if symb in self.w2i:
            idx = self.w2i[symb]
        else:
            idx = len(self.i2w)
            self.i2w.append(symb)
            self.w2i[symb] = idx
        return idx

    def decode(self, idx: int) -> actions.Edit:
        """Returns symbol referenced by given index in encoding table.

        Args:
            idx (int): encoding index that references symbol in encoding table.

        Returns:
            actions.Edit: action in encoding table referenced by given index.
        """
        return self.i2w[idx]

    def encode_actions(self, target: Any) -> None:
        """Encodes all possible edit actions for given target.

        Args:
            target (Any): target symbol produced by encoded edits.
        """
        for c in target:
            if c in self.target_characters:
                continue
            self.encode(actions.ConditionalSub(c))
            self.encode(actions.ConditionalIns(c))
            self.target_characters.add(c)

    def encode_unseen_action(self, action: actions.Edit) -> int:
        """Encodes action unseen in training.

        Operates same as encode_action but does not expand lookup table.

        Args:
            actions (actions.Edit).

        Returns:
            int.
        """
        if action in self.w2i:
            return self.lookup(action)
        else:
            raise ActionError(f"Action {action} is out of vocabulary")

    def __len__(self) -> int:
        return len(self.i2w)

    def __repr__(self) -> str:
        return f"Vocabulary({str(self.w2i)})"

    def lookup(self, word: str) -> int:
        return self.w2i[word]

    def to_i2w(self) -> List[str]:
        return self.i2w[len(self.start_vocab_idx) :]  # noqa: E203

    @property
    def substitutions(
        self,
    ) -> List[Tuple[int, actions.ConditionalEdit]]:
        return [
            i
            for i, a in enumerate(self.i2w)
            if isinstance(a, actions.ConditionalSub)
        ]

    @property
    def insertions(self) -> List[Tuple[int, actions.ConditionalEdit]]:
        return [
            i
            for i, a in enumerate(self.i2w)
            if isinstance(a, actions.ConditionalIns)
        ]

    @property
    def beg_idx(self) -> int:
        return self.i2w.index(actions.Start())

    @property
    def end_idx(self) -> int:
        return self.i2w.index(actions.End())

    @property
    def del_idx(self) -> int:
        return self.i2w.index(actions.ConditionalDel())

    @property
    def copy_idx(self) -> int:
        return self.i2w.index(actions.ConditionalCopy())


@dataclasses.dataclass
class Prefix:
    """Represents overlapping target and prediction strings."""

    prediction: Sequence[Any]
    target: Sequence[Any]
    alignment: int

    @property
    def suffix(self) -> Sequence[Any]:
        return self.target[self.alignment :]  # noqa: E203

    @property
    def leftmost_of_suffix(self) -> Any:
        try:
            return self.target[self.alignment]
        except IndexError:
            return None


@dataclasses.dataclass
class ActionPrefix:
    """Class for wrapping possible actions associated with given prefix."""

    action: Set[actions.Edit]
    prefix: Prefix


def edit_distance(
    x: Sequence[Any],
    y: Sequence[Any],
    del_cost=1.0,
    ins_cost=1.0,
    sub_cost=1.0,
    x_offset=0,
    y_offset=0,
) -> numpy.ndarray:
    """Generates edit distance matrix from source (x) to target (y).

    Args:
        x (Sequence[Any]): source sequence.
        y (Sequence[Any]): target sequence.
        del_cost (float): weight to deletion actions.
        ins_cost (float): weight to insertion actions.
        sub_cost (float): weight to substitution actions.
        x_offset (int): starting index of source sequence.
        y_offset (int): starting index of target sequence.

    Returns:
        numpy.ndarray: edit sequence matrix.
    """
    x_size = len(x) - x_offset + 1
    y_size = len(y) - y_offset + 1
    prefix_matrix = numpy.full(
        (x_size, y_size), numpy.inf, dtype=numpy.float32
    )
    for i in range(x_size):
        prefix_matrix[i, 0] = i * del_cost
    for j in range(y_size):
        prefix_matrix[0, j] = j * ins_cost
    for i in range(1, x_size):
        for j in range(1, y_size):
            if x[i - 1 + x_offset] == y[j - 1 + y_offset]:
                substitution = 0.0
            else:
                substitution = sub_cost
            prefix_matrix[i, j] = min(
                prefix_matrix[i - 1, j] + del_cost,
                prefix_matrix[i, j - 1] + ins_cost,
                prefix_matrix[i - 1, j - 1] + substitution,
            )
    return prefix_matrix


class Expert(abc.ABC):
    """Oracle scores possible edit actions between prediction and target.

    Args:
        actions (ActionVocabulary): vocabulary of possible edit actions.
        aligner (StochasticEditDistance): an alignment object to score edits
            between source and target strings.
        oracle_factor (int): a scaling factor for scheduling predictions used
            in transducer training.
    """

    actions: ActionVocabulary
    aligner: sed.StochasticEditDistance
    oracle_factor: int
    roll_in: int

    def __init__(self, actions, aligner, oracle_factor=defaults.ORACLE_FACTOR):
        self.actions = actions
        self.oracle_factor = oracle_factor
        self.roll_in = 1
        self.aligner = aligner

    def find_valid_actions(
        self,
        source: Sequence[Any],
        alignment: int,
        prefixes: Iterable[Prefix],
    ) -> List[ActionPrefix]:
        """Provides edit actions for source symbol and prefix.

        Args:
            source (Sequence[Any]): source string to perform edit actions over.
            alignment (int): index for current symbol to edit in source string.
            prefixes (Iterable[Prefix]): prefix objects aligning current prefix
                overlap between target and source strings. (Can be multiple.)

        Returns:
            List[ActionPrefix]: edit actions and current prefix
                overlap between target and prediction associated with action.
        """
        input_not_empty = alignment < len(source)
        attention = source[alignment] if input_not_empty else None
        action_prefixes = []
        for prefix in prefixes:
            prefix_insert = (
                prefix.leftmost_of_suffix
            )  # First symbol in target string not overlapping with prediction.
            if prefix_insert is None:  # No remaining symbols in target.
                valid_action = {actions.End()}
            else:  # More symbols to go. Insertion is always valid.
                valid_action = {actions.Ins(prefix_insert)}
            if input_not_empty:
                if prefix_insert is not None:
                    # The target symbol and source symbol are same. Copy.
                    if prefix_insert == attention:
                        # TODO: These actions are in maxwell. Remove?
                        valid_action.add(
                            actions.Copy(attention, prefix_insert)
                        )
                    # Target and source symbol are different. Sub.
                    else:
                        valid_action.add(
                            actions.Sub(old=attention, new=prefix_insert)
                        )
                # Source symbol deleted. Further actions can occur.
                valid_action.add(actions.Del(attention))
            action_prefix = ActionPrefix(valid_action, prefix)
            action_prefixes.append(action_prefix)
        return action_prefixes

    def roll_in_schedule(self, epoch: int) -> float:
        """Gets probability of sampling from oracle given current epoch."""
        if self.aligner is None:
            raise AlignerError(
                """Expert called `roll_in_schedule` but there
                               is no aligner present to allow oracle
                               predictions"""
            )
        self.roll_in = 1 - self.oracle_factor / (
            self.oracle_factor + numpy.exp(epoch / self.oracle_factor)
        )

    def explore(self) -> bool:
        """Randomly determines whether expert should advise model."""
        return numpy.random.rand() <= self.roll_in

    def roll_out(
        self,
        source: Sequence[Any],
        target: Sequence[Any],
        alignment: int,
        action_prefixes: Iterable[ActionPrefix],
    ) -> Dict[actions.Edit, float]:
        """Scores potential actions by a predicted 'cost to go' for target.

        Score sums potential edit sequence with cost of action.

        Args:
            source (Sequence[Any]): source string to perform edit actions over.
            target (Sequence[Any]): target string for edit actions.
            alignment (int): index for current symbol to edit in source string.
            action_prefixes (Iterable[ActionsPrefix]): set of edit
                actions-prefix pairs to evaluate.

        Returns:
            Dict[Edit, float]: edit actions and their respective scores.
        """
        costs_to_go = {}
        for action_prefix in action_prefixes:
            suffix_begin = action_prefix.prefix.alignment
            for action in action_prefix.action:
                if isinstance(action, actions.Del):
                    s_offset = alignment + 1
                    t_offset = suffix_begin
                elif isinstance(action, actions.Ins):
                    s_offset = alignment
                    t_offset = suffix_begin + 1
                elif isinstance(action, actions.Sub):
                    s_offset = alignment + 1
                    t_offset = suffix_begin + 1
                elif isinstance(action, actions.End):
                    s_offset = alignment
                    t_offset = suffix_begin
                else:
                    raise ActionError(f"Unknown action: {action}")
                sequence_cost = self.aligner.action_sequence_cost(
                    source, target, s_offset, t_offset
                )
                action_cost = self.aligner.action_cost(action)
                cost = action_cost + sequence_cost
                if action not in costs_to_go or costs_to_go[action] > cost:
                    costs_to_go[action] = cost
        return costs_to_go

    def score(
        self,
        source: Sequence[Any],
        target: Sequence[Any],
        alignment: int,
        prediction: Sequence[Any],
        max_action_seq_len=150,
    ) -> Dict[actions.Edit, float]:
        """Provides potential actions given source, target, and prediction.

        Args:
            source (Sequence[Any]): source string to perform edit actions over.
            target (Sequence[Any]): target string for edit actions.
            alignment (int): index for current symbol to edit in source string.
            prediction (Sequence[Any]): current prediction from previous edits.

        Returns:
            Dict[Edit, float]: edit actions and their respective scores.
        """
        prefixes = self.find_prefixes(prediction, target)
        valid_actions = (
            self.find_valid_actions(source, alignment, prefixes)
            if len(prediction) <= max_action_seq_len
            else {actions.End()}
        )
        valid_action_scores = self.roll_out(
            source, target, alignment, valid_actions
        )
        return valid_action_scores

    @staticmethod
    def find_prefixes(
        prediction: Sequence[Any], target: Sequence[Any]
    ) -> List[Prefix]:
        """Creates prefix objects for prediction and target.

        Args:
            prediction (Sequence[Any]): current prediction from prior edits.
            target (Sequence[Any]): target string for edit actions.

        Returns:
            List[Prefix]: prefix objects that index overlap between
               prediction and target with minimal edit distance.
        """
        prefix_matrix = edit_distance(prediction, target)
        prediction_row = prefix_matrix[-1]
        return [
            Prefix(prediction, target, i)
            for i in numpy.where(prediction_row == prediction_row.min())[0]
        ]


def get_expert(
    dataset: data.Dataset,
    index: data.Index,
    epochs: int = defaults.ORACLE_EM_EPOCHS,
    oracle_factor: int = defaults.ORACLE_FACTOR,
    sed_params_path: str = None,
    read_from_file: bool = False,
) -> Expert:
    """Generates expert object for training transducer.

    Args:
        dataset (data.Dataset): dataset for generating expert vocabulary.
        index (data.Index): index for mapping symbols to indices.
        epochs (int): number of EM epochs.
        oracle_factor (float): scaling factor to determine rate of
            expert rollout sampling.
        sed_params_path (str): path to read/write location of sed parameters.
        read_from_file (bool): bool for whether to use existing parameters.
            If True, reads from sed_params_path. If False, writes to
            sed_params_path.

    Returns:
        expert.Expert.
    """

    def _generate_data(
        dataset: data.Dataset,
        index: data.Index,
    ) -> Iterator[Tuple[List[int], List[int]]]:
        """Helper function to manage data encoding for SED.

        We want encodings without padding. This encodes only raw source-target
        text for the Maxwell library.

        Args:
            dataset (data.Dataset): dataset for generating expert vocabulary.
            index (data.Index): index for mapping symbols to indices.

        Yields:
            Tuple[List[int, List[int]]]: lists of source and target entries.
        """
        if not dataset.has_target:
            raise Error("Dataset has no target")
        for sample in dataset.samples:
            source, *_, target = sample
            yield (
                [index(symbol) for symbol in source],
                [index(symbol) for symbol in target],
            )

    actions = ActionVocabulary(index)
    if read_from_file:
        sed_params = sed.ParamDict.read_params(sed_params_path)
        sed_aligner = sed.StochasticEditDistance(sed_params)
    else:
        sed_aligner = sed.StochasticEditDistance.fit_from_data(
            _generate_data(dataset, index),
            epochs=epochs,
        )
        sed_aligner.params.write_params(sed_params_path)
    return Expert(actions, aligner=sed_aligner, oracle_factor=oracle_factor)


def add_argparse_args(parser: argparse.ArgumentParser) -> None:
    """Adds expert configuration options to the argument parser.

    These are only needed at training time.

    Args:
        parser (argparse.ArgumentParser).
    """
    parser.add_argument(
        "--oracle_em_epochs",
        type=int,
        default=defaults.ORACLE_EM_EPOCHS,
        help="Number of EM epochs "
        "(transducer architecture only). Default: %(default)s.",
    )
    parser.add_argument(
        "--oracle_factor",
        type=int,
        default=defaults.ORACLE_FACTOR,
        help="Roll-in schedule parameter "
        "(transducer architecture only). Default: %(default)s.",
    )
    parser.add_argument(
        "--sed_params",
        type=str,
        help="Path to input SED parameters (transducer architecture only).",
    )
