# -*- coding: utf-8 -*-
"""(working) Creating Federated Data- Copy of Classification.ipynb

Automatically generated by Colaboratory.

Original file is located at
    https://colab.research.google.com/drive/1CUctqVr_-V3bf-mMAyphQICxkxUU4r0Z
"""

#import nest_asyncio

#nest_asyncio.apply()

# import attr
# import abc

import numpy as np
import tensorflow as tf
import tensorflow_federated as tff

import tensorflow_datasets as tfds
import collections
# from tensorflow.keras import layers
# from collections import OrderedDict

np.random.seed(0)

tff.federated_computation(lambda: 'Hello, World!')()

import tensorflow as tf
import os
import glob
import itertools
import csv
import numpy as np
import pandas as pd
from PIL import Image
import os.path


class Reader:
    def __init__(self, data_type, path):
        self.data_type = data_type.lower()
        self.path = path

    def _path_checker_(self):
        return True if os.path.splitext(self.path)[-1][-3:] == 'csv' or os.path.isdir(self.path) else False

    def _data_type_checker_(self):
        return True if self.data_type in ['image', 'text', 'csv'] else False

    def _read_image_(self) -> object:
        sub_folders = [f.path for f in os.scandir(self.path) if f.is_dir()]

        all_files = []
        if sub_folders:
            all_files = [f for s in sub_folders for f in glob.glob(s + "/*.png")]

        print(f'Total number of image files: {len(sub_folders)} folders and {len(all_files)} image files')

        images = [np.array(Image.open(filename)) for filename in all_files]
        labels = [label.split('/')[-2] for label in all_files]

        return images, labels

    def _read_text_folder_(self):
        sub_folders = [f.path for f in os.scandir(self.path) if f.is_dir()]

        all_files = []
        if sub_folders:
            all_files = [f for s in sub_folders for f in glob.glob(s + "/*.txt")]

        print(f'Total number of text files: {len(sub_folders)} folders and {len(all_files)} text files')

        texts = ["".join(open(filename, "r").readlines()) for filename in all_files]
        labels = [label.split('/')[-2] for label in all_files]

        return texts, labels

    def _read_text_csv_(self, columns):

        dataset = pd.read_csv(self.path, dtype=str)
        dataset.fillna('', inplace=True)

        print(f'Total number of rows:{dataset.shape[0]}')

        return dataset[list(dataset.iloc[:, :columns - 1])].values.tolist(), list(dataset.iloc[:, columns - 1])

    def _read_csv_(self):
        if self._data_type_checker_:
            if self.data_type == 'csv':
                f = open(self.path, 'r')
                reader1, reader2 = itertools.tee(csv.reader(f, delimiter=','))
                columns = len(next(reader1))
                print(f"Last csv column is considered as the label. Total number of columns:{columns}")
                return self._read_text_csv_(columns)

            elif self.data_type in ['image', 'text']:
                raise ValueError(
                    f'CSV file cannot contain an image or text file, check the input type and try again.')
            else:
                raise ValueError(
                    f'Given path: {self.path} is not pointing to text file, check the path again.')

        else:
            raise ValueError(
                f'Data type: "{self.data_type}" is incorrect. It should be "text", "image" or "csv" instead')

    def _read_folder_(self):
        if self._data_type_checker_:
            return self._read_text_folder_() if self.data_type == 'text' else self._read_image_()
        else:
            raise ValueError(
                f'Data type: "{self.data_type}" is incorrect. It should be "text","image", or "csv" instead')

    def read_data(self):
        print('Reading data...')
        if self._path_checker_:
            return self._read_folder_() if os.path.isdir(self.path) else self._read_csv_()
        else:
            raise ValueError(f'Given path: {self.path} is not a csv file or a folder, check the path and try again.')


import random
from random import randrange
import time
from sklearn.model_selection import train_test_split
import numpy as np
import warnings
from collections import Counter


def partition_list(list_in, n):
    random.shuffle(list_in)
    return [list_in[i::n] for i in range(n)]


