# This is a sample Python script.
import math
import random

import torch
import torch.nn as nn
import fasttext
import numpy as np
from torch import optim
import pickle
import time
torch.backends.cudnn.enabled = False
# Global Variables
N_VECTOR = 20


# Use GPU or not
device = torch.device('cuda') if torch.cuda.is_available() else torch.device('cpu')

# Special Words
special_embedding_tensor = pickle.load(open('./first_stage_state/special_embedding','rb'))



class Vocab:

    def __init__(self, word_embedding_path, special_embedding_tensor):
        self.word_vector = fasttext.load_model(word_embedding_path)
        self.special_vector = special_embedding_tensor

    def get_sentence_tensor(self, sentence, source=False):
        words = sentence.strip().split(' ')
        length = len(words)
        sentence_tensor = torch.zeros(length, N_VECTOR) if not source else torch.zeros(length+1, N_VECTOR)
        if source:
            sentence_tensor[0] = special_embedding_tensor[1]
        for idx, word in enumerate(words):
            index = idx if not source else idx+1
            word_tensor = torch.Tensor(self.word_vector.get_word_vector(word))
            if torch.equal(word_tensor, torch.zeros((N_VECTOR))):
                word_tensor = self.special_vector[2]
            sentence_tensor[index] = word_tensor

        # sentence_tensor[-1] = special_embedding_tensor[-1]
        return sentence_tensor

# vocab = Vocab('english20', special_embedding_tensor)
# abc = vocab.get_sentence_tensor('hello there', source = True)
def get_padded_tensor(tensor, n_padding, padding_index, special_embedding, front =True):
    """
    Args:
        tensor (torch.Tensor): tensor to be padded
        n_padding (int): padding nums
        padding_index (int): which chars to use:  (0-SOS 1-NULL 2-UNK 3-EOS)
        special_embedding (torch.Tensor): embedding tensor of 4 special chars
        front (bool): pad front or end

    Returns:
         padded_tensor (torch.Tensor)
    """
    if tensor == 'NULL':
        return special_embedding[1].repeat(n_padding,1)

    padding_tensor = special_embedding[padding_index].repeat(n_padding,1)
    padded_tensor = torch.cat((padding_tensor, tensor), dim=0) if front \
        else torch.cat((tensor, padding_tensor), dim = 0)

    return padded_tensor

def judge_pad(index, n_window, sentence_tensor, special_embedding, src, sen_li):
    n_sen, radius = sentence_tensor.size()[0] if not sen_li else sen_li, n_window//2
    if src and index == 0:
        return get_padded_tensor("NULL", n_window, 1, special_embedding)


    #fist half (including index)
    if index - radius < 0:
        n_pad = radius - index
        tensor_f = get_padded_tensor(sentence_tensor[:index+1], n_pad, 0, special_embedding, True)
    else:
        tensor_f = sentence_tensor[index-radius:index+1]

    # second half (No index)
    if index + radius >= n_sen:
        n_pad = radius + index - n_sen + 1
        tensor_s = get_padded_tensor(sentence_tensor[index+1:n_sen], n_pad, 3, special_embedding, False)
    else:
        tensor_s = sentence_tensor[index+1:index + radius + 1]

    tensor_window = torch.cat((tensor_f,tensor_s), dim=0)
    return tensor_window




# E <NULL> the commissionners went on :
# F je poursuis la lecture de les recommandations de les commissaires :
# 0-4 0-5 0-7 2-0 2-1 2-2 2-3 2-6 2-8 2-9 5-10
def find_ga_pairs(word_alignments, n_window, src, tgt, special_embedding):
    alignments = set([tuple(map(int, x.split("-"))) for x in word_alignments.strip().split()])
    concated_vector = torch.zeros((len(alignments), 2*n_window, N_VECTOR))
    for idx, alignment in enumerate(alignments):
        tgt_index, src_index, radius = alignment[1], alignment[0], n_window//2
        tgt_window = judge_pad(tgt_index, n_window, tgt, special_embedding, False)
        src_window = judge_pad(src_index, n_window, src, special_embedding, True)
        concated_vector[idx] = torch.cat((src_window,tgt_window), dim=0)

    return concated_vector

def find_pairs(alignment, n_window, src, tgt, special_embedding, sen_li= (None,None)):
    tgt_index, src_index, radius = alignment[1], alignment[0], n_window//2
    tgt_window = judge_pad(tgt_index, n_window, tgt, special_embedding, False, sen_li[1])
    src_window = judge_pad(src_index, n_window, src, special_embedding, True, sen_li[0])
    return torch.cat((src_window,tgt_window), dim = 0)

