# -*- 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', (['(0)'], {}), True, 'import numpy as np\n'), (99, 'pytest.main', 'pytest.main', (['[__file__]'], {}), 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', (['nb_entities', '(embedding_size * 2)'], {'sparse': '(True)'}), False, 'from torch import nn\n'), (46, 'torch.nn.Embedding', 'nn.Embedding', (['nb_predicates', '(embedding_size * 2)'], {'sparse': '(True)'}), False, 'from torch import nn\n'), (53, 'kbcr.smart.NeuralKB', 'NeuralKB', ([], {'entity_embeddings': 'entity_embeddings', 'predicate_embeddings': 'predicate_embeddings', 'kernel': 'kernel', 'facts': 'facts'}), False, 'from kbcr.smart import NeuralKB\n'), (68, 'torch.LongTensor', 'torch.LongTensor', (['xs_np'], {}), False, 'import torch\n'), (69, 'torch.LongTensor', 'torch.LongTensor', (['xp_np'], {}), False, 'import torch\n'), (70, 'torch.LongTensor', 'torch.LongTensor', (['xo_np'], {}), False, 'import torch\n'), (48, 'numpy.array', 'np.array', (['[predicate_to_index[p] for _, p, _ in triples]'], {}), True, 'import numpy as np\n'), (49, 'numpy.array', 'np.array', (['[entity_to_index[s] for s, _, _ in triples]'], {}), True, 'import numpy as np\n'), (50, 'numpy.array', 'np.array', (['[entity_to_index[o] for _, _, o in triples]'], {}), True, 'import numpy as np\n')]