import os

from datasets.kitti_360.kitti_360_dataset import Kitti360Dataset
from datasets.kitti_odom.kitti_odometry_dataset import KittiOdometryDataset
from datasets.kitti_raw.kitti_raw_dataset import KittiRawDataset
from datasets.nyu_depth_v2.nyu_depth_v2_dataset import NYUDepthV2Dataset
from datasets.realestate10k.realestate10k_dataset import RealEstate10kDataset
from datasets.waymo.waymo_dataset import WaymoDataset


def make_datasets(config):
    type = config.get("type", "KITTI_Raw")
    if type == "KITTI_Odometry":
        train_dataset = KittiOdometryDataset(
            base_path=config["data_path"],
            frame_count=config.get("data_fc", 1),
            target_image_size=config.get("image_size", (128, 256)),
            return_stereo=config.get("data_stereo", False),
            sequences=config.get("train_sequences", ("00",)),
            custom_pose_path=config.get("custom_pose_path", None),
            keyframe_offset=0 #-(config.get("data_fc", 1) // 2)
        )
        test_dataset = KittiOdometryDataset(
            base_path=config["data_path"],
            frame_count=config.get("data_fc", 1),
            target_image_size=config.get("image_size", (128, 256)),
            return_stereo=config.get("data_stereo", False),
            sequences=config.get("val_sequences", ("00",)),
            custom_pose_path=config.get("custom_pose_path", None),
            keyframe_offset=0 #-(config.get("data_fc", 1) // 2)
        )
        return train_dataset, test_dataset

    elif type == "KITTI_Raw":
        train_dataset = KittiRawDataset(
            data_path=config["data_path"],
            pose_path=config["pose_path"],
            split_path=os.path.join(config["split_path"], "train_files.txt"),
            target_image_size=config.get("image_size", (192, 640)),
            frame_count=config.get("data_fc", 1),
            return_stereo=config.get("data_stereo", False),
            keyframe_offset=config.get("keyframe_offset", 0),
            dilation=config.get("dilation", 1),
            color_aug=config.get("color_aug", False)
        )
        test_dataset = KittiRawDataset(
            data_path=config["data_path"],
            pose_path=config["pose_path"],
            split_path=os.path.join(config["split_path"], "val_files.txt"),
            target_image_size=config.get("image_size", (192, 640)),
            frame_count=config.get("data_fc", 1),
            return_stereo=config.get("data_stereo", False),
            keyframe_offset=config.get("keyframe_offset", 0),
            dilation=config.get("dilation", 1),
        )
        return train_dataset, test_dataset

    elif type == "KITTI_360":
        if config.get("split_path", None) is None:
            train_split_path = None
            test_split_path = None
        else:
            train_split_path = os.path.join(config["split_path"], "train_files.txt")
            test_split_path = os.path.join(config["split_path"], "val_files.txt")

        train_dataset = Kitti360Dataset(
            data_path=config["data_path"],
            pose_path=config["pose_path"],
            split_path=train_split_path,
            target_image_size=tuple(config.get("image_size", (192, 640))),
            frame_count=config.get("data_fc", 3),
            return_stereo=config.get("data_stereo", True),
            return_fisheye=config.get("data_fisheye", True),
            return_3d_bboxes=config.get("data_3d_bboxes", False),
            return_segmentation=config.get("data_segmentation", False),
            keyframe_offset=config.get("keyframe_offset", 0),
            dilation=config.get("dilation", 1),
            fisheye_rotation=config.get("fisheye_rotation", 0),
            fisheye_offset=config.get("fisheye_offset", 1),
            color_aug=config.get("color_aug", False),
            is_preprocessed=config.get("is_preprocessed", False)
        )
        test_dataset = Kitti360Dataset(
            data_path=config["data_path"],
            pose_path=config["pose_path"],
            split_path=test_split_path,
            target_image_size=tuple(config.get("image_size", (192, 640))),
            frame_count=config.get("data_fc", 3),
            return_stereo=config.get("data_stereo", True),
            return_fisheye=config.get("data_fisheye", True),
            return_3d_bboxes=config.get("data_3d_bboxes", False),
            return_segmentation=config.get("data_segmentation", False),
            keyframe_offset=config.get("keyframe_offset", 0),
            fisheye_rotation=config.get("fisheye_rotation", 0),
            fisheye_offset=config.get("fisheye_offset", 1),
            dilation=config.get("dilation", 1),
            is_preprocessed=config.get("is_preprocessed", False)
        )
        return train_dataset, test_dataset

    elif type == "RealEstate10k":
        train_dataset = RealEstate10kDataset(
            data_path=config["data_path"],
            split_path=None,
            target_image_size=config.get("image_size", (256, 384)),
            frame_count=config.get("data_fc", 2),
            keyframe_offset=0, #-(config.get("data_fc", 1) // 2),
            dilation=config.get("dilation", 10),
            color_aug=config.get("color_aug", False)
        )
        test_dataset = RealEstate10kDataset(
            data_path=config["data_path"],
            split_path=os.path.join(config["split_path"], "val_files.txt"),
            target_image_size=config.get("image_size", (256, 384)),
            frame_count=config.get("data_fc", 2),
            keyframe_offset=0, #-(config.get("data_fc", 1) // 2),
            dilation=config.get("dilation", 10),
            color_aug=False
        )
        return train_dataset, test_dataset

    elif type == "Waymo":
        if config.get("split_path", None) is None:
            train_split_path = None
            test_split_path = None
        else:
            train_split_path = os.path.join(config["split_path"], "train_files.txt")
            test_split_path = os.path.join(config["split_path"], "val_files.txt")

        train_dataset = WaymoDataset(
            data_path=config["data_path"],
            mode="training",
            split_path=train_split_path,
            target_image_size=tuple(config.get("image_size", (320, 480))),
            frame_count=config.get("data_fc", 2),
            keyframe_offset=config.get("keyframe_offset", 0),
            return_45=config.get("return_45", True),
            return_90=config.get("return_90", True),
            offset_45=config.get("offset_45", 5),
            offset_90=config.get("offset_90", 10),
            dilation=config.get("dilation", 1),
            color_aug=config.get("color_aug", True),
            correct_exposure=config.get("correct_exposure", True),
        )
        test_dataset = WaymoDataset(
            data_path=config["data_path"],
            mode="validation",
            split_path=test_split_path,
            target_image_size=tuple(config.get("image_size", (320, 480))),
            frame_count=config.get("data_fc", 2),
            keyframe_offset=config.get("keyframe_offset", 0),
            return_45=config.get("return_45", True),
            return_90=config.get("return_90", True),
            offset_45=config.get("offset_45", 5),
            offset_90=config.get("offset_90", 10),
            dilation=config.get("dilation", 1),
            color_aug=False,
            return_depth=True,
            correct_exposure=config.get("correct_exposure", True),
        )
        return train_dataset, test_dataset

    elif type == "KITTI_Raw_DFT":
        train_dataset = KittiRawDataset(
            data_path=config["data_path"],
            pose_path=config["pose_path"],
            split_path=os.path.join(config["split_path"], "train_files.txt"),
            target_image_size=config.get("image_size", (192, 640)),
            frame_count=config.get("data_fc", 1),
            return_stereo=config.get("data_stereo", False),
            keyframe_offset=config.get("keyframe_offset", 0),
            dilation=config.get("dilation", 1),
            color_aug=config.get("color_aug", False)
        )
        test_dataset = KittiRawDataset(
            data_path=config["data_path"],
            pose_path=config["pose_path"],
            split_path=os.path.join(config["split_path"], "val_files.txt"),
            target_image_size=config.get("image_size", (192, 640)),
            frame_count=config.get("data_fc", 1),
            return_stereo=config.get("data_stereo", False),
            keyframe_offset=config.get("keyframe_offset", 0),
            dilation=config.get("dilation", 1),
        )
        return train_dataset, test_dataset

    elif type == "KITTI_360_DFT":
        if config.get("split_path", None) is None:
            train_split_path = None
            test_split_path = None
        else:
            train_split_path = os.path.join(config["split_path"], "train_files.txt")
            test_split_path = os.path.join(config["split_path"], "val_files.txt")

        train_dataset = Kitti360Dataset(
            data_path=config["data_path"],
            pose_path=config["pose_path"],
            split_path=train_split_path,
            target_image_size=tuple(config.get("image_size", (192, 640))),
            frame_count=config.get("data_fc", 3),       
            return_stereo=config.get("data_stereo", True),
            return_fisheye=config.get("data_fisheye", True),
            return_3d_bboxes=config.get("data_3d_bboxes", False),
            return_segmentation=config.get("data_segmentation", False),
            keyframe_offset=config.get("keyframe_offset", 0),
            dilation=config.get("dilation", 1),
            fisheye_rotation=config.get("fisheye_rotation", 0),
            fisheye_offset=config.get("fisheye_offset", 1),
            stereo_offset=config.get("stereo_offset", 1),
            color_aug=config.get("color_aug", False),
            is_preprocessed=config.get("is_preprocessed", False)
        )
        test_dataset = Kitti360Dataset(
            data_path=config["data_path"],
            pose_path=config["pose_path"],
            split_path=test_split_path,
            target_image_size=tuple(config.get("image_size", (192, 640))),
            frame_count=config.get("data_fc", 3),
            return_stereo=config.get("data_stereo", True),
            return_fisheye=config.get("data_fisheye", True),
            return_3d_bboxes=config.get("data_3d_bboxes", False),
            return_segmentation=config.get("data_segmentation", False),
            keyframe_offset=config.get("keyframe_offset", 0),
            fisheye_rotation=config.get("fisheye_rotation", 0),
            fisheye_offset=config.get("fisheye_offset", [10])[0],        ## this modifies the offsets for all datasets including the training dataset
            stereo_offset= config.get("stereo_offset", [1])[0],         ## This is to set consistent evaluation with test and viz
            dilation=config.get("dilation", 1),
            is_preprocessed=config.get("is_preprocessed", False)
        )
        return train_dataset, test_dataset

    else:
        raise NotImplementedError(f"Unsupported dataset type: {type}")