def get_ns_index(ori_index, sentence_length):
    if ori_index == 0:
        return random.randint(min(1,sentence_length-1), sentence_length-1)
    elif ori_index == sentence_length -1:
        return random.randint(0,max(0,sentence_length- 2))
    else:
        choice = random.randint(0,1)
        return choice * random.randint(0, ori_index-1) + (1-choice) * random.randint(ori_index+1, sentence_length-1)


def find_ns_pairs(word_alignments, n_window, src, tgt, special_embedding):
    alignments = set([tuple(map(int, x.split("-"))) for x in word_alignments.strip().split()])
    concated_vector = torch.zeros((3 * len(alignments), 2 * n_window, N_VECTOR))
    for idx, alignment in enumerate(alignments):
        tgt_index, src_index, radius = alignment[1], alignment[0], n_window // 2
        # choice_src, choice_tgt = random.randint(0,1), random.randint(0,1)
        ns_src_idx = get_ns_index(src_index, src.size()[0])
        ns_tgt_idx = get_ns_index(tgt_index, tgt.size()[0])
        concated_vector[3*idx] = find_pairs(alignment,n_window,src,tgt,special_embedding)
        concated_vector[3*idx+1] = find_pairs((ns_src_idx, tgt_index), n_window, src, tgt, special_embedding)
        concated_vector[3*idx+2] = find_pairs((src_index, ns_tgt_idx), n_window, src, tgt, special_embedding)

    return concated_vector, len(alignments)


def find_ns_batch_pairs(word_alignments, n_window, src_all, tgt_all, special_embedding, sentence_length):
    alignments, batch_size = [set([tuple(map(int, x.split("-"))) for x in ali.strip().split()]) for ali in word_alignments], len(word_alignments)
    total_length = sum([len(alignment) for alignment in alignments])
    concated_vector, count = torch.zeros((3, total_length, 2 * n_window, N_VECTOR)),0
    for realization, src, tgt, src_li, tgt_li in zip(alignments, src_all, tgt_all, sentence_length[0], sentence_length[1]):
        sen_li = (src_li, tgt_li)
        for alignment in realization:

            tgt_index, src_index, radius = alignment[1], alignment[0], n_window // 2
            # choice_src, choice_tgt = random.randint(0,1), random.randint(0,1)
            ns_src_idx = get_ns_index(src_index, src_li)
            ns_tgt_idx = get_ns_index(tgt_index, tgt_li)
            concated_vector[0,count] = find_pairs(alignment,n_window,src,tgt,special_embedding,sen_li)
            concated_vector[1,count] = find_pairs((ns_src_idx, tgt_index), n_window, src, tgt, special_embedding,sen_li)
            concated_vector[2,count] = find_pairs((src_index, ns_tgt_idx), n_window, src, tgt, special_embedding,sen_li)
            count += 1
    return concated_vector
#
# E = "<NULL> the commissionners went on :"
# F = "je poursuis la lecture de les recommandations de les commissaires :"
# VocabE = Vocab('english20', special_embedding_tensor)
# VocabF = Vocab('french20', special_embedding_tensor)
# E = VocabE.get_sentence_tensor(E, 1)
# F = VocabF.get_sentence_tensor(F)
# A = "0-4 0-5 0-7 2-0 2-1 2-2 2-3 2-6 2-8 2-9 5-10"
# res = find_ns_pairs(A, 11, E, F, special_embedding_tensor)
# print(res)

def htanh(input):
    input[input > 1] = 1
    input[input < -1] = -1
    return input

class Htanh(nn.Module):
    def __init__(self):
        super().__init__()

    def forward(self, input):
        return htanh(input)

# def train(input_data, optimizer, model):
#     length = input_data[0].size()[0]
#     optimizer.zero_grad()
#
#     loss =  torch.max(torch.zeros((length,1)), 1 - model(input_data[0].flatten(1)) + model(input_data[1].flatten(1))) +\
#             torch.max(torch.zeros((length,1)), 1 - model(input_data[0].flatten(1)) + model(input_data[2].flatten(1)))
#     # loss += torch.max(torch.zeros((length,1),requires_grad=True), 1 - model(input_data[0].flatten(1)) + model(input_data[2].flatten(1)))
#     loss = loss.sum()
#
#     loss.backward()
#     optimizer.step()
#
#     return loss.item()/length

