import numpy as np
import torch
from torch.utils import data
import warnings
import json, csv
import traceback
from collections import Counter
from sybil.augmentations import get_augmentations
from tqdm import tqdm 
from sybil.serie import Serie
from sybil.datasets.utils import order_slices, METAFILE_NOTFOUND_ERR, LOAD_FAIL_MSG
from sybil.loaders.image_loaders import OpenCVLoader, DicomLoader 



class CSVDataset(data.Dataset):
    """
    Dataset used for large validations
    """
    def __init__(self, args, split_group):
        '''
        params: args - config.
        params: split_group - ['train'|'dev'|'test'].

        constructs: standard pytorch Dataset obj, which can be fed in a DataLoader for batching
        '''
        super(CSVDataset, self).__init__()
        
        self.split_group = split_group
        self.args = args
        self._num_images = args.num_images # number of slices in each volume
        self._max_followup = args.max_followup

        try:
            self.dataset_dicts = self.parse_csv_dataset(args.dataset_file_path)
        except Exception as e:
            raise Exception(METAFILE_NOTFOUND_ERR.format(args.dataset_file_path, e))

        augmentations = get_augmentations(split_group, args)
        if args.img_file_type == 'dicom':
            self.input_loader = DicomLoader(args.cache_path, augmentations, args)  
        else:
            self.input_loader = OpenCVLoader(args.cache_path, augmentations, args)
            
        self.dataset = self.create_dataset(split_group)
        if len(self.dataset) == 0:
            return
        
        print(self.get_summary_statement(self.dataset, split_group))
        
        dist_key = 'y'
        label_dist = [d[dist_key] for d in self.dataset]
        label_counts = Counter(label_dist)
        weight_per_label = 1./ len(label_counts)
        label_weights = {
            label: weight_per_label/count for label, count in label_counts.items()
            }
        
        print("Class counts are: {}".format(label_counts))
        print("Label weights are {}".format(label_weights))
        self.weights = [ label_weights[d[dist_key]] for d in self.dataset]


    def parse_csv_dataset(self, file_path):
        """
        Convert a CSV file into a list of dictionaries for each patient like:
        [
            {
                'patient_id': str, 
                'split': str, 
                'exam_id': str,
                'series_id': str,
                'ever_has_future_cancer': bool
                'years_to_cancer': int,
                'years_to_last_negative_followup': int,
                'paths': [str],
                'slice_locations': [str]
            }
        ]

        Parameters
        ----------
        file_path : str
            path to csv file

        Returns
        -------
        list
            list patient cases in the above structure
        """
        dataset_dicts = {}
        _reader = csv.DictReader(open(file_path,'r'))
        for _row in _reader:
            row = {k.encode('ascii', 'ignore').decode(): v.encode('ascii', 'ignore').decode() for k,v in _row.items()}
            patient_id, exam_id, series_id = row['patient_id'], row['exam_id'], row['series_id']
            unique_id = '{}_{}_{}'.format(patient_id, exam_id, series_id)
            if unique_id in dataset_dicts:
                dataset_dicts[unique_id]['paths'].append(row['file_path'])
                dataset_dicts[unique_id]['slice_locations'].append(row['fileslice_position_path'])
            else:
                dataset_dicts[unique_id] = {
                    'unique_id': unique_id,
                    'patient_id': patient_id, 
                    'exam_id': exam_id,
                    'series_id': series_id,
                    'split': row['split'],
                    'ever_has_future_cancer': row['ever_has_future_cancer'],
                    'years_to_cancer': row['years_to_cancer'],
                    'years_to_last_negative_followup': row['years_to_last_negative_followup'],
                    'paths': [ row['file_path'] ],
                    'slice_locations': [ row['slice_position'] ]
                }
        
        dataset_dicts = list(dataset_dicts.values())
        return dataset_dicts

    def create_dataset(self, split_group):
        """
        Gets the dataset from the paths and labels in the json.
        Arguments:
            split_group(str): One of ['train'|'dev'|'test'].
        Returns:
            The dataset as a dictionary with img paths, label, 
            and additional information regarding exam or participant
        """
        
        if self.args.assign_splits:
            np.random.seed(self.args.cross_val_seed)
            self.assign_splits(self.metadata_json)

        dataset = []
        
        for mrn_row in tqdm(self.dataset_dicts, position = 0):
            
            label = mrn_row['ever_has_future_cancer']
            censor_time = mrn_row['years_to_cancer'] if label else mrn_row['years_to_last_negative_followup']    
            paths = order_slices(mrn_row['paths'], mrn_row['slice_locations'])
            try:
                series_object = Serie(
                    paths, 
                    label, 
                    censor_time,
                    self.args.img_file_type,
                    mrn_row['split']
                    )
            except Exception:
                continue 

            if self.skip_sample(series_object, mrn_row, split_group):
                continue
            
            dataset.append({
                'serie': series_object,
                'exam': mrn_row['unique_id']
                })

        return dataset

    def skip_sample(self, series_object, row, split_group ):
        if row['split'] != split_group:
            return True
        
        if not series_object.has_label():
            return True
        
        return False

    def get_summary_statement(self, dataset, split_group):
        summary = "Contructed Sybil Cancer Risk {} dataset with {} records, {} exams, {} patients, and the following class balance \n {}"
        class_balance = Counter([d['y'] for d in dataset])
        exams = set([d['exam'] for d in dataset])
        patients = set([d['patient_id'] for d in dataset])
        statement = summary.format(split_group, len(dataset), len(exams), len(patients), class_balance)
        statement += "\n" + "Censor Times: {}".format( Counter([d['time_at_event'] for d in dataset]))
        return statement

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

    def __getitem__(self, index):
        sample = self.dataset[index]
        serie = sample['serie']
        try:
            labels = serie.get_label()
            item = {
                'x': serie.get_volume(),
                'y': labels.y,
                'y_seq': labels.y_seq,
                'y_mask': labels.y_mask,
                'time_at_event': labels.censor_time,
                'exam': sample['exam']
            }
            return item
        except Exception:
            warnings.warn(LOAD_FAIL_MSG.format(sample['paths'], traceback.print_exc()))