# -*- coding: utf-8 -*-
# emacs: -*- mode: python; py-indent-offset: 4; indent-tabs-mode: nil -*-
# vi: set ft=python sts=4 ts=4 sw=4 et:
"""
Test the resource management of MultiProc
"""
import os
import pytest
from nipype.pipeline import engine as pe
from nipype.interfaces import base as nib


class InputSpec(nib.TraitedSpec):
    input1 = nib.traits.Int(desc='a random int')
    input2 = nib.traits.Int(desc='a random int')


class OutputSpec(nib.TraitedSpec):
    output1 = nib.traits.List(nib.traits.Int, desc='outputs')


class MultiprocTestInterface(nib.BaseInterface):
    input_spec = InputSpec
    output_spec = OutputSpec

    def _run_interface(self, runtime):
        runtime.returncode = 0
        return runtime

    def _list_outputs(self):
        outputs = self._outputs().get()
        outputs['output1'] = [1, self.inputs.input1]
        return outputs


def test_run_multiproc(tmpdir):
    tmpdir.chdir()

    pipe = pe.Workflow(name='pipe')
    mod1 = pe.Node(MultiprocTestInterface(), name='mod1')
    mod2 = pe.MapNode(
        MultiprocTestInterface(), iterfield=['input1'], name='mod2')
    pipe.connect([(mod1, mod2, [('output1', 'input1')])])
    pipe.base_dir = os.getcwd()
    mod1.inputs.input1 = 1
    pipe.config['execution']['poll_sleep_duration'] = 2
    execgraph = pipe.run(plugin="MultiProc")
    names = [node.fullname for node in execgraph.nodes()]
    node = list(execgraph.nodes())[names.index('pipe.mod1')]
    result = node.get_output('output1')
    assert result == [1, 1]


class InputSpecSingleNode(nib.TraitedSpec):
    input1 = nib.traits.Int(desc='a random int')
    input2 = nib.traits.Int(desc='a random int')


class OutputSpecSingleNode(nib.TraitedSpec):
    output1 = nib.traits.Int(desc='a random int')


class SingleNodeTestInterface(nib.BaseInterface):
    input_spec = InputSpecSingleNode
    output_spec = OutputSpecSingleNode

    def _run_interface(self, runtime):
        runtime.returncode = 0
        return runtime

    def _list_outputs(self):
        outputs = self._outputs().get()
        outputs['output1'] = self.inputs.input1
        return outputs


def test_no_more_memory_than_specified(tmpdir):
    tmpdir.chdir()
    pipe = pe.Workflow(name='pipe')
    n1 = pe.Node(SingleNodeTestInterface(), name='n1', mem_gb=1)
    n2 = pe.Node(SingleNodeTestInterface(), name='n2', mem_gb=1)
    n3 = pe.Node(SingleNodeTestInterface(), name='n3', mem_gb=1)
    n4 = pe.Node(SingleNodeTestInterface(), name='n4', mem_gb=1)

    pipe.connect(n1, 'output1', n2, 'input1')
    pipe.connect(n1, 'output1', n3, 'input1')
    pipe.connect(n2, 'output1', n4, 'input1')
    pipe.connect(n3, 'output1', n4, 'input2')
    n1.inputs.input1 = 1

    max_memory = 0.5
    with pytest.raises(RuntimeError):
        pipe.run(
            plugin='MultiProc',
            plugin_args={
                'memory_gb': max_memory,
                'n_procs': 2
            })


def test_no_more_threads_than_specified(tmpdir):
    tmpdir.chdir()

    pipe = pe.Workflow(name='pipe')
    n1 = pe.Node(SingleNodeTestInterface(), name='n1', n_procs=2)
    n2 = pe.Node(SingleNodeTestInterface(), name='n2', n_procs=2)
    n3 = pe.Node(SingleNodeTestInterface(), name='n3', n_procs=4)
    n4 = pe.Node(SingleNodeTestInterface(), name='n4', n_procs=2)

    pipe.connect(n1, 'output1', n2, 'input1')
    pipe.connect(n1, 'output1', n3, 'input1')
    pipe.connect(n2, 'output1', n4, 'input1')
    pipe.connect(n3, 'output1', n4, 'input2')
    n1.inputs.input1 = 4

    max_threads = 2
    with pytest.raises(RuntimeError):
        pipe.run(plugin='MultiProc', plugin_args={'n_procs': max_threads})


def test_hold_job_until_procs_available(tmpdir):
    tmpdir.chdir()

    pipe = pe.Workflow(name='pipe')
    n1 = pe.Node(SingleNodeTestInterface(), name='n1', n_procs=2)
    n2 = pe.Node(SingleNodeTestInterface(), name='n2', n_procs=2)
    n3 = pe.Node(SingleNodeTestInterface(), name='n3', n_procs=2)
    n4 = pe.Node(SingleNodeTestInterface(), name='n4', n_procs=2)

    pipe.connect(n1, 'output1', n2, 'input1')
    pipe.connect(n1, 'output1', n3, 'input1')
    pipe.connect(n2, 'output1', n4, 'input1')
    pipe.connect(n3, 'output1', n4, 'input2')
    n1.inputs.input1 = 4

    max_threads = 2
    pipe.run(plugin='MultiProc', plugin_args={'n_procs': max_threads})