class Distribute:
    def __init__(self, data, label):

        self.data = data
        self.label = label
        self.selected_feature = label
        self.type = 'iid'
        self.client_no = 10
        self.data_sample_fraction = 0.1
        self.min_user_number = 10
        self.max_user_number = 20
        self.train_data_fraction = 0.9
        self.random_sampling_seed = 4
        self.random_split_seed = 1
        self.max_seen_labels = 10
        self.min_seen_labels = 1
        self.split_type = 'sample'

    def __shuffle(self, data, label):
        random.Random(self.random_sampling_seed).shuffle(data)
        random.Random(self.random_sampling_seed).shuffle(label)

    def _iid_no_clint(self):
        size = random.randrange(2, len(self.data))
        self.__shuffle(self.data, self.label)

        glist = []
        glabel = []
        group_size = int(len(self.data) / size)
        for i in range(size):
            glist.append(self.data[group_size * i: group_size * (i + 1)])
            glabel.append(self.label[group_size * i: group_size * (i + 1)])

        return glist, glabel

    def _iid_clint(self, number_of_clients):

        self.__shuffle(self.data, self.label)

        glist = []
        glabel = []
        group_size = int(len(self.data) / number_of_clients)

        for i in range(number_of_clients):
            glist.append(self.data[group_size * i: group_size * (i + 1)])
            glabel.append(self.label[group_size * i: group_size * (i + 1)])

        return glist, glabel

    def _iid(self, **kwargs):
        number_of_clients = kwargs.get('number_of_clients')
        if number_of_clients:
            return self._iid_clint(number_of_clients)
        else:
            return self._iid_no_clint()

    def _niid(self, **kwargs):

        selected_feature = kwargs.get('selected_feature', self.selected_feature)
        min_seen_labels = kwargs.get('min_seen_labels', self.min_seen_labels)
        max_seen_labels = kwargs.get('max_seen_labels', self.max_seen_labels)
        min_user_number = kwargs.get('min_user_number', self.min_user_number)
        max_user_number = kwargs.get('max_user_number', self.max_user_number)
        number_of_clients = kwargs.get('number_of_clients')

        data_type = kwargs.get('data_type')

        if data_type == 'image':
            if number_of_clients:
                if number_of_clients > len(self.data):
                    raise ValueError('Total number of data:', len(self.data),
                                     'is less than total number of clients specified:', number_of_clients)
                else:
                    data, label = self.__select_feature_image_client(min_user_number, max_user_number,
                                                                     number_of_clients,
                                                                     min_seen_labels, max_seen_labels)
            else:
                data, label = self.__select_feature_image_no_client(min_user_number, max_user_number,
                                                                    min_seen_labels, max_seen_labels)
        elif data_type == 'text':
            if number_of_clients:
                if number_of_clients > len(self.data):
                    raise ValueError('Total number of data:', len(self.data),
                                     'is less than total number of clients specified:', number_of_clients)
                else:
                    data, label = self.__select_feature_text_client(min_user_number, max_user_number,
                                                                    number_of_clients,
                                                                    min_seen_labels, max_seen_labels)
            else:
                data, label = self.__select_feature_text_no_client(min_user_number, max_user_number,
                                                                   min_seen_labels, max_seen_labels)

        elif data_type == 'csv':
            if number_of_clients:
                if number_of_clients > len(self.data):
                    raise ValueError('Total number of data:', len(self.data),
                                     'is less than total number of clients specified:', number_of_clients)
                else:
                    data, label = self.__select_feature_csv_client(selected_feature,
                                                                   min_seen_labels, max_seen_labels,
                                                                   min_user_number, max_user_number,
                                                                   number_of_clients)
            else:
                data, label = self.__select_feature_csv_no_client(selected_feature,
                                                                  min_seen_labels, max_seen_labels,
                                                                  min_user_number, max_user_number)
        else:
            raise ValueError(
                f'Given data type: "{data_type}" is not correct, choose between options "text" or "image".')

        return data, label

    def distribute_data(self, **kwargs):
        if kwargs.get('dist_type', self.type) == 'iid':
            return self._iid(**kwargs)
        else:
            return self._niid(**kwargs)

    def __select_feature_image_no_client(self, min_user_number, max_user_number, min_seen_labels, max_seen_labels):

        remained_data = self.data
        remained_label = self.label

        unique_features = list(set(self.label))
        max_feature_len = len(unique_features)

        if max_seen_labels < min_seen_labels:
            raise ValueError(
                f'Max seen labels value: ({max_seen_labels}) cannot be less'
                f' than min seen labels: ({min_seen_labels})')

        if min_seen_labels > min(max_seen_labels, max_feature_len):
            raise ValueError(
                f'Total number of unique features: ({max_feature_len}) is more'
                f' than what is set for min seen labels: ({min_seen_labels})')
        else:
            unique_feature_size = random.randint(min_seen_labels, min(max_seen_labels, max_feature_len))

        grouped_data = []
        grouped_data_label = []

        while remained_data:
            random_selected_features = random.sample(unique_features, k=unique_feature_size)

            data_index = [i for i, e in enumerate(remained_label) if e in random_selected_features]

            selected_data = [remained_data[i] for i in data_index]

            if selected_data:
                rng = randrange(min_user_number, max_user_number)
                user_size = len(selected_data) if rng > len(selected_data) else rng

                random_selected_index = random.sample(data_index, k=user_size)

                random_selected_data = [remained_data[ind] for ind in random_selected_index]
                random_selected_labels = [remained_label[ind] for ind in random_selected_index]

                grouped_data.append(random_selected_data)
                grouped_data_label.append(random_selected_labels)

                remained_data = [x for i, x in enumerate(remained_data) if i not in random_selected_index]
                remained_label = [x for i, x in enumerate(remained_label) if i not in random_selected_index]

        return grouped_data, grouped_data_label

    def __select_feature_image_client(self, min_user_number, max_user_number,
                                      number_of_clients, min_seen_labels, max_seen_labels):
        remained_data = self.data
        remained_label = self.label

        unique_features = list(set(self.label))
        max_feature_len = len(unique_features)

        if max_seen_labels < min_seen_labels:
            raise ValueError(
                f'Max seen labels value: ({max_seen_labels}) cannot be less'
                f' than min seen labels: ({min_seen_labels})')

        if min_seen_labels > min(max_seen_labels, max_feature_len):
            raise ValueError(
                f'Given total number of unique features: ({max_feature_len}) is less'
                f' than what is set for min seen labels: ({min_seen_labels})')
        else:
            unique_feature_size = random.randint(min_seen_labels, min(max_seen_labels, max_feature_len))

        selected_data = []

        grouped_data = []
        grouped_data_label = []

        distributed_labels = []

        keep = set()
        for i in range(number_of_clients):
            while i == number_of_clients - 1 and len(keep) != max_feature_len:
                selected_labels = random.sample(unique_features, k=unique_feature_size)
                keep.update(selected_labels)
                distributed_labels.append(selected_labels)

            selected_labels = random.sample(unique_features, k=unique_feature_size)
            keep.update(selected_labels)
            distributed_labels.append(selected_labels)

        flatted_distributed_labels = [item for sublist in distributed_labels for item in sublist]

        data_dict = {}

        for label, frequency in Counter(flatted_distributed_labels).items():
            datapoints = [i for i, e in enumerate(self.label) if e == label]
            data_dict[label] = partition_list(datapoints, frequency)

        for client in distributed_labels:
            selected_index = []
            selected_data = []
            selected_label = []
            for client_label in client:
                index = random.sample(data_dict[client_label], k=1)

                data_dict[client_label].remove(index[0])
                selected_index.append(index[0])

                for i in range(len(index[0])):
                    selected_label.append(client_label)

            flatted_selected_index = [item for sublist in selected_index for item in sublist]
            selected_data = [self.data[i] for i in flatted_selected_index]

            grouped_data_label.append(selected_label)
            grouped_data.append(selected_data)

        return grouped_data, grouped_data_label

    def __select_feature_text_no_client(self, min_user_number, max_user_number, min_seen_labels, max_seen_labels):

        remained_data = self.data
        remained_label = self.label

        unique_features = list(set(self.label))
        max_feature_len = len(unique_features)

        if max_seen_labels < min_seen_labels:
            raise ValueError(
                f'Max seen labels value: ({max_seen_labels}) cannot be less'
                f' than min seen labels: ({min_seen_labels})')

        if min_seen_labels > min(max_seen_labels, max_feature_len):
            raise ValueError(
                f'Total number of unique features: ({max_feature_len}) is more'
                f' than what is set for min seen labels: ({min_seen_labels})')
        else:
            unique_feature_size = random.randint(min_seen_labels, min(max_seen_labels, max_feature_len))

        grouped_data = []
        grouped_data_label = []

        selected_data = []

        while remained_data:

            random_selected_features = random.sample(unique_features, k=unique_feature_size)

            data_index = [i for i, e in enumerate(remained_label) if e in random_selected_features]

            selected_data = [remained_data[i] for i in data_index]

            if selected_data:
                rng = randrange(min_user_number, max_user_number)
                user_size = len(selected_data) if rng > len(selected_data) else rng

                random_selected_index = random.sample(data_index, k=user_size)

                random_selected_data = [remained_data[ind] for ind in random_selected_index]
                random_selected_labels = [remained_label[ind] for ind in random_selected_index]

                grouped_data.append(random_selected_data)
                grouped_data_label.append(random_selected_labels)

                remained_data = [x for i, x in enumerate(remained_data) if i not in random_selected_index]
                remained_label = [x for i, x in enumerate(remained_label) if i not in random_selected_index]

        return grouped_data, grouped_data_label

    def __select_feature_text_client(self, min_user_number, max_user_number,
                                     number_of_clients,
                                     min_seen_labels, max_seen_labels):

        unique_features = list(set(self.label))
        max_feature_len = len(unique_features)

        if max_seen_labels < min_seen_labels:
            raise ValueError(
                f'Max seen labels value: ({max_seen_labels}) cannot be less'
                f' than min seen labels: ({min_seen_labels})')

        if min_seen_labels > min(max_seen_labels, max_feature_len):
            raise ValueError(
                f'Given total number of unique features: ({max_feature_len}) is less'
                f' than what is set for min seen labels: ({min_seen_labels})')
        else:
            unique_feature_size = random.randint(min_seen_labels, min(max_seen_labels, max_feature_len))

        selected_data = []

        grouped_data = []
        grouped_data_label = []

        distributed_labels = []

        keep = set()
        for i in range(number_of_clients):
            while i == number_of_clients - 1 and len(keep) != max_feature_len:
                selected_labels = random.sample(unique_features, k=unique_feature_size)
                keep.update(selected_labels)
                distributed_labels.append(selected_labels)

            selected_labels = random.sample(unique_features, k=unique_feature_size)
            keep.update(selected_labels)
            distributed_labels.append(selected_labels)

        flatted_distributed_labels = [item for sublist in distributed_labels for item in sublist]

        data_dict = {}

        for label, frequency in Counter(flatted_distributed_labels).items():
            datapoints = [i for i, e in enumerate(self.label) if e == label]
            data_dict[label] = partition_list(datapoints, frequency)

        for client in distributed_labels:
            selected_index = []
            selected_data = []
            selected_label = []
            for client_label in client:
                index = random.sample(data_dict[client_label], k=1)

                data_dict[client_label].remove(index[0])
                selected_index.append(index[0])

                for i in range(len(index[0])):
                    selected_label.append(client_label)

            flatted_selected_index = [item for sublist in selected_index for item in sublist]
            selected_data = [self.data[i] for i in flatted_selected_index]

            grouped_data_label.append(selected_label)
            grouped_data.append(selected_data)
        return grouped_data, grouped_data_label

    def __select_feature_csv_no_client(self, feature_column,
                                       min_seen_labels, max_seen_labels,
                                       min_user_number, max_user_number):

        if '' in [item[feature_column] for item in self.data]:
            raise ValueError('Selected feature should not have any null values. Select another feature column')

        unique_features = list(set([item[feature_column] for item in self.data]))
        max_feature_len = len(unique_features)

        if max_seen_labels < min_seen_labels:
            raise ValueError(
                f'Max seen labels value: ({max_seen_labels}) cannot be less'
                f' than min seen labels: ({min_seen_labels})')

        if min_seen_labels > min(max_seen_labels, max_feature_len):
            raise ValueError(
                f'Total number of unique features: ({max_feature_len}) for column ({feature_column}) is more'
                f' than what is set for min seen labels: ({min_seen_labels})')
        else:
            unique_feature_size = random.randint(min_seen_labels, min(max_seen_labels, max_feature_len))

        grouped_data = []
        grouped_data_label = []
        remained_data = self.data
        remained_label = self.label

        while remained_data:
            random_selected_features = random.sample(unique_features, k=unique_feature_size)

            data_index = [i for i, x in enumerate(remained_data) if x[feature_column] in random_selected_features]
            selected_data = [remained_data[i] for i in data_index]

            if selected_data:
                rng = randrange(min_user_number, max_user_number)
                user_size = len(selected_data) if rng > len(selected_data) else rng
                random_selected_index = random.sample(data_index, k=user_size)

                random_selected_data = [remained_data[ind] for ind in random_selected_index]
                random_selected_labels = [remained_label[ind] for ind in random_selected_index]

                grouped_data.append(random_selected_data)
                grouped_data_label.append(random_selected_labels)

                remained_data = [x for i, x in enumerate(remained_data) if i not in random_selected_index]
                remained_label = [x for i, x in enumerate(remained_label) if i not in random_selected_index]

        return grouped_data, grouped_data_label

    def __select_feature_csv_client(self, feature_column,
                                    min_seen_labels, max_seen_labels,
                                    min_user_number, max_user_number,
                                    number_of_clients):

        if '' in [item[feature_column] for item in self.data]:
            raise ValueError('Selected feature should not have any null values. Select another feature column')

        unique_features = list(set([item[feature_column] for item in self.data]))
        max_feature_len = len(unique_features)

        if max_seen_labels < min_seen_labels:
            raise ValueError(
                f'Max seen labels value: ({max_seen_labels}) cannot be less'
                f' than min seen labels: ({min_seen_labels})')

        if min_seen_labels > min(max_seen_labels, max_feature_len):
            raise ValueError(
                f'Total number of unique features: ({max_feature_len}) for column ({feature_column}) is more'
                f' than what is set for min seen labels: ({min_seen_labels})')
        else:
            unique_feature_size = random.randint(min_seen_labels, min(max_seen_labels, max_feature_len))

        selected_data = []

        grouped_data = []
        grouped_data_label = []

        distributed_labels = []

        for i in range(number_of_clients):
            distributed_labels.append(random.sample(unique_features, k=unique_feature_size))

        flatted_distributed_labels = [item for sublist in distributed_labels for item in sublist]

        data_dict = {}

        for label, frequency in Counter(flatted_distributed_labels).items():
            datapoints = [i for i, x in enumerate(self.data) if x[feature_column] == label]
            data_dict[label] = partition_list(datapoints, frequency)

        for client in distributed_labels:
            selected_index = []
            selected_data = []
            selected_label = []
            for client_label in client:
                index = random.sample(data_dict[client_label], k=1)

                data_dict[client_label].remove(index[0])
                selected_index.append(index[0])

                for i in range(len(index[0])):
                    selected_label.append(client_label)

            flatted_selected_index = [item for sublist in selected_index for item in sublist]

            selected_data = [self.data[i] for i in flatted_selected_index]

            grouped_data_label.append(selected_label)
            grouped_data.append(selected_data)

        return grouped_data, grouped_data_label

    def split_data(self, x, y, **kwargs):
        train_data_fraction = kwargs.get('train_data_fraction', self.train_data_fraction)
        if kwargs.get('type', self.type) == 'sample':
            return self._sample_split(x, y, train_data_fraction)
        else:
            return self._user_split(train_data_fraction)

    def _user_split(self, train_data_fraction):
        rng_seed = (self.random_split_seed if (self.random_split_seed is not None and self.random_split_seed >= 0)
                    else int(time.time()))
        rng = random.Random(rng_seed)
        # randomly sample from user_files to pick training set users
        num_users = self.client_no
        num_train_users = int(train_data_fraction * num_users)
        indices = [i for i in range(num_users)]
        train_indices = rng.sample(indices, num_train_users)
        train_blist = [False for i in range(num_users)]
        for i in train_indices:
            train_blist[i] = True
        train_user_files = []
        test_user_files = []
        train_labels = []
        test_labels = []

        for i in range(num_users):
            if train_blist[i]:
                train_user_files.append(self.data[i])
                train_labels.append(self.label[i])
            else:
                test_user_files.append(self.data[i])
                test_labels.append(self.label[i])

        return train_user_files, test_user_files, train_labels, test_labels

    def _sample_split(self, x, y, train_data_fraction):
        x_train, x_test, y_train, y_test = train_test_split(x, y, train_size=train_data_fraction)
        return x_train, x_test, y_train, y_test


