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, 'os.path.dirname', 'os.path.dirname', False, 'import os\n'), (14, 'hivemind.moe.server.background_server', 'background_server', False, 'from hivemind.moe.server import background_server\n'), (23, 'hivemind.RemoteExpert', 'RemoteExpert', False, 'from hivemind import RemoteExpert\n'), (24, 'hivemind.RemoteExpert', 'RemoteExpert', False, 'from hivemind import RemoteExpert\n'), (40, 'hivemind.moe.server.background_server', 'background_server', False, 'from hivemind.moe.server import background_server\n'), (49, 'hivemind.RemoteExpert', 'RemoteExpert', False, 'from hivemind import RemoteExpert\n'), (50, 'hivemind.RemoteExpert', 'RemoteExpert', False, 'from hivemind import RemoteExpert\n'), (27, 'torch.randn', 'torch.randn', False, 'import torch\n'), (54, 'torch.randn', 'torch.randn', False, 'import torch\n'), (55, 'torch.randn', 'torch.randn', False, 'import torch\n'), (56, 'torch.randn', 'torch.randn', False, 'import torch\n')]