import numpy as np
import torch
from depth_anything_v2.dpt import DepthAnythingV2


class DepthMopdel:
    def __init__(self):
        self.input_size = 518
        self.initModel()

    def initModel(self):
        DEVICE = 'cuda' if torch.cuda.is_available() else 'mps' if torch.backends.mps.is_available() else 'cpu'
        
        model_configs = {
            'vits': {'encoder': 'vits', 'features': 64, 'out_channels': [48, 96, 192, 384]},
            'vitb': {'encoder': 'vitb', 'features': 128, 'out_channels': [96, 192, 384, 768]},
            'vitl': {'encoder': 'vitl', 'features': 256, 'out_channels': [256, 512, 1024, 1024]},
            'vitg': {'encoder': 'vitg', 'features': 384, 'out_channels': [1536, 1536, 1536, 1536]}
        }
        
        depth_anything = DepthAnythingV2(**model_configs['vits'])
        depth_anything.load_state_dict(torch.load(f'checkpoints/depth_anything_v2_vits.pth', map_location='cpu', weights_only=True))
        self.depth_anything = depth_anything.to(DEVICE).eval()

    def inference(self, raw_image):
        depth = self.depth_anything.infer_image(raw_image, self.input_size)
    
        depth = (depth - depth.min()) / (depth.max() - depth.min()) * 255.0
        depth = depth.astype(np.uint8)
        depth = np.repeat(depth[..., np.newaxis], 3, axis=-1)
        
        return depth