def convert_to_client_data(data, labels, data_type, **kwargs):
    distributor_obj = Distribute(data, labels)

    distributed_data, distributed_label = distributor_obj.distribute_data(data_type=data_type, **kwargs)

    client_train_dataset = collections.OrderedDict()

    for i in range(len(distributed_data)):
        client_name = "client_" + str(i)
        data = collections.OrderedDict((('label', distributed_label[i]), ('data', distributed_data[i])))
        client_train_dataset[client_name] = data

    print(f'Converting data to {len(distributed_data)} client data...')
    train_dataset=tff.simulation.FromTensorSlicesClientData(client_train_dataset)

    # # train_dataset = tff.simulation.datasets.TestClientData(client_train_dataset)
    # train_dataset = tff.simulation.TensorSlicesClientData(client_train_dataset)

    print(f'Data successfully converted to {len(distributed_data)} client data.')

    return train_dataset


"""# **Data**"""

# data_type = 'csv'
# input_path = '/content/drive/MyDrive/Vesal-Yasaman/v2/data/csv/fake_job_postings_newLabel.csv'  # accepts only a csv file

# data_type = 'text'
# input_path = '/content/drive/MyDrive/Vesal-Yasaman/v2/data/text/topics_sample'  # accepts either folder or csv file

# data_type = 'image'
# input_path = '/content/drive/MyDrive/Divya-Yasaman/v2/data/image/fish_sample'  # accepts folder only
#
# obj = Reader(data_type, input_path)
#
# # Commented out IPython magic to ensure Python compatibility.
# # %%time
# data, labels = obj.read_data()
#
# print(len(labels))
#
# # Commented out IPython magic to ensure Python compatibility.
# # %%time
# NUM_CLIENTS = 10
# SELECTED_FEATURE = 17
# dataset = convert_to_client_data(data, labels,  data_type,
#                                                 dist_type = 'niid',
#                                                 number_of_clients = NUM_CLIENTS,
#                                                 selected_feature=SELECTED_FEATURE,
#                                                 min_seen_labels=2,
#                                                 max_seen_labels=3
#                                                 )

"""# **Counting unique labels**"""


def count_unique_labels(example_dataset):
    labels = []
    it = iter(example_dataset)
    for x in example_dataset:
        example_element = next(it)
        labels.append(example_element['label'].numpy())  # for image & text
        # labels.append(example_element['data'][SELECTED_FEATURE].numpy()) #for CSV

    labels_set = set(labels)
    print(len(labels))
    return labels_set, labels


"""# **Counting**"""

def print_counting(NUM_CLIENTS,dataset):
    all_data_point = 0
    for i in range(NUM_CLIENTS):
        example_dataset = dataset.create_tf_dataset_for_client(dataset.client_ids[i])
        print(example_dataset)

        labels_set, labels = count_unique_labels(example_dataset)
        all_data_point = all_data_point + len(labels)
        print('client ' + str(i) + ' has: ')
        print(str(len(labels)) + ' data points')
        print('with ' + str(len(labels_set)) + ' unique labels')
        print('The labels are:')
        print(labels_set)
        print('In total, this clients has ' + str(len(labels)) + ' data points')
        print('so far, all clients have ' + str(all_data_point) + ' data points')
        print('\r\n')
        print('In total, all clients have ' + str(all_data_point) + ' data points')
        print(len(labels))