import loadTextDoc
from pickle import load
from keras.preprocessing.text import Tokenizer
from keras.preprocessing.sequence import pad_sequences
from keras.utils import to_categorical
from numpy import array

#load a pre defined list of photo identifiers
def load_set(filename):
    doc = loadTextDoc.load_doc(filename)
    dataset = list()
    #process line by line
    for line in doc.split('\n'):
        #skip empty lines
        if len(line) < 1:
            continue
        #get the image identifier
        identifier = line.split('.')[0]
        dataset.append(identifier)
    return set(dataset)


def load_clean_descriptions(filename,dataset):
    #load document
    doc = loadTextDoc.load_doc(filename)
    descriptions = dict()
    for line in doc.split('\n'):
        #split tokens
        tokens = line.split()
        #split image id from descriptions
        image_id,image_desc = tokens[0],tokens[1:]
        #skip images not in the set
        if image_id in dataset:
            #create list
            if image_id not in descriptions:
                descriptions[image_id] = list()
            #wrap description in tokens
            desc = 'startseq ' + ' '.join(image_desc) + ' endseq'
            #store
            descriptions[image_id].append(desc)
    return descriptions  

#load interested photo features
def load_photo_features(filename,dataset):
    #load all features
    all_features = load(open(filename,'rb'))
    #filter features
    features = {k : all_features[k] for k in dataset}
    return features

#covnert dictionary of descriptions to a list of descriptions
def to_lines(descriptions):
    all_desc = list()
    for key in descriptions.keys():
        [all_desc.append(d) for d in descriptions[key]]
    return all_desc
    
#fit a tokenizer to given caption descriptions
def create_tokenizer(descriptions):
    lines = to_lines(descriptions)
    tokenizer = Tokenizer()
    tokenizer.fit_on_texts(lines)
    return tokenizer

#calcuate the length of the description with the most words
def max_length(descriptions):
    lines = to_lines(descriptions)
    return max(len(d.split()) for d in lines)

#create sequences of images, input sequences and output words for an image
def create_sequences(tokenizer,max_length,descriptions,photos,vocab_size):
    X1,X2,y = list(),list(),list()
    #walk through each image identifier
    for key,desc_list in descriptions.items():
        #walk through each description for the image
        for desc in desc_list:
            #encode the sequence
            seq = tokenizer.texts_to_sequences([desc])[0]
            #split one sequence into multiple X,y pairs
            for i in range(1,len(seq)):
                #split into input and output pairs
                in_seq,out_seq = seq[:i],seq[i]
                #pad input sequence
                in_seq = pad_sequences([in_seq],maxlen=max_length)[0]
                #encode output sequence
                out_seq = to_categorical([out_seq],num_classes=vocab_size)[0]
                #store
                X1.append(photos[key][0])
                X2.append(in_seq)
                y.append(out_seq)
    
    return array(X1),array(X2),array(y)
    

    
    
    
    
    
    
    
    
    