def train_sgd(input_ga, input_ns, optimizer, model):
    optimizer.zero_grad()

    loss = torch.max(torch.tensor(0.0, requires_grad=True), 1 - model(input_ga.flatten()) + model(input_ns.flatten()))

    loss.backward(retain_graph = True)
    optimizer.step()

    return loss.item()

class FFNN_word_alignment(nn.Module):

    def __init__(self, input_size, f1_size, f2_size):
        super(FFNN_word_alignment, self).__init__()
        self.fc1 = nn.Linear(input_size, f1_size)
        self.fc2 = nn.Linear(f1_size,f2_size)
        self.fc3 = nn.Linear(f2_size, 1)
        self.htanh = Htanh()

    def forward(self, x):
        output = self.fc1(x)

        output =  self.htanh(output)

        output = self.fc2(output)

        output = self.fc3(output)

        return output

def read_in_chunk(file_name, chunk_size):
    file_object = open(file_name, 'r')
    while True:
        data = [0 for _ in range(chunk_size)]
        for i in range(chunk_size):
            line = file_object.readline()
            if not line:
                data = data[:i]
                break
            data[i] = line
        if not data:
            break
        yield data

# gen_t = read_in_chunk('./data/text.txt',2)
# for li in gen_t:
#     print(li)
def get_chunk_sentence_tensor(sentence_list, vocab, N_VECTOR, source):
    n_sentence = len(sentence_list)
    # sentence_tensor = torch.zeros((n_sentence, max_len,N_VECTOR))
    li_length,li_tensor = [0 for _ in range(n_sentence)], [0 for _ in range(n_sentence)]
    for idx, sentence in enumerate(sentence_list):

        temp_tensor = vocab.get_sentence_tensor(sentence, source)
        li_length[idx], li_tensor[idx] = temp_tensor.size()[0], temp_tensor

    final_sentence_tensor = torch.nn.utils.rnn.pad_sequence(li_tensor, batch_first=True)

    return final_sentence_tensor, li_length

def mini_batch(batch_size, train_list, alignment_list, sentence_length, shuffle=True):
    length = len(train_list)
    batch_num = math.ceil(length / batch_size)
    index_array = list(range(length))
    if shuffle:
        np.random.shuffle(index_array)
    for i in range(batch_num):
        indices = index_array[i*batch_size:(i+1)*batch_size]
        curr_data = [train_list[index] for index in indices]
        align_data = [alignment_list[index] for index in indices]
        src_length,tgt_length = [sentence_length[0][index] for index in indices], \
                                [sentence_length[1][index] for index in indices]
        src_sent = [e[0] for e in curr_data]
        tgt_sent = [f[1] for f in curr_data]
        yield src_sent, tgt_sent, align_data,(src_length,tgt_length)

