import os
import random
import numpy as np
from torch.utils import data
from torchvision import datasets, transforms
from sklearn.model_selection import train_test_split
from collections import defaultdict

from config import *


# データの前処理
def data_preprocessing(args):
    # データの前処理を設定
    data_transforms = {
        'known': transforms.Compose([
            transforms.Resize([224, 224]),
            transforms.Grayscale(num_output_channels=3),
            transforms.ToTensor(),
            transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
        ]),
        'unknown': transforms.Compose([
            transforms.Resize([224, 224]),
            transforms.Grayscale(num_output_channels=3),
            transforms.ToTensor(),
            transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
        ]),
    }

    image_datasets = {x: datasets.ImageFolder(os.path.join(DATA_DIR, x), data_transforms[x])
                      for x in ['known', 'unknown']}

    # calculate target weights
    targets_counts = np.unique(np.array(image_datasets['known'].targets), return_counts=True)[1]

    print(targets_counts)

    # weights for each sample
    w = targets_counts.sum() / (targets_counts.size * targets_counts)

    targets_weights = np.apply_along_axis(lambda x: w[x], 0, np.array(image_datasets['known'].targets))

    # split dataset and weights
    indices = dict()
    indices['train'], indices['tmp'] = train_test_split(np.arange(len(image_datasets['known'].targets)), test_size=0.4,
                                                        stratify=image_datasets['known'].targets)
    indices['val'], indices['test'] = train_test_split(np.array(indices['tmp']), test_size=0.5,
                                                       stratify=np.array(image_datasets['known'].targets)[indices['tmp']])

    known_datasets = {x: data.Subset(image_datasets['known'], indices=indices[x]) for x in ['train', 'val', 'test']}
    known_weights = {x: targets_weights[indices[x]] for x in ['train', 'val', 'test']}

    # Split Dataset
    num_split = {
        'train': int(targets_counts.min() * 0.6) * len(targets_counts),
        'val': int(targets_counts.min() * 0.2) * len(targets_counts),
        'test': int(targets_counts.min() * 0.2) * len(targets_counts)
        # 'train': 100 * len(targets_counts),
        # 'val': 25 * len(targets_counts),
        # 'test': 25 * len(targets_counts)
    }

    samplers = {x: data.WeightedRandomSampler(known_weights[x], num_samples=num_split[x], replacement=True)
                for x in ['train', 'val', 'test']}

    data_loaders = {x: data.DataLoader(known_datasets[x], batch_size=args.batch_size, sampler=samplers[x])
                    for x in ['train', 'val', 'test']}

    unknown_loader = torch.utils.data.DataLoader(image_datasets['unknown'], batch_size=args.batch_size, shuffle=True)

    data_loaders['unknown'] = unknown_loader

    datasets_size = {"train": num_split['train'], "val": num_split['val'], "test": num_split['test'],
                     "unknown": len(image_datasets['unknown'])}

    for x in ['train', 'val', 'test']:
        count = defaultdict(int)
        for da, label in data_loaders[x]:
            for la in label:
                count[int(la)] += 1
        print('{} num_samples: {}'.format(x, num_split[x]))
        print([(k, count[k]) for k in sorted(count.keys())])

    # 訓練(train)，検証(val)，評価(test)のデータ分割サイズを定義
    train_ratio = 0.6
    val_ratio = 0.2
    train_size = int(train_ratio * len(image_datasets['known']))
    val_size = int(val_ratio * len(image_datasets['known']))
    test_size = len(image_datasets['known']) - train_size - val_size
    datasets_size = {"train": train_size, "val": val_size, "test": test_size,
                     "unknown": len(image_datasets['unknown'])}

    # 既知のマルウェアファミリのデータセットを3種類に分割
    train_datasets, val_datasets, test_datasets = torch.utils.data.random_split(image_datasets['known'],
                                                                                    [train_size, val_size, test_size])

    # データローダを作成
    train_loader = torch.utils.data.DataLoader(known_datasets['train'], batch_size=args.batch_size, shuffle=True)
    val_loader = torch.utils.data.DataLoader(known_datasets['val'], batch_size=args.batch_size, shuffle=True)
    test_loader = torch.utils.data.DataLoader(known_datasets['test'], batch_size=args.batch_size, shuffle=True)
    unknown_loader = torch.utils.data.DataLoader(image_datasets['unknown'], batch_size=args.batch_size, shuffle=True)
    data_loaders = {"train": train_loader, "val": val_loader, "test": test_loader,
                    "unknown": unknown_loader}

    # 既知クラスと未知クラスのファミリ名を保存
    class_names = {"known": image_datasets['known'].classes, "unknown": image_datasets['unknown'].classes}

    return data_loaders, datasets_size, class_names
