from ppd.data.depth_estimation import Dataset as BaseDataset
from ppd.data.depth_estimation import *
import os
import json
from ppd.utils.logger import Log


def load_in_extrinsic(path):
    with open(path, 'r') as f:
        lines = f.readlines()
        lines = [line.strip() for line in lines]
        lines = [line.split(' ') for line in lines]
        lines = [[float(x) for x in line] for line in lines]

    intrinsic = np.array(lines[0]).reshape(3,3)
    extrinsic_ = np.array(lines[1]).reshape(3,4)
    extrinsic = np.eye(4)
    extrinsic[:3,:3] = extrinsic_[:3,:3]
    extrinsic[:3,3] = extrinsic_[:3,3]
    return intrinsic, extrinsic

def get_baseline_focal(left_path, right_path):
    intrinsic0, extrinsic0 = load_in_extrinsic(left_path)
    intrinsic1, extrinsic1 = load_in_extrinsic(right_path)
    cam_extrinsic = extrinsic1 @ np.linalg.inv(extrinsic0) # ref:https://github.com/fabiotosi92/SMD-Nets/issues/8
    baseline = abs(cam_extrinsic[0, 3])
    return baseline, (intrinsic0[0,0] + intrinsic0[1,1]) / 2

def get_depth(disp, focal, baseline):
    """
    get depth from reference frame disparity and camera intrinsics
    """
    return focal * baseline / disp

class Dataset(BaseDataset):

    def build_metas(self):
        self.focal = 1920
        self.dataset_name = 'unrealstereo4k'
        # Auto-discover scene folders (mirror layout uses 00000..0000N).
        if "scenes" in self.cfg:
            folder_names = list(self.cfg.scenes)
        else:
            folder_names = sorted(
                d for d in os.listdir(self.cfg.data_root)
                if os.path.isdir(os.path.join(self.cfg.data_root, d, 'Image0'))
            )
        self.rgb_files = []
        self.depth_files = []
        # The fabiotosi92/UnrealStereo4K mirror ships .jpg under Image0/; the
        # original release uses .png. Accept either.
        exts = ('.png', '.jpg', '.jpeg')
        for folder_name in folder_names:
            folder_path = os.path.join(self.cfg.data_root, folder_name, 'Image0')
            if not os.path.isdir(folder_path):
                continue
            for file_name in os.listdir(folder_path):
                base, ext = os.path.splitext(file_name)
                if ext.lower() not in exts:
                    continue
                rgb_path = os.path.join(folder_path, file_name)
                dpt_path = os.path.join(self.cfg.data_root, folder_name, 'Disp0', base + '.npy')
                if os.path.isfile(rgb_path) and os.path.isfile(dpt_path):
                    self.rgb_files.append(rgb_path)
                    self.depth_files.append(dpt_path)

        assert len(self.rgb_files) == len(self.depth_files), \
            f"{len(self.rgb_files)} rgbs vs {len(self.depth_files)} depths"


    def read_depth(self, index, depth=None):
        disp = np.load(self.depth_files[index])
        invalid_mask = np.isnan(disp) | (disp == 0)
        disp[invalid_mask] = 1e-4
        left_extrin_path = self.depth_files[index].replace('Disp0', 'Extrinsics0').replace('.npy', '.txt')
        right_extrin_path = self.depth_files[index].replace('Disp0', 'Extrinsics1').replace('.npy', '.txt')
        if os.path.isfile(left_extrin_path) and os.path.isfile(right_extrin_path):
            baseline, focal = get_baseline_focal(left_extrin_path, right_extrin_path)
        else:
            # Mirrors that ship only Image0/Image1/Disp0/Disp1 (e.g. fabiotosi92's HF
            # mirror) lack the Extrinsics .txt files. Fall back to the published
            # SMD-Nets defaults — UnrealStereo4K is rendered with a 30 cm baseline
            # at 3840×2160 with focal length 1920.
            baseline, focal = 0.3, 1920.0
        depth = get_depth(disp, focal, baseline).astype(np.float32)
        
        min_val, max_val = 0.1, 80
        tiankong_mask = depth > 200.

        valid_mask = np.logical_and(
            depth > 0.1, ~np.isnan(depth)) & (~np.isinf(depth))
        valid_mask = np.logical_and(valid_mask, depth < max_val)
        if valid_mask.sum() == 0:
            Log.warn('No valid mask in the depth map of {}'.format(
                self.depth_files[index]))
        if valid_mask.sum() != 0 and np.isnan(depth).sum() != 0:
            depth[np.isnan(depth)] = depth[valid_mask].max()
        if valid_mask.sum() != 0 and np.isinf(depth).sum() != 0:
            depth[np.isinf(depth)] = depth[valid_mask].max()

        depth = np.clip(depth, min_val, max_val)

        depth[tiankong_mask] = depth.max() + 1.0

        valid_mask = np.logical_or(valid_mask, tiankong_mask)

        return depth, valid_mask.astype(np.uint8)