import torch
import torch.nn as nn
import src.utils as utils
from torch.utils.data import Dataset, DataLoader
import os
SAVE_INFO_PATH = "saveInfo"


class PassageEncoder(nn.Module):
    def __init__(self, input_size, output_size):
        '''
        initial encoder
        :param input_size: input size of lstm
        :param output_size:  output size of lstm
        '''
        super(PassageEncoder,self).__init__()
        self.bilstm = nn.LSTM(input_size=input_size,hidden_size=output_size,
                              dropout=0.33,bidirectional=True,num_layers=2)

    def forward(self, feature, input_length, hidden=None):
        '''
        :param feature: torch tensor of shape(batch_size, )
        :param input_length: length of input, torch tensor of dtype int
        :param hidden: initial hidden state, torch tensor shape(output_size)
        :return: the output of lstm of shape(batch_size, length, output_size)
        '''
        # feature = feature.float()
        packed = nn.utils.rnn.pack_padded_sequence(feature, batch_first=True, lengths=input_length)
        lstm_out,hidden = self.bilstm(packed, hidden)
        out = nn.utils.rnn.pad_packed_sequence(lstm_out,batch_first=True)
        return out


def load_model(path, hidden_size=10, version="last"):
    model = PassageEncoder(hidden_size,hidden_size)
    if not os.path.exists(path):
        print("not exist")
        return model, 0
    else:
        saveinfo_path = os.path.join(path, SAVE_INFO_PATH)
        if not os.path.exists(saveinfo_path):
            return model, 0
        saveinfo = torch.load(saveinfo_path)[version]
        model.load_state_dict(torch.load(saveinfo)["model_state_dict"])
        cur_epoch = torch.load(saveinfo)["epoch"]
        return model, cur_epoch


def aggregate_title_content(self, title, content):
    '''
    aggregate title and content to form passage vector
    here just cat content after title
    :param title: torch tensor of shape (title_len, d)
    :param content: torch tensor of shape (content_len, d)
    :return: aggregated passage tensor of shape(title_len+content_len, d)
    '''
    passage = torch.cat([title,content],dim=0)
    return passage


def loss_fn(encoded_passage,label, entity_candidate, device):
    '''
    triplet loss with anchor as encoded_passage, positive as label, negtive as entity_candidate
    :param encoded_passage: output of lstm with shape (batch_size, seq_length, num_layer, output_size)
    :param label: a batch_size length list of dict contain the passage label of 3 entity
    :param entity_candidate: a batch_size length list of dict contain negative entity generated by tf-idf
    :return: loss
    '''
    total_loss = torch.zeros(1,dtype=torch.float).to(device=device)
    batch_size = encoded_passage.shape[0]
    for i_batch in range(batch_size):
        for entity in label[i_batch]:
            for candidate in entity_candidate[i_batch]:
                total_loss += nn.functional.triplet_margin_loss(
                    anchor=encoded_passage[i_batch][-1,:], positive=entity["encoding"].squeeze(),
                    negative=candidate["encoding"].squeeze()
                )
    return total_loss/batch_size


def getdata(dataloader):
    return dataloader


def train(dataloader, model, optimizer, device ,num_epcoh=1, save_path=None, cur_epoch=0):
    model.train()
    for step in range(num_epcoh):
        print("epoch ",cur_epoch)
        # for idx, data in enumerate(dataloader):
        optimizer.zero_grad()
        data = getdata(dataloader)
        passage, length, label, candidate = data["passage"], data["length"], data["entity"], data["candidate"]
        passage = passage.float().to(device=device)
        encoded_passage, _ = model(passage, length)
        shape = encoded_passage.shape
        encoded_passage = encoded_passage.view(shape[0],shape[1],2,-1)
        loss = loss_fn(encoded_passage, label, candidate, device=device)
        loss.backward()
        print("loss at epoch {}: {}".format(cur_epoch,loss.item()))
        optimizer.step()
        if save_path:
            saveinfo_path = os.path.join(save_path,SAVE_INFO_PATH)
            check_point_path = os.path.join(save_path,"model_{}".format(cur_epoch))
            torch.save({"last":check_point_path},saveinfo_path)
            torch.save({"model_state_dict": model.state_dict(),"epoch": cur_epoch},check_point_path)


def eval(dataloader, model):
    model.eval()
    data = getdata(dataloader)
    passage, length, label, candidate = data["passage"], data["length"], data["entity"], data["candidate"]
    encoded_passage, _ = model(passage, length)
    shape = encoded_passage.shape
    encoded_passage = encoded_passage.view(shape[0], shape[1], 2, -1)
    print(encoded_passage.shape)
    # loss = loss_fn(encoded_passage, label, candidate)
    batch_top3 = []
    for i_batch in range(encoded_passage.shape[0]):
        hidden = encoded_passage[i_batch,-1,:,:]
        bi_passage = torch.sum(hidden, dim=0)
        _, top3 = utils.top3entity(bi_passage,candidate[i_batch])
        batch_top3.append({"passage": label[i_batch],
                           "y_hat":top3})
        print(i_batch)
        print("label: ", [x["entity"] for x in label[i_batch]])
        print("top3: ", top3)
    f1_score = utils.F1Score(batch_top3)
    return f1_score


def stackpassage(passages):
    '''

    :param passages:
    :return: passages sorted by length using device
    '''
    passages.sort(key=lambda x:x["length"],reverse=True)
    dataset = {}
    for key in passages[0]:
        dataset[key] = []
        for p in passages:
            dataset[key].append(p[key])
    dataset["passage"] = nn.utils.rnn.pad_sequence(dataset["passage"],batch_first=True)
    return dataset


def test():
    nd = 10
    passages = []
    for i in range(3):
        passage = {}
        # passage["title"] = torch.randn(size=(1+i,nd))
        # passage["content"] = torch.randn(size=(2+i,nd))
        passage["passage"] = torch.randn(size=(2+i, nd))
        passage["length"] = 2+i
        passage["entity"] = [{"entity":"a", "encoding": torch.randn(size=(nd,))}]
        passage["candidate"] = [{"entity":"a", "encoding": torch.randn(size=(nd,))},
                                {"entity": "c", "encoding": torch.randn(size=(nd,))}]
        passages.append(passage)
    # dataset = PassageDataset(passages)
    # dataloader = DataLoader(dataset,2, True,collate_fn=collate_wrapper, pin_memory=True)

    dataloader = stackpassage(passages)
    model,epoch = load_model(SAVE_INFO_PATH)
    train(dataloader,model,cur_epoch=epoch)


if __name__ == "__main__":
    test()
    # load_model(SAVE_INFO_PATH,10)
