# Copyright 2021 MosaicML. All Rights Reserved.

import math
from dataclasses import asdict, dataclass
from typing import Dict, Mapping, Optional

import torch
import yahp as hp

from composer.algorithms import AlgorithmHparams
from composer.core.types import Algorithm, Batch, Event, Logger, State, Tensor
from composer.models.transformer_shared import MosaicTransformer
from composer.utils import ensure_tuple


def apply_curriculum(batch: Dict[str, Tensor], curr_seq_len: int, truncate: bool) -> Batch:
    """
    Progressively increases the sequence length during training.

    Changes the sequence length of all tensors in the provided dictionary
    to ``curr_seq_len``, by either truncating the tensors (``truncate=True``)
    or reshaping the tensors to create new examples from the extra tokens
    (``truncate=False``).

    The schedule for ``curr_seq_len`` over training time should be managed
    out of this function.

    Args:
        batch: The input batch to the model, must be a dictionary.
        curr_seq_length (int): The desired sequence length to apply.
        truncate (bool): Truncate sequences early, or reshape tensors
                         to create new examples out of the extra tokens.

    Returns:
        batch: a Mapping of input tensors to the model,
               where all tensors have curr_seq_len in the second dimension.
    """

    assert isinstance(batch, Mapping)

    if truncate:
        for k in batch.keys():
            batch[k] = batch[k][:, :curr_seq_len]
    else:
        # ensure new tensor shape is divisible by curr_seq_len
        input_ids = batch['input_ids'].view(-1)
        tensor_len = (input_ids.shape[0] // curr_seq_len) * curr_seq_len

        input_ids = input_ids[:tensor_len]
        input_ids = input_ids.view(-1, curr_seq_len)
        batch['input_ids'] = input_ids

        for k, v in batch.items():
            if k == "input_ids":
                continue
            v = v.view(-1)
            v = v[:tensor_len]
            batch[k] = v.view(-1, curr_seq_len)

    return batch


@dataclass
class CurriculumLearningHparams(AlgorithmHparams):

    duration: float = hp.optional("Fraction of total training time to apply curriculum learning.", default=0.3)
    min_seq_length: int = hp.optional("Starting sequence length.", default=8)
    max_seq_length: int = hp.optional("End sequence length", default=1024)
    step_size: int = hp.optional("Sequence length step size", default=8)
    truncate: bool = hp.optional("Truncate tensors or reshape extra tokens to new examples.", default=True)

    def validate(self):
        if self.duration < 0 or self.duration > 1:
            raise ValueError(f'Duration must be getween 0 and 1, got: {self.duration}')

        if self.max_seq_length < self.min_seq_length:
            raise ValueError(f'max_seq_length={self.max_seq_length} must be '
                             f'greater than min_seq_length={self.min_seq_length}')

    def initialize_object(self) -> "CurriculumLearning":
        return CurriculumLearning(**asdict(self))


class CurriculumLearning(Algorithm):
    """Progressively increases the sequence length during training.

    Changes the sequence length of all tensors in the input batch. The
    sequence length increases from ``min_seq_length`` to ``max_seq_length``
    in steps of ``step_size`` during the first ``duration`` fraction of
    training.

    The sequence length is then kept at ``max_seq_length``
    for the rest of training.

    Tensors are either truncated (``truncate=True``) or reshaped to
    create new examples from the extra tokens (``truncate=False``).

    .. note::

        ``step_size`` should be a multiple of eight for GPUs

    .. note::

        Variable input lengths can create CUDA OOM errors. To avoid this,
        we follow pytorch notes and pre-allocate the memory with a blank
        forward and backward pass.

    Args:
        duration (float): fraction of total training for curriculum learning.
        min_seq_length (int): Minimum sequence length to start the curriculum.
        max_seq_length (int): Maximum sequence length to stop the curriculum.
        step_size (int): Step size of sequence length.

        truncate (bool): Truncate tensors or reshape extra tokens to new examples
    """

    def __init__(
        self,
        duration: float = 0.3,
        min_seq_length: int = 8,
        max_seq_length: int = 1024,
        step_size: int = 8,
        truncate: bool = True,
    ):
        self.hparams = CurriculumLearningHparams(duration=duration,
                                                 min_seq_length=min_seq_length,
                                                 max_seq_length=max_seq_length,
                                                 step_size=step_size,
                                                 truncate=truncate)
        self.hparams.validate()

    def match(self, event: Event, state: State) -> bool:
        return event in (Event.TRAINING_START, Event.AFTER_DATALOADER)

    def apply(self, event: Event, state: State, logger: Logger) -> Optional[int]:

        # in order to avoid OOMs, we do a forward and a backward pass on a dummy input.
        if event == Event.TRAINING_START:
            # ensure that input_ids is a valid model input. since we don't need the
            # results, we don't use all inputs.

            original_model = state.model.module
            assert isinstance(original_model, MosaicTransformer)
            model_inputs = original_model.get_model_inputs()  # type: ignore
            assert 'input_ids' in model_inputs
            assert 'labels' in model_inputs

            # create fake inputs
            vocab_size = len(original_model.tokenizer)  # type: ignore

            # simplifying assumption: Composer doesn't support model-parallelism,
            # so the first parameter's device is likely the same device for
            # all of the parameters
            device = next(state.model.parameters()).device

            assert (state.train_batch_size % state.world_size) == 0
            per_gpu_batch = math.ceil(state.train_batch_size / (state.world_size * state.grad_accum))
            input_ids = torch.randint(low=0,
                                      high=vocab_size - 1,
                                      size=(per_gpu_batch, self.hparams.max_seq_length),
                                      device=device).long()
            labels = input_ids.clone()
            attn_mask = torch.ones_like(labels)
            model_inputs = {
                "input_ids": input_ids,
                "labels": labels,
                "attention_mask": attn_mask,
            }

            # start by running a forward and backward pass
            # of the maximum sequence length to allocate cache.
            with state.precision_context(state.precision):
                outputs = state.model.forward(model_inputs)
                loss = original_model.loss(outputs, model_inputs)

            # since use_grad_scaling is in the Trainer, and we
            # don't care about the loss values, skip scaling
            for loss_item in ensure_tuple(loss):
                loss_item.backward()

            # zero out gradients and proceed to normal training
            assert state.optimizers is not None, \
                "optimizers are set before TRAINING_START"

            for optimizer in ensure_tuple(state.optimizers):
                optimizer.zero_grad()
        else:
            num_optimization_steps = state.steps_per_epoch * state.max_epochs
            num_curriculum_steps = int(num_optimization_steps * self.hparams.duration)

            # assume the full sequence length is the unaltered sequence length
            num_update_steps = (self.hparams.max_seq_length - self.hparams.min_seq_length) // self.hparams.step_size
            update_every_n_steps = num_curriculum_steps // num_update_steps

            curr_seq_len = self.hparams.step_size * (state.step // update_every_n_steps)
            curr_seq_len = max(curr_seq_len, self.hparams.min_seq_length)
            curr_seq_len = min(curr_seq_len, self.hparams.max_seq_length)

            state.batch = apply_curriculum(state.batch_dict, curr_seq_len, self.hparams.truncate)

            batch_size = state.batch_dict['input_ids'].shape[0]
            logger.metric_batch({
                'curriculum_learning/curr_seq_len': curr_seq_len,
                'curriculum_learning/curr_bs': batch_size,
            })
