import os

import logging

from sklearn.model_selection import GroupShuffleSplit, GroupKFold

from rocaseg.datasets import (index_from_path_oai_imo,

                              index_from_path_okoa,

                              index_from_path_maknee)

logging.basicConfig()

logger = logging.getLogger('datasets')

logger.setLevel(logging.DEBUG)

def sources_from_path(path_data_root,

                      selection=None,

                      with_folds=False,

                      fold_num=5,

                      seed_trainval_test=0):

    """

    Args:

        path_data_root: str

        selection: iterable or str or None

        with_folds: bool

            Whether to split trainval subset into the folds.

        fold_num: int

            Number of folds.

        seed_trainval_test: int

            Random state for the trainval/test splitting.

    Returns:

    """

    if selection is None:

        selection = ('oai_imo', 'okoa', 'maknee')

    elif isinstance(selection, str):

        selection = (selection, )

    sources = dict()

    for name in selection:

        if name == 'oai_imo':

            logger.info('--- OAI iMorphics dataset ---')

            tmp = dict()

            tmp['path_root'] = os.path.join(path_data_root,

                                            '91_OAI_iMorphics_full_meta')

            if not os.path.exists(tmp['path_root']):

                logger.warning(f"Dataset {name} is not found in {tmp['path_root']}")

                continue

            tmp['full_df'] = index_from_path_oai_imo(tmp['path_root'])

            logger.info(f"Total number of samples: "

                        f"{len(tmp['full_df'])}")

            # Select the specific subset

            # Remove two series from the dataset as they are completely missing

            # information on patellar cartilage:

            #        /0.C.2/9674570/20040913/10699609/

            #        /1.C.2/9674570/20050829/10488714/

            tmp['sel_df'] = tmp['full_df'][tmp['full_df']['patient'] != '9674570']

            logger.info(f"Selected number of samples: "

                        f"{len(tmp['sel_df'])}")

            if with_folds:

                # Get trainval/test split

                tmp_groups = tmp['sel_df'].loc[:, 'patient'].values

                tmp_grades = tmp['sel_df'].loc[:, 'KL'].values

                tmp_gss = GroupShuffleSplit(n_splits=1, test_size=0.2,

                                            random_state=seed_trainval_test)

                tmp_idcs_trainval, tmp_idcs_test = next(tmp_gss.split(X=tmp['sel_df'],

                                                                      y=tmp_grades,

                                                                      groups=tmp_groups))

                tmp['trainval_df'] = tmp['sel_df'].iloc[tmp_idcs_trainval]

                tmp['test_df'] = tmp['sel_df'].iloc[tmp_idcs_test]

                logger.info(f"Made trainval-test split, number of samples: "

                            f"{len(tmp['trainval_df'])}, "

                            f"{len(tmp['test_df'])}")

                # Make k folds

                tmp_gkf = GroupKFold(n_splits=fold_num)

                tmp_groups = tmp['trainval_df'].loc[:, 'patient'].values

                tmp_grades = tmp['trainval_df'].loc[:, 'KL'].values

                tmp['trainval_folds'] = tmp_gkf.split(X=tmp['trainval_df'],

                                                      y=tmp_grades, groups=tmp_groups)

            sources['oai_imo'] = tmp

        elif name == 'okoa':

            logger.info('--- OKOA dataset ---')

            tmp = dict()

            tmp['path_root'] = os.path.join(path_data_root,

                                            '32_OKOA_full_meta_rescaled')

            if not os.path.exists(tmp['path_root']):

                logger.warning(f"Dataset {name} is not found in {tmp['path_root']}")

                continue

            tmp['full_df'] = index_from_path_okoa(tmp['path_root'])

            logger.info(f"Total number of samples: "

                        f"{len(tmp['full_df'])}")

            # Select the specific subset

            tmp['sel_df'] = tmp['full_df']

            logger.info(f"Selected number of samples: "

                        f"{len(tmp['sel_df'])}")

            if with_folds:

                # Get trainval/test split

                tmp['trainval_df'] = tmp['sel_df'][tmp['sel_df']['subset'] == 'training']

                tmp['test_df'] = tmp['sel_df'][tmp['sel_df']['subset'] == 'evaluation']

                logger.info(f"Made trainval-test split, number of samples: "

                            f"{len(tmp['trainval_df'])}, "

                            f"{len(tmp['test_df'])}")

                # Make k folds

                tmp_gkf = GroupKFold(n_splits=fold_num)

                tmp_groups = tmp['trainval_df'].loc[:, 'patient'].values

                tmp['trainval_folds'] = tmp_gkf.split(X=tmp['trainval_df'],

                                                      groups=tmp_groups)

            sources['okoa'] = tmp

        elif name == 'maknee':

            logger.info('--- MAKNEE dataset ---')

            tmp = dict()

            tmp['path_root'] = os.path.join(path_data_root,

                                            '42_MAKNEE_full_meta_rescaled')

            if not os.path.exists(tmp['path_root']):

                logger.warning(f"Dataset {name} is not found in {tmp['path_root']}")

                continue

            tmp['full_df'] = index_from_path_maknee(tmp['path_root'])

            logger.info(f"Total number of samples: "

                        f"{len(tmp['full_df'])}")

            # Select the specific subset

            tmp['sel_df'] = tmp['full_df']

            logger.info(f"Selected number of samples: "

                        f"{len(tmp['sel_df'])}")

            # Get trainval/test split

            tmp_groups = tmp['sel_df'].loc[:, 'patient'].values

            tmp_gss = GroupShuffleSplit(n_splits=1, test_size=0.2,

                                        random_state=seed_trainval_test)

            tmp_idcs_trainval, tmp_idcs_test = next(tmp_gss.split(X=tmp['sel_df'],

                                                                  groups=tmp_groups))

            tmp['trainval_df'] = tmp['sel_df'].iloc[tmp_idcs_trainval]

            tmp['test_df'] = tmp['sel_df'].iloc[tmp_idcs_test]

            logger.info(f"Made trainval-test split, number of samples: "

                        f"{len(tmp['trainval_df'])}, "

                        f"{len(tmp['test_df'])}")

            if with_folds:

                # Make k folds

                tmp_gkf = GroupKFold(n_splits=fold_num)

                tmp_groups = tmp['trainval_df'].loc[:, 'patient'].values

                tmp['trainval_folds'] = tmp_gkf.split(X=tmp['trainval_df'],

                                                      groups=tmp_groups)

            sources['maknee'] = tmp

        else:

            raise ValueError(f'Unknown dataset `{name}`')

    return sources