import numpy as np

import torch

import os

import sys

def set_seed(seed=0xD153A53):

    np.random.seed(seed)

    torch.manual_seed(seed)

    torch.cuda.manual_seed(seed)

    os.environ['PYTHONHASHSEED'] = str(seed)

def get_metric(cfg):

    return getattr(sys.modules['custom.metrics'], cfg.metric)

def get_criterion(cfg):

    return getattr(sys.modules['custom.losses'], cfg.criterion)

def get_model(cfg):

    return getattr(sys.modules['custom.models'], cfg.model)

def get_optimizer(cfg):

    return getattr(sys.modules['torch.optim'], cfg.optimizer)

def get_scheduler(cfg):

    if cfg.scheduler:

        return getattr(sys.modules['torch.optim.lr_scheduler'], cfg.scheduler)

    return FakeScheduler

def get_paths_1_dataset(data_folder, dataset_name):

    paths_folder = os.path.join(data_folder, dataset_name)

    last_number = 0

    paths, _paths = [], []

    # MedSeg has only one NIftI for many patients

    if 'MedSeg' in dataset_name:

        for i, name in enumerate(sorted(os.listdir(paths_folder))):

            path = os.path.join(paths_folder, name)

            if i % 5 == 0:

                paths.append(_paths)

                _paths = []

            _paths.append(path)

        paths.append(_paths)

        return paths

    # adding paths by patients

    for name in sorted(os.listdir(paths_folder)):

        path = os.path.join(paths_folder, name)

        number_of_patient = int(name.split('_')[0])

        if last_number != number_of_patient:

            paths.append(_paths)

            _paths = []

            last_number = number_of_patient

        _paths.append(path)

    paths.append(_paths)

    return paths

def get_paths(cfg):

    # if list for adding paths from every dataset

    if isinstance(cfg.dataset_name, list):

        paths = []

        for name in cfg.dataset_name:

            paths.extend(get_paths_1_dataset(cfg.data_folder, name))

        return paths

    # if solo dataset

    paths = get_paths_1_dataset(cfg.data_folder, cfg.dataset_name)

    return paths

class OneHotEncoder:

    def __init__(self, cfg):

        self.zeros = [0] * cfg.num_classes

    def encode_num(self, x):

        zeros = self.zeros.copy()

        zeros[int(x)] = 1

        return zeros

    def __call__(self, y):

        y = np.array(y)

        y = np.expand_dims(y, -1)

        y = np.apply_along_axis(self.encode_num, -1, y)

        y = np.swapaxes(y, -1, 1)

        y = np.ascontiguousarray(y)

        return torch.Tensor(y)

class FakeScheduler:

    def __init__(self, *args, **kwargs):

        self.lr = kwargs['lr']

    def step(self, *args, **kwargs):

        pass

    def get_last_lr(self, *args, **kwargs):

        return [self.lr]