import numpy as np
import torch
from torch.autograd import Variable
from torch.utils.data import DataLoader
from torch import optim, nn
import data_preprocess
import os

torch.manual_seed(1)

os.environ['CUDA_VISIBLE_DEVICES'] = '0'
use_cuda = torch.cuda.is_available()

word2index, index2word, tag2index, index2tag = data_preprocess.get_dic()
test_x_cut, test_y_cut, test_mask, test_x_len, test_x_cut_word, test_x_fenge = data_preprocess.getTest_xy(
    './data/test_data')
testDataSet = data_preprocess.TextDataSet(test_x_cut, test_y_cut, test_mask)

testDataLoader = DataLoader(testDataSet, batch_size=16, shuffle=False)

MAXLEN = 100
vcab_size = len(word2index)
emb_dim = 128
hidden_dim = 256
num_epoches = 20
batch_size = 16


class BILSTM_CRF(nn.Module):
    def __init__(self,vcab_size,tag2index,emb_dim,hidden_dim,batch_size):
        super(BILSTM_CRF,self).__init__()
        self.vcab_size=vcab_size
        self.tag2index=tag2index
        self.num_tags=len(tag2index)
        self.emb_dim=emb_dim
        self.hidden_dim=hidden_dim
        self.batch_size=batch_size
        self.use_cuda=torch.cuda.is_available()
        self.embed=nn.Embedding(num_embeddings=vcab_size,embedding_dim=emb_dim)#b,100,128
        #->100,b,128
        self.bilstm=nn.LSTM(input_size=emb_dim,hidden_size=hidden_dim,num_layers=1,bidirectional=True,dropout=0.1)#100,b,256*2
        self.conv1 = nn.Sequential(
            #b,1,100,128
            nn.Conv2d(1, 128, (1,emb_dim),padding=0),  # b,128,100,1
            nn.BatchNorm2d(128),
            nn.ReLU(True),
        )
        self.conv2 = nn.Sequential(
            nn.Conv2d(1, 128, (3,emb_dim+2), padding=1),  # b,128,100,1
            nn.BatchNorm2d(128),
            nn.ReLU(True),
        )
        self.conv3 = nn.Sequential(
            nn.Conv2d(1, 128, (5,emb_dim+4), padding=2),  # b,128,100,1
            nn.BatchNorm2d(128),
            nn.ReLU(True),
        )
        #b,128*3,100,1->100,b,128*3
        self.linear1 = nn.Linear(hidden_dim * 2+128*3,hidden_dim)
        self.drop=nn.Dropout(0.2)
        self.classfy=nn.Linear(hidden_dim,self.num_tags)#100*b,10
        #->100,b,10
        # init transitions
        self.start_transitions = nn.Parameter(torch.Tensor(self.num_tags))#
        self.end_transitions = nn.Parameter(torch.Tensor(self.num_tags))#
        self.transitions = nn.Parameter(torch.Tensor(self.num_tags, self.num_tags))#
        nn.init.uniform(self.start_transitions, -0.1, 0.1)
        nn.init.uniform(self.end_transitions, -0.1, 0.1)
        nn.init.uniform(self.transitions, -0.1, 0.1)

    def init_hidden(self,batch_size):#
        h_h=Variable(torch.randn(2,batch_size,self.hidden_dim))
        h_c=Variable(torch.randn(2,batch_size,self.hidden_dim))
        if use_cuda:
            h_h=h_h.cuda()
            h_c=h_c.cuda()
        return (h_h,h_c)

    def get_bilstm_out(self,x):#
        batch_size = x.size(0)
        emb=self.embed(x)

        #cnn输出
        emb_cnn=emb.unsqueeze(1)
        cnn1=self.conv1(emb_cnn)
        cnn2=self.conv2(emb_cnn)
        cnn3=self.conv3(emb_cnn)
        cnn_cat=torch.cat((cnn1,cnn2,cnn3),1)
        cnn_out=cnn_cat.squeeze().permute(2,0,1)#100,b,128*3

        emb_rnn=emb.permute(1,0,2)
        init_hidden=self.init_hidden(batch_size)
        lstm_out,hidden=self.bilstm(emb_rnn,init_hidden)

        cat_out=torch.cat((cnn_out,lstm_out),2)#100,b,128*3+256*2
        s,b,h=cat_out.size()
        cat_out=cat_out.view(s*b,h)
        cat_out=self.linear1(cat_out)
        cat_out=self.drop(cat_out)
        cat_out=self.classfy(cat_out)
        cat_out=cat_out.view(s,b,-1)
        # out=out.permute(1,0,2)
        return cat_out

    def _log_sum_exp(self,tensor,dim):
        # Find the max value along `dim`
        offset, _ = tensor.max(dim)#b,m
        # Make offset broadcastable
        broadcast_offset = offset.unsqueeze(dim)#b,1,m
        # Perform log-sum-exp safely
        safe_log_sum_exp = torch.log(torch.sum(torch.exp(tensor - broadcast_offset), dim))#b,m
        # Add offset back
        return offset + safe_log_sum_exp

    def get_all_score(self,emissions,mask):#
        # emissions: (seq_length, batch_size, num_tags)
        # mask: (batch_size,seq_length)
        seq_length = emissions.size(0)
        mask = mask.permute(1,0).contiguous().float()

        log_prob = self.start_transitions.view(1, -1) + emissions[0]  # b,m,

        for i in range(1, seq_length):
            broadcast_log_prob = log_prob.unsqueeze(2)  # b,m,1
            broadcast_transitions = self.transitions.unsqueeze(0)  #1,m,m
            broadcast_emissions = emissions[i].unsqueeze(1)  # b,1,m

            score = broadcast_log_prob + broadcast_transitions \
                    + broadcast_emissions  # b,m,m

            score = self._log_sum_exp(score, 1)  # 

            log_prob = score * mask[i].unsqueeze(1) + log_prob * (1. - mask[i]).unsqueeze(
                1)  # 

        # End transition score
        log_prob += self.end_transitions.view(1, -1)
        # Sum (log-sum-exp) over all possible tags
        return self._log_sum_exp(log_prob, 1)  # (batch_size,)

    def get_real_score(self,emissions,mask,tags):#
        # emissions: (seq_length, batch_size, num_tags)
        # tags: (batch_size,seq_length)
        # mask: (batch_size,seq_length)
        seq_length = emissions.size(0)#s
        mask = mask.permute(1,0).contiguous().float()
        tags=tags.permute(1,0).contiguous()

        # Start transition score
        llh = self.start_transitions[tags[0]]  # (batch_size,),T(start->firstTag)

        for i in range(seq_length - 1):
            cur_tag, next_tag = tags[i], tags[i+1]
            # Emission score for current tag
            llh += emissions[i].gather(1, cur_tag.view(-1, 1)).squeeze(1) * mask[i]#(b,1)->b->b*mask，
            # Transition score to next tag
            transition_score = self.transitions[cur_tag.data, next_tag.data]#
            # Only add transition score if the next tag is not masked (mask == 1)
            llh += transition_score * mask[i+1]#

        # Find last tag index
        last_tag_indices = mask.long().sum(0) - 1  # (batch_size,)
        last_tags = tags.gather(0, last_tag_indices.view(1, -1)).squeeze(0)#b

        # End transition score
        llh += self.end_transitions[last_tags]#
        # Emission score for the last tag, if mask is valid (mask == 1)
        llh += emissions[-1].gather(1, last_tags.view(-1, 1)).squeeze(1) * mask[-1]#

        return llh#b

    def neg_log_likelihood(self,feats,tags,mask):
        #feats:  
        batch_size=feats.size(1)
        all_score=self.get_all_score(feats,mask)#
        real_score=self.get_real_score(feats,mask,tags)#
        loss=(all_score.view(batch_size,1)-real_score.view(batch_size,1)).sum()/batch_size
        return loss #

    def viterbi_decode(self, emissions,mask):
        # emissions: (seq_length, batch_size, num_tags)
        # mask: (batch_size,seq_length)
        seq_length=emissions.size(0)
        batch_size=emissions.size(1)
        num_tags=emissions.size(2)
        length_mask = torch.sum(mask, dim=1).view(batch_size, 1).long()  #
        mask=mask.permute(1,0).contiguous().float()#s,b

        viterbi_history=[]
        viterbi_score = self.start_transitions.view(1, -1) + emissions[0]  # 

        for i in range(1, seq_length):
            broadcast_viterbi_score = viterbi_score.unsqueeze(2)  # b,m,1
            broadcast_transitions = self.transitions.unsqueeze(0)  #1,m,m
            broadcast_emissions = emissions[i].unsqueeze(1)  # b,1,m

            score = broadcast_viterbi_score + broadcast_transitions \
                    + broadcast_emissions  # b,m,m

            best_score,best_path = torch.max(score, 1)  # 
            viterbi_history.append(best_path*mask[i].long().unsqueeze(1))#
            viterbi_score = best_score * mask[i].unsqueeze(1) + viterbi_score * (1. - mask[i]).unsqueeze(
                1)  # 
        viterbi_score+=self.end_transitions.view(1,-1)#b,m
        best_score,last_path=torch.max(viterbi_score,1)#b
        last_path=last_path.view(-1,1)#b,1
        last_position = (length_mask.contiguous().view(batch_size, 1, 1).expand(batch_size, 1, num_tags) - 1).contiguous()  #
        pad_zero = Variable(torch.zeros(batch_size, num_tags)).long()
        if use_cuda:
            pad_zero = pad_zero.cuda()
        viterbi_history.append(pad_zero)#(s-1,b,m)->(s,b,m)
        viterbi_history = torch.cat(viterbi_history).view(-1, batch_size, num_tags)  # s,b,m
        insert_last = last_path.view(batch_size, 1, 1).expand(batch_size, 1, num_tags) #
        viterbi_history = viterbi_history.transpose(1, 0).contiguous()  # b,s,m
        viterbi_history.scatter_(1, last_position, insert_last)  # 
        viterbi_history = viterbi_history.transpose(1, 0).contiguous()  # s,b,m
        decode_idx = Variable(torch.LongTensor(seq_length, batch_size))#
        if use_cuda:
            decode_idx = decode_idx.cuda()
        # decode_idx[-1] = 0
        for idx in range(len(viterbi_history)-2,-1,-1):
            last_path=torch.gather(viterbi_history[idx],1,last_path)
            decode_idx[idx]=last_path.data
        decode_idx=decode_idx.transpose(1,0)#b,s
        return decode_idx

    def forward(self, feats,mask):
        #feats    #bilstm的输出#100.b.10
        best_path=self.viterbi_decode(feats,mask)#最佳路径b,s
        return best_path

