from __future__ import annotations

import logging
from typing import List

import torch

from s3prl import Logs, Output
from s3prl.nn.predictor_mockingjay import PredictorMockingjay as PredictorExample
from s3prl.nn.transformer_mockingjay import TransformerMockingjay as UpstreamExample

from . import Task

logger = logging.getLogger(__name__)


class FeatReconstructionTask(Task):
    """
    Attributes:
        upstream (torch.nn.Module): The upstream encoder (transformers, rnn, etc) that outputs `hidden_states`
        predictor (torch.nn.Module): The pre-training predictor that takes `hidden_states` as input and maps to the task target
        loss (torch.nn Loss Functions): The reconstruction loss (torch.nn.L1Loss, torch.nn.MSELoss, etc)
    """

    def __init__(
        self,
        upstream: UpstreamExample,
        predictor: PredictorExample,
        loss: torch.nn.L1Loss,
        loss_config: dict = {},
        **kwargs,
    ):
        """
        The input feature does not necessary have to be the same as the target feature.

        Args:
            upstream (Encoder)
            predictor (Projection NN)
            loss (reconstruction loss)
                feat_A -> upstream -> predictor -> feat_B
                loss(feat_A, feat_B)
        """

        super().__init__()
        self.upstream = upstream
        self.predictor = predictor
        self.loss = loss(**loss_config)

    def predict(
        self,
        x: torch.Tensor,
        label: torch.Tensor,
        label_mask: torch.BoolTensor = None,
        position_encoding: torch.Tensor = None,
        attention_mask: torch.LongTensor = None,
    ):
        """
        Args:
            x (torch.Tensor): source_feat - (batch_size, timestamps, input_size)
            label (torch.Tensor): target_feat - (batch_size, timestamps, output_size)
            label_mask (torch.BoolTensor): (batch_size, timestamps, output_size)
            position_encoding (torch.Tensor): (batch_size, timestamps, input_size)
            attention_mask (torch.LongTensor): (batch_size, timestamps)

        Return:
            hidden_states (torch.Tensor): (batch_size, timestamps, hidden_size)
            loss (torch.Tensor): scalar.
            prediction (torch.Tensor): (batch_size, timestamps, output_size)
        """
        if position_encoding is None and attention_mask is None:
            upstream_output: torch.Tensor = self.upstream(x)
        else:
            upstream_output: torch.Tensor = self.upstream(
                x, position_encoding, attention_mask
            )
        prediction: torch.Tensor = self.predictor(upstream_output).prediction

        if label_mask is None:
            reconstruction_loss = self.loss(prediction, label)
        else:
            assert label_mask.sum() > 0, "Without any masking, loss might go NaN."
            reconstruction_loss = self.loss(
                prediction.masked_select(label_mask), label.masked_select(label_mask)
            )

        return Output(
            loss=reconstruction_loss,
            hidden_states=upstream_output.hidden_states,
            prediction=prediction,
        )

    def _general_forward(
        self,
        x: torch.Tensor,
        label: torch.Tensor,
        label_mask: torch.BoolTensor = None,
        position_encoding: torch.Tensor = None,
        attention_mask: torch.LongTensor = None,
        unique_name: List[str] = None,
    ):
        loss, hidden_states, prediction = self.predict(
            x, label, label_mask, position_encoding, attention_mask
        ).slice(3)

        logs = Logs()
        logs.add_hidden_state("hidden_states", hidden_states)
        logs.add_hidden_state("prediction", prediction)

        return Output(
            loss=loss,
            prediction=prediction,
            label=label,
            unique_name=unique_name,
            logs=logs,
        )

    def _general_reduction(self, batch_results: list, on_epoch_end: bool = None):
        losses = []
        for batch_result in batch_results:
            losses.append(batch_result.loss)

        loss = (sum(losses) / len(losses)).item()

        logs = Logs()
        logs.add_scalar("loss", loss)

        return Output(
            logs=logs,
        )

    def train_step(
        self,
        x: torch.Tensor,
        label: torch.Tensor,
        label_mask: torch.BoolTensor = None,
        position_encoding: torch.Tensor = None,
        attention_mask: torch.LongTensor = None,
        unique_name: List[str] = None,
        **kwargs,
    ):
        """
        Each forward step in the training loop

        Args:
            source_feat (torch.Tensor): (batch_size, timestamps, input_size)
            target_feat (torch.Tensor): (batch_size, timestamps, output_size)
            label_mask (torch.BoolTensor): (batch_size, timestamps, output_size)
            pos_enc (torch.Tensor): (batch_size, timestamps, input_size)
            attn_mask (torch.LongTensor): (batch_size, timestamps)

        Return:
            hidden_states (torch.Tensor): (batch_size, timestamps, hidden_size)
            loss (torch.Tensor): scalar.
            prediction (torch.Tensor): (batch_size, timestamps, output_size)
        """
        return self._general_forward(
            x, label, label_mask, position_encoding, attention_mask, unique_name
        )

    def train_reduction(self, batch_results: list, on_epoch_end: bool = False):
        """
        After several forward steps, outputs should be collected untouched (but detaching the Tensors)
        into a list and passed as batch_results. This function examine the collected items and compute
        metrics across these batches. This function might be called in the middle of an epoch for quick
        logging, or after exactly an epoch to know the epoch level performance.

        Args:
            batch_results (List[cacheable version of the output of self.train_step])
            on_epoch_end (bool):
                usually you should keep the same behavior between sub-epoch and epoch level
                this parameter is here in case you need specific postprocessing which must
                only be done right on the end of an epoch

        Return:
            logs (List[Log]):
                a list of content to log onto any logger
                each content should be in the Log class format
        """
        return self._general_reduction(batch_results, on_epoch_end)

    def valid_step(
        self,
        x: torch.Tensor,
        label: torch.Tensor,
        label_mask: torch.BoolTensor = None,
        position_encoding: torch.Tensor = None,
        attention_mask: torch.LongTensor = None,
        unique_name: List[str] = None,
        **kwargs,
    ):
        return self._general_forward(
            x, label, label_mask, position_encoding, attention_mask, unique_name
        )

    def test_step(
        self,
        x: torch.Tensor,
        label: torch.Tensor,
        label_mask: torch.BoolTensor = None,
        position_encoding: torch.Tensor = None,
        attention_mask: torch.LongTensor = None,
        unique_name: List[str] = None,
        **kwargs,
    ):
        return self._general_forward(
            x, label, label_mask, position_encoding, attention_mask, unique_name
        )

    def valid_reduction(self, batch_results: list, on_epoch_end: bool = True):
        return self._general_reduction(batch_results, on_epoch_end)

    def test_reduction(self, batch_results: list, on_epoch_end: bool = True):
        return self._general_reduction(batch_results, on_epoch_end)