WINDOW_SIZE = 11
Learning_rate= 0.01
N_epoch = 1000
def main():
    model_FFNN = FFNN_word_alignment(2*WINDOW_SIZE*N_VECTOR, 120, 10)
    optimizer = optim.SGD(model_FFNN.parameters(), lr=Learning_rate)
    vocab_e, vocab_f = Vocab('english20', special_embedding_tensor), Vocab('french20', special_embedding_tensor)
    src = open('./data/hansards.f','r').readlines()[:37]
    tgt = open('./data/hansards.e','r').readlines()[:37]
    golden_alignment = open('./data/alignmentfe','r').readlines()[:37]
    shuffle_indeies = list(range(len(golden_alignment)))
    # np.random.shuffle(shuffle_indeies)
    # gen_src, gen_tgt = read_in_chunk('./data/hansards.e', 10000), read_in_chunk('./data/hansards.f', 10000)
    # gen_alignment = read_in_chunk('./data/alignment',10000)
    # count = 0
    # for src,tgt,ali in zip(gen_src,gen_tgt,gen_alignment):
    #     src_tensor, li_src_len = get_chunk_sentence_tensor(src, vocab_e, N_VECTOR, True)
    #     tgt_tensor, li_tgt_len= get_chunk_sentence_tensor(tgt, vocab_f, N_VECTOR, False)
    #     sentence_length = (li_src_len, li_tgt_len)
    #     train_data = [(src_tensor[i],tgt_tensor[i]) for i in range(src_tensor.size()[0])]
    #     beg = time.time()
    #     for src_sent, tgt_sent, ali_sent, sentence_length, in mini_batch(BATCH_SIZE, train_data, ali, sentence_length, shuffle=False):
    #         train_data = find_ns_batch_pairs(ali_sent,WINDOW_SIZE, src_sent, tgt_sent,special_embedding_tensor, sentence_length)
    #         train_ga, train_ns_src, train_ns_tgt = train_data[0], train_data[1], train_data[2]
    #         loss = train(train_data,optimizer,model_FFNN)
    #         count += 1
    #         if count % 10 == 0:
    #             end = time.time()
    #             print(end-beg)
    #             beg = time.time()



    loss, epoch, print_per = 0, 0, 1000
    loss_his, beg = [], time.time()
    # check_point = torch.load('./first_stage_state/state2198000.pth')
    # model_FFNN.load_state_dict(check_point)
    model_FFNN.train()
    while epoch < N_epoch:
        print_n_alignment = 0
        for i in range(len(src)):
            wv_src, wv_tgt = vocab_f.get_sentence_tensor(src[shuffle_indeies[i]], True), \
                             vocab_e.get_sentence_tensor(tgt[shuffle_indeies[i]],False)

            training_ns,n_alignment = find_ns_pairs(golden_alignment[shuffle_indeies[i]], WINDOW_SIZE, wv_src, wv_tgt,
                                                    special_embedding_tensor)
            training_ns = training_ns.view(n_alignment, 3, 2*WINDOW_SIZE, N_VECTOR)
            for idx in range(n_alignment):
                loss_f = train_sgd(training_ns[idx][0], training_ns[idx][1], optimizer,model_FFNN)
                loss_s = train_sgd(training_ns[idx][0], training_ns[idx][2], optimizer,model_FFNN)
                loss += (loss_f + loss_s) /2
            print_n_alignment += n_alignment
            if i % print_per == 36 and i != 0 :
                end = time.time()
                print("Loss per {} iters: {}".format(print_per, loss/print_n_alignment))
                print("Time per {} iters: {}".format(print_per, end-beg))
                torch.save(model_FFNN.state_dict(),'./first_stage_state/statefe37'  + '.pth')
                loss_his.append(loss)
                loss,print_n_alignment = 0,0
                beg = time.time()
        epoch += 1
        print('epoch:{}'.format(epoch))
    # pickle.dump(loss_his, open('./first_stage_state/loss','wb'))






if __name__ == '__main__':
    main()

# See PyCharm help at https://www.jetbrains.com/help/pycharm/

# This is a sample Python script.
import math
import random

import torch
import torch.nn as nn
import fasttext
import numpy as np
from torch import optim
import pickle
import time
torch.backends.cudnn.enabled = False
# Global Variables
N_VECTOR = 20


# Use GPU or not
device = torch.device('cuda') if torch.cuda.is_available() else torch.device('cpu')

# Special Words
special_embedding_tensor = pickle.load(open('./first_stage_state/special_embedding','rb'))



class Vocab:

    def __init__(self, word_embedding_path, special_embedding_tensor):
        self.word_vector = fasttext.load_model(word_embedding_path)
        self.special_vector = special_embedding_tensor

    def get_sentence_tensor(self, sentence, source=False):
        words = sentence.strip().split(' ')
        length = len(words)
        sentence_tensor = torch.zeros(length, N_VECTOR) if not source else torch.zeros(length+1, N_VECTOR)
        if source:
            sentence_tensor[0] = special_embedding_tensor[1]
        for idx, word in enumerate(words):
            index = idx if not source else idx+1
            word_tensor = torch.Tensor(self.word_vector.get_word_vector(word))
            if torch.equal(word_tensor, torch.zeros((N_VECTOR))):
                word_tensor = self.special_vector[2]
            sentence_tensor[index] = word_tensor

        # sentence_tensor[-1] = special_embedding_tensor[-1]
        return sentence_tensor

# vocab = Vocab('english20', special_embedding_tensor)
# abc = vocab.get_sentence_tensor('hello there', source = True)
def get_padded_tensor(tensor, n_padding, padding_index, special_embedding, front =True):
    """
    Args:
        tensor (torch.Tensor): tensor to be padded
        n_padding (int): padding nums
        padding_index (int): which chars to use:  (0-SOS 1-NULL 2-UNK 3-EOS)
        special_embedding (torch.Tensor): embedding tensor of 4 special chars
        front (bool): pad front or end

    Returns:
         padded_tensor (torch.Tensor)
    """
    if tensor == 'NULL':
        return special_embedding[1].repeat(n_padding,1)

    padding_tensor = special_embedding[padding_index].repeat(n_padding,1)
    padded_tensor = torch.cat((padding_tensor, tensor), dim=0) if front \
        else torch.cat((tensor, padding_tensor), dim = 0)

    return padded_tensor

