# -*- coding: utf-8 -*-
"""
Created on Sun Jun 13 09:34:51 2021

@author: MARCELLOCHIESA
"""

import os,sys
import numpy as np
import torch
from torchvision import datasets,transforms

########################################################################################################################

# SVHN divided into 5 sets.

def get(data_path,seed,fixed_order=False,pc_valid=0):
    data={}
    taskcla=[]
    size=[3,32,32]

    # MNIST
    mean=(0.4514,)
    std=(0.1993,)
    dat={}
    dat['train']=datasets.SVHN(data_path,split='train',download=True,transform=transforms.Compose([transforms.ToTensor(),transforms.Normalize(mean,std)]))
    dat['test']=datasets.SVHN(data_path,split='test',download=True,transform=transforms.Compose([transforms.ToTensor(),transforms.Normalize(mean,std)]))
    data[0]={}
    data[0]['name']='SVHN-0-1'
    data[0]['ncla']=2
    data[1]={}
    data[1]['name']='SVHN-2-3'
    data[1]['ncla']=2
    data[2]={}
    data[2]['name']='SVHN-4-5'
    data[2]['ncla']=2
    data[3]={}
    data[3]['name']='SVHN-6-7'
    data[3]['ncla']=2
    data[4]={}
    data[4]['name']='SVHN-8-9'
    data[4]['ncla']=2
    for s in ['train','test']:
        loader=torch.utils.data.DataLoader(dat[s],batch_size=1,shuffle=False)
        data[0][s]={'x': [],'y': []}
        data[1][s]={'x': [],'y': []}
        data[2][s]={'x': [],'y': []}
        data[3][s]={'x': [],'y': []}
        data[4][s]={'x': [],'y': []}
        for image,target in loader:
            label=target.numpy()-1
            # image = torch.mean(sample[0], dim = 0)
            if label==0:
                data[0][s]['x'].append(image)
                data[0][s]['y'].append(0)
            elif label==1:
                data[0][s]['x'].append(image)
                data[0][s]['y'].append(1)
                
            elif label==2:
                data[1][s]['x'].append(image)
                data[1][s]['y'].append(0)
            elif label==3:
                data[1][s]['x'].append(image)
                data[1][s]['y'].append(1)

            elif label==4:
                data[2][s]['x'].append(image)
                data[2][s]['y'].append(0)
            elif label==5:
                data[2][s]['x'].append(image)
                data[2][s]['y'].append(1)

            elif label==6:
                data[3][s]['x'].append(image)
                data[3][s]['y'].append(0)
            elif label==7:
                data[3][s]['x'].append(image)
                data[3][s]['y'].append(1)

            elif label==8:
                data[4][s]['x'].append(image)
                data[4][s]['y'].append(0)
            elif label==9:
                data[4][s]['x'].append(image)
                data[4][s]['y'].append(1)


    # "Unify" and save
    for n in range(5):
        for s in ['train','test']:
            data[n][s]['x']=torch.stack(data[n][s]['x']).view(-1,size[0],size[1],size[2])
            data[n][s]['y']=torch.LongTensor(np.array(data[n][s]['y'],dtype=int)).view(-1)

    # Validation
    for t in data.keys():
        data[t]['valid']={}
        data[t]['valid']['x']=data[t]['train']['x'].clone()
        data[t]['valid']['y']=data[t]['train']['y'].clone()

    # Others
    n=0
    for t in data.keys():
        taskcla.append((t,data[t]['ncla']))
        n+=data[t]['ncla']
    data['ncla']=n

    return data,taskcla,size

########################################################################################################################
