import functools
import warnings
from typing import Dict, Tuple

import pytorch_lightning as pl
import torch
from pytorch_lightning.callbacks.finetuning import BaseFinetuning

from aixd.mlmodel.architecture.cond_ae_model import CondAEModel

warnings.filterwarnings("ignore", ".*does not have many workers.*")


class InverseModel:
    """
    Wrapper class for taking a Cond(V)AEModel and fine-tuning the decoder using the encoder as surrogate model.
    It was observed that the decoder may neglect the conditional features y, which is why a second-stage training was suggested,
    where the order of the decoder and encoder are swapped, the encoder frozen and the model subsequently trained to reconstruct the correct y.

    This class copies the complete functionality of the input model, except for the `forward` and `fit` methods, which are redefined for fine-tuning.

    Parameters
    ----------
    model : CondAEModel)
        Model whose decoder should be fine-tuned using its (frozen) encoder as surrogate.
    gen_z_strategy : str, default='encode'
        Which strategy should be employed for optaining the latent vectors z.
        One of 'encode', 'sample' or 'sample_around'.
        -   if 'encode', the latent variables z are generated by the (frozen) encoder
        -   if 'sample', the latent variables z are sampled normally
        -   if 'sample_around', the latent variables z are generated by the (frozen) encoder and added to random gaussian noise with std `sample_around_std`.
    sample_around_std : float, default=0.1
        If `gen_z_strategy`is 'sample_around', the random gaussion noise is sampled using this value as standard deviation.
    """

    def __init__(self, model: CondAEModel, gen_z_strategy: str = "encode", sample_around_std: float = 0.1, **kwargs):
        self.model = model
        self.gen_z_strategy = gen_z_strategy
        self.sample_around_std = sample_around_std

        # get all functions of `model` and add them to this instance except for `forward` and `fit`, which are redefined.
        for name in dir(model):
            attr = getattr(model, name)
            if callable(attr) and name not in dir(self):
                setattr(self, name, self.__wrap_method(attr))

    def __wrap_method(self, method):
        """
        A helper function that wraps the given method and modifies its behavior according to the method name.
        If the method name is `forward`, it returns the result of `self.forward` method.
        If the method name is `fit`, it returns the result of `self.fit` method.
        Otherwise, it returns the result of the given method directly.

        Parameters
        ----------
        method : callable
            The method to be wrapped.

        Returns
        -------
        wrapper : callable
            The wrapped method with modified behavior.
        """

        @functools.wraps(method)
        def wrapper(*args, **kwargs):
            if method.__name__ == "forward":
                return self.forward(*args, **kwargs)
            elif method.__name__ == "fit":
                return self.fit(*args, **kwargs)
            else:
                return method(*args, **kwargs)

        return wrapper

    def forward(self, data: Tuple[torch.Tensor, torch.Tensor]) -> Dict[str, torch.Tensor]:
        """
        Obtains latent vectors z according to `self.gen_z_strategy`, then decodes the data to obtain x_hat
        and finally encodes x_hat using the (frozen) encoder to obtain y_hat.

        Parameters
        ----------
        data : Tuple[torch.Tensor, torch.Tensor]
            The training data containing x and y.

        Returns
        -------
        Dict[str, torch.Tensor]
            Prediction
        """
        x, y = data

        if self.gen_z_strategy in ["encode", "sample_around"]:
            z = self.encoder(x)["z"]

            if self.gen_z_strategy == "sample_around":
                z = z + torch.normal(mean=0.0, std=self.sample_around_std, size=(y.shape[0], self.model.latent_dim))

        elif self.gen_z_strategy == "sample":
            z = torch.normal(mean=0.0, std=1.0, size=(y.shape[0], self.model.latent_dim))

        pred = self.decoder({"z": z, "y": y})
        pred.update(self.encoder(pred["x"]))

        return pred

    def fit(
        self, datamodule: pl.LightningDataModule, max_epochs: int = 100, callbacks: list = None, loggers: list = None, accelerator: str = "auto", flag_wandb=False, **kwargs
    ) -> None:
        """
        Launches fine-tuning of the decoder in `self.model`.
        Freezes the encoder prior to calling `self.model.fit()` and unfreezes it upon completion.

        Parameters
        ----------
        data_module : pl.LightningDataModule
            DataModule object that provides the training, validation, and test data.
        max_epochs : int, optional, default=100
            The maximum number of epochs to train for.
        callbacks : list, optional, default=None
            A list of PyTorch Lightning Callback objects to use during training.
        loggers : list, optional, default=None
            A list of PyTorch Lightning Logger objects to use during training.
        accelerator : str, optional, default='auto'
            Which accelerator should be used (e.g. cpu, gpu, mps, etc.)
        """
        freeze_callback = FreezeEncoder()

        if not callbacks:
            callbacks = []
        callbacks.append(freeze_callback)

        self.model.fit(datamodule, max_epochs, callbacks, loggers, accelerator, flag_wandb, **kwargs)

        for optimizer in self.model.trainer.optimizers:
            freeze_callback.unfreeze_and_add_param_group(
                modules=self.model.encoder,
                optimizer=optimizer,
                train_bn=True,
            )


class FreezeEncoder(BaseFinetuning):
    """
    A PyTorch Lightning callback class that extends the `BaseFinetuning` class and
    freezes the encoder of the `CondAEModel` before training.
    """

    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)

    def freeze_before_training(self, module: CondAEModel):
        """
        Freezes the encoder of the `CondAEModel` before training.

        Parameters
        ----------
        module : `CondAEModel`
            The model whose encoder should be frozen.
        """
        self.freeze(module.encoder, train_bn=False)

    def finetune_function(self, pl_module, current_epoch, optimizer, optimizer_idx):
        pass
