import os
import glob
import logging
from collections import defaultdict

import numpy as np
import pandas as pd

import cv2
from torch.utils.data.dataset import Dataset


logging.basicConfig()
logger = logging.getLogger('dataset')
logger.setLevel(logging.DEBUG)


def index_from_path_maknee(path_root, force=False):
    fname_meta_dyn = os.path.join(path_root, 'meta_dynamic.csv')
    fname_meta_base = os.path.join(path_root, 'meta_base.csv')

    if not os.path.exists(fname_meta_dyn) or force:
        fnames_image = glob.glob(
            os.path.join(path_root, '**', 'images', '*.png'), recursive=True)
        logger.info('{} images found'.format(len(fnames_image)))
        fnames_mask = glob.glob(
            os.path.join(path_root, '**', 'masks', '*.png'), recursive=True)
        logger.info('{} masks found'.format(len(fnames_mask)))

        df_meta = pd.read_csv(fname_meta_base,
                              dtype={'patient': str,
                                     'release': str,
                                     'sequence': str,
                                     'side': str,
                                     'slice_idx': int,
                                     'pixel_spacing_0': float,
                                     'pixel_spacing_1': float,
                                     'slice_thickness': float,
                                     'KL': int},
                              index_col=False)

        if len(fnames_image) != len(df_meta):
            raise ValueError("Number of images doesn't match with the metadata")

        df_meta['path_image'] = [os.path.join(path_root, e)
                                 for e in df_meta['path_rel_image']]

        # Sort the records
        df_meta_sorted = (df_meta
                          .sort_values(['patient', 'sequence', 'slice_idx'])
                          .reset_index()
                          .drop('index', axis=1))

        df_meta_sorted.to_csv(fname_meta_dyn, index=False)
    else:
        df_meta_sorted = pd.read_csv(fname_meta_dyn,
                                     dtype={'patient': str,
                                            'release': str,
                                            'sequence': str,
                                            'side': str,
                                            'slice_idx': int,
                                            'pixel_spacing_0': float,
                                            'pixel_spacing_1': float,
                                            'slice_thickness': float,
                                            'KL': int},
                                     index_col=False)

    return df_meta_sorted


def read_image(path_file):
    image = cv2.imread(path_file, cv2.IMREAD_GRAYSCALE)
    return image.reshape((1, *image.shape))


class DatasetMAKNEESagittal2d(Dataset):
    def __init__(self, df_meta, mask_mode=None, name=None, transforms=None,
                 sample_mode='x_y', **kwargs):
        logger.warning('Redundant dataset init arguments:\n{}'
                       .format(repr(kwargs)))

        self.df_meta = df_meta
        self.mask_mode = mask_mode
        self.name = name
        self.transforms = transforms
        self.sample_mode = sample_mode

    def __len__(self):
        return len(self.df_meta)

    def _getitem_x_y(self, idx):
        image = read_image(self.df_meta['path_image'].iloc[idx])
        mask = np.zeros_like(image)

        # Apply transformations
        if self.transforms is not None:
            for t in self.transforms:
                if hasattr(t, 'randomize'):
                    t.randomize()
                image, mask = t(image, mask)

        tmp = dict(self.df_meta.iloc[idx])
        tmp['image'] = image
        tmp['mask'] = mask

        tmp['xs'] = tmp['image']
        tmp['ys'] = tmp['mask']
        return tmp

    def __getitem__(self, idx):
        if self.sample_mode == 'x_y':
            return self._getitem_x_y(idx)
        else:
            raise ValueError('Invalid `sample_mode`')

    def describe(self):
        summary = defaultdict(float)
        for i in range(len(self)):
            if self.sample_mode == 'x_y':
                _, mask = self.__getitem__(i)
            else:
                mask = self.__getitem__(i)['mask']
            summary['num_class_pixels'] += mask.numpy().sum(axis=(1, 2))
        summary['class_importance'] = \
            np.sum(summary['num_class_pixels']) / summary['num_class_pixels']
        summary['class_importance'] /= np.sum(summary['class_importance'])
        logger.info('Dataset statistics:')
        logger.info(sorted(summary.items()))