import json
import pickle
import random

import dgl

import numpy as np
import torch

NODE_TYPE = {"entity": 0, "root": 1, "relation": 2}


def write_txt(batch, seqs, w_file, args):
    # converting the prediction to real text.
    ret = []
    for b, seq in enumerate(seqs):
        txt = []
        for token in seq:
            # copy the entity
            if token >= len(args.text_vocab):
                ent_text = batch["raw_ent_text"][b][
                    token - len(args.text_vocab)
                ]
                ent_text = filter(lambda x: x != "<PAD>", ent_text)
                txt.extend(ent_text)
            else:
                if int(token) not in [
                    args.text_vocab(x) for x in ["<PAD>", "<BOS>", "<EOS>"]
                ]:
                    txt.append(args.text_vocab(int(token)))
            if int(token) == args.text_vocab("<EOS>"):
                break
        w_file.write(" ".join([str(x) for x in txt]) + "\n")
        ret.append([" ".join([str(x) for x in txt])])
    return ret


def replace_ent(x, ent, V):
    # replace the entity
    mask = x >= V
    if mask.sum() == 0:
        return x
    nz = mask.nonzero()
    fill_ent = ent[nz, x[mask] - V]
    x = x.masked_scatter(mask, fill_ent)
    return x


def len2mask(lens, device):
    max_len = max(lens)
    mask = (
        torch.arange(max_len, device=device)
        .unsqueeze(0)
        .expand(len(lens), max_len)
    )
    mask = mask >= torch.LongTensor(lens).to(mask).unsqueeze(1)
    return mask


def pad(var_len_list, out_type="list", flatten=False):
    if flatten:
        lens = [len(x) for x in var_len_list]
        var_len_list = sum(var_len_list, [])
    max_len = max([len(x) for x in var_len_list])
    if out_type == "list":
        if flatten:
            return [
                x + ["<PAD>"] * (max_len - len(x)) for x in var_len_list
            ], lens
        else:
            return [x + ["<PAD>"] * (max_len - len(x)) for x in var_len_list]
    if out_type == "tensor":
        if flatten:
            return (
                torch.stack(
                    [
                        torch.cat(
                            [
                                x,
                                torch.zeros(
                                    [max_len - len(x)] + list(x.shape[1:])
                                ).type_as(x),
                            ],
                            0,
                        )
                        for x in var_len_list
                    ],
                    0,
                ),
                lens,
            )
        else:
            return torch.stack(
                [
                    torch.cat(
                        [
                            x,
                            torch.zeros(
                                [max_len - len(x)] + list(x.shape[1:])
                            ).type_as(x),
                        ],
                        0,
                    )
                    for x in var_len_list
                ],
                0,
            )


class Vocab(object):
    def __init__(
        self,
        max_vocab=2**31,
        min_freq=-1,
        sp=["<PAD>", "<BOS>", "<EOS>", "<UNK>"],
    ):
        self.i2s = []
        self.s2i = {}
        self.wf = {}
        self.max_vocab, self.min_freq, self.sp = max_vocab, min_freq, sp

    def __len__(self):
        return len(self.i2s)

    def __str__(self):
        return "Total " + str(len(self.i2s)) + str(self.i2s[:10])

    def update(self, token):
        if isinstance(token, list):
            for t in token:
                self.update(t)
        else:
            self.wf[token] = self.wf.get(token, 0) + 1

    def build(self):
        self.i2s.extend(self.sp)
        sort_kv = sorted(self.wf.items(), key=lambda x: x[1], reverse=True)
        for k, v in sort_kv:
            if (
                len(self.i2s) < self.max_vocab
                and v >= self.min_freq
                and k not in self.sp
            ):
                self.i2s.append(k)
        self.s2i.update(list(zip(self.i2s, range(len(self.i2s)))))

    def __call__(self, x):
        if isinstance(x, int):
            return self.i2s[x]
        else:
            return self.s2i.get(x, self.s2i["<UNK>"])

    def save(self, fname):
        pass

    def load(self, fname):
        pass


