"""
Utterance Classification Tasks

Authors
  * Leo 2022
"""

import logging
from typing import List

import numpy as np
import torch
import torch.nn.functional as F

from s3prl.dataio.encoder.category import CategoryEncoder, CategoryEncoders
from s3prl.metric import accuracy

from . import Task

logger = logging.getLogger(__name__)

__all__ = [
    "UtteranceClassifierExample",
    "UtteranceClassificationTask",
]


class UtteranceClassifierExample(torch.nn.Module):
    """
    Attributes:
        input_size: int
        output_size: int
    """

    def __init__(self, input_size=3, output_size=4):
        super().__init__()
        self._input_size = input_size
        self._output_size = output_size

    @property
    def input_size(self):
        return self._input_size

    @property
    def output_size(self):
        return self._output_size

    def forward(self, x, x_len):
        """
        Args:
            x (torch.Tensor): (batch_size, timestemps, input_size)
            x_len (torch.LongTensor): (batch_size, )

        Return:
            output (torch.Tensor): (batch_size, output_size)
        """
        assert x.size(-1) == self.input_size
        output = torch.randn(x.size(0), self.output_size)
        assert output


class UtteranceClassificationTask(Task):
    """
    Attributes:
        input_size (int): defined by model.input_size
        output_size (int): defined by len(categories)
    """

    def __init__(self, model: UtteranceClassifierExample, category: CategoryEncoder):
        """
        model.output_size should match len(categories)

        Args:
            model (UtteranceClassifier)
            category:
                encode: str -> int
                decode: int -> str
                __len__: -> int
        """

        super().__init__()
        self.model = model
        self.category = category
        assert self.model.output_size == len(category)

    def predict(self, x: torch.Tensor, x_len: torch.LongTensor):
        """
        Args:
            x (torch.Tensor): (batch_size, timestamps, input_size)
            x_len (torch.LongTensor): (batch_size, )

        Return:
            logits (torch.Tensor): (batch_size, output_size)
            prediction (list): prediction strings
        """
        logits: torch.Tensor = self.model(x, x_len)
        predictions = [
            self.category.decode(index)
            for index in logits.argmax(dim=-1).detach().cpu().tolist()
        ]
        return logits, predictions

    def forward(
        self,
        _mode: str,
        x: torch.Tensor,
        x_len: torch.LongTensor,
        class_id: torch.LongTensor,
        label: List[str],
        unique_name: List[str],
        _dump_dir: str = None,
    ):
        logits, prediction = self.predict(x, x_len)
        loss = F.cross_entropy(logits, class_id)

        cacheable = dict(
            loss=loss.detach().cpu(),
            prediction=prediction,
            label=[self.category.decode(idx) for idx in class_id],
            unique_name=unique_name,
        )

        return loss, cacheable

    def reduction(self, _mode: str, cached_results: List[dict], _dump_dir: str = None):
        results = self.parse_cached_results(cached_results)
        predictions = results["prediction"]
        labels = results["label"]
        losses = results["loss"]

        acc = accuracy(predictions, labels)
        loss = (sum(losses) / len(losses)).item()

        return dict(
            loss=loss,
            accuracy=acc,
        )


class UtteranceMultiClassClassificationTask(Task):
    def __init__(self, model: UtteranceClassifierExample, categories: CategoryEncoders):
        super().__init__()
        self.model = model
        self.categories = categories
        assert self.model.output_size == len(categories)

    def predict(self, x: torch.Tensor, x_len: torch.LongTensor):
        """
        Args:
            x (torch.Tensor): (batch_size, timestamps, input_size)
            x_len (torch.LongTensor): (batch_size, )

        Return:
            logit (torch.Tensor): List[(batch_size, sub_output_size)]
            prediction (np.array): (batch_size, num_category)
        """
        logits: torch.Tensor = self.model(x, x_len)

        logit_start = 0
        sub_logits, sub_predictions = [], []
        for category in self.categories:
            logit_end = logit_start + len(category)
            sub_logit = logits[:, logit_start:logit_end]
            sub_logits.append(sub_logit)
            sub_predictions.append(
                [
                    category.decode(index)
                    for index in sub_logit.argmax(dim=-1).detach().cpu().tolist()
                ]
            )
            logit_start = logit_end
        prediction = np.array(sub_predictions, dtype="object").T

        return sub_logits, prediction

    def forward(
        self,
        _mode: str,
        x: torch.Tensor,
        x_len: torch.LongTensor,
        class_ids: torch.LongTensor,
        labels: np.ndarray,
        unique_name: List[str],
        _dump_dir: str = None,
    ):
        """
        Args:
            x: torch.Tensor, (batch_size, timestamps, input_size)
            x_len: torch.LongTensor, (batch_size)
            class_ids: torch.LongTensor, (batch_size, num_category)
            labels: np.ndarray, (batch_size, num_category)

        Return:
            loss: torch.Tensor
            prediction: np.ndarray
            label: np.ndarray
        """
        logit, prediction = self.predict(x, x_len)
        loss = sum(
            [
                F.cross_entropy(sub_logit, class_id)
                for sub_logit, class_id in zip(logit, class_ids.T)
            ]
        )

        cacheable = dict(
            loss=loss.detach().cpu(),
            prediction=prediction.tolist(),
            label=labels.tolist(),
            unique_name=unique_name,
        )

        return loss, cacheable

    def reduction(self, _mode: str, cached_results: List[dict], _dump_dir: str = None):
        results = self.parse_cached_results(cached_results)
        losses = results["loss"]
        predictions = results["prediction"]
        labels = results["label"]

        acc = accuracy(predictions, labels)
        loss = (sum(losses) / len(losses)).item()

        return dict(
            loss=loss,
            accuracy=acc,
        )