'''
Class for estimating relative inverse depth map by MiDaS
'''

import rospy

import torch
from cv2 import INTER_CUBIC
import numpy as np
import time
from pathlib import Path

# midas modules
from torchvision.transforms import Compose
from midas.dpt_depth import DPTDepthModel
from midas.midas_net import MidasNet
from midas.midas_net_custom import MidasNet_small
from midas.transforms import Resize, NormalizeImage, PrepareForNet

FILE = Path(__file__).resolve()
ROOT = FILE.parents[2]

class MidasDepth:

    def __init__(self, device):
        
        ### launch parameters
        try:
            MIDAS_MODEL_PATH = str(rospy.get_param('midas_model_path'))
        except Exception:
            print("Can't load Midas parameters from ROS parameter server, using defaults.")
            MIDAS_MODEL_PATH = 'data/midas_small.pt' # midas_hybrid.pt / midas_small.pt

        ### init atributes
        self.DO_DEBUGING = False

        ### midas initialization

        self.optimize = False # use half precision on cuda - True nefunguje na Jetsonu pro MidasSmall
        model_path = ROOT / MIDAS_MODEL_PATH  # model.pt path(s)
        # Tested model types: midas_small.pt / midas_hybrid.pt
        model_type = Path(model_path).stem
        self.device = torch.device(
            "cuda" if torch.cuda.is_available() and device != 'cpu' else "cpu")

        # load network
        if model_type == "dpt_large": # DPT-Large
            self.model = DPTDepthModel(
                path=model_path,
                backbone="vitl16_384",
                non_negative=True,
            )
            net_w, net_h = 384, 384
            resize_mode = "minimal"
            normalization = NormalizeImage(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5])
        elif model_type == "midas_hybrid": #DPT-Hybrid
            self.model = DPTDepthModel(
                path=model_path,
                backbone="vitb_rn50_384",
                non_negative=True,
            )
            net_w, net_h = 384, 384
            resize_mode="upper_bound" # originally "minimal"
            normalization = NormalizeImage(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5])
        elif model_type == "midas_v21":
            self.model = MidasNet(model_path, non_negative=True)
            net_w, net_h = 384, 384
            resize_mode="upper_bound"
            normalization = NormalizeImage(
                mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]
            )
        elif model_type == "midas_small":
            self.model = MidasNet_small(model_path, features=64, backbone="efficientnet_lite3",
                exportable=True, non_negative=True, blocks={'expand': True})
            net_w, net_h = 256, 256
            resize_mode="minimal" # originally "upper_bound"
            normalization = NormalizeImage(
                mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]
            )   
        else:
            print(f"Model_type '{model_type}' not implemented.")
            assert False

        # compose transformer
        self.transform = Compose(
            [
                Resize(
                    net_w,
                    net_h,
                    resize_target=None,
                    keep_aspect_ratio=True,
                    ensure_multiple_of=32,
                    resize_method=resize_mode,
                    image_interpolation_method=INTER_CUBIC,
                ),
                normalization,
                PrepareForNet(),
            ]
        )

        # switch to inference mode
        self.model.eval()
        
        if self.optimize == True:
            if self.device == torch.device("cuda"):
                self.model = self.model.to(memory_format=torch.channels_last)  
                self.model = self.model.half()

        self.model.to(self.device)

    '''
    Estimate relative inverse depth map using MiDaS
    '''
    def estimate_depth_map(self,img):

        t0 = time.time()
        
        ### PREDICTION
        print(img.shape)
        img_input = self.transform({"image": img})["image"]
        print(img_input.shape)

        t1 = time.time()
        dt1 = np.round(1000*(t1-t0))

        # predict depth map
        with torch.no_grad():
            sample = torch.from_numpy(img_input).to(self.device).unsqueeze(0)
            if self.optimize==True and self.device == torch.device("cuda"):
                sample = sample.to(memory_format=torch.channels_last)  
                sample = sample.half()
            self.depth_map_rel = self.model.forward(sample)
            self.depth_map_rel = (
                torch.nn.functional.interpolate(
                    self.depth_map_rel.unsqueeze(1),
                    size=img.shape[:2],
                    mode="bicubic",
                    align_corners=False,
                )
                .squeeze()
                .cpu()
                .numpy()
            )
        
        t2 = time.time()
        dt2 = np.round(1000*(t2-t1))

        if self.DO_DEBUGING:
            rospy.loginfo(f'Image transform: {dt1} ms; Depth estimation: {dt2} ms.')


