import os
from scipy.misc import imread
import numpy as np
import pickle
import pandas as pd

def load_sub_omniglot(path,n=0):

    '''
    path => Path of train directory or test directory
    '''
    X=[]
    y = []
    cat_dict = {}
    lang_dict = {}
    curr_y = n

    # we load every alphabet seperately so we can isolate them later
    for alphabet in os.listdir(path):
        if alphabet == '.DS_Store':
            continue
        print("loading alphabet: " + alphabet)
        lang_dict[alphabet] = [curr_y,None]
        alphabet_path = os.path.join(path,alphabet)

        # every letter/category has it's own column in the array, so  load seperately
        for letter in os.listdir(alphabet_path):
            if letter == '.DS_Store':
                continue
            cat_dict[curr_y] = (alphabet, letter)
            category_images=[]
            letter_path = os.path.join(alphabet_path, letter)

            # read all the images in the current category
            for filename in os.listdir(letter_path):
                image_path = os.path.join(letter_path, filename)
                image = imread(image_path)
                category_images.append(image)
                y.append(curr_y)
            try:
                #X.append(np.stack(category_images))
                X += list(np.stack(category_images))
            # edge case  - last one
            except ValueError as e:
                print(e)
                print("error - category_images:", category_images)
            curr_y += 1
            lang_dict[alphabet][1] = curr_y - 1
    y = np.vstack(y)
    X = np.stack(X)
    print(X.shape,y.shape,len(np.unique(y)))
    return X,y


def load_omniglot(path=None):
    if path is None:
        train_path = './data/omniglot/images_background'
        test_path = './data/omniglot/images_evaluation'
    else:
        train_path,test_path = path


    return load_sub_omniglot(train_path),load_sub_omniglot(test_path)




def load_tihm(path=None):
    if path is None:
        path = './data/timh1.5'

    data = []
    label = []

    for filename in os.listdir(path):
        if '_lag_new.csv' in filename:
            file_path = os.path.join(path, filename)
            d = pd.read_csv(file_path)
            if 'halway' in list(d.columns):
                d = d.drop(columns='halway')
            if 'd_front' not in list(d.columns):
                continue
            d = np.array(d.values)
            d = d[:,2:]
            data += d.tolist()
            f_label = filename.split('_lag_new.csv')[0]
            f_label += '_class.csv'
            file_path = os.path.join(path, f_label)
            d = np.array(pd.read_csv(file_path,header=None))

            label += d.reshape(-1).tolist()

    data = np.array(data)
    label = np.array(label)

    return (data[:10000],label[:10000]),(data[10000:],label[10000:])



def load_agitation(path=None):
    if path is None:
        path = './data/aigitation'

    data = []
    label = []





























