# -*- coding: utf-8 -*-

import numpy as np

import torch
from torch import nn

from kbcr.kernels import GaussianKernel
from kbcr.smart import NeuralKB

import pytest


@pytest.mark.light
def test_smart_v1():
    embedding_size = 50

    rs = np.random.RandomState(0)

    for _ in range(32):
        with torch.no_grad():
            triples = [
                ('a', 'p', 'b'),
                ('c', 'q', 'd'),
                ('e', 'q', 'f'),
                ('g', 'q', 'h'),
                ('i', 'q', 'l'),
                ('m', 'q', 'n'),
                ('o', 'q', 'p'),
                ('q', 'q', 'r'),
                ('s', 'q', 't'),
                ('u', 'q', 'v')
            ]

            entity_lst = sorted({s for (s, _, _) in triples} | {o for (_, _, o) in triples})
            predicate_lst = sorted({p for (_, p, _) in triples})

            nb_entities, nb_predicates = len(entity_lst), len(predicate_lst)

            entity_to_index = {e: i for i, e in enumerate(entity_lst)}
            predicate_to_index = {p: i for i, p in enumerate(predicate_lst)}

            kernel = GaussianKernel()

            entity_embeddings = nn.Embedding(nb_entities, embedding_size * 2, sparse=True)
            predicate_embeddings = nn.Embedding(nb_predicates, embedding_size * 2, sparse=True)

            fact_rel = torch.LongTensor(np.array([predicate_to_index[p] for (_, p, _) in triples]))
            fact_arg1 = torch.LongTensor(np.array([entity_to_index[s] for (s, _, _) in triples]))
            fact_arg2 = torch.LongTensor(np.array([entity_to_index[o] for (_, _, o) in triples]))
            facts = [fact_rel, fact_arg1, fact_arg2]

            model = NeuralKB(entity_embeddings=entity_embeddings, predicate_embeddings=predicate_embeddings,
                             kernel=kernel, facts=facts)

            xs_np = rs.randint(nb_entities, size=32)
            xp_np = rs.randint(nb_predicates, size=32)
            xo_np = rs.randint(nb_entities, size=32)

            xs_np[0] = 0
            xp_np[0] = 0
            xo_np[0] = 1

            xs_np[1] = 2
            xp_np[1] = 1
            xo_np[1] = 3

            xs = torch.LongTensor(xs_np)
            xp = torch.LongTensor(xp_np)
            xo = torch.LongTensor(xo_np)

            xs_emb = entity_embeddings(xs)
            xp_emb = predicate_embeddings(xp)
            xo_emb = entity_embeddings(xo)

            print('xp_emb', xp_emb.shape)

            res_sp, res_po = model.forward(xp_emb, xs_emb, xo_emb)
            inf = model.score(xp_emb, xs_emb, xo_emb)

            assert inf[0] > 0.9
            assert inf[1] > 0.9

            scores_sp, emb_sp = res_sp
            scores_po, emb_po = res_po

            print(scores_sp.shape, emb_sp.shape)
            print(scores_po.shape, emb_po.shape)

            inf = inf.cpu().numpy()
            scores_sp = scores_sp.cpu().numpy()
            scores_po = scores_po.cpu().numpy()

            print('AAA', inf)
            print('BBB', scores_sp)


if __name__ == '__main__':
    pytest.main([__file__])
    # test_smart_v1()

[(18, 'numpy.random.RandomState', 'np.random.RandomState', True, 'import numpy as np\n'), (99, 'pytest.main', 'pytest.main', False, 'import pytest\n'), (21, 'torch.no_grad', 'torch.no_grad', False, 'import torch\n'), (43, 'kbcr.kernels.GaussianKernel', 'GaussianKernel', False, 'from kbcr.kernels import GaussianKernel\n'), (45, 'torch.nn.Embedding', 'nn.Embedding', False, 'from torch import nn\n'), (46, 'torch.nn.Embedding', 'nn.Embedding', False, 'from torch import nn\n'), (53, 'kbcr.smart.NeuralKB', 'NeuralKB', False, 'from kbcr.smart import NeuralKB\n'), (68, 'torch.LongTensor', 'torch.LongTensor', False, 'import torch\n'), (69, 'torch.LongTensor', 'torch.LongTensor', False, 'import torch\n'), (70, 'torch.LongTensor', 'torch.LongTensor', False, 'import torch\n'), (48, 'numpy.array', 'np.array', True, 'import numpy as np\n'), (49, 'numpy.array', 'np.array', True, 'import numpy as np\n'), (50, 'numpy.array', 'np.array', True, 'import numpy as np\n')]