"""
ScanNet20 / ScanNet200 / ScanNet Data Efficient Dataset

Author: Xiaoyang Wu (xiaoyang.wu.cs@gmail.com)
Please cite our work if the code is helpful to you.
"""

import os
import glob
import numpy as np
import torch
from copy import deepcopy
from torch.utils.data import Dataset
from collections.abc import Sequence

from pointcept.utils.logger import get_root_logger
from pointcept.utils.cache import shared_dict
from .builder import DATASETS
from .defaults import DefaultDataset
from .transform import Compose, TRANSFORMS
from .preprocessing.scannet.meta_data.scannet200_constants import (
    VALID_CLASS_IDS_20,
    VALID_CLASS_IDS_200,
)


@DATASETS.register_module()
class ScanNetDataset(DefaultDataset):
    VALID_ASSETS = [
        "coord",
        "color",
        "normal",
        "segment20",
        "instance",
    ]
    class2id = np.array(VALID_CLASS_IDS_20)

    def __init__(
        self,
        lr_file=None,
        la_file=None,
        split="train",
        data_root=None,  # <-- 确保接收这个参数
        split_file=None,
        **kwargs,
    ):
        self.lr = np.loadtxt(lr_file, dtype=str) if lr_file is not None else None
        self.la = torch.load(la_file) if la_file is not None else None
        self.split_file = split_file
        self.split = split
        # ========== 关键修复：在调用父类构造函数前，打印 data_root ==========
        print(f"[DEBUG ScanNetDataset __init__] data_root received: {data_root}")
        # ===============================================

        # ========== 关键修复：显式将 data_root 传递给父类 ==========
        super().__init__(split=split, data_root=data_root, **kwargs)
        # ===============================================

        # ========== 关键修复：在调用父类构造函数后，再次打印 ==========
        print(f"[DEBUG ScanNetDataset __init__] data_root after super(): {self.data_root}")
        # ===============================================

    def get_data_list(self):
        # ========== 修复：使用被强制覆盖的 self.data_root ==========
        if self.split_file is not None and os.path.exists(self.split_file):
            with open(self.split_file, 'r') as f:
                data_list = [
                    os.path.join(self.data_root, self.split, line.strip())
                    for line in f.readlines()
                ]
            print(f"[INFO ScanNetDataset] Loaded {len(data_list)} scenes from split file: {self.split_file}")
            print(f"[INFO ScanNetDataset] Sample path: {data_list[0] if data_list else 'N/A'}")
            return data_list
        # ===============================================

        # 如果没有 split_file，则使用原有的 lr_file 逻辑
        if self.lr is None:
            return super().get_data_list()
        else:
            return [
                os.path.join(self.data_root, "train", name) for name in self.lr
            ]

    def get_data(self, idx):
        data_name = self.get_data_name(idx)
        data_path = os.path.join(self.data_root, data_name)
        data_dict = {}

        # 加载所有 .npy 文件
        for asset in self.VALID_ASSETS:
            asset_path = os.path.join(data_path, f"{asset}.npy")
            if os.path.exists(asset_path):
                data_dict[asset] = np.load(asset_path)
            else:
                # 如果是关键字段缺失，打印警告
                if asset in ["coord", "segment20"]:
                    print(f"❌ Error: {asset}.npy not found in {data_path}, this may cause issues.")

        # 🚨 关键修复：强制检查所有数组长度
        required_keys = ["coord", "segment20"]
        for key in required_keys:
            if key not in data_dict:
                print(f"❌ Error: {key} not found in {data_name}, using dummy data...")
                if key == "coord":
                    data_dict[key] = np.zeros((1, 3), dtype=np.float32)
                else:
                    data_dict[key] = np.array([self.ignore_index], dtype=np.int64)

        # 🚨 关键修复：确保 coord 和 segment20 长度一致
        if len(data_dict["coord"]) != len(data_dict["segment20"]):
            print(f"❌ Error: coord and segment20 length mismatch in {data_name}")
            print(f"   coord: {len(data_dict['coord'])}, segment20: {len(data_dict['segment20'])}")
            min_len = min(len(data_dict["coord"]), len(data_dict["segment20"]))
            data_dict["coord"] = data_dict["coord"][:min_len]
            data_dict["segment20"] = data_dict["segment20"][:min_len]
            if "color" in data_dict:
                data_dict["color"] = data_dict["color"][:min_len]
            if "normal" in data_dict:
                data_dict["normal"] = data_dict["normal"][:min_len]
            if "instance" in data_dict:
                data_dict["instance"] = data_dict["instance"][:min_len]

        # 🚨 关键修复：确保点数 > 0
        if len(data_dict["coord"]) == 0:
            print(f"❌ Error: Empty point cloud in {data_name}, using dummy point...")
            data_dict["coord"] = np.zeros((1, 3), dtype=np.float32)
            data_dict["segment20"] = np.array([self.ignore_index], dtype=np.int64)
            if "color" in data_dict:
                data_dict["color"] = np.zeros((1, 3), dtype=np.uint8)
            if "normal" in data_dict:
                data_dict["normal"] = np.zeros((1, 3), dtype=np.float32)
            if "instance" in data_dict:
                data_dict["instance"] = np.array([self.ignore_index], dtype=np.int64)
        
        # 🚨 关键修复：确保所有数组长度一致
        keys_to_check = ["coord", "color", "normal", "segment20", "segment200", "instance"]
        lengths = [len(data_dict[k]) for k in keys_to_check if k in data_dict]
        if len(lengths) > 0:
            min_len = min(lengths)
            if min_len == 0:
                print(f"⚠️  Warning: Inconsistent array lengths in {data_name}, truncating...")
                for k in keys_to_check:
                    if k in data_dict:
                        data_dict[k] = data_dict[k][:min_len]
        
        data_dict["name"] = data_name
        if "segment20" in data_dict:
            data_dict["segment"] = data_dict["segment20"]
        return data_dict


@DATASETS.register_module()
class ScanNet200Dataset(ScanNetDataset):
    VALID_ASSETS = [
        "coord",
        "color",
        "normal",
        "segment200",
        "instance",
    ]
    class2id = np.array(VALID_CLASS_IDS_200)