# Copyright (c) Facebook, Inc. and its affiliates.
#
# This source code is licensed under the MIT license found in the
# LICENSE file in the root directory of this source tree.
import torch

from .. import tasks
from .. import models
from .. import losses
from ..datasets import MMDataset
from .. import processors


class Task(object):
    """
    A task refers to one generic training task (e.g., training one model).
    """

    @classmethod
    def config_task(cls, config):
        """
        determine whether to load a hard-coded task or config from a generic one.
        via if a task string is available in config.
        """
        if config.task is not None:
            # TODO (huxu): expand the search scope.
            task_cls = getattr(tasks, config.task)
            return task_cls(config)
        else:
            return Task(config)

    def __init__(self, config):
        self.config = config
        self.train_data = None
        self.val_data = None
        self.test_data = None

        self.model = None
        self.loss_fn = None
        self.eval_fn = None

    def build_dataset(self):
        """TODO (huxu): move processor breakdown to MMDataset."""
        """fill-in `self.train_data`, `self.val_data` and `self.test_data`."""

        meta_processor_cls = getattr(
            processors, self.config.dataset.meta_processor)
        video_processor_cls = getattr(
            processors, self.config.dataset.video_processor)
        text_processor_cls = getattr(
            processors, self.config.dataset.text_processor)
        aligner_cls = getattr(
            processors, self.config.dataset.aligner)

        if self.config.dataset.train_path is not None:
            self.config.dataset.split = "train"
            # may be used by meta processor.
            # meta_processor controls different dataset.
            meta_processor = meta_processor_cls(self.config.dataset)
            video_processor = video_processor_cls(self.config.dataset)
            text_processor = text_processor_cls(self.config.dataset)
            aligner = aligner_cls(self.config.dataset)
            self.train_data = MMDataset(
                meta_processor, video_processor, text_processor, aligner
            )
            print("train_len", len(self.train_data))
            output = self.train_data[0]
            self.train_data.print_example(output)
        if self.config.dataset.val_path is not None:
            self.config.dataset.split = "valid"
            # may be used by meta processor.
            meta_processor = meta_processor_cls(self.config.dataset)
            video_processor = video_processor_cls(self.config.dataset)
            text_processor = text_processor_cls(self.config.dataset)
            aligner = aligner_cls(self.config.dataset)
            self.val_data = MMDataset(
                meta_processor, video_processor, text_processor, aligner
            )
            print("val_len", len(self.val_data))
            output = self.val_data[0]
            self.val_data.print_example(output)

        if self.config.dataset.split == "test":
            # the following is run via lauching fairseq-validate.
            meta_processor = meta_processor_cls(self.config.dataset)
            video_processor = video_processor_cls(self.config.dataset)
            text_processor = text_processor_cls(self.config.dataset)

            self.test_data = MMDataset(
                meta_processor, video_processor, text_processor, aligner
            )
            print("test_len", len(self.test_data))
            output = self.test_data[0]
            self.test_data.print_example(output)

    def build_model(self, checkpoint=None):
        if self.model is None:
            model_cls = getattr(models, self.config.model.model_cls)
            self.model = model_cls(self.config)
        if checkpoint is not None:
            self.load_checkpoint(checkpoint)
        return self.model

    def load_checkpoint(self, checkpoint):
        if self.model is None:
            raise ValueError("model is not initialized.")
        state_dict = torch.load(checkpoint)
        state_dict = self._trim_state_dict(state_dict)
        self.model.load_state_dict(state_dict, strict=False)
        # if it's a fp16 model, turn it back.
        if next(self.model.parameters()).dtype == torch.float16:
            self.model = self.model.float()
        return self.model

    def _trim_state_dict(self, state_dict):
        from collections import OrderedDict

        if "state_dict" in state_dict:
            state_dict = state_dict["state_dict"]
        if "model" in state_dict:  # fairseq checkpoint format.
            state_dict = state_dict["model"]
        ret_state_dict = OrderedDict()
        for (
            key,
            value,
        ) in state_dict.items():
            # remove fairseq wrapper since this is a task.
            if key.startswith("mmmodel"):
                key = key[len("mmmodel."):]
            ret_state_dict[key] = value
        return ret_state_dict

    def build_loss(self):
        if self.loss_fn is None and self.config.loss is not None:
            loss_cls = getattr(losses, self.config.loss.loss_cls)
            self.loss_fn = loss_cls()
        return self.loss_fn

    def flat_subsample(self, tensor):
        size = tensor.size()
        if len(size) >= 2:
            batch_size = size[0] * size[1]
            expanded_size = (
                (batch_size,) + size[2:] if len(size) > 2
                else (batch_size,)
            )
            tensor = tensor.view(expanded_size)
        return tensor

    def reshape_subsample(self, sample):
        if (
            hasattr(self.config.dataset, "subsampling")
            and self.config.dataset.subsampling is not None
            and self.config.dataset.subsampling > 1
        ):
            for key in sample:
                if torch.is_tensor(sample[key]):
                    sample[key] = self.flat_subsample(sample[key])
        return sample

    def __call__(self, model, sample):
        loss = None
        loss_scalar = float("inf")

        sample = self.reshape_subsample(sample)
        outputs = self.model(**sample)
        sample.update(outputs)
        if self.loss_fn is not None:
            loss = self.loss_fn(**sample)
            loss_scalar = loss.item()

        batch_size = sample["caps"].size(0)
        sample_size = 1
        return {
            "loss": loss,
            "loss_scalar": loss_scalar,
            "max_len": self.config.dataset.max_len,
            "batch_size": batch_size,
            "sample_size": sample_size,
        }

    def build_dataloader(self):
        """only used for trainer that lacks building loaders."""
        raise NotImplementedError