import random

from .segdataloader import *

from .utils import *

import csv

import glob 

import cv2

fold = 1

def readCSV(filename):

    lines = []

    with open(filename, "r") as f:

        csvreader = csv.reader(f)

        for line in csvreader:

            lines.append(line[0])

    return lines

def get_dataloader(config, mode='train', batchsize=64, width=64, height=64):

    train_datas = []

    train_masks = []

    for index in config.training_fold_index:

        tempdata = readCSV(os.path.join(config.csvPath, 'data_fold' + str(index) + '.csv'))

        tempmask = readCSV(os.path.join(config.csvPath, 'mask_fold' + str(index) + '.csv'))

        train_datas += tempdata

        train_masks += tempmask

    test_datas = readCSV(os.path.join(config.csvPath, 'data_fold' + str(config.test_fold_index[0]) + '.csv'))

    test_masks = readCSV(os.path.join(config.csvPath, 'mask_fold' + str(config.test_fold_index[0]) + '.csv'))

    if mode=='train':

        # remove features labels

        temp_train_datas = []

        for one in train_datas:

            one_temp = one.split('/')[-1]

            one_list = one_temp.split('_')

            temp_train_datas.append(one_list[0] + '_' + one_list[1] + '_' + one_list[2])

        temp_test_datas = []

        for one in test_datas:

            one_temp = one.split('/')[-1]

            one_list = one_temp.split('_')

            temp_test_datas.append(one_list[0] + '_' + one_list[1] + '_' + one_list[2])

        mid_files = os.listdir(config.maskPath2)

        temp2_train_inter = []

        temp2_train_union = []

        temp2_train_data = []

        temp2_train_lung = []

        temp2_train_media = []

        temp2_train_mask = []

        for one_train_data in temp_train_datas:

            imagename = one_train_data + '.png'

            if imagename in mid_files:

                innertemp0 = config.midPath + one_train_data + '.npy'

                innertemp1 = config.lungPath + one_train_data + '_lung.npy'

                innertemp2 = config.mediaPath + one_train_data + '_mediastinal.npy'

                innertemp3 = config.maskPath2 + one_train_data + '_red.png'

                innertemp4 = config.maskPath2 + one_train_data + '_blue.png'

                innertemp5 = config.maskPath1 + 'mid_' + one_train_data + '_mask.png'

                temp2_train_data.append(innertemp0)

                temp2_train_lung.append(innertemp1)

                temp2_train_media.append(innertemp2)

                temp2_train_union.append(innertemp3) 

                temp2_train_inter.append(innertemp4)

                temp2_train_mask.append(innertemp5)

        temp2_test_data = []

        temp2_test_lung = []

        temp2_test_media = []

        temp2_test_inter = []

        temp2_test_union = []

        temp2_test_mask = []

        for one_test_data in temp_test_datas:

            imagename = one_test_data + '.png'       

            if imagename in mid_files:

                innertemp0 = config.midPath + one_test_data + '.npy' 

                innertemp1 = config.lungPath + one_test_data + '_lung.npy'

                innertemp2 = config.mediaPath + one_test_data + '_mediastinal.npy'

                innertemp3 = config.maskPath2 + one_test_data + '_red.png'

                innertemp4 = config.maskPath2 + one_test_data + '_blue.png'

                innertemp5 = config.maskPath1 + 'mid_' + one_test_data + '_mask.png'

                temp2_test_data.append(innertemp0)

                temp2_test_lung.append(innertemp1)

                temp2_test_media.append(innertemp2)

                temp2_test_union.append(innertemp3) 

                temp2_test_inter.append(innertemp4)

                temp2_test_mask.append(innertemp5)

        print('***********')

        print('the length of train data: ', len(temp2_train_data))

        print('the length of test data: ', len(temp2_test_data))

        print('-----------')

        dataloader = loader(Dataset(temp2_train_data, temp2_train_lung, temp2_train_media, temp2_train_inter, temp2_train_union, temp2_train_mask,  width=width, height=height), batchsize)

        dataloader_val = loader(Dataset(temp2_test_data, temp2_test_lung, temp2_test_media, temp2_test_inter, temp2_test_union, temp2_test_mask, width=width, height=height), batchsize)

        return dataloader, dataloader_val

    if mode=='row':

        # remove features labels

        temp_train_datas = []

        for one in train_datas:

            one_temp = one.split('/')[-1]

            one_list = one_temp.split('_')

            temp_train_datas.append('mid_' + one_list[0] + '_' + one_list[1] + '_' + one_list[2])

        temp_test_datas = []

        for one in test_datas:

            one_temp = one.split('/')[-1]

            one_list = one_temp.split('_')

            temp_test_datas.append('mid_' + one_list[0] + '_' + one_list[1] + '_' + one_list[2])

        temp2_train_datas = []

        temp2_train_masks = []

        temp2_test_datas = []

        temp2_test_masks = []

        row_files = os.listdir(config.rowPath)

        for one_train_data in temp_train_datas:

            imagename = one_train_data + '.png'

            if imagename in row_files:

                innertemp0 = config.rowPath + one_train_data + '.png'

                innertemp1 = config.rowPath + one_train_data + '_mask.png'

                temp2_train_datas.append(innertemp0)

                temp2_train_masks.append(innertemp1)

        for one_test_data in temp_test_datas:

            imagename = one_test_data + '.png'

            if imagename in row_files:

                innertemp0 = config.rowPath + one_test_data + '.png'

                innertemp1 = config.rowPath + one_test_data + '_mask.png'

                temp2_test_datas.append(innertemp0)

                temp2_test_masks.append(innertemp1)

        dataloader = loader(RowDataset(temp2_train_datas, temp2_train_masks, width=width, height=height), batchsize)

        dataloader_val = loader(RowDataset(temp2_test_datas, temp2_test_masks, width=width, height=height), batchsize)

        return dataloader, dataloader_val