import argparse
import cv2
import numpy as np
import os
import torch
import torch.nn as nn
import torch.nn.functional as F
from torchvision.transforms import Compose

from models.monoD.depth_anything.dpt import DPT_DINOv2
from models.monoD.depth_anything.util.transform import (
    Resize, NormalizeImage, PrepareForNet
)


def build(config):
    """
        Build the model from the config
        NOTE: the config should contain the following
        - encoder: the encoder type of the model
        - load_from: the path to the pretrained model
    """
    args = config
    assert args.encoder in ['vits', 'vitb', 'vitl']
    if args.encoder == 'vits':
        depth_anything = DPT_DINOv2(encoder='vits', features=64,
                                           out_channels=[48, 96, 192, 384],
                                            localhub=args.localhub).cuda()
    elif args.encoder == 'vitb':
        depth_anything = DPT_DINOv2(encoder='vitb', features=128,
                                     out_channels=[96, 192, 384, 768],
                                       localhub=args.localhub).cuda()
    else:
        depth_anything = DPT_DINOv2(encoder='vitl', features=256,
                                     out_channels=[256, 512, 1024, 1024],
                                       localhub=args.localhub).cuda()
    depth_anything.load_state_dict(torch.load(args.load_from,
                                               map_location='cpu'), strict=True)
    total_params = sum(param.numel() for param in depth_anything.parameters())
    print('Total parameters: {:.2f}M'.format(total_params / 1e6))    
    depth_anything.eval()

    return depth_anything

class DepthAnything(nn.Module):
    def __init__(self, args):
        super(DepthAnything, self).__init__()

        # build the chosen model
        self.dpAny = build(args)

    def infer(self, rgbs):
        """
            Infer the depth map from the input RGB image
        
        Args:
            rgbs: the input RGB image B x 3 x H x W (Cuda Tensor)
        
        Asserts:
            the input should be a cuda tensor
        """
        assert (rgbs.is_cuda)&(len(rgbs.shape) == 4)
        T, C, H, W = rgbs.shape
        # prepare the input
        Resizer = Resize(
                width=518,
                height=518,
                resize_target=False,
                keep_aspect_ratio=True,
                ensure_multiple_of=14,
                resize_method='lower_bound',
                image_interpolation_method=cv2.INTER_CUBIC,
            )       
        #NOTE: step 1 Resize 
        width, height = Resizer.get_size( 
            rgbs.shape[2], rgbs.shape[3]
        )
        rgbs = F.interpolate(
            rgbs, (int(height), int(width)), mode='bicubic', align_corners=False
        )
        #NOTE: step 2 NormalizeImage
        mean_ = torch.tensor([0.485, 0.456, 0.406], 
                             device=rgbs.device).view(1, 3, 1, 1)
        std_ = torch.tensor([0.229, 0.224, 0.225],
                            device=rgbs.device).view(1, 3, 1, 1)
        rgbs = (rgbs - mean_)/std_
        #NOTE: step 3 PrepareForNet

        # get the depth map

        disp = self.dpAny(rgbs)
        disp = F.interpolate(
            disp[:,None], (H, W), 
            mode='bilinear', align_corners=False
        )
        # clamping the farthest depth to 100x of the nearest
        depth_map = disp        

        return depth_map