"""Transducer model class."""

from typing import Callable, Dict, List, Optional, Tuple

import numpy
import torch
from maxwell import actions
from torch import nn

from .. import data, defaults, special, util
from . import expert, modules, rnn


class TransducerRNNModel(rnn.RNNModel):
    """Base class for transducer models.

    This uses a trained oracle for imitation learning edits.

    After:
        Makarov, P., and Clematide, S. 2018. Imitation learning for neural
        morphological string transduction. In Proceedings of the 2018
        Conference on Empirical Methods in Natural Language Processing, pages
        2877–2882.

     Args:
        expert (expert.Expert): oracle that guides training for transducer.
        *args: passed to superclass.
        **kwargs: passed to superclass.
    """

    expert: expert.Expert

    def __init__(
        self,
        expert,
        *args,
        **kwargs,
    ):
        """Initializes transducer model.

        Args:
            expert (expert.Expert): oracle that guides training for transducer.
            *args: passed to superclass.
            **kwargs: passed to superclass.
        """
        super().__init__(*args, **kwargs)
        # Model specific variables.
        self.vocab_offset = self.vocab_size - self.target_vocab_size
        self.expert = expert  # Oracle to train model.
        self.actions = self.expert.actions
        self.substitutions = self.actions.substitutions
        self.insertions = self.actions.insertions

    def forward(
        self,
        batch: data.PaddedBatch,
    ) -> Tuple[List[List[int]], torch.Tensor]:
        """Runs the encoder-decoder model.

        Args:
            batch (data.PaddedBatch).

        Returns:
            Tuple[List[List[int]], torch.Tensor]: encoded prediction values
                and loss tensor; due to transducer setup, prediction is
                performed during training, so these are returned.
        """
        encoder_out = self.source_encoder(batch.source)
        encoded = encoder_out.output[:, 1:, :]  # Ignores start symbol.
        source_padded = batch.source.padded[:, 1:]
        source_mask = batch.source.mask[:, 1:]
        # Start of decoding.
        if self.has_features_encoder:
            features_encoder_out = self.features_encoder(batch.features)
            features_encoded = features_encoder_out.output
            if features_encoder_out.has_hiddens:
                h_features, c_features = features_encoder_out.hiddens
                h_features = h_features.mean(dim=0, keepdim=True).expand(
                    self.decoder_layers, -1, -1
                )
                c_features = c_features.mean(dim=0, keepdim=True).expand(
                    self.decoder_layers, -1, -1
                )
                last_hiddens = h_features, c_features
            else:
                last_hiddens = self.init_hiddens(source_mask.shape[0])
            features_encoded = features_encoded.mean(dim=1, keepdim=True)
            encoded = torch.cat(
                (
                    encoded,
                    features_encoded.expand(-1, encoded.shape[1], -1),
                ),
                dim=2,
            )
        else:
            last_hiddens = self.init_hiddens(source_mask.shape[0])
        if self.beam_width > 1:
            # Will raise a NotImplementedError.
            prediction = self.beam_decode(
                encoder_out=encoded,
                mask=batch.source.mask,
                beam_width=self.beam_width,
            )
        else:
            prediction, loss = self.decode(
                encoded,
                last_hiddens,
                source_padded,
                source_mask,
                teacher_forcing=(
                    self.teacher_forcing if self.training else False
                ),
                target=batch.target.padded if batch.target else None,
                target_mask=batch.target.mask if batch.target else None,
            )
        return prediction, loss

    def decode(
        self,
        encoder_out: torch.Tensor,
        last_hiddens: Tuple[torch.Tensor, torch.Tensor],
        source: torch.Tensor,
        source_mask: torch.Tensor,
        teacher_forcing: bool,
        target: Optional[torch.Tensor] = None,
        target_mask: Optional[torch.Tensor] = None,
    ) -> Tuple[List[List[int]], torch.Tensor]:
        """Decodes a sequence given the encoded input.

        This essentially serves as a wrapper for looping decode_step.

        Args:
            encoder_out (torch.Tensor): input symbols of shape
                B x seq_len x emb_size.
            source (torch.Tensor): encoded source input.
            source_mask (torch.Tensor): mask for source input.
            teacher_forcing (bool): whether or not to decode
                with teacher forcing; determines whether or not to rollout
                optimal actions.
            target (torch.Tensor, optional): encoded target input.
            target_mask (torch.Tensor, optional): mask for target input.

        Returns:
            Tuple[List[List[int]], torch.Tensor]: encoded prediction values
                and loss tensor; due to transducer setup, prediction is
                performed during training, so these are returned.
        """
        batch_size = source_mask.size(dim=0)
        input_length = (~source_mask).sum(dim=1)
        # Initializing values.
        alignment = torch.zeros(
            batch_size, device=self.device, dtype=torch.int64
        )
        action_count = torch.zeros(
            batch_size, device=self.device, dtype=torch.int64
        )
        last_action = torch.full(
            (batch_size,), self.actions.beg_idx, device=self.device
        )
        loss = torch.zeros(batch_size, device=self.device)
        prediction = [[] for _ in range(batch_size)]
        # Converting encodings for prediction.
        if target is not None:
            # Target and source need to be integers for SED values.
            # Clips EOW (idx = -1) for source and target.
            source = [
                s[~smask].tolist()[:-1]
                for s, smask in zip(source, source_mask)
            ]
            target = [
                t[~tmask].tolist()[:-1]
                for t, tmask in zip(target, target_mask)
            ]
        for _ in range(self.max_target_length):
            # Checks if completed all sequences.
            not_complete = last_action != self.actions.end_idx
            if not any(not_complete):
                break
            # Proceeds to make new edit; new action for all current decoding.
            action_count = torch.where(
                not_complete.to(self.device),
                action_count + 1,
                action_count,
            )
            # Decoding.
            # We offset the action idx by the symbol vocab size so that we
            # can index into the shared embeddings matrix.
            decoder_output = self.decoder(
                last_action.unsqueeze(dim=1) + self.vocab_offset,
                last_hiddens,
                encoder_out,
                # Accomodates RNNDecoder; see encoder_mask behavior.
                ~(alignment.unsqueeze(1) + 1),
            )
            decoded, last_hiddens = (
                decoder_output.output,
                decoder_output.hiddens,
            )
            logits = self.classifier(decoded).squeeze(dim=1)
            # If given targets, asks expert for optimal actions.
            optim_actions = (
                self.batch_expert_rollout(
                    source,
                    target,
                    alignment,
                    prediction,
                    not_complete,
                )
                if target is not None
                else None
            )
            last_action = self.decode_action_step(
                logits,
                alignment,
                input_length,
                not_complete,
                optim_actions=(optim_actions if teacher_forcing else None),
            )
            alignment = self.update_prediction(
                last_action, source, alignment, prediction
            )
            # If target, validation or training step loss required.
            if target is not None:
                log_sum_loss = self.log_sum_softmax_loss(logits, optim_actions)
                loss = torch.where(not_complete, log_sum_loss + loss, loss)
        avg_loss = torch.mean(loss / action_count)
        return prediction, -avg_loss

    def decode_action_step(
        self,
        logits: torch.Tensor,
        alignment: torch.Tensor,
        input_length: torch.Tensor,
        not_complete: torch.Tensor,
        optim_actions: Optional[torch.Tensor] = None,
    ) -> torch.Tensor:
        """Decodes logits to find edit action.

        Finds possible actions given remaining size of source input and masks
        logits for edit action decoding.

        Args:
            logits (torch.Tensor): logit values from decode_step of shape
                B x num_actions.
            alignment (torch.Tensor): index of encoding symbols for decoding,
                per item in batch of shape B x seq_len.
            input_length (torch.Tensor): length of each item in batch.
            not_complete (torch.Tensor): boolean values designating which items
                have not terminated edits.
            optim_actions (List[List[int]], optional): optimal actions
                determined by expert, present when loss is being calculated.

        Returns:
            torch.Tensor: chosen edit action.
        """
        # Finds valid actions given remaining input length.
        end_of_input = (input_length - alignment) <= 1  # 1 -> Last char.
        valid_actions = [
            (self.compute_valid_actions(eoi) if nc else [self.actions.end_idx])
            for eoi, nc in zip(end_of_input, not_complete)
        ]
        # Masks invalid actions.
        logits = self.action_probability_mask(logits, valid_actions)
        return self.choose_action(logits, not_complete, optim_actions)

    def compute_valid_actions(self, end_of_input: bool) -> List[int]:
        """Gives all possible actions for remaining length of edits.

        Args:
            end_of_input (bool): indicates if this is the last input from
                string; if true, only insertions are available.

        Returns:
            List[actions.Edit]: actions known by transducer.
        """
        valid_actions = [self.actions.end_idx]
        valid_actions.extend(self.insertions)
        if not end_of_input:
            valid_actions.extend([self.actions.copy_idx, self.actions.del_idx])
            valid_actions.extend(self.substitutions)
        return valid_actions

    def action_probability_mask(
        self, logits: torch.Tensor, valid_actions: List[int]
    ) -> torch.Tensor:
        """Masks non-valid actions in logits."""
        with torch.no_grad():
            mask = torch.full(
                logits.shape, defaults.NEG_INF, device=self.device
            )
            for row, action in zip(mask, valid_actions):
                row[action] = 0.0
            logits = mask + logits
        return logits

    def choose_action(
        self,
        logits: torch.Tensor,
        not_complete: torch.Tensor,
        optim_actions: Optional[List[List[int]]] = None,
    ) -> torch.Tensor:
        """Chooses transducer action from log_prob distribution.

        If training, uses dynamic oracle for selection.

        Args:
            log_probs (torch.Tensor): probability distribution of actions.
            not_complete (torch.Tensor): boolean tensor of batch length to
                indicate if each item in batch is complete.
            optim_actions (Optional[List[List[int]]]): optional encoded actions
                to use for action selection.

        Returns:
            torch.Tensor: action encodings.
        """
        # TODO: Merge logic into PyTorch methods.
        log_probs = nn.functional.log_softmax(logits, dim=1)
        if optim_actions is None:
            # Argmax decoding.
            next_action = [
                (torch.argmax(probs, dim=0) if nc else self.actions.end_idx)
                for probs, nc in zip(log_probs, not_complete)
            ]
        else:
            # Training with dynamic oracle; chooses from optimal actions.
            with torch.no_grad():
                if self.expert.explore():
                    # Action is picked by random exploration.
                    next_action = [
                        (self.sample(probs) if nc else self.actions.end_idx)
                        for probs, nc in zip(log_probs, not_complete)
                    ]
                else:
                    # Action is picked from optim_actions.
                    next_action = []
                    for action, probs, nc in zip(
                        optim_actions, log_probs, not_complete
                    ):
                        if nc:
                            optim_logs = probs[action]
                            idx = int(torch.argmax(optim_logs, 0))
                            next_action.append(action[idx])
                        else:  # Already complete, so skip.
                            next_action.append(self.actions.end_idx)
        return torch.tensor(next_action, device=self.device, dtype=torch.int)

    # TODO: Merge action classes to remove need for this method.
    @staticmethod
    def remap_actions(
        action_scores: Dict[actions.Edit, float]
    ) -> Dict[actions.Edit, float]:
        """Maps generative oracle's edit to conditional counterpart.

        Oracle edits are a distinct subclass from edits learned from samples.

        This will eventually be removed.

        Args:
            action_scores (Dict[actions.Edit, float]): weights for each action.

        Returns:
            Dict[actions.Edit, float]: edit action-weight pairs.
        """
        remapped_action_scores = {}
        for action, score in action_scores.items():
            if isinstance(action, actions.GenerativeEdit):
                remapped_action = action.conditional_counterpart()
            elif isinstance(action, actions.Edit):
                remapped_action = action
            else:
                raise expert.ActionError(
                    f"Unknown action: {action}, {score}, "
                    f"action_scores: {action_scores}"
                )
            remapped_action_scores[remapped_action] = score
        return remapped_action_scores

    def expert_rollout(
        self,
        source: List[int],
        target: List[int],
        alignment: int,
        prediction: List[int],
    ) -> List[int]:
        """Rolls out with optimal expert policy.

        Args:
            source (List[int]): input string.
            target (List[int]): target string.
            alignment (int): position in source to edit.
            prediction (List[str]): current prediction.

        Returns:
            List[int]: optimal action encodings.
        """
        raw_action_scores = self.expert.score(
            source,
            target,
            alignment,
            prediction,
            max_action_seq_len=self.max_target_length,
        )
        action_scores = self.remap_actions(raw_action_scores)
        optimal_value = min(action_scores.values())
        optimal_action = sorted(
            [
                self.actions.encode_unseen_action(action)
                for action, value in action_scores.items()
                if value == optimal_value
            ]
        )
        return optimal_action

    def batch_expert_rollout(
        self,
        source: List[List[int]],
        target: List[List[int]],
        alignment: torch.Tensor,
        prediction: List[List[int]],
        not_complete: torch.Tensor,
    ) -> List[List[int]]:
        """Performs expert rollout over batch."""
        return [
            (
                self.expert_rollout(s, t, align, pred)
                if nc
                else self.actions.end_idx
            )
            for s, t, align, pred, nc in zip(
                source, target, alignment, prediction, not_complete
            )
        ]

    def update_prediction(
        self,
        action: List[actions.Edit],
        source: List[int],
        alignment: torch.Tensor,
        prediction: List[List[str]],
    ) -> torch.Tensor:
        """Batch updates prediction and alignment information from actions.

        Args:
           action (List[actions.Edit]): valid actions, one per item in
                batch.
           source (List[int]): source strings, one per item in batch.
           alignment (torch.Tensor): index of current symbol for each item
                in batch.
           prediction (List[List[str]]): current predictions for each item
                in batch, one list of symbols per item.

        Return:
            torch.Tensor: new alignments for transduction.
        """
        alignment_update = torch.zeros(
            alignment.shape, device=self.device, dtype=torch.int64
        )
        for i in range(len(source)):
            a = self.actions.decode(action[i])
            if isinstance(a, actions.ConditionalCopy):
                symb = source[i][alignment[i]]
                prediction[i].append(symb)
                alignment_update[i] += 1
            elif isinstance(a, actions.ConditionalDel):
                alignment_update[i] += 1
            elif isinstance(a, actions.ConditionalIns):
                prediction[i].append(a.new)
            elif isinstance(a, actions.ConditionalSub):
                alignment_update[i] += 1
                prediction[i].append(a.new)
            elif isinstance(a, actions.End):
                prediction[i].append(special.END_IDX)
            else:
                raise expert.ActionError(f"Unknown action: {action[i]}")
        return alignment + alignment_update

    @staticmethod
    def log_sum_softmax_loss(
        logits: torch.Tensor, optimal_actions: List[int]
    ) -> torch.Tensor:
        """Computes log loss.

        After:
            Riezler, S. Prescher, D., Kuhn, J. and Johnson, M. 2000.
            Lexicalized stochastic modeling of constraint-based grammars using
            log-linear measures and EM training. In Proceedings of the 38th
            Annual Meeting of the Association for Computational
            Linguistics, pages 480–487.
        """
        opt_act = [
            log[actions] for log, actions in zip(logits, optimal_actions)
        ]
        log_sum_exp_terms = torch.stack(
            [torch.logsumexp(act, dim=-1) for act in opt_act]
        )
        normalization_term = torch.logsumexp(logits, -1)
        return log_sum_exp_terms - normalization_term

    def _get_loss_func(
        self,
    ) -> Callable[[torch.Tensor, torch.Tensor], torch.Tensor]:
        # Prevents base construction of unused loss function.
        return None

    def training_step(self, batch: data.PaddedBatch, batch_idx: int) -> Dict:
        """Runs one step of training.

        This is called by the PL Trainer.

        Args:
            batch (data.PaddedBatch)
            batch_idx (int).

        Returns:
            torch.Tensor: loss.
        """
        # Forward pass produces loss by default.
        _, loss = self(batch)
        self.log(
            "train_loss",
            loss,
            batch_size=len(batch),
            on_step=False,
            on_epoch=True,
        )
        return loss

    def validation_step(self, batch: data.PaddedBatch, batch_idx: int) -> Dict:
        predictions, loss = self(batch)
        # Evaluation requires prediction as a tensor.
        predictions = self.convert_predictions(predictions)
        # Gets a dict of all eval metrics for this batch.
        val_eval_items_dict = {
            evaluator.name: evaluator.evaluate(
                predictions,
                batch.target.padded,
                predictions_finalized=True,
            )
            for evaluator in self.evaluators
        }
        val_eval_items_dict.update({"val_loss": loss})
        return val_eval_items_dict

    def predict_step(self, batch: Tuple[torch.tensor], batch_idx: int) -> Dict:
        predictions, _ = self.forward(
            batch,
        )
        # Evaluation requires prediction tensor.
        return self.convert_predictions(predictions)

    def convert_predictions(
        self, predictions: List[List[int]]
    ) -> torch.Tensor:
        """Converts prediction values to tensor for evaluator compatibility."""
        # FIXME: the two steps below may be partially redundant.
        # TODO: Clean this up and make it more efficient.
        max_len = len(max(predictions, key=len))
        for i, pred in enumerate(predictions):
            pad = [self.actions.end_idx] * (max_len - len(pred))
            pred.extend(pad)
            predictions[i] = torch.tensor(pred, dtype=torch.int)
        predictions = torch.stack(predictions)
        # This turns all symbols after the first EOS into PAD so prediction
        # tensors match gold tensors.
        return util.pad_tensor_after_eos(predictions)

    def on_train_epoch_start(self) -> None:
        """Scheduler for oracle."""
        self.expert.roll_in_schedule(self.current_epoch)

    @staticmethod
    def sample(log_probs: torch.Tensor) -> torch.Tensor:
        """Samples an action from a log-probability distribution."""
        dist = torch.exp(log_probs)
        rand = numpy.random.rand()
        for action, p in enumerate(dist):
            rand -= p
            if rand <= 0:
                break
        return action

    def get_decoder(self):
        raise NotImplementedError

    @property
    def name(self) -> str:
        raise NotImplementedError


