# Infinigen dataset adapter for PPD eval. Mirrors ppd/data/nyu.py.
# Scene layout (EvalMDE convention):
#   <data_root>/<scene>/rgb.png
#   <data_root>/<scene>/gt_depth.npz  keys: depth (H,W) f32, intr (4,) px, valid (H,W) bool
import os
from os.path import join
import numpy as np
from ppd.data.depth_estimation import Dataset as BaseDataset
from ppd.utils.logger import Log


class Dataset(BaseDataset):
    def build_metas(self):
        self.dataset_name = 'infinigen'
        # split_path: optional newline-delimited scene names. If absent, auto-discover.
        sp = self.cfg.get('split_path', None)
        if sp and os.path.exists(sp):
            scenes = [s.strip() for s in open(sp).read().splitlines() if s.strip()]
        else:
            root = self.cfg.data_root
            scenes = sorted(d for d in os.listdir(root)
                            if os.path.isdir(join(root, d))
                            and os.path.exists(join(root, d, 'rgb.png'))
                            and os.path.exists(join(root, d, 'gt_depth.npz')))
        self.rgb_files   = [join(self.cfg.data_root, s, 'rgb.png')     for s in scenes]
        self.depth_files = [join(self.cfg.data_root, s, 'gt_depth.npz') for s in scenes]
        self.scene_names = scenes

    def read_depth(self, index, depth=None):
        # Infinigen gt_depth.npz uses keys 'depth' + 'valid' (not 'data' as
        # depth_estimation.py default assumes). Override here.
        data = np.load(self.depth_files[index])
        depth = data['depth'].astype(np.float32)
        valid_mask = data['valid'].astype(bool) & np.isfinite(depth) & (depth > 0.01)
        if valid_mask.sum() == 0:
            Log.warn(f'No valid pixels in {self.depth_files[index]}')
        # Replace invalid pixels with a finite value so downstream isn't NaN-poisoned.
        if (~np.isfinite(depth)).any():
            fill = depth[valid_mask].max() if valid_mask.sum() > 0 else 1.0
            depth = np.where(np.isfinite(depth), depth, fill)
        return depth, valid_mask.astype(np.uint8)

    def read_rgb_name(self, index):
        # Use scene dir name as identifier (mirrors NYU two-level join, but Infinigen
        # has single scene-dir layer).
        return self.scene_names[index] + '.png'