# -*- coding: utf-8 -*-
r"""
Polos Ranker Model
======================
    The goal of this model is to rank good translations closer to the reference and source text
    and bad translations further by a small margin.

    https://pytorch.org/docs/stable/nn.html#tripletmarginloss
"""
from argparse import Namespace
from typing import Dict, List, Tuple, Union

import torch
import torch.nn.functional as F
from tqdm import tqdm

from polos.models.ranking.ranking_base import RankingBase
from polos.models.utils import move_to_cuda
from torchnlp.utils import collate_tensors


class PolosRanker(RankingBase): # extends ptl.LightningModule
    """
    Polos Ranker class that uses a pretrained encoder to extract features
    from the sequences and then passes those features through a Triplet Margin Loss.

    :param hparams: Namespace containing the hyperparameters.
    """

    def __init__(self, hparams: Namespace) -> None:
        super().__init__(hparams)

    def compute_metrics(self, outputs: List[Dict[str, torch.Tensor]]) -> dict:
        """  Computes WMT19 shared task kendall tau like metric. """
        distance_pos, distance_neg = [], []
        for minibatch in outputs:
            minibatch = minibatch["val_prediction"]
            src_embedding = minibatch["src_sentemb"]
            ref_embedding = minibatch["ref_sentemb"]
            pos_embedding = minibatch["pos_sentemb"]
            neg_embedding = minibatch["neg_sentemb"]

            distance_src_pos = F.pairwise_distance(pos_embedding, src_embedding)
            distance_ref_pos = F.pairwise_distance(pos_embedding, ref_embedding)
            harmonic_distance_pos = (2 * distance_src_pos * distance_ref_pos) / (
                distance_src_pos + distance_ref_pos
            )
            distance_pos.append(harmonic_distance_pos)

            distance_src_neg = F.pairwise_distance(neg_embedding, src_embedding)
            distance_ref_neg = F.pairwise_distance(neg_embedding, ref_embedding)
            harmonic_distance_neg = (2 * distance_src_neg * distance_ref_neg) / (
                distance_src_neg + distance_ref_neg
            )
            distance_neg.append(harmonic_distance_neg)

        return {
            "kendall": self.metrics.compute(
                torch.cat(distance_pos), torch.cat(distance_neg)
            )
        }

    def compute_loss(self, model_out: Dict[str, torch.Tensor], *args) -> torch.Tensor: 
        """
        # forwardの結果がmodel_outに入っているのでlossを計算
        Computes Triplet Margin Loss for both the reference and the source.

        :param model_out: model specific output with src_anchor, ref_anchor, pos and neg
            sentence embeddings.
        """

        # 参考
        #     "src_sentemb": self.get_sentence_embedding(src_tokens, src_lengths),
        #     "ref_sentemb": self.get_sentence_embedding(ref_tokens, ref_lengths),
        #     "pos_sentemb": self.get_sentence_embedding(pos_tokens, pos_lengths),
        #     "neg_sentemb": self.get_sentence_embedding(neg_tokens, neg_lengths),
        
        ref_anchor = model_out["ref_sentemb"]
        src_anchor = model_out["src_sentemb"]
        positive = model_out["pos_sentemb"]
        negative = model_out["neg_sentemb"]

        return self.loss(src_anchor, positive, negative) + self.loss(
            ref_anchor, positive, negative
        )

    def predict(
        self,
        samples: Dict[str, str],
        cuda: bool = False,
        show_progress: bool = False,
        batch_size: int = -1,
    ) -> (Dict[str, Union[str, float]], List[float]):
        """Function that runs a model prediction,

        :param samples: List of dictionaries with 'mt' and 'ref' keys.
        :param cuda: Flag that runs inference using 1 single GPU.
        :param show_progress: Flag to show progress during inference of multiple examples.
        :para batch_size: Batch size used during inference. By default uses the same batch size used during training.

        :return: Dictionary with model outputs
        """
        if self.training:
            self.eval()

        if cuda and torch.cuda.is_available():
            self.to("cuda")

        batch_size = self.hparams.batch_size if batch_size < 1 else batch_size
        with torch.no_grad():
            batches = [
                samples[i : i + batch_size] for i in range(0, len(samples), batch_size)
            ]
            model_inputs = []
            if show_progress:
                pbar = tqdm(
                    total=len(batches), desc="Preparing batches....", dynamic_ncols=True
                )
            for batch in batches:
                model_inputs.append(self.prepare_sample(batch, inference=True))
                if show_progress:
                    pbar.update(1)

            if show_progress:
                pbar.close()

            if show_progress:
                pbar = tqdm(
                    total=len(batches), desc="Scoring hypothesis...", dynamic_ncols=True
                )

            distance_weighted, distance_src, distance_ref = [], [], []
            for k, model_input in enumerate(model_inputs):
                src_input, mt_input, ref_input, alt_input = model_input
                if cuda and torch.cuda.is_available():
                    src_embeddings = self.get_sentence_embedding(
                        **move_to_cuda(src_input)
                    )
                    mt_embeddings = self.get_sentence_embedding(
                        **move_to_cuda(mt_input)
                    )
                    ref_embeddings = self.get_sentence_embedding(
                        **move_to_cuda(ref_input)
                    )
                    ref_distances = F.pairwise_distance(
                        mt_embeddings, ref_embeddings
                    ).cpu()
                    src_distances = F.pairwise_distance(
                        mt_embeddings, src_embeddings
                    ).cpu()

                    # When 2 references are given the distance to the reference is the Min between
                    # both references.
                    if alt_input is not None:
                        alt_embeddings = self.get_sentence_embedding(
                            **move_to_cuda(alt_input)
                        )
                        alt_distances = F.pairwise_distance(
                            mt_embeddings, alt_embeddings
                        ).cpu()
                        ref_distances = torch.stack([ref_distances, alt_distances])
                        ref_distances = ref_distances.min(dim=0).values

                else:
                    src_embeddings = self.get_sentence_embedding(**src_input)
                    mt_embeddings = self.get_sentence_embedding(**mt_input)
                    ref_embeddings = self.get_sentence_embedding(**ref_input)
                    ref_distances = F.pairwise_distance(mt_embeddings, ref_embeddings)
                    src_distances = F.pairwise_distance(mt_embeddings, src_embeddings)

                # Harmonic mean between the distances:
                distances = (2 * ref_distances * src_distances) / (
                    ref_distances + src_distances
                )
                src_distances = ref_distances.numpy().tolist()
                ref_distances = ref_distances.numpy().tolist()
                distances = distances.numpy().tolist()

                for i in range(len(distances)):
                    distance_weighted.append(1 / (1 + distances[i]))
                    distance_src.append(1 / (1 + src_distances[i]))
                    distance_ref.append(1 / (1 + ref_distances[i]))

                if show_progress:
                    pbar.update(1)

            if show_progress:
                pbar.close()

        assert len(distance_weighted) == len(samples)
        scores = []
        for i in range(len(samples)):
            scores.append(distance_weighted[i])
            samples[i]["predicted_score"] = scores[-1]
            samples[i]["reference_distance"] = distance_ref[i]
            samples[i]["source_distance"] = distance_src[i]

        return samples, scores

    def prepare_sample(
        self, sample: List[Dict[str, Union[str, float]]], inference: bool = False
    ) -> Union[Tuple[Dict[str, torch.Tensor], None], List[Dict[str, torch.Tensor]]]:
        """
        Function that prepares a sample to input the model.

        :param sample: list of dictionaries.
        :param inference: If set to to False, then the model expects
            a MT and reference instead of anchor, pos, and neg segments.

        :return: Tuple with a dictionary containing the model inputs and None OR List
            with source, MT and reference tokenized and vectorized.
        """
        sample = collate_tensors(sample)
        if inference:
            src_inputs = self.encoder.prepare_sample(sample["src"])
            mt_inputs = self.encoder.prepare_sample(sample["mt"])
            ref_inputs = self.encoder.prepare_sample(sample["ref"])
            alt_inputs = (
                self.encoder.prepare_sample(sample["alt"]) if "alt" in sample else None
            )
            return src_inputs, mt_inputs, ref_inputs, alt_inputs

        ref_inputs = self.encoder.prepare_sample(sample["ref"])
        src_inputs = self.encoder.prepare_sample(sample["src"])
        pos_inputs = self.encoder.prepare_sample(sample["pos"])
        neg_inputs = self.encoder.prepare_sample(sample["neg"])

        ref_inputs = {"ref_" + k: v for k, v in ref_inputs.items()}
        src_inputs = {"src_" + k: v for k, v in src_inputs.items()}
        pos_inputs = {"pos_" + k: v for k, v in pos_inputs.items()}
        neg_inputs = {"neg_" + k: v for k, v in neg_inputs.items()}

        return {**ref_inputs, **src_inputs, **pos_inputs, **neg_inputs}, torch.empty(0)

    def forward(
        self,
        src_tokens: torch.tensor,
        ref_tokens: torch.tensor,
        pos_tokens: torch.tensor,
        neg_tokens: torch.tensor,
        src_lengths: torch.tensor,
        ref_lengths: torch.tensor,
        pos_lengths: torch.tensor,
        neg_lengths: torch.tensor,
        **kwargs
    ) -> Dict[str, torch.Tensor]:
        """
        Function that encodes the anchor, positive samples and negative samples
        and returns embeddings for the triplet.

        :param src_tokens: anchor sequences [batch_size x anchor_seq_len]
        :param ref_tokens: anchor sequences [batch_size x anchor_seq_len]
        :param pos_tokens: positive sequences [batch_size x pos_seq_len]
        :param neg_tokens: negative sequences [batch_size x neg_seq_len]
        :param src_lengths: anchor lengths [batch_size]
        :param ref_lengths: anchor lengths [batch_size]
        :param pos_lengths: positive lengths [batch_size]
        :param neg_lengths: negative lengths [batch_size]

        :return: Dictionary with model outputs to be passed to the loss function.
        """
        return {
            "src_sentemb": self.get_sentence_embedding(src_tokens, src_lengths),
            "ref_sentemb": self.get_sentence_embedding(ref_tokens, ref_lengths),
            "pos_sentemb": self.get_sentence_embedding(pos_tokens, pos_lengths),
            "neg_sentemb": self.get_sentence_embedding(neg_tokens, neg_lengths),
        }