# Copyright (c) OpenMMLab. All rights reserved.
from typing import List, Optional, Sequence, Union

import mmengine
import torch
from mmedit.structures import EditDataSample, PixelData
from mmengine import Config
from mmengine.model.base_model.data_preprocessor import BaseDataPreprocessor
from mmengine.registry import Registry
from mmengine.structures import BaseDataElement
from torch import nn

from mmdeploy.codebase.base import BaseBackendModel
from mmdeploy.utils import (Backend, get_backend, get_codebase_config,
                            get_root_logger, load_config)

__BACKEND_MODEL = Registry('backend_models')


@__BACKEND_MODEL.register_module('end2end')
class End2EndModel(BaseBackendModel):
    """End to end model for inference of super resolution.

    Args:
        backend (Backend): The backend enum, specifying backend type.
        backend_files (Sequence[str]): Paths to all required backend files(e.g.
            '.onnx' for ONNX Runtime, '.param' and '.bin' for ncnn).
        device (str): A string represents device type.
        model_cfg(mmengine.Config): Input model config object.
        deploy_cfg(str | mmengine.Config):Deployment config file or loaded
            Config object.
        data_preprocessor (BaseDataPreprocessor): The data preprocessor
                of the model. Default to `None`.
    """

    def __init__(self,
                 backend: Backend,
                 backend_files: Sequence[str],
                 device: str,
                 model_cfg: mmengine.Config,
                 deploy_cfg: Union[str, mmengine.Config] = None,
                 data_preprocessor: Optional[Union[dict, nn.Module]] = None,
                 **kwargs):
        super().__init__(
            deploy_cfg=deploy_cfg, data_preprocessor=data_preprocessor)
        self.deploy_cfg = deploy_cfg
        self.test_cfg = model_cfg.test_cfg
        self.device = device
        self._init_wrapper(
            backend=backend,
            backend_files=backend_files,
            device=device,
            **kwargs)

    def _init_wrapper(self, backend: Backend, backend_files: Sequence[str],
                      device: str, **kwargs):
        output_names = self.output_names
        self.wrapper = BaseBackendModel._build_wrapper(
            backend=backend,
            backend_files=backend_files,
            device=device,
            input_names=[self.input_name],
            output_names=output_names,
            deploy_cfg=self.deploy_cfg,
            **kwargs)

    def forward(self,
                inputs: torch.Tensor,
                data_samples: Optional[List[BaseDataElement]] = None,
                mode: str = 'predict',
                **kwargs) -> Sequence[EditDataSample]:
        """Run test inference for restorer.

        We want forward() to output an image or a evaluation result.
        When test_mode is set, the output is evaluation result. Otherwise
        it is an image.

        Args:
            inputs (torch.Tensor): The input tensors
            data_samples (List[BaseDataElement], optional): The data samples.
                Defaults to None.
            mode (str, optional): forward mode, only support `predict`.
            **kwargs: Other key-pair arguments.

        Returns:
            list | dict: High resolution image or a evaluation results.
        """
        assert mode == 'predict', \
            'Backend model only support mode==predict,' f' but get {mode}'
        lq = inputs
        if lq.device != torch.device(self.device):
            get_root_logger().warning(f'expect input device {self.device}'
                                      f' but get {lq.device}.')
        lq = lq.to(self.device)
        batch_outputs = self.wrapper({self.input_name:
                                      lq})[self.output_names[0]].to('cpu')
        if hasattr(self.data_preprocessor, 'destructor'):
            batch_outputs = self.data_preprocessor.destructor(
                batch_outputs.to(self.data_preprocessor.outputs_std.device))
        predictions = []

        for sr_pred, data_sample in zip(batch_outputs, data_samples):
            pred = EditDataSample()
            pred.set_data(dict(pred_img=PixelData(**dict(data=sr_pred))))
            data_sample.set_data(dict(output=pred))
            '''
            data_sample.set_data(
                dict(pred_img=PixelData(**dict(data=sr_pred))))
            '''
            predictions.append(data_sample)
        return predictions


@__BACKEND_MODEL.register_module('sdk')
class SDKEnd2EndModel(End2EndModel):
    """SDK inference class, converts SDK output to mmedit format."""

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

    def forward(self,
                inputs: torch.Tensor,
                data_samples: Optional[List[BaseDataElement]] = None,
                mode: str = 'predict',
                *args,
                **kwargs) -> list:
        """Run test inference for restorer.

        We want forward() to output an image or a evaluation result.
        When test_mode is set, the output is evaluation result. Otherwise
        it is an image.

        Args:
            inputs (torch.Tensor): The input tensors
            data_samples (List[BaseDataElement], optional): The data samples.
                Defaults to None.
            mode (str, optional): forward mode, only support `predict`.
            *args: Other arguments.
            **kwargs: Other key-pair arguments.

        Returns:
            list | dict: High resolution image or a evaluation results.
        """
        if hasattr(self.data_preprocessor, 'destructor'):
            inputs = self.data_preprocessor.destructor(
                inputs.to(self.data_preprocessor.input_std.device))
        outputs = []
        for i in range(inputs.shape[0]):
            output = self.wrapper.invoke(inputs[i].permute(
                1, 2, 0).contiguous().detach().cpu().numpy())
            outputs.append(
                torch.from_numpy(output).permute(2, 0, 1).contiguous())
        outputs = torch.stack(outputs, 0) / 255.
        if hasattr(self.data_preprocessor, 'destructor'):
            outputs = self.data_preprocessor.destructor(
                outputs.to(self.data_preprocessor.outputs_std.device))

        for i, sr_pred in enumerate(outputs):
            pred = EditDataSample()
            pred.set_data(dict(pred_img=PixelData(**dict(data=sr_pred))))
            data_samples[i].set_data(dict(output=pred))
        return data_samples


def build_super_resolution_model(
        model_files: Sequence[str],
        model_cfg: Union[str, mmengine.Config],
        deploy_cfg: Union[str, mmengine.Config],
        device: str,
        data_preprocessor: Optional[Union[Config,
                                          BaseDataPreprocessor]] = None,
        **kwargs):
    """Build super resolution model for different backends.

    Args:
        model_files (Sequence[str]): Input model file(s).
        model_cfg (str | Config): Input model config file or Config
            object.
        deploy_cfg (str | Config): Input deployment config file or
            Config object.
        device (str):  Device to input model
        data_preprocessor (BaseDataPreprocessor | Config): The data
            preprocessor of the model.
    Returns:
        End2EndModel: Super Resolution model for a configured backend.
    """
    model_cfg = load_config(model_cfg)[0]
    deploy_cfg = load_config(deploy_cfg)[0]

    backend = get_backend(deploy_cfg)
    model_type = get_codebase_config(deploy_cfg).get('model_type', 'end2end')
    backend_model = __BACKEND_MODEL.build(
        dict(
            type=model_type,
            backend=backend,
            backend_files=model_files,
            device=device,
            model_cfg=model_cfg,
            deploy_cfg=deploy_cfg,
            data_preprocessor=data_preprocessor,
            **kwargs))

    return backend_model
