# -*- coding: utf-8 -*-
"""patchdata

Automatically generated by Colaboratory.

Original file is located at
    https://colab.research.google.com/drive/1BUVGX_Ymu9kxdmP0gB3rwVAlBJMJQ_QI
"""

import torchvision
import torchvision.transforms as transforms
from torch.utils.data import DataLoader
import matplotlib.pyplot as plt
import numpy as np
import torch

# transforms.Compose에 들어갈 전처리 루틴을 제작하는 방식
class PatchGenerator:

    def __init__(self, patch_size):
        self.patch_size = patch_size

    # 전처리기에는 Data가 1장만 들어간다고 가정된다.
    # 그래서 img는 (channel, height, width) 차원을 가지고 있다
    def __call__(self, img):
        num_channels = img.size(0) # channel수

        # patch로 만드는 과정 : dimension 1방향으로 p조각 분해 > dimension 2방향으로 p조각 분해 하면 (num_channel, h/p, p , w/p, p)
        # 여기서 reshape통해서 (num_channels, N , p, p) 형태로 바꾼다
        # N : 생성된 patch의 개수
        # p : patch의 가로 세로 size
        patches = img.unfold(1, self.patch_size, self.patch_size).unfold(2, self.patch_size, self.patch_size).reshape(num_channels, -1, self.patch_size, self.patch_size)
        
        # (N, num_channels, p,p) 형태가 된다.
        patches = patches.permute(1,0,2,3)
        num_patch = patches.size(0)

        # 모든 patch들을 일렬로 펴서 반환 (num_patch, num_channels*p*p)
        # Encoder 입력 전 Linear Projection을 하기 직전의 상태가 된다
        return patches.reshape(num_patch,-1)


class Flattened2Dpatches:
    # dataname인자를 넣어서 다양한 Dataset에 대한 작업을 수행할 수 있도록 하였다
    def __init__(self, patch_size=16, dataname='imagenet', img_size=256, batch_size=64):
        self.patch_size = patch_size
        self.dataname = dataname
        self.img_size = img_size
        self.batch_size = batch_size

    # 각 class의 원소에다가 weight를 곱해주어 동일한 분포의 class 학습 데이터가 나올 수 있도록 한다
    def make_weights(self, labels, nclasses):
        labels = np.array(labels)
        weight_arr = np.zeros_like(labels)
        _, counts = np.unique(labels, return_counts=True)
        for cls in range(nclasses):
            weight_arr = np.where(labels == cls, 1/counts[cls], weight_arr) 
    
        return weight_arr 

    def patchdata(self):
        # CIFAR-10 Dataset에 대한 RGB채널의 평균과 표준편차값이다
        # CIFAR-10 Data로 Pre-trained된 model을 사용하기 때문에, test data도 모두 학습 데이터를 기준으로 정규화 되어야 한다
        mean = (0.4914, 0.4822, 0.4465) 
        std = (0.2023, 0.1994, 0.2010)

        # Data의 Pre-processor 정의
        # ToTensor() : CIFAR10을 불러오면 PIL이미지 객체로 구성되어 있다. 이를 Pytorch에서 처리하기 위하여 Tensor형으로 만든다
        # PatchGenerator() : 입력 이미지를 patch단위로 분해하는 transforms이다
        train_transform = transforms.Compose([transforms.Resize(self.img_size), transforms.RandomCrop(self.img_size, padding=2),
                                              transforms.RandomHorizontalFlip(), transforms.ToTensor(), transforms.Normalize(mean, std),
                                              PatchGenerator(self.patch_size)])
        test_transform = transforms.Compose([transforms.Resize(self.img_size), transforms.ToTensor(),
                                             transforms.Normalize(mean, std), PatchGenerator(self.patch_size)])

        # CIFAR-10 Dataset에 대한 Classification을 수행하고 싶다면
        if self.dataname == 'cifar10':
            # 훈련데이터 / 테스트 데이터 불러오기
            trainset = torchvision.datasets.CIFAR10(root='./data', train=True, download=True, transform=train_transform)
            testset = torchvision.datasets.CIFAR10(root='./data', train=False, download=True, transform=test_transform)
            
            # Test Data를 Test Data / Validation Data으로 나누기
            evens = list(range(0, len(testset), 2)) # 짝수
            odds = list(range(1, len(testset), 2)) # 홀수
            valset = torch.utils.data.Subset(testset, evens) # 짝수 index의 data는 validation Dataset
            testset = torch.utils.data.Subset(testset, odds) # 홀수 index의 data는 test Dataset
          
        elif self.dataname == 'imagenet':
            pass

        weights = self.make_weights(trainset.targets, len(trainset.classes))  # 가중치 계산
        weights = torch.DoubleTensor(weights) # 가중치는 numpy이므로 이를 Tensor로 변환하기

        # 학습 시 하나의 batch당 동일한 class분포의 데이터가 들어올 수 있도록 sampler를 사용한다
        sampler = torch.utils.data.sampler.WeightedRandomSampler(weights, len(weights))
        trainloader = DataLoader(trainset, batch_size=self.batch_size, sampler=sampler)

        valloader = DataLoader(valset, batch_size=self.batch_size, shuffle=False)
        testloader = DataLoader(testset, batch_size=self.batch_size, shuffle=False)

        return trainloader, valloader, testloader