# TODO: Implement beam decoding.


class TransducerGRUModel(TransducerRNNModel, rnn.GRUModel):
    """Transducer with GRU backend."""

    # h0: nn.Parameter

    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)
        # self.h0 = nn.Parameter(torch.rand(self.hidden_size))

    def forward(
        self,
        batch: data.PaddedBatch,
    ) -> Tuple[List[List[int]], torch.Tensor]:
        """Runs the encoder-decoder model.

        Args:
            batch (data.PaddedBatch).

        Returns:
            Tuple[List[List[int]], torch.Tensor]: encoded prediction values
                and loss tensor; due to transducer setup, prediction is
                performed during training, so these are returned.
        """
        encoder_out = self.source_encoder(batch.source)
        encoded = encoder_out.output[:, 1:, :]  # Ignores start symbol.
        source_padded = batch.source.padded[:, 1:]
        source_mask = batch.source.mask[:, 1:]
        # Start of decoding.
        if self.has_features_encoder:
            features_encoder_out = self.features_encoder(batch.features)
            features_encoded = features_encoder_out.output
            if features_encoder_out.has_hiddens:
                h_features = features_encoder_out.hiddens
                last_hiddens = h_features.mean(dim=0, keepdim=True).expand(
                    self.decoder_layers, -1, -1
                )
                last_hiddens = h_features
            else:
                last_hiddens = self.init_hiddens(source_mask.shape[0])
            features_encoded = features_encoded.mean(dim=1, keepdim=True)
            encoded = torch.cat(
                (
                    encoded,
                    features_encoded.expand(-1, encoded.shape[1], -1),
                ),
                dim=2,
            )
        else:
            last_hiddens = self.init_hiddens(source_mask.shape[0])
        prediction, loss = self.decode(
            encoded,
            last_hiddens,
            source_padded,
            source_mask,
            teacher_forcing=(self.teacher_forcing if self.training else False),
            target=batch.target.padded if batch.target else None,
            target_mask=batch.target.mask if batch.target else None,
        )
        return prediction, loss

    def get_decoder(self) -> modules.GRUDecoder:
        return modules.GRUDecoder(
            bidirectional=False,
            decoder_input_size=(
                self.source_encoder.output_size
                + self.features_encoder.output_size
                if self.has_features_encoder
                else self.source_encoder.output_size
            ),
            dropout=self.dropout,
            embeddings=self.embeddings,
            embedding_size=self.embedding_size,
            layers=self.decoder_layers,
            hidden_size=self.hidden_size,
            num_embeddings=self.vocab_size,
        )

    @property
    def name(self) -> str:
        return "transducer GRU"