def judge_pad(index, n_window, sentence_tensor, special_embedding, src, sen_li):
    n_sen, radius = sentence_tensor.size()[0] if not sen_li else sen_li, n_window//2
    if src and index == 0:
        return get_padded_tensor("NULL", n_window, 1, special_embedding)


    #fist half (including index)
    if index - radius < 0:
        n_pad = radius - index
        tensor_f = get_padded_tensor(sentence_tensor[:index+1], n_pad, 0, special_embedding, True)
    else:
        tensor_f = sentence_tensor[index-radius:index+1]

    # second half (No index)
    if index + radius >= n_sen:
        n_pad = radius + index - n_sen + 1
        tensor_s = get_padded_tensor(sentence_tensor[index+1:n_sen], n_pad, 3, special_embedding, False)
    else:
        tensor_s = sentence_tensor[index+1:index + radius + 1]

    tensor_window = torch.cat((tensor_f,tensor_s), dim=0)
    return tensor_window




# E <NULL> the commissionners went on :
# F je poursuis la lecture de les recommandations de les commissaires :
# 0-4 0-5 0-7 2-0 2-1 2-2 2-3 2-6 2-8 2-9 5-10
def find_ga_pairs(word_alignments, n_window, src, tgt, special_embedding):
    alignments = set([tuple(map(int, x.split("-"))) for x in word_alignments.strip().split()])
    concated_vector = torch.zeros((len(alignments), 2*n_window, N_VECTOR))
    for idx, alignment in enumerate(alignments):
        tgt_index, src_index, radius = alignment[1], alignment[0], n_window//2
        tgt_window = judge_pad(tgt_index, n_window, tgt, special_embedding, False)
        src_window = judge_pad(src_index, n_window, src, special_embedding, True)
        concated_vector[idx] = torch.cat((src_window,tgt_window), dim=0)

    return concated_vector

def find_pairs(alignment, n_window, src, tgt, special_embedding, sen_li= (None,None)):
    tgt_index, src_index, radius = alignment[1], alignment[0], n_window//2
    tgt_window = judge_pad(tgt_index, n_window, tgt, special_embedding, False, sen_li[1])
    src_window = judge_pad(src_index, n_window, src, special_embedding, True, sen_li[0])
    return torch.cat((src_window,tgt_window), dim = 0)

def get_ns_index(ori_index, sentence_length):
    if ori_index == 0:
        return random.randint(min(1,sentence_length-1), sentence_length-1)
    elif ori_index == sentence_length -1:
        return random.randint(0,max(0,sentence_length- 2))
    else:
        choice = random.randint(0,1)
        return choice * random.randint(0, ori_index-1) + (1-choice) * random.randint(ori_index+1, sentence_length-1)


def find_ns_pairs(word_alignments, n_window, src, tgt, special_embedding):
    alignments = set([tuple(map(int, x.split("-"))) for x in word_alignments.strip().split()])
    concated_vector = torch.zeros((3 * len(alignments), 2 * n_window, N_VECTOR))
    for idx, alignment in enumerate(alignments):
        tgt_index, src_index, radius = alignment[1], alignment[0], n_window // 2
        # choice_src, choice_tgt = random.randint(0,1), random.randint(0,1)
        ns_src_idx = get_ns_index(src_index, src.size()[0])
        ns_tgt_idx = get_ns_index(tgt_index, tgt.size()[0])
        concated_vector[3*idx] = find_pairs(alignment,n_window,src,tgt,special_embedding)
        concated_vector[3*idx+1] = find_pairs((ns_src_idx, tgt_index), n_window, src, tgt, special_embedding)
        concated_vector[3*idx+2] = find_pairs((src_index, ns_tgt_idx), n_window, src, tgt, special_embedding)

    return concated_vector, len(alignments)


