import os

import pytest
import torch

from hivemind import RemoteExpert
from hivemind.moe.server import background_server

CUSTOM_EXPERTS_PATH = os.path.join(os.path.dirname(__file__), "test_utils", "custom_networks.py")


@pytest.mark.forked
def test_custom_expert(hid_dim=16):
    with background_server(
        expert_cls="perceptron",
        num_experts=2,
        device="cpu",
        hidden_dim=hid_dim,
        num_handlers=2,
        no_dht=True,
        custom_module_path=CUSTOM_EXPERTS_PATH,
    ) as (server_endpoint, _):
        expert0 = RemoteExpert("expert.0", server_endpoint)
        expert1 = RemoteExpert("expert.1", server_endpoint)

        for batch_size in (1, 4):
            batch = torch.randn(batch_size, hid_dim)

            output0 = expert0(batch)
            output1 = expert1(batch)

            loss = output0.sum()
            loss.backward()
            loss = output1.sum()
            loss.backward()


@pytest.mark.forked
def test_multihead_expert(hid_dim=16):
    with background_server(
        expert_cls="multihead",
        num_experts=2,
        device="cpu",
        hidden_dim=hid_dim,
        num_handlers=2,
        no_dht=True,
        custom_module_path=CUSTOM_EXPERTS_PATH,
    ) as (server_endpoint, _):
        expert0 = RemoteExpert("expert.0", server_endpoint)
        expert1 = RemoteExpert("expert.1", server_endpoint)

        for batch_size in (1, 4):
            batch = (
                torch.randn(batch_size, hid_dim),
                torch.randn(batch_size, 2 * hid_dim),
                torch.randn(batch_size, 3 * hid_dim),
            )

            output0 = expert0(*batch)
            output1 = expert1(*batch)

            loss = output0.sum()
            loss.backward()
            loss = output1.sum()
            loss.backward()

[((9, 35, 9, 60), 'os.path.dirname', 'os.path.dirname', ({(9, 51, 9, 59): '__file__'}, {}), '(__file__)', False, 'import os\n'), ((14, 9, 22, 5), 'hivemind.moe.server.background_server', 'background_server', (), '', False, 'from hivemind.moe.server import background_server\n'), ((23, 18, 23, 59), 'hivemind.RemoteExpert', 'RemoteExpert', ({(23, 31, 23, 41): '"""expert.0"""', (23, 43, 23, 58): 'server_endpoint'}, {}), "('expert.0', server_endpoint)", False, 'from hivemind import RemoteExpert\n'), ((24, 18, 24, 59), 'hivemind.RemoteExpert', 'RemoteExpert', ({(24, 31, 24, 41): '"""expert.1"""', (24, 43, 24, 58): 'server_endpoint'}, {}), "('expert.1', server_endpoint)", False, 'from hivemind import RemoteExpert\n'), ((40, 9, 48, 5), 'hivemind.moe.server.background_server', 'background_server', (), '', False, 'from hivemind.moe.server import background_server\n'), ((49, 18, 49, 59), 'hivemind.RemoteExpert', 'RemoteExpert', ({(49, 31, 49, 41): '"""expert.0"""', (49, 43, 49, 58): 'server_endpoint'}, {}), "('expert.0', server_endpoint)", False, 'from hivemind import RemoteExpert\n'), ((50, 18, 50, 59), 'hivemind.RemoteExpert', 'RemoteExpert', ({(50, 31, 50, 41): '"""expert.1"""', (50, 43, 50, 58): 'server_endpoint'}, {}), "('expert.1', server_endpoint)", False, 'from hivemind import RemoteExpert\n'), ((27, 20, 27, 52), 'torch.randn', 'torch.randn', ({(27, 32, 27, 42): 'batch_size', (27, 44, 27, 51): 'hid_dim'}, {}), '(batch_size, hid_dim)', False, 'import torch\n'), ((54, 16, 54, 48), 'torch.randn', 'torch.randn', ({(54, 28, 54, 38): 'batch_size', (54, 40, 54, 47): 'hid_dim'}, {}), '(batch_size, hid_dim)', False, 'import torch\n'), ((55, 16, 55, 52), 'torch.randn', 'torch.randn', ({(55, 28, 55, 38): 'batch_size', (55, 40, 55, 51): '(2 * hid_dim)'}, {}), '(batch_size, 2 * hid_dim)', False, 'import torch\n'), ((56, 16, 56, 52), 'torch.randn', 'torch.randn', ({(56, 28, 56, 38): 'batch_size', (56, 40, 56, 51): '(3 * hid_dim)'}, {}), '(batch_size, 3 * hid_dim)', False, 'import torch\n')]