def at_least(x):
    # handling the illegal data
    if len(x) == 0:
        return ["<UNK>"]
    else:
        return x


class Example(object):
    def __init__(self, title, ent_text, ent_type, rel, text):
        # one object corresponds to a data sample
        self.raw_title = title.split()
        self.raw_ent_text = [at_least(x.split()) for x in ent_text]
        assert min([len(x) for x in self.raw_ent_text]) > 0, str(
            self.raw_ent_text
        )
        self.raw_ent_type = ent_type.split()  # <method> .. <>
        self.raw_rel = []
        for r in rel:
            rel_list = r.split()
            for i in range(len(rel_list)):
                if (
                    i > 0
                    and i < len(rel_list) - 1
                    and rel_list[i - 1] == "--"
                    and rel_list[i] != rel_list[i].lower()
                    and rel_list[i + 1] == "--"
                ):
                    self.raw_rel.append(
                        [
                            rel_list[: i - 1],
                            rel_list[i - 1] + rel_list[i] + rel_list[i + 1],
                            rel_list[i + 2 :],
                        ]
                    )
                    break
        self.raw_text = text.split()
        self.graph = self.build_graph()

    def __str__(self):
        return "\n".join(
            [str(k) + ":\t" + str(v) for k, v in self.__dict__.items()]
        )

    def __len__(self):
        return len(self.raw_text)

    @staticmethod
    def from_json(json_data):
        return Example(
            json_data["title"],
            json_data["entities"],
            json_data["types"],
            json_data["relations"],
            json_data["abstract"],
        )

    def build_graph(self):
        graph = dgl.DGLGraph()
        ent_len = len(self.raw_ent_text)
        rel_len = len(
            self.raw_rel
        )  # treat the repeated relation as different nodes, refer to the author's code

        graph.add_nodes(
            ent_len, {"type": torch.ones(ent_len) * NODE_TYPE["entity"]}
        )
        graph.add_nodes(1, {"type": torch.ones(1) * NODE_TYPE["root"]})
        graph.add_nodes(
            rel_len * 2,
            {"type": torch.ones(rel_len * 2) * NODE_TYPE["relation"]},
        )
        graph.add_edges(ent_len, torch.arange(ent_len))
        graph.add_edges(torch.arange(ent_len), ent_len)
        graph.add_edges(
            torch.arange(ent_len + 1 + rel_len * 2),
            torch.arange(ent_len + 1 + rel_len * 2),
        )
        adj_edges = []
        for i, r in enumerate(self.raw_rel):
            assert len(r) == 3, str(r)
            st, rt, ed = r
            st_ent, ed_ent = self.raw_ent_text.index(
                st
            ), self.raw_ent_text.index(ed)
            # according to the edge_softmax operator, we need to reverse the graph
            adj_edges.append([ent_len + 1 + 2 * i, st_ent])
            adj_edges.append([ed_ent, ent_len + 1 + 2 * i])
            adj_edges.append([ent_len + 1 + 2 * i + 1, ed_ent])
            adj_edges.append([st_ent, ent_len + 1 + 2 * i + 1])

        if len(adj_edges) > 0:
            graph.add_edges(*list(map(list, zip(*adj_edges))))
        return graph

    def get_tensor(
        self, ent_vocab, rel_vocab, text_vocab, ent_text_vocab, title_vocab
    ):
        if hasattr(self, "_cached_tensor"):
            return self._cached_tensor
        else:
            title_data = ["<BOS>"] + self.raw_title + ["<EOS>"]
            title = [title_vocab(x) for x in title_data]
            ent_text = [
                [ent_text_vocab(y) for y in x] for x in self.raw_ent_text
            ]
            ent_type = [
                text_vocab(x) for x in self.raw_ent_type
            ]  # for inference
            rel_data = ["--root--"] + sum(
                [[x[1], x[1] + "_INV"] for x in self.raw_rel], []
            )
            rel = [rel_vocab(x) for x in rel_data]

            text_data = ["<BOS>"] + self.raw_text + ["<EOS>"]
            text = [text_vocab(x) for x in text_data]
            tgt_text = []
            # the input text and decoding target are different since the consideration of the copy mechanism.
            for i, str1 in enumerate(text_data):
                if str1[0] == "<" and str1[-1] == ">" and "_" in str1:
                    a, b = str1[1:-1].split("_")
                    text[i] = text_vocab("<" + a + ">")
                    tgt_text.append(len(text_vocab) + int(b))
                else:
                    tgt_text.append(text[i])
            self._cached_tensor = {
                "title": torch.LongTensor(title),
                "ent_text": [torch.LongTensor(x) for x in ent_text],
                "ent_type": torch.LongTensor(ent_type),
                "rel": torch.LongTensor(rel),
                "text": torch.LongTensor(text[:-1]),
                "tgt_text": torch.LongTensor(tgt_text[1:]),
                "graph": self.graph,
                "raw_ent_text": self.raw_ent_text,
            }
            return self._cached_tensor

    def update_vocab(
        self, ent_vocab, rel_vocab, text_vocab, ent_text_vocab, title_vocab
    ):
        ent_vocab.update(self.raw_ent_type)
        ent_text_vocab.update(self.raw_ent_text)
        title_vocab.update(self.raw_title)
        rel_vocab.update(
            ["--root--"]
            + [x[1] for x in self.raw_rel]
            + [x[1] + "_INV" for x in self.raw_rel]
        )
        text_vocab.update(self.raw_ent_type)
        text_vocab.update(self.raw_text)