def find_ns_batch_pairs(word_alignments, n_window, src_all, tgt_all, special_embedding, sentence_length):
    alignments, batch_size = [set([tuple(map(int, x.split("-"))) for x in ali.strip().split()]) for ali in word_alignments], len(word_alignments)
    total_length = sum([len(alignment) for alignment in alignments])
    concated_vector, count = torch.zeros((3, total_length, 2 * n_window, N_VECTOR)),0
    for realization, src, tgt, src_li, tgt_li in zip(alignments, src_all, tgt_all, sentence_length[0], sentence_length[1]):
        sen_li = (src_li, tgt_li)
        for alignment in realization:

            tgt_index, src_index, radius = alignment[1], alignment[0], n_window // 2
            # choice_src, choice_tgt = random.randint(0,1), random.randint(0,1)
            ns_src_idx = get_ns_index(src_index, src_li)
            ns_tgt_idx = get_ns_index(tgt_index, tgt_li)
            concated_vector[0,count] = find_pairs(alignment,n_window,src,tgt,special_embedding,sen_li)
            concated_vector[1,count] = find_pairs((ns_src_idx, tgt_index), n_window, src, tgt, special_embedding,sen_li)
            concated_vector[2,count] = find_pairs((src_index, ns_tgt_idx), n_window, src, tgt, special_embedding,sen_li)
            count += 1
    return concated_vector
#
# E = "<NULL> the commissionners went on :"
# F = "je poursuis la lecture de les recommandations de les commissaires :"
# VocabE = Vocab('english20', special_embedding_tensor)
# VocabF = Vocab('french20', special_embedding_tensor)
# E = VocabE.get_sentence_tensor(E, 1)
# F = VocabF.get_sentence_tensor(F)
# A = "0-4 0-5 0-7 2-0 2-1 2-2 2-3 2-6 2-8 2-9 5-10"
# res = find_ns_pairs(A, 11, E, F, special_embedding_tensor)
# print(res)

def htanh(input):
    input[input > 1] = 1
    input[input < -1] = -1
    return input

class Htanh(nn.Module):
    def __init__(self):
        super().__init__()

    def forward(self, input):
        return htanh(input)

# def train(input_data, optimizer, model):
#     length = input_data[0].size()[0]
#     optimizer.zero_grad()
#
#     loss =  torch.max(torch.zeros((length,1)), 1 - model(input_data[0].flatten(1)) + model(input_data[1].flatten(1))) +\
#             torch.max(torch.zeros((length,1)), 1 - model(input_data[0].flatten(1)) + model(input_data[2].flatten(1)))
#     # loss += torch.max(torch.zeros((length,1),requires_grad=True), 1 - model(input_data[0].flatten(1)) + model(input_data[2].flatten(1)))
#     loss = loss.sum()
#
#     loss.backward()
#     optimizer.step()
#
#     return loss.item()/length

def train_sgd(input_ga, input_ns, optimizer, model):
    optimizer.zero_grad()

    loss = torch.max(torch.tensor(0.0, requires_grad=True), 1 - model(input_ga.flatten()) + model(input_ns.flatten()))

    loss.backward(retain_graph = True)
    optimizer.step()

    return loss.item()

class FFNN_word_alignment(nn.Module):

    def __init__(self, input_size, f1_size, f2_size):
        super(FFNN_word_alignment, self).__init__()
        self.fc1 = nn.Linear(input_size, f1_size)
        self.fc2 = nn.Linear(f1_size,f2_size)
        self.fc3 = nn.Linear(f2_size, 1)
        self.htanh = Htanh()

    def forward(self, x):
        output = self.fc1(x)

        output =  self.htanh(output)

        output = self.fc2(output)

        output = self.fc3(output)

        return output

def read_in_chunk(file_name, chunk_size):
    file_object = open(file_name, 'r')
    while True:
        data = [0 for _ in range(chunk_size)]
        for i in range(chunk_size):
            line = file_object.readline()
            if not line:
                data = data[:i]
                break
            data[i] = line
        if not data:
            break
        yield data

# gen_t = read_in_chunk('./data/text.txt',2)
# for li in gen_t:
#     print(li)
def get_chunk_sentence_tensor(sentence_list, vocab, N_VECTOR, source):
    n_sentence = len(sentence_list)
    # sentence_tensor = torch.zeros((n_sentence, max_len,N_VECTOR))
    li_length,li_tensor = [0 for _ in range(n_sentence)], [0 for _ in range(n_sentence)]
    for idx, sentence in enumerate(sentence_list):

        temp_tensor = vocab.get_sentence_tensor(sentence, source)
        li_length[idx], li_tensor[idx] = temp_tensor.size()[0], temp_tensor

    final_sentence_tensor = torch.nn.utils.rnn.pad_sequence(li_tensor, batch_first=True)

    return final_sentence_tensor, li_length