if use_cuda:
    model = BILSTM_CRF(vcab_size, tag2index, emb_dim, hidden_dim, batch_size).cuda()
else:
    model = BILSTM_CRF(vcab_size, tag2index, emb_dim, hidden_dim, batch_size)

model.load_state_dict(torch.load('./model/best_model.pth'))

# model.eval()
test_loss = 0
test_acc = 0
batch_len_all = 0
prepath_all=[]#
for i, data in enumerate(testDataLoader):
    x, y, mask = data
    batch_len = len(x)
    batch_len_all += batch_len
    if use_cuda:
        x = Variable(x, volatile=True).cuda()
        y = Variable(y, volatile=True).cuda()
        mask = Variable(mask, volatile=True).cuda()
    else:
        x = Variable(x, volatile=True)
        y = Variable(y, volatile=True)
        mask = Variable(mask, volatile=True)
    feats = model.get_bilstm_out(x)
    loss = model.neg_log_likelihood(feats, y, mask)
    test_loss += loss.data[0]
    prepath = model(feats, mask)  # b,s
    prepath_all.append(prepath)
    pre_y = prepath.masked_select(mask)
    true_y = y.masked_select(mask)
    acc_num = (pre_y == true_y).data.sum()
    acc_pro = float(acc_num) / len(pre_y)
    test_acc += acc_pro
print('test loss is:{:.6f},test acc is:{:.6f}'.format(test_loss / (len(testDataLoader)),test_acc / (len(testDataLoader))))

#写入结果文件
prepath_all=torch.cat(prepath_all).data
data_preprocess.write_result_to_file('./data/result_data',prepath_all,test_x_len,test_x_cut_word,test_x_fenge)
