import torch
from torch import optim

from first_stage import FFNN_word_alignment
import torch.nn as nn
import utils.get_tensor as gt
import pickle
import numpy as np
import random
import time
model = FFNN_word_alignment(440,120,10)
model.load_state_dict(torch.load('./first_stage_state/statefe37.pth'))
model.eval()
special_embedding = pickle.load(open('./first_stage_state/special_embedding','rb'))
class distortion_model(nn.Module):
    def __init__(self):
        super().__init__()
        temp = torch.tensor([.10 for i in range(15)])
        # temp = torch.cat((torch.tensor([.05 for _ in range(5)]),torch.tensor([.3 for _ in range(5)]), torch.tensor([.05 for _ in range(5)])))
        self.li_paras_nomral = torch.nn.parameter.Parameter(temp)
        self.li_paras_outlier = torch.nn.parameter.Parameter(torch.tensor((.1,.1)))

    def find_index(self,n_col,index):
        # even = [1/n_row for _ in range(n_row)]
        # return torch.tensor(even).unsqueeze(1)
        window_list = torch.zeros((n_col))
        for i in range(n_col):
            if 7 - index + i < 0:
                window_list[i] = self.li_paras_outlier[0]
            elif 7 - index + i > 14:
                window_list[i] = self.li_paras_outlier[1]
            else:
                window_list[i] = self.li_paras_nomral[7-index+i]
        return window_list
    #
    def forward(self, lex_score, alignments, jump_dist):
        if not alignments:
            return 'NEXT'
        s_nn = 1
        distort_params = self.get_distort_params(jump_dist)
        init_state = alignments.pop(0)
        s_nn *= torch.tensor(lex_score[init_state[0], init_state[1]])
        for distort, alignment in zip(distort_params,alignments):
            src_index, tgt_index = alignment[0], alignment[1]
            s_nn *= lex_score[src_index, tgt_index] * distort

        return s_nn
    # def forward(self, lex_score, alignment, jump_dist):
    #     distort_params = self.get_distort_word_params(jump_dist)
    #     src_index, tgt_index = alignment[0], alignment[1]
    #     return lex_score[src_index,tgt_index] * distort_params

    def get_jump_dist(self, alignment):
        sorted_index = [0] + [x[1] for x in alignment]
        difference = np.diff(sorted_index).astype(int)
        jump_dist = list(difference)
        return jump_dist
    def check_params(self,num_normal, num_out):
        normal = self.li_paras_nomral
        outlier = self.li_paras_outlier
        normal.data[normal.data < 0] = num_normal
        outlier.data[outlier.data < 0] = num_out
        outlier.data[outlier.data > 1] = num_out
        normal.data[normal.data > 1] = num_normal

    def decode_naive(self, lex_score):
        def compute_jump(align_list):
            numpy_align = align_list.numpy()
            jump_dist = np.diff(numpy_align).astype(int)
            return jump_dist

        n_row, n_col = lex_score.size()[0], lex_score.size()[1]
        path = torch.zeros((n_row, n_col))
        align_list = torch.zeros((n_row))
        for r_idx in range(n_row):
            if r_idx == 0:
                # dist_window = self.find_index(n_col, 0)
                # lex_score[0] = lex_score[0] * dist_window
                continue
            temp_tensor = []
            for idx in range(n_col):
                dist_window = self.find_index(n_col, idx)
                # temp_tensor[idx] = (lex_score[r_idx - 1, idx] * dist_window * lex_score[r_idx, idx]).squeeze(1)
                temp_tensor.append((lex_score[r_idx - 1, idx] * dist_window * lex_score[r_idx]).unsqueeze(0))
                # temp_tensor.append((lex_score[r_idx - 1, idx]  * dist_window * lex_score[r_idx]).unsqueeze(0))
            temp_tensor = torch.cat(temp_tensor, dim = 0)
            max_row = torch.argmax(temp_tensor, dim=0)
            max_tensor = torch.max(temp_tensor, dim=0)
            lex_score[r_idx] = max_tensor.values
            path[r_idx - 1] = max_row
        final_align = torch.argmax(lex_score[-1])
        align_list[-1] = final_align
        for i in range(n_row-2,-1,-1):
            align_list[i] = path[i,final_align]
        jump_dist = compute_jump(align_list)
        align_list = [(i+1,int(align_list[i])) for i in range(len(align_list))]

        return align_list,jump_dist

    def decode(self, lex_score, ga):
        def compute_jump(align_list):
            numpy_align = align_list.numpy()
            jump_dist = np.diff(numpy_align).astype(int)
            return jump_dist

        def get_ga(ga):
            ga_li = list(set([x[0] for x in ga]))
            if 0 in ga_li:
                ga_li.remove(0)
            ga_li.sort()
            return ga_li
        ga_li = get_ga(ga)
        n_row, n_col = len(ga_li), lex_score.size()[1]
        lex_slice = lex_score[ga_li]
        path = torch.zeros((n_row, n_col))
        align_list = torch.zeros((n_row))
        for r_idx in range(n_row):
            if r_idx == 0:
                # dist_window = self.find_index(n_col, 0)
                # lex_score[0] = lex_score[0] * dist_window
                continue
            temp_tensor = []
            for idx in range(n_col):
                dist_window = self.find_index(n_col, idx)
                # temp_tensor[idx] = (lex_score[r_idx - 1, idx] * dist_window * lex_score[r_idx, idx]).squeeze(1)
                temp_tensor.append((lex_slice[r_idx - 1, idx] * dist_window * lex_slice[r_idx]).unsqueeze(0))
                # temp_tensor.append((lex_score[r_idx - 1, idx]  * dist_window * lex_score[r_idx]).unsqueeze(0))
            temp_tensor = torch.cat(temp_tensor, dim=0)
            max_row = torch.argmax(temp_tensor, dim=0)
            max_tensor = torch.max(temp_tensor, dim=0)
            lex_slice[r_idx] = max_tensor.values
            path[r_idx - 1] = max_row
        final_align = torch.argmax(lex_slice[-1])
        align_list[-1] = final_align
        for i in range(n_row - 2, -1, -1):
            align_list[i] = path[i, final_align]
        jump_dist = compute_jump(align_list)
        align_list = [(ga_li[i], int(align_list[i])) for i in range(len(align_list))]

        return align_list, jump_dist

    def choose_golden_alignment(self, golden_alignments):
        rand_chosen_index = [[] for _ in range(len(golden_alignments))]
        jump_dist = [[] for _ in range(len(golden_alignments))]
        for idx, alignment in enumerate(golden_alignments):
            align = list(set([tuple(map(int, x.split("-"))) for x in alignment.strip().split()]))
            src_index = set([x[0] for x in align]) - set([0])
            src_map = {x:True for x in src_index}
            random.shuffle(align)

            for ali in align:
                temp_index = ali[0]
                if temp_index == 0:
                    continue
                if src_map[temp_index] == True:
                    src_map[temp_index] = False
                    rand_chosen_index[idx].append(ali)

            rand_chosen_index[idx].sort()

            sorted_index = [x[1] for x in rand_chosen_index[idx]]
            difference = np.diff(sorted_index).astype(int)
            jump_dist[idx] = list(difference)



        return rand_chosen_index, jump_dist

    def get_distort_params(self, just_dist):
        li_distort_params = torch.zeros(len(just_dist))
        for idx,dist in enumerate(just_dist):
            if dist < -7:
                li_distort_params[idx] = self.li_paras_outlier[0]
            elif dist > 7:
                li_distort_params[idx] = self.li_paras_outlier[1]
            else:
                li_distort_params[idx] = self.li_paras_nomral[7+dist]

        return li_distort_params

    def get_distort_word_params(self, dist):
        if dist < -7:
            return self.li_paras_outlier[0]
        elif dist > 7:
            return self.li_paras_outlier[1]
        else:
            return self.li_paras_nomral[7+dist]




