#!/usr/bin/env python
# -*- coding: utf-8 -*-

"""
Program to preprocess datasets and prepare weight embedding vectors
from word2vec, dependency and fatcual embeddings

Copyright (C) 2016
Ubiquitous Knowledge Processing (UKP) Lab
Technische Universität Darmstadt

Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

    http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
"""

import numpy as np
import cPickle
from collections import defaultdict
import sys, re
import pandas as pd

def preprocess_data(data_folder):
    
    #trainset
    train_file = data_folder[0]
    #testset
    test_file = data_folder[1]

    claims = []  

    #vocabulary from training and test data 
    vocab = defaultdict(float)

    labelshash = {}
    labelList = []	

    # Read labels list from training data and convert labels into indices 0 ... n
    with open(train_file, "rb") as fl:
        for line in fl:
	 line = line.strip()	
         #print line
	 label_index = line.split("\t")
	 label = label_index[1]
	 labelList.append(label)
	 labelshash[label_index[0]]=label_index[1]
    fl.close()	

    labelSet = set(labelList)
    labelSortedList = sorted(list(labelSet))
    indices = range(0,len(labelSortedList))
    labelshash = dict(zip(labelSortedList,indices))
		
    with open(train_file, "rb") as ftr:
        for line in ftr:       
            claim = []
	    line = line.strip()	
		
	    claim_and_label = line.split("\t")
           
	    claim = claim_and_label[0]
            clabel = claim_and_label[1]
           
            orig_claim = clean_str(claim)

	    if len(orig_claim.split()) >500:
                orig_claim = orig_claim[0:2500]

            words = set(orig_claim.split())
	    
            for word in words:
                vocab[word] += 1

            claim_meta  = {"y":labelshash[clabel],
                      "text": orig_claim,                             
                      "num_words": len(orig_claim.split()),
                      "split": "train"
			}
            if len(orig_claim.split()) <=500:
	     claims.append(claim_meta)
             slen = len(orig_claim.split())
	    #print slen

    ftr.close()		

    c = 0 
    with open(test_file, "rb") as fts:
        for line in fts:
	    claim = []	
	    line = line.strip()	
		
	    claim_and_label = line.split("\t")

            claim = claim_and_label[0]
            clabel = claim_and_label[1]

	    orig_claim = clean_str(claim)
	   	
            if len(orig_claim.split()) >500:
                orig_claim = orig_claim[0:2500]
		c = c + 1

            words = set(orig_claim.split())

            claim_meta  = {"y":labelshash[clabel],
                      "text": orig_claim,      
                      "num_words": len(orig_claim.split()),
                      "split": "test"
                        }		
	    #if len(orig_claim.split()) <=500:
            claims.append(claim_meta)	
	    slen = len(orig_claim.split())
            #print slen

    #print c
    fts.close()
    #print claims
    return claims, labelSortedList, vocab

#Return weight matrix and word word index map    
def get_W(word_vecs, k=300):
    """
    Get word matrix. W[i] is the vector for word indexed by i
    """
    vocab_size = len(word_vecs)
    word_idx_map = dict()
    W = np.zeros(shape=(vocab_size+1, k))            
    W[0] = np.zeros(k)
    i = 1
    for word in word_vecs:
        W[i] = word_vecs[word]
        word_idx_map[word] = i
        i += 1
    return W, word_idx_map

def load_bin_vec(fname, vocab):
    """
    Loads 300x1 word vecs from Google (Mikolov) word2vec
    """
    word_vecs = {}
    with open(fname, "rb") as f:
        header = f.readline()
        vocab_size, layer1_size = map(int, header.split())
        binary_len = np.dtype('float32').itemsize * layer1_size
        for line in xrange(vocab_size):
            word = []
            while True:
                ch = f.read(1)
                if ch == ' ':
                    word = ''.join(word)
                    break
                if ch != '\n':
                    word.append(ch)   
            if word in vocab:
               word_vecs[word] = np.fromstring(f.read(binary_len), dtype='float32')  
            else:
                f.read(binary_len)
    return word_vecs

def load_levy_vec(fname, vocab, k):

    #deps.words is a filename  from Omer Levy WordEmbeddings

    word_vecs = {}
    with open(fname, "rb") as f:
      for line in f:

        wordvector = line.strip().split(' ')
        word = wordvector[0]
        #print ('word vector',word,wordvector)
        del wordvector[0]
        #print ('word',word)
        if word in vocab:
          vec = np.array(wordvector, dtype='float32')
          word_vecs[word]=vec[0:k]
         # print('word vec',word,word_vecs[word])
    return word_vecs


