# -*- coding: utf-8 -*-
r"""
Quality Estimator Model
=======================
    Quality Estimator Model estimates a quality score for the hyphotesis (e.g: the MT text)
    by looking only at source and MT.
"""
from argparse import Namespace
from typing import Dict, List, Tuple, Union

import pandas as pd
import torch

from polos.models.estimators import PolosEstimator, Estimator
from polos.modules.feedforward import FeedForward
from polos.modules.scalar_mix import ScalarMixWithDropout
from torchnlp.utils import collate_tensors


class QualityEstimator(PolosEstimator):
    """
    Estimator class that uses a pretrained encoder to extract features from
    the sequences and then passes those features to a feed forward estimator.

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

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

    def _build_model(self) -> Estimator:
        """
        Initializes the estimator architecture.
        """
        super()._build_model()
        if self.hparams.encoder_model != "LASER":
            self.layer = (
                int(self.hparams.layer)
                if self.hparams.layer != "mix"
                else self.hparams.layer
            )

            self.scalar_mix = (
                ScalarMixWithDropout(
                    mixture_size=self.encoder.num_layers,
                    dropout=self.hparams.scalar_mix_dropout,
                    do_layer_norm=True,
                )
                if self.layer == "mix" and self.hparams.pool != "default"
                else None
            )

        self.ff = FeedForward(
            in_dim=self.encoder.output_units * 4,
            hidden_sizes=self.hparams.hidden_sizes,
            activations=self.hparams.activations,
            dropout=self.hparams.dropout,
        )

    def read_csv(self, path: str) -> List[dict]:
        """Reads a comma separated value file.

        :param path: path to a csv file.

        :return: List of records as dictionaries
        """
        df = pd.read_csv(path)
        df = df[["mt", "src", "score"]]
        df["mt"] = df["mt"].astype(str)
        df["src"] = df["src"].astype(str)
        df["score"] = df["score"].astype(float)
        return df.to_dict("records")

    def prepare_sample(
        self, sample: List[Dict[str, Union[str, float]]], inference: bool = False
    ) -> Union[
        Tuple[Dict[str, torch.Tensor], Dict[str, torch.Tensor]], Dict[str, torch.Tensor]
    ]:
        """
        Function that prepares a sample to input the model.
        :param sample: list of dictionaries.
        :param inference: If set to true prepares only the model inputs.

        :returns: Tuple with 2 dictionaries (model inputs and targets).
            If `inference=True` returns only the model inputs.
        """
        sample = collate_tensors(sample)
        mt_inputs = self.encoder.prepare_sample(sample["mt"])
        src_inputs = self.encoder.prepare_sample(sample["src"])

        mt_inputs = {"mt_" + k: v for k, v in mt_inputs.items()}
        src_inputs = {"src_" + k: v for k, v in src_inputs.items()}

        inputs = {**mt_inputs, **src_inputs}

        if inference:
            return inputs

        targets = {"score": torch.tensor(sample["score"], dtype=torch.float)}
        return inputs, targets

    def forward(
        self,
        mt_tokens: torch.tensor,
        src_tokens: torch.tensor,
        mt_lengths: torch.tensor,
        src_lengths: torch.tensor,
        **kwargs
    ) -> Dict[str, torch.Tensor]:
        """
        Function that encodes both Source, MT and returns a quality score.

        :param mt_tokens: MT sequences [batch_size x mt_seq_len]
        :param src_tokens: SRC sequences [batch_size x src_seq_len]
        :param mt_lengths: MT lengths [batch_size]
        :param src_lengths: SRC lengths [batch_size]

        :return: Dictionary with model outputs to be passed to the loss function.
        """
        mt_sentemb = self.get_sentence_embedding(mt_tokens, mt_lengths)
        src_sentemb = self.get_sentence_embedding(src_tokens, src_lengths)

        diff_src = torch.abs(mt_sentemb - src_sentemb)
        prod_src = mt_sentemb * src_sentemb

        embedded_sequences = torch.cat(
            (mt_sentemb, src_sentemb, prod_src, diff_src), dim=1
        )
        return {"score": self.ff(embedded_sequences)}