def main():
    Learning_rate =0.001
    N_sentence = 100000
    N_window = 11
    #each of them is very complex , but the link between the two is even more complex which makes the whole situation for most people understandably confusing .
    #chacun en lui - même est très complexe et le lien entre les deux le est encore davantage de sorte que pour beaucoup la situation présente est confuse .
    dist_model = distortion_model()
    # dist_model.load_state_dict(torch.load('./second_stage/state1.pth'))
    # dist_model.eval()
    optimizer = optim.SGD(dist_model.parameters(), lr=Learning_rate)
    vocab_e, vocab_f = gt.Vocab('english20', special_embedding), gt.Vocab('french20', special_embedding)
    src = open('./data/hansards.f','r').readlines()
    tgt = open('./data/hansards.e','r').readlines()
    golden_alignment = open('./data/alignmentfe','r').readlines()
    golden_alignment, jump_ga = dist_model.choose_golden_alignment(golden_alignment)
    shuffle_indeies = list(range(len(golden_alignment)))
    # shuffle_indeies = list(range(661,700))
    np.random.shuffle(shuffle_indeies)
    state = 0
    # # vocab_e, vocab_f = gt.Vocab('english20', special_embedding), gt.Vocab('french20', special_embedding)
    # # lex_score = gt.get_lex_score(e, f, special_embedding, model, (vocab_e, vocab_f),11)
    # # print(lex_score)
    print('begin sgd')
    for i in range(N_sentence):
        optimizer.zero_grad()
        total_loss, beg= 0, time.time()

        src_train,tgt_train, ga_train, jump_dist_ga_train = src[shuffle_indeies[i]], tgt[shuffle_indeies[i]], \
                                                            golden_alignment[shuffle_indeies[i]], jump_ga[shuffle_indeies[i]]
        # src_train, tgt_train, ga_train, jump_dist_ga_train = src[0], tgt[0], golden_alignment[
        #     0], jump_ga[0]
        lex_score = gt.get_lex_score(src_train, tgt_train, special_embedding, model, (vocab_f,vocab_e), N_window)
        lex_score = lex_score.detach()

        align_infer, jump_dist_infer = dist_model.decode(lex_score.clone(), ga_train)
        # align_ga, jump_ga = dist_model.choose_golden_alignment(ga_train)


        s_nn_infer = dist_model(lex_score.detach().numpy(), align_infer, jump_dist_infer)
        s_nn_ga = dist_model(lex_score.detach().numpy(), ga_train, jump_dist_ga_train)
        # if s_nn_infer.isnan() or s_nn_ga.isnan() or s_nn_infer.isinf() or s_nn_ga.isinf():
        #     print('nan detected, i: {}'.format(i))
        #     print('s_nn_ga: {}'.format(s_nn_ga))
        #     print('s_nn_infer: {}'.format(s_nn_infer))
        #
        #     if i % 50 == 0 and i != 0:
        #         torch.save(dist_model.state_dict(), './second_stage/state' + str(state) + '.pth')
        #         print('save to {}'.format(state))
        #         state = (state +1) % 5
        #     continue
        # if s_nn_infer / s_nn_ga >= 10:
        #     print("TOO LARGE")
        #     if i % 50 == 0 and i != 0:
        #         torch.save(dist_model.state_dict(), './second_stage/state' + str(state) + '.pth')
        #         print('save to {}'.format(state))
        #         state = (state +1) % 5
        #     continue
        if s_nn_ga =='NEXT' or s_nn_infer == 'NEXT':
            print('no alignment, continue')
            continue
        # denom = torch.max(s_nn_ga,s_nn_infer).detach()
        # s_nn_infer, s_nn_ga = s_nn_infer / denom, s_nn_ga/ denom
        loss = torch.max(torch.tensor(0.00,requires_grad=True), 1 - s_nn_ga + s_nn_infer)
        # total_loss += loss
        loss.backward()
        # torch.nn.utils.clip_grad_norm_(dist_model.parameters(), 5)
        optimizer.step()
        dist_model.check_params(.1,.1)
        # #
        # ga_src = [a[0] for a in ga_train]
        # infer_train = []
        # for a_infer in align_infer:
        #     if a_infer[0] in ga_src:
        #         infer_train.append(a_infer)
        # jump_dist_infer_train = dist_model.get_jump_dist(infer_train)
        # for j in range(len(jump_dist_infer_train)):
        #     optimizer.zero_grad()
        #     s_nn_infer = dist_model(lex_score.detach().numpy(), infer_train[j], jump_dist_infer_train[j])
        #     s_nn_ga = dist_model(lex_score.detach().numpy(), ga_train[j], jump_dist_ga_train[j])
        #
        #     loss = torch.max(torch.tensor(0.00, requires_grad=True), 1- s_nn_ga + s_nn_infer)
        #     # total_loss += loss
        #     loss.backward(retain_graph=True)
        #     optimizer.step()

        if i % 50 == 0:
            print(i)
            for param in dist_model.parameters():
                print(param)
            print(loss)

        if i != 0 and i % 100 == 0:
            end = time.time()
            print(loss)
            print(end-beg)
            torch.save(dist_model.state_dict(), './second_stage/statefe37' +   '.pth')
            total_loss, beg = 0, time.time()
            state = (state + 1) % 5
    # #
    #

    # for i in range(len(src)):
    #     lex_score = gt.get_lex_score(src[i], tgt[i], special_embedding, model, (vocab_e, vocab_f), N_window)
    #     align_infer, jump_dist_infer = dist_model.decode(lex_score.clone())
    # a = distortion_model(model)
    # a.decode(lex_score)
    # a.choose_golden_alignment(golden_alignment)
    # print(golden_alignment)

if __name__ == '__main__':
    main()
