# Copyright (c) Alibaba, Inc. and its affiliates.
import os
import os.path as osp
from glob import glob

import cv2
import numpy as np
import torch
import tqdm

from modelscope.metainfo import Models
from modelscope.models.base.base_torch_model import TorchModel
from modelscope.models.builder import MODELS
from modelscope.models.cv.video_depth_estimation.models.model_wrapper import \
    ModelWrapper
from modelscope.models.cv.video_depth_estimation.utils.augmentations import (
    resize_image, to_tensor)
from modelscope.models.cv.video_depth_estimation.utils.config import \
    parse_test_file
from modelscope.models.cv.video_depth_estimation.utils.depth import (
    inv2depth, viz_inv_depth, write_depth)
from modelscope.models.cv.video_depth_estimation.utils.image import (
    get_intrinsics, load_image, parse_video)
from modelscope.outputs import OutputKeys
from modelscope.utils.constant import ModelFile, Tasks


@MODELS.register_module(
    Tasks.video_depth_estimation,
    module_name=Models.dro_resnet18_depth_estimation)
class DROEstimation(TorchModel):

    def __init__(self, model_dir: str, **kwargs):
        """str -- model file root."""
        super().__init__(model_dir, **kwargs)
        self.check_trust_remote_code(model_dir=model_dir)

        model_path = osp.join(model_dir, ModelFile.TORCH_MODEL_FILE)

        # Parse arguments
        config, state_dict = parse_test_file(model_path)

        # If no image shape is provided, use the checkpoint one
        self.image_shape = config.datasets.augmentation.image_shape
        print(f'== input image shape:{self.image_shape}')

        # Initialize model wrapper from checkpoint arguments
        self.model_wrapper = ModelWrapper(config, load_datasets=False)
        # Restore monodepth_model state
        self.model_wrapper.load_state_dict(state_dict)

        # Send model to GPU if available
        if torch.cuda.is_available():
            model_wrapper = self.model_wrapper.to('cuda')
        else:
            model_wrapper = self.model_wrapper
            print('cuda is not available, use cpu')

        # Set to eval mode
        model_wrapper.eval()

    def forward(self, Inputs):
        return self.model_wrapper(Inputs)

    def postprocess(self, Inputs):
        return Inputs

    def inference(self, data):

        print('processing video input:.........')
        input_type = 'video'
        sample_rate = 1
        data_type = 'indoor'

        assert osp.splitext(data['video_path'])[1] in [
            '.mp4', '.avi', '.mov', '.mpeg', '.flv', '.wmv'
        ]
        input_video_images = os.path.join('tmp/input_video_images')
        parse_video(data['video_path'], input_video_images, sample_rate)
        # update input
        input = input_video_images

        files = []
        for ext in ['png', 'jpg', 'bmp']:
            files.extend(glob((os.path.join(input, '*.{}'.format(ext)))))

        if input_type == 'folder':
            print('processing folder input:...........')
            print(f'folder total frames num: {len(files)}')
            files = files[::sample_rate]

        files.sort()
        print('Found total {} files'.format(len(files)))
        assert len(files) > 2

        # Process each file
        list_of_files = list(zip(files[:-2], files[1:-1], files[2:]))

        depth_list = []
        pose_list = []
        vis_depth_list = []
        depth_upsample_list = []
        vis_depth_upsample_list = []

        print(f'*********************data_type:{data_type}')
        print('inference start.....................')
        for fn1, fn2, fn3 in tqdm.tqdm(list_of_files):
            depth, vis_depth, depth_upsample, vis_depth_upsample, pose21, pose23, intr, rgb = self.infer_and_save_pose(
                [fn1, fn3], fn2, self.model_wrapper, self.image_shape,
                data_type)
            pose_list.append((pose21, pose23))
            depth_list.append(depth)
            vis_depth_list.append(vis_depth.astype(np.uint8))
            depth_upsample_list.append(depth_upsample)
            vis_depth_upsample_list.append(vis_depth_upsample.astype(np.uint8))

        return {
            'depths': depth_list,
            'depths_color': vis_depth_upsample_list,
            'poses': pose_list
        }

    @torch.no_grad()
    def infer_and_save_pose(self, input_file_refs, input_file, model_wrapper,
                            image_shape, data_type):
        """
        Process a single input file to produce and save visualization

        Parameters
        ----------
        input_file_refs : list(str)
            Reference image file paths
        input_file : str
            Image file for pose estimation
        model_wrapper : nn.Module
            Model wrapper used for inference
        image_shape : Image shape
            Input image shape
        half: bool
            use half precision (fp16)
        save: str
            Save format (npz or png)
        """

        image_raw_wh = load_image(input_file).size

        # Load image
        def process_image(filename):
            image = load_image(filename)
            # Resize and to tensor
            intr = get_intrinsics(image.size, image_shape, data_type)  # (3, 3)
            image = resize_image(image, image_shape)
            image = to_tensor(image).unsqueeze(0)
            intr = torch.from_numpy(intr).unsqueeze(0)  # (1, 3, 3)
            # Send image to GPU if available
            if torch.cuda.is_available():
                image = image.to('cuda')
                intr = intr.to('cuda')
            return image, intr

        image_ref = [
            process_image(input_file_ref)[0]
            for input_file_ref in input_file_refs
        ]
        image, intrinsics = process_image(input_file)

        batch = {
            'rgb': image,
            'rgb_context': image_ref,
            'intrinsics': intrinsics
        }

        output = self.forward(batch)
        inv_depth = output['inv_depths'][0]  # (1, 1, h, w)
        depth = inv2depth(inv_depth)[0, 0].detach().cpu().numpy()  # (h, w)

        pose21 = output['poses'][0].mat[0].detach().cpu().numpy(
        )  # (4, 4)  # TODO check: targe -> ref[0]
        pose23 = output['poses'][1].mat[0].detach().cpu().numpy(
        )  # (4, 4)  # TODO check: targe -> ref[0]

        vis_depth = viz_inv_depth(inv_depth[0]) * 255

        vis_depth_upsample = cv2.resize(
            vis_depth, image_raw_wh, interpolation=cv2.INTER_LINEAR)
        depth_upsample = cv2.resize(
            depth, image_raw_wh, interpolation=cv2.INTER_NEAREST)

        return depth, vis_depth, depth_upsample, vis_depth_upsample, pose21, pose23, intrinsics[
            0].detach().cpu().numpy(), image[0].permute(
                1, 2, 0).detach().cpu().numpy() * 255
