from typing import Optional

from torch.utils.data import DataLoader

from move.models.vae import VAE

TrainingLoopOutput = tuple[list[float], list[float], list[float], list[float], float]


def dilate_batch(dataloader: DataLoader) -> DataLoader:
    """
    Increase the batch size of a dataloader.

    Args:
        dataloader (DataLoader): An object feeding data to the VAE

    Returns:
        DataLoader: An object feeding data to the VAE
    """
    assert dataloader.batch_size is not None
    dataset = dataloader.dataset
    batch_size = int(dataloader.batch_size * 1.5)
    return DataLoader(dataset, batch_size, shuffle=True, drop_last=True)


BATCH_DILATION_STEPS = []
KLD_WARMUP_STEPS = []


def training_loop(
    model: VAE,
    train_dataloader: DataLoader,
    valid_dataloader: Optional[DataLoader] = None,
    lr: float = 1e-4,
    num_epochs: int = 100,
    batch_dilation_steps: list[int] = BATCH_DILATION_STEPS,
    kld_warmup_steps: list[int] = KLD_WARMUP_STEPS,
    early_stopping: bool = False,
    patience: int = 0,
) -> TrainingLoopOutput:
    """
    Trains a VAE model with batch dilation and KLD warm-up. Optionally,
    enforce early stopping.

    Args:
        model (VAE): trained VAE model object
        train_dataloader (DataLoader):  An object feeding data to the VAE
                                        with training data
        valid_dataloader (Optional[DataLoader], optional): An object feeding data to the
                                            VAE with validation data. Defaults to None.
        lr (float, optional): learning rate. Defaults to 1e-4.
        num_epochs (int, optional): number of epochs. Defaults to 100.
        batch_dilation_steps (list[int], optional): a list with integers corresponding
                                to epochs when batch size is increased. Defaults to [].
        kld_warmup_steps (list[int], optional):  a list with integers corresponding to
                    epochs when kld is decreased by the selected rate. Defaults to [].
        early_stopping (bool, optional): boolean if use early stopping.
                                         Defaults to False.

        patience (int, optional): number of epochs to wait before early stop
                                  if no progress on the validation set. Defaults to 0.

    Returns:
        (tuple): a tuple containing:
            *outputs (*list): lists containing information of epoch loss, BCE loss,
                              SSE loss, KLD loss
            kld_weight (float): final KLD after dilations during the training
    """

    outputs = [[] for _ in range(4)]
    min_likelihood = float("inf")
    counter = 0

    kld_weight = 0.0

    for epoch in range(1, num_epochs + 1):
        if epoch in kld_warmup_steps:
            kld_weight += 1 / len(kld_warmup_steps)

        if epoch in batch_dilation_steps:
            train_dataloader = dilate_batch(train_dataloader)

        for i, output in enumerate(
            model.encoding(train_dataloader, epoch, lr, kld_weight)
        ):
            outputs[i].append(output)

        if early_stopping and valid_dataloader is not None:
            output = model.latent(valid_dataloader, kld_weight)
            valid_likelihood = output[-1]
            if valid_likelihood > min_likelihood and counter < patience:
                counter += 1
                if counter % 5 == 0:
                    lr *= 0.9
            elif counter == patience:
                break
            else:
                min_likelihood = valid_likelihood
                counter = 0

    return *outputs, kld_weight