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()