class TransducerLSTMModel(TransducerRNNModel):
    """Transducer with LSTM backend."""

    h0: nn.Parameter
    c0: nn.Parameter

    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)
        self.h0 = nn.Parameter(torch.rand(self.hidden_size))
        self.c0 = nn.Parameter(torch.rand(self.hidden_size))

    def forward(
        self,
        batch: data.PaddedBatch,
    ) -> Tuple[List[List[int]], torch.Tensor]:
        """Runs the encoder-decoder model.

        Args:
            batch (data.PaddedBatch).

        Returns:
            Tuple[List[List[int]], torch.Tensor]: encoded prediction values
                and loss tensor; due to transducer setup, prediction is
                performed during training, so these are returned.
        """
        encoder_out = self.source_encoder(batch.source)
        encoded = encoder_out.output[:, 1:, :]  # Ignores start symbol.
        source_padded = batch.source.padded[:, 1:]
        source_mask = batch.source.mask[:, 1:]
        # Start of decoding.
        if self.has_features_encoder:
            features_encoder_out = self.features_encoder(batch.features)
            features_encoded = features_encoder_out.output
            if features_encoder_out.has_hiddens:
                h_features, c_features = features_encoder_out.hiddens
                h_features = h_features.mean(dim=0, keepdim=True).expand(
                    self.decoder_layers, -1, -1
                )
                c_features = c_features.mean(dim=0, keepdim=True).expand(
                    self.decoder_layers, -1, -1
                )
                last_hiddens = h_features, c_features
            else:
                last_hiddens = self.init_hiddens(source_mask.shape[0])
            features_encoded = features_encoded.mean(dim=1, keepdim=True)
            encoded = torch.cat(
                (
                    encoded,
                    features_encoded.expand(-1, encoded.shape[1], -1),
                ),
                dim=2,
            )
        else:
            last_hiddens = self.init_hiddens(source_mask.shape[0])
        prediction, loss = self.decode(
            encoded,
            last_hiddens,
            source_padded,
            source_mask,
            teacher_forcing=(self.teacher_forcing if self.training else False),
            target=batch.target.padded if batch.target else None,
            target_mask=batch.target.mask if batch.target else None,
        )
        return prediction, loss

    def init_hiddens(
        self, batch_size: int
    ) -> Tuple[torch.Tensor, torch.Tensor]:
        """Initializes the hidden state to pass to the RNN.

        We treat the initial value as a model parameter.

        Args:
            batch_size (int).

        Returns:
            Tuple[torch.Tensor, torch.Tensor].
        """
        return (
            self.h0.repeat(self.decoder_layers, batch_size, 1),
            self.c0.repeat(self.decoder_layers, batch_size, 1),
        )

    def get_decoder(self) -> modules.LSTMDecoder:
        return modules.LSTMDecoder(
            bidirectional=False,
            decoder_input_size=(
                self.source_encoder.output_size
                + self.features_encoder.output_size
                if self.has_features_encoder
                else self.source_encoder.output_size
            ),
            dropout=self.dropout,
            embeddings=self.embeddings,
            embedding_size=self.embedding_size,
            layers=self.decoder_layers,
            hidden_size=self.hidden_size,
            num_embeddings=self.vocab_size,
        )

    @property
    def name(self) -> str:
        return "transducer LSTM"
