from dataclasses import fields
from typing import Callable
from torch.utils.data import Dataset, ConcatDataset
import bisect

from ..misc.step_tracker import StepTracker
from .types import Stage
from .view_sampler import get_view_sampler
from .dataset_dl3dv import DatasetDL3DV, DatasetDL3DVCfgWrapper
from .dataset_scannetpp import DatasetScannetpp, DatasetScannetppCfgWrapper
from .dataset_co3d import DatasetCo3d, DatasetCo3dCfgWrapper

DATASETS: dict[str, Dataset] = {
    "co3d": DatasetCo3d,
    "scannetpp": DatasetScannetpp,
    "dl3dv": DatasetDL3DV,
}

DatasetCfgWrapper = DatasetDL3DVCfgWrapper | DatasetScannetppCfgWrapper | DatasetCo3dCfgWrapper

class TestDatasetWarpper(Dataset):
    def __init__(self, dataset: Dataset):
        self.dataset = dataset

    def __getitem__(self, idx):

        return self.dataset[(idx, self.dataset.view_sampler.num_context_views, self.dataset.cfg.input_image_shape[1] // 14)] # fake parameters here, to fit the input of dataset
    
    def __len__(self):
        return len(self.dataset)

        
    
class CustomConcatDataset(ConcatDataset):

    def __getitem__(self, idx_tuple):

        if isinstance(idx_tuple, list):
            idx_tuple = idx_tuple[0]

        idx = idx_tuple[0]
        if idx < 0:
            if -idx > len(self):
                raise ValueError("absolute value of index should not exceed dataset length")
            idx = len(self) + idx
        dataset_idx = bisect.bisect_right(self.cumulative_sizes, idx)
        if dataset_idx == 0:
            sample_idx = idx
        else:
            sample_idx = idx - self.cumulative_sizes[dataset_idx - 1]
        return self.datasets[dataset_idx][(sample_idx, idx_tuple[1], idx_tuple[2])]


def get_dataset(
    cfgs: list[DatasetCfgWrapper],
    stage: Stage,
    step_tracker: StepTracker | None,
    dataset_shim: Callable[[Dataset, str], Dataset]
) -> list[Dataset]:
    datasets = []
    if stage != "test":
        if stage == "val":
            cfgs = [cfgs[0]]
        for cfg in cfgs:
            (field,) = fields(type(cfg))
            cfg = getattr(cfg, field.name)
            
            view_sampler = get_view_sampler(
                cfg.view_sampler,
                stage,
                cfg.overfit_to_scene is not None,
                cfg.cameras_are_circular,
                step_tracker,
            )
            dataset = DATASETS[cfg.name](cfg, stage, view_sampler)
            dataset = dataset_shim(dataset, stage)
            datasets.append(dataset)

        return CustomConcatDataset(datasets), datasets
    elif stage == "test":
        assert len(cfgs) == 1
        cfg = cfgs[0]
        (field,) = fields(type(cfg))
        cfg = getattr(cfg, field.name)
        
        view_sampler = get_view_sampler(
            cfg.view_sampler,
            stage,
            cfg.overfit_to_scene is not None,
            cfg.cameras_are_circular,
            step_tracker,
        )
        dataset = DATASETS[cfg.name](cfg, stage, view_sampler)
        dataset = dataset_shim(dataset, stage)

        return TestDatasetWarpper(dataset)
    else:
        NotImplementedError(f"Stage {stage} is not supported")