class BucketSampler(torch.utils.data.Sampler):
    def __init__(self, data_source, batch_size=32, bucket=3):
        self.data_source = data_source
        self.bucket = bucket
        self.batch_size = batch_size

    def __iter__(self):
        # the magic number comes from the author's code
        perm = torch.randperm(len(self.data_source))
        lens = torch.Tensor([len(x) for x in self.data_source])
        lens = lens[perm]
        t1 = []
        t2 = []
        t3 = []
        for i, l in enumerate(lens):
            if l < 100:
                t1.append(perm[i])
            elif l > 100 and l < 220:
                t2.append(perm[i])
            else:
                t3.append(perm[i])
        datas = [t1, t2, t3]
        random.shuffle(datas)
        idxs = sum(datas, [])
        batch = []

        lens = torch.Tensor([len(x) for x in self.data_source])
        for idx in idxs:
            batch.append(idx)
            mlen = max([0] + [lens[x] for x in batch])
            if (
                (mlen < 100 and len(batch) == 32)
                or (mlen > 100 and mlen < 220 and len(batch) >= 24)
                or (mlen > 220 and len(batch) >= 8)
                or len(batch) == 32
            ):
                yield batch
                batch = []
        if len(batch) > 0:
            yield batch

    def __len__(self):
        return (len(self.data_source) + self.batch_size - 1) // self.batch_size