def make_test_dataset(config):
    type = config.get("type", "KITTI_Raw")
    if type == "KITTI_Raw":
        test_dataset = KittiRawDataset(
            data_path=config["data_path"],
            pose_path=config["pose_path"],
            split_path=os.path.join(config["split_path"], "test_files.txt"),
            target_image_size=config.get("image_size", (192, 640)),
            return_depth=True,
            frame_count=config.get("data_fc", 1),
            return_stereo=config.get("data_stereo", False),
            keyframe_offset=0
        )
        return test_dataset
    elif type == "KITTI_360":
        test_dataset = Kitti360Dataset(
            data_path=config["data_path"],
            pose_path=config["pose_path"],
            split_path=os.path.join(config.get("split_path", None), "test_files.txt"),
            target_image_size=tuple(config.get("image_size", (192, 640))),
            frame_count=config.get("data_fc", 1),
            return_stereo=config.get("data_stereo", False),
            return_fisheye=config.get("data_fisheye", False),
            return_3d_bboxes=config.get("data_3d_bboxes", False),
            return_segmentation=config.get("data_segmentation", False),
            keyframe_offset=0,
            fisheye_rotation=config.get("fisheye_rotation", 0),
            fisheye_offset=config.get("fisheye_offset", 1),
            dilation=config.get("dilation", 1),
            is_preprocessed=config.get("is_preprocessed", False)
        )
        return test_dataset
    elif type == "KITTI_360_DFT":
        test_dataset = Kitti360Dataset(
            data_path=config["data_path"],
            pose_path=config["pose_path"],
            split_path=os.path.join(config.get("split_path", None), "test_files.txt"),
            target_image_size=tuple(config.get("image_size", (192, 640))),
            frame_count=config.get("data_fc", 1),
            return_stereo=config.get("data_stereo", False),
            return_fisheye=config.get("data_fisheye", False),
            return_3d_bboxes=config.get("data_3d_bboxes", False),
            return_segmentation=config.get("data_segmentation", False),
            keyframe_offset=0,
            fisheye_rotation=config.get("fisheye_rotation", 0),
            fisheye_offset=config.get("fisheye_offset", [10])[0],        ## this modifies the offsets for all datasets including the training dataset
            stereo_offset= config.get("stereo_offset", [1])[0],         ## This is to set consistent evaluation with test and viz
            dilation=config.get("dilation", 1),
            is_preprocessed=config.get("is_preprocessed", False),

            return_depth=True
        )
        return test_dataset
    elif type == "RealEstate10k":
        test_dataset = RealEstate10kDataset(
            data_path=config["data_path"],
            split_path=os.path.join(config["split_path"], "test_files.txt"),
            target_image_size=config.get("image_size", (256, 384)),
            frame_count=config.get("data_fc", 2),
            keyframe_offset=0,
            dilation=config.get("dilation", 10),
            color_aug=False
        )
        return test_dataset
    elif type == "NYU_Depth_V2":
        test_dataset = NYUDepthV2Dataset(
            data_path=config["data_path"],
            target_image_size=config.get("image_size", (256, 384)),
        )
        return test_dataset
    else:
        raise NotImplementedError(f"Unsupported dataset type: {type}")