def clean_str(string):
    """
    Tokenization/string cleaning for all datasets
   
    """
    #print string
    string = re.sub(r"[^A-Za-z0-9(),!?\'\`]", " ", string)
    string = re.sub(r"\'s", " \'s", string)
    string = re.sub(r"\'ve", " \'ve", string)
    string = re.sub(r"n\'t", " n\'t", string)
    string = re.sub(r"\'re", " \'re", string)
    string = re.sub(r"\'d", " \'d", string)
    string = re.sub(r"\'ll", " \'ll", string)
    string = re.sub(r",", " , ", string)
    string = re.sub(r"!", " ! ", string)
    string = re.sub(r"\(", " \( ", string)
    string = re.sub(r"\)", " \) ", string)
    string = re.sub(r"\?", " \? ", string)
    string = re.sub(r"\s{2,}", " ", string)
    return string.strip().lower()

def add_unknown_words(word_vecs, vocab, min_df=1, k=300):
    """
    For words that occur in at least min_df sentences, create a separate word vector.    
    0.25 is chosen so the unknown vectors have (approximately) same variance as pre-trained ones
    """
    for word in vocab:
        if word not in word_vecs and vocab[word] >= min_df:
            word_vecs[word] = np.random.uniform(-0.25,0.25,k)  


if __name__=="__main__":    
   
   # Labelled train and test set      
    data_folder = ["../data/data-fact-feel-v2/train/train.txt","../data/data-fact-feel-v2/test/test.txt"]
    
   # Labelled train and dev set      
   # data_folder = ["../data/data-fact-feel-v2/train/train.txt","../data/data-fact-feel-v2/dev/dev.txt"]  
 
    #Read word2vec, dependency and factual embeddings
    w2v_file = sys.argv[1]
    d2v_file = sys.argv[2]
    f2v_file = sys.argv[3]	

    print "preprocessing data..."        

    claims, labels, vocab = preprocess_data(data_folder)

    max_l = np.max(pd.DataFrame(claims)["num_words"])

    print "data loaded!"
    print "number of sentences: " + str(len(claims))
    print " labels: ", labels
    print "vocab size: " + str(len(vocab))
    print "max sentence length: " + str(max_l)


    #loading word2vec embeddings        
    print "loading word2vec vectors...",
    w2v = load_bin_vec(w2v_file, vocab)
    print "word2vec loaded!"
    print "num words already in word2vec: " + str(len(w2v))

    #load omerlevy word embeddings
    print "loading dependency word2vec vectors...",
    d2v = load_levy_vec(d2v_file, vocab, 300)
    print "dep2vec loaded!"
    print "num words already in dep2vec: " + str(len(d2v))

    #loading factual embeddings 
    print "loading fact2vec vectors...",
    f2v = load_bin_vec(f2v_file, vocab)
    print "fact2vec loaded!"
    print "num words already in fact2vec: " + str(len(f2v))

    #Add distributions for unseen words in w2v
    add_unknown_words(w2v, vocab)
    # get weight vector using w2v       
    W_w2v, word_idx_map_w2v = get_W(w2v)

    #Add distributions for unseen words in d2v 
    add_unknown_words(d2v, vocab)
    # get weight vector using d2v
    W_d2v, word_idx_map_d2v = get_W(d2v)

    #Add distributions for unseen words in f2v
    add_unknown_words(f2v, vocab)
    # get weight vector using f2v       
    W_f2v, word_idx_map_f2v = get_W(f2v)

    #random vectors and weights
    rand_vecs = {}
    add_unknown_words(rand_vecs, vocab)
    W2, _ = get_W(rand_vecs)

    # consider 100 dimensions from each embeddings to form 300 dimnesion stacked embedding vector
    k = 100
    #concatenated embeddings - word2vec, de2vec, fact2vec embeddings
    cvecs = {}

    #concatenate all vectors into final vector
    for voc in w2v:
        if voc in d2v and voc in f2v:
                cvecs[voc] = np.concatenate([w2v[voc][0:k],d2v[voc][0:k],f2v[voc][0:k]])
        elif voc not in d2v and voc in f2v:
                cvecs[voc] = np.concatenate([w2v[voc][0:2*k],f2v[voc][0:k]])
        elif voc in d2v and voc not in f2v:
                cvecs[voc] = np.concatenate([w2v[voc][0:2*k],d2v[voc][0:k]])
        else:
                cvecs[voc] = w2v[voc]

    #Add distributions for unseen words in cvecs
    add_unknown_words(cvecs, vocab)
    # get weight vector using cvecs
    W_concat, word_idx_map_concat = get_W(cvecs)

    cPickle.dump([claims, W_w2v, W2, word_idx_map_w2v, vocab, max_l], open("claims_factfeel_word2vec.p", "wb"))
    cPickle.dump([claims, W_d2v, W2, word_idx_map_d2v, vocab, max_l], open("claims_factfeel_dep2vec.p", "wb"))
    cPickle.dump([claims, W_f2v, W2, word_idx_map_f2v, vocab, max_l], open("claims_factfeel_fact2vec.p", "wb"))
    cPickle.dump([claims, W_concat, W2, word_idx_map_concat, vocab, max_l], open("claims_factfeel_concat.p", "wb"))

    print ("word2vec, dep2vec, fact2vec, stacked pickle datasets are created!")