class GWdataset(torch.utils.data.Dataset):
    def __init__(
        self,
        exs,
        ent_vocab=None,
        rel_vocab=None,
        text_vocab=None,
        ent_text_vocab=None,
        title_vocab=None,
        device=None,
    ):
        super(GWdataset, self).__init__()
        self.exs = exs
        (
            self.ent_vocab,
            self.rel_vocab,
            self.text_vocab,
            self.ent_text_vocab,
            self.title_vocab,
            self.device,
        ) = (
            ent_vocab,
            rel_vocab,
            text_vocab,
            ent_text_vocab,
            title_vocab,
            device,
        )

    def __iter__(self):
        return iter(self.exs)

    def __getitem__(self, index):
        return self.exs[index]

    def __len__(self):
        return len(self.exs)

    def batch_fn(self, batch_ex):
        (
            batch_title,
            batch_ent_text,
            batch_ent_type,
            batch_rel,
            batch_text,
            batch_tgt_text,
            batch_graph,
        ) = ([], [], [], [], [], [], [])
        batch_raw_ent_text = []
        for ex in batch_ex:
            ex_data = ex.get_tensor(
                self.ent_vocab,
                self.rel_vocab,
                self.text_vocab,
                self.ent_text_vocab,
                self.title_vocab,
            )
            batch_title.append(ex_data["title"])
            batch_ent_text.append(ex_data["ent_text"])
            batch_ent_type.append(ex_data["ent_type"])
            batch_rel.append(ex_data["rel"])
            batch_text.append(ex_data["text"])
            batch_tgt_text.append(ex_data["tgt_text"])
            batch_graph.append(ex_data["graph"])
            batch_raw_ent_text.append(ex_data["raw_ent_text"])
        batch_title = pad(batch_title, out_type="tensor")
        batch_ent_text, ent_len = pad(
            batch_ent_text, out_type="tensor", flatten=True
        )
        batch_ent_type = pad(batch_ent_type, out_type="tensor")
        batch_rel = pad(batch_rel, out_type="tensor")
        batch_text = pad(batch_text, out_type="tensor")
        batch_tgt_text = pad(batch_tgt_text, out_type="tensor")
        batch_graph = dgl.batch(batch_graph)
        batch_graph.to(self.device)
        return {
            "title": batch_title.to(self.device),
            "ent_text": batch_ent_text.to(self.device),
            "ent_len": ent_len,
            "ent_type": batch_ent_type.to(self.device),
            "rel": batch_rel.to(self.device),
            "text": batch_text.to(self.device),
            "tgt_text": batch_tgt_text.to(self.device),
            "graph": batch_graph,
            "raw_ent_text": batch_raw_ent_text,
        }


def get_datasets(
    fnames,
    min_freq=-1,
    sep=";",
    joint_vocab=True,
    device=None,
    save="tmp.pickle",
):
    # min_freq : not support now since it's very sensitive to the final results, but you can set it via passing min_freq to the Vocab class.
    # sep : not support now
    # joint_vocab : not support now
    ent_vocab = Vocab(sp=["<PAD>", "<UNK>"])
    title_vocab = Vocab(min_freq=5)
    rel_vocab = Vocab(sp=["<PAD>", "<UNK>"])
    text_vocab = Vocab(min_freq=5)
    ent_text_vocab = Vocab(sp=["<PAD>", "<UNK>"])
    datasets = []
    for fname in fnames:
        exs = []
        json_datas = json.loads(open(fname).read())
        for json_data in json_datas:
            # construct one data example
            ex = Example.from_json(json_data)
            if fname == fnames[0]:  # only training set
                ex.update_vocab(
                    ent_vocab,
                    rel_vocab,
                    text_vocab,
                    ent_text_vocab,
                    title_vocab,
                )
            exs.append(ex)
        datasets.append(exs)
    ent_vocab.build()
    rel_vocab.build()
    text_vocab.build()
    ent_text_vocab.build()
    title_vocab.build()
    datasets = [
        GWdataset(
            exs,
            ent_vocab,
            rel_vocab,
            text_vocab,
            ent_text_vocab,
            title_vocab,
            device,
        )
        for exs in datasets
    ]
    with open(save, "wb") as f:
        pickle.dump(datasets, f)
    return datasets


if __name__ == "__main__":
    ds = get_datasets(
        [
            "data/unprocessed.val.json",
            "data/unprocessed.val.json",
            "data/unprocessed.test.json",
        ]
    )
    print(ds[0].exs[0])
    print(
        ds[0]
        .exs[0]
        .get_tensor(
            ds[0].ent_vocab,
            ds[0].rel_vocab,
            ds[0].text_vocab,
            ds[0].ent_text_vocab,
            ds[0].title_vocab,
        )
    )