def mini_batch(batch_size, train_list, alignment_list, sentence_length, shuffle=True):
    length = len(train_list)
    batch_num = math.ceil(length / batch_size)
    index_array = list(range(length))
    if shuffle:
        np.random.shuffle(index_array)
    for i in range(batch_num):
        indices = index_array[i*batch_size:(i+1)*batch_size]
        curr_data = [train_list[index] for index in indices]
        align_data = [alignment_list[index] for index in indices]
        src_length,tgt_length = [sentence_length[0][index] for index in indices], \
                                [sentence_length[1][index] for index in indices]
        src_sent = [e[0] for e in curr_data]
        tgt_sent = [f[1] for f in curr_data]
        yield src_sent, tgt_sent, align_data,(src_length,tgt_length)

WINDOW_SIZE = 11
Learning_rate= 0.01
N_epoch = 1000
def main():
    model_FFNN = FFNN_word_alignment(2*WINDOW_SIZE*N_VECTOR, 120, 10)
    optimizer = optim.SGD(model_FFNN.parameters(), lr=Learning_rate)
    vocab_e, vocab_f = Vocab('english20', special_embedding_tensor), Vocab('french20', special_embedding_tensor)
    src = open('./data/hansards.f','r').readlines()[:37]
    tgt = open('./data/hansards.e','r').readlines()[:37]
    golden_alignment = open('./data/alignmentfe','r').readlines()[:37]
    shuffle_indeies = list(range(len(golden_alignment)))
    # np.random.shuffle(shuffle_indeies)
    # gen_src, gen_tgt = read_in_chunk('./data/hansards.e', 10000), read_in_chunk('./data/hansards.f', 10000)
    # gen_alignment = read_in_chunk('./data/alignment',10000)
    # count = 0
    # for src,tgt,ali in zip(gen_src,gen_tgt,gen_alignment):
    #     src_tensor, li_src_len = get_chunk_sentence_tensor(src, vocab_e, N_VECTOR, True)
    #     tgt_tensor, li_tgt_len= get_chunk_sentence_tensor(tgt, vocab_f, N_VECTOR, False)
    #     sentence_length = (li_src_len, li_tgt_len)
    #     train_data = [(src_tensor[i],tgt_tensor[i]) for i in range(src_tensor.size()[0])]
    #     beg = time.time()
    #     for src_sent, tgt_sent, ali_sent, sentence_length, in mini_batch(BATCH_SIZE, train_data, ali, sentence_length, shuffle=False):
    #         train_data = find_ns_batch_pairs(ali_sent,WINDOW_SIZE, src_sent, tgt_sent,special_embedding_tensor, sentence_length)
    #         train_ga, train_ns_src, train_ns_tgt = train_data[0], train_data[1], train_data[2]
    #         loss = train(train_data,optimizer,model_FFNN)
    #         count += 1
    #         if count % 10 == 0:
    #             end = time.time()
    #             print(end-beg)
    #             beg = time.time()



    loss, epoch, print_per = 0, 0, 1000
    loss_his, beg = [], time.time()
    # check_point = torch.load('./first_stage_state/state2198000.pth')
    # model_FFNN.load_state_dict(check_point)
    model_FFNN.train()
    while epoch < N_epoch:
        print_n_alignment = 0
        for i in range(len(src)):
            wv_src, wv_tgt = vocab_f.get_sentence_tensor(src[shuffle_indeies[i]], True), \
                             vocab_e.get_sentence_tensor(tgt[shuffle_indeies[i]],False)

            training_ns,n_alignment = find_ns_pairs(golden_alignment[shuffle_indeies[i]], WINDOW_SIZE, wv_src, wv_tgt,
                                                    special_embedding_tensor)
            training_ns = training_ns.view(n_alignment, 3, 2*WINDOW_SIZE, N_VECTOR)
            for idx in range(n_alignment):
                loss_f = train_sgd(training_ns[idx][0], training_ns[idx][1], optimizer,model_FFNN)
                loss_s = train_sgd(training_ns[idx][0], training_ns[idx][2], optimizer,model_FFNN)
                loss += (loss_f + loss_s) /2
            print_n_alignment += n_alignment
            if i % print_per == 36 and i != 0 :
                end = time.time()
                print("Loss per {} iters: {}".format(print_per, loss/print_n_alignment))
                print("Time per {} iters: {}".format(print_per, end-beg))
                torch.save(model_FFNN.state_dict(),'./first_stage_state/statefe37'  + '.pth')
                loss_his.append(loss)
                loss,print_n_alignment = 0,0
                beg = time.time()
        epoch += 1
        print('epoch:{}'.format(epoch))
    # pickle.dump(loss_his, open('./first_stage_state/loss','wb'))






if __name__ == '__main__':
    main()

# See PyCharm help at https://www.jetbrains.com/help/pycharm/

# This is a sample Python script. # Global Variables # Use GPU or not # Special Words # sentence_tensor[-1] = special_embedding_tensor[-1] # vocab = Vocab('english20', special_embedding_tensor) # abc = vocab.get_sentence_tensor('hello there', source = True) Args:
         tensor (torch.Tensor): tensor to be padded
         n_padding (int): padding nums
         padding_index (int): which chars to use:  (0-SOS 1-NULL 2-UNK 3-EOS)
         special_embedding (torch.Tensor): embedding tensor of 4 special chars
         front (bool): pad front or end
 
     Returns:
          padded_tensor (torch.Tensor) #fist half (including index) # second half (No index) # E <NULL> the commissionners went on : # F je poursuis la lecture de les recommandations de les commissaires : # 0-4 0-5 0-7 2-0 2-1 2-2 2-3 2-6 2-8 2-9 5-10 # choice_src, choice_tgt = random.randint(0,1), random.randint(0,1) # choice_src, choice_tgt = random.randint(0,1), random.randint(0,1) # # E = "<NULL> the commissionners went on :" # F = "je poursuis la lecture de les recommandations de les commissaires :" # VocabE = Vocab('english20', special_embedding_tensor) # VocabF = Vocab('french20', special_embedding_tensor) # E = VocabE.get_sentence_tensor(E, 1) # F = VocabF.get_sentence_tensor(F) # A = "0-4 0-5 0-7 2-0 2-1 2-2 2-3 2-6 2-8 2-9 5-10" # res = find_ns_pairs(A, 11, E, F, special_embedding_tensor) # print(res) # def train(input_data, optimizer, model): #     length = input_data[0].size()[0] #     optimizer.zero_grad() # #     loss =  torch.max(torch.zeros((length,1)), 1 - model(input_data[0].flatten(1)) + model(input_data[1].flatten(1))) +\ #             torch.max(torch.zeros((length,1)), 1 - model(input_data[0].flatten(1)) + model(input_data[2].flatten(1))) #     # loss += torch.max(torch.zeros((length,1),requires_grad=True), 1 - model(input_data[0].flatten(1)) + model(input_data[2].flatten(1))) #     loss = loss.sum() # #     loss.backward() #     optimizer.step() # #     return loss.item()/length # gen_t = read_in_chunk('./data/text.txt',2) # for li in gen_t: #     print(li) # sentence_tensor = torch.zeros((n_sentence, max_len,N_VECTOR)) # np.random.shuffle(shuffle_indeies) # gen_src, gen_tgt = read_in_chunk('./data/hansards.e', 10000), read_in_chunk('./data/hansards.f', 10000) # gen_alignment = read_in_chunk('./data/alignment',10000) # count = 0 # for src,tgt,ali in zip(gen_src,gen_tgt,gen_alignment): #     src_tensor, li_src_len = get_chunk_sentence_tensor(src, vocab_e, N_VECTOR, True) #     tgt_tensor, li_tgt_len= get_chunk_sentence_tensor(tgt, vocab_f, N_VECTOR, False) #     sentence_length = (li_src_len, li_tgt_len) #     train_data = [(src_tensor[i],tgt_tensor[i]) for i in range(src_tensor.size()[0])] #     beg = time.time() #     for src_sent, tgt_sent, ali_sent, sentence_length, in mini_batch(BATCH_SIZE, train_data, ali, sentence_length, shuffle=False): #         train_data = find_ns_batch_pairs(ali_sent,WINDOW_SIZE, src_sent, tgt_sent,special_embedding_tensor, sentence_length) #         train_ga, train_ns_src, train_ns_tgt = train_data[0], train_data[1], train_data[2] #         loss = train(train_data,optimizer,model_FFNN) #         count += 1 #         if count % 10 == 0: #             end = time.time() #             print(end-beg) #             beg = time.time() # check_point = torch.load('./first_stage_state/state2198000.pth') # model_FFNN.load_state_dict(check_point) # pickle.dump(loss_his, open('./first_stage_state/loss','wb')) # See PyCharm help at https://www.jetbrains.com/help/pycharm/