# pylint: disable=missing-docstring,no-self-use,line-too-long, unused-argument, redefined-outer-name
# type: ignore

import unittest
from unittest.mock import patch
from soil import decorator, modulify, task
from soil.data_structure import DataStructure
import soil


@modulify(output_types=lambda *input_types, **args: [input_types[0]])
def train(*data, module=None, model_params=None, **kwargs):
    res, = task(module)(*data, **model_params)
    return [res]


@decorator(depth=3)
def trainable(module):
    def decorator(**model_params):
        def train_fn(**train_params):
            def inner_module(*data):
                return train(*data, module=module, model_params=model_params, **train_params)
            return inner_module
        return train_fn
    return decorator


@trainable
@modulify(output_types=lambda *input_types, **args: [input_types[0]])
def my_model(*data, **my_model_params):
    if 'p3' in my_model_params:
        return task(my_model_params['p3'](cross_validation=True))(*data)
    return data


@modulify(output_types=lambda *input_types, **args: [input_types[0]])
def my_model2(*data, **my_model_params):
    return data


def test_fn1():
    pass


def test_fn2():
    pass


def mock_soil_data_fn(data):
    return DataStructure('test_id')


class TestUtils(unittest.TestCase):
    @patch('soil.data')
    def test_build_arguments_1(self, mock_soil_data):
        mock_soil_data.side_effect = mock_soil_data_fn
        data = soil.data([{'test': [1, 2, 3]}])
        res, = my_model(test_param=1)(cross_validation=True)(data)
        plan = res.pipeline.plan
        first = plan[0]
        args = first['args']
        expected_args = {
            'module': {
                '__soil_arg_type': 'decorated',
                'module': 'soil.modulify',
                'name': 'modulify',
                'calls': [
                    {
                        'args': [],
                        'kwargs': {'output_types': {'__soil_arg_type': 'lambda'}}
                    }, {
                        'args': [{
                            '__soil_arg_type': 'function',
                            'module': 'unit.soil.test_utils', 'name': 'my_model'
                        }],
                        'kwargs': {}}
                ]
            },
            'model_params': {'test_param': 1},
            'cross_validation': True
        }
        assert first['module'] == 'unit.soil.test_utils.train'
        assert args == expected_args

    @patch('soil.data')
    def test_build_arguments_2(self, mock_soil_data):
        mock_soil_data.side_effect = mock_soil_data_fn
        data = soil.data([{'test': 'test'}])
        res, = my_model(test_param=[test_fn1, test_fn2])(cross_validation=True)(data)
        plan = res.pipeline.plan
        first = plan[0]
        args = first['args']
        expected = {
            "cross_validation": True,
            "model_params": {
                "test_param": [
                    {
                        "__soil_arg_type": "function",
                        "module": "unit.soil.test_utils",
                        "name": "test_fn1"
                    },
                    {
                        "__soil_arg_type": "function",
                        "module": "unit.soil.test_utils",
                        "name": "test_fn2"
                    }
                ]
            },
            "module": {
                "__soil_arg_type": "decorated",
                "calls": [
                    {
                        "args": [],
                        "kwargs": {
                            "output_types": {
                                "__soil_arg_type": "lambda"
                            }
                        }
                    },
                    {
                        "args": [
                            {
                                "__soil_arg_type": "function",
                                "module": "unit.soil.test_utils",
                                "name": "my_model"
                            }
                        ],
                        "kwargs": {}
                    }
                ],
                "module": "soil.modulify",
                "name": "modulify"
            }
        }
        assert args == expected

    @patch('soil.data')
    def test_build_arguments_3(self, mock_soil_data):
        mock_soil_data.side_effect = mock_soil_data_fn
        data = soil.data([{'test': 'test'}])
        res, = my_model(test_param=my_model2)(cross_validation=True)(data)
        plan = res.pipeline.plan
        first = plan[0]
        args = first['args']
        expected = {
            "cross_validation": True,
            "model_params": {
                "test_param": {
                    "__soil_arg_type": "decorated",
                    "calls": [
                        {
                            "args": [],
                            "kwargs": {
                                "output_types": {
                                    "__soil_arg_type": "lambda"
                                }
                            }
                        },
                        {
                            "args": [
                                {
                                    "__soil_arg_type": "function",
                                    "module": "unit.soil.test_utils",
                                    "name": "my_model2"
                                }
                            ],
                            "kwargs": {}
                        }
                    ],
                    "module": "soil.modulify",
                    "name": "modulify"
                }
            },
            "module": {
                "__soil_arg_type": "decorated",
                "calls": [
                    {
                        "args": [],
                        "kwargs": {
                            "output_types": {
                                "__soil_arg_type": "lambda"
                            }
                        }
                    },
                    {
                        "args": [
                            {
                                "__soil_arg_type": "function",
                                "module": "unit.soil.test_utils",
                                "name": "my_model"
                            }
                        ],
                        "kwargs": {}
                    }
                ],
                "module": "soil.modulify",
                "name": "modulify"
            }
        }
        assert args == expected

    @patch('soil.data')
    def test_build_arguments_4(self, mock_soil_data):
        mock_soil_data.side_effect = mock_soil_data_fn
        data = soil.data([{'test': 'test'}])
        res, = my_model(test_param=trainable(my_model2))(cross_validation=True)(data)
        plan = res.pipeline.plan
        first = plan[0]
        args = first['args']
        expected = {
            "cross_validation": True,
            "model_params": {
                "test_param": {
                    "__soil_arg_type": "decorated",
                    "calls": [
                        {
                            "args": [
                                {
                                    "__soil_arg_type": "decorated",
                                    "calls": [
                                        {
                                            "args": [],
                                            "kwargs": {
                                                "output_types": {
                                                    "__soil_arg_type": "lambda"
                                                }
                                            }
                                        },
                                        {
                                            "args": [
                                                {
                                                    "__soil_arg_type": "function",
                                                    "module": "unit.soil.test_utils",
                                                    "name": "my_model2"
                                                }
                                            ],
                                            "kwargs": {}
                                        }
                                    ],
                                    "module": "soil.modulify",
                                    "name": "modulify"
                                }
                            ],
                            "kwargs": {}
                        }
                    ],
                    "module": "unit.soil.test_utils",
                    "name": "trainable"
                }
            },
            "module": {
                "__soil_arg_type": "decorated",
                "calls": [
                    {
                        "args": [],
                        "kwargs": {
                            "output_types": {
                                "__soil_arg_type": "lambda"
                            }
                        }
                    },
                    {
                        "args": [
                            {
                                "__soil_arg_type": "function",
                                "module": "unit.soil.test_utils",
                                "name": "my_model"
                            }
                        ],
                        "kwargs": {}
                    }
                ],
                "module": "soil.modulify",
                "name": "modulify"
            }
        }
        assert args == expected

    @patch('soil.data')
    def test_build_arguments_5(self, mock_soil_data):
        mock_soil_data.side_effect = mock_soil_data_fn
        data = soil.data([{'test': 'test'}])
        res, = my_model(test_param=my_model)(cross_validation=True)(data)
        plan = res.pipeline.plan
        first = plan[0]
        args = first['args']
        expected = {
            "cross_validation": True,
            "model_params": {
                "test_param": {
                    "__soil_arg_type": "decorated",
                    "calls": [
                        {
                            "args": [
                                {
                                    "__soil_arg_type": "decorated",
                                    "calls": [
                                        {
                                            "args": [],
                                            "kwargs": {
                                                "output_types": {
                                                    "__soil_arg_type": "lambda"
                                                }
                                            }
                                        },
                                        {
                                            "args": [
                                                {
                                                    "__soil_arg_type": "function",
                                                    "module": "unit.soil.test_utils",
                                                    "name": "my_model"
                                                }
                                            ],
                                            "kwargs": {}
                                        }
                                    ],
                                    "module": "soil.modulify",
                                    "name": "modulify"
                                }
                            ],
                            "kwargs": {}
                        }
                    ],
                    "module": "unit.soil.test_utils",
                    "name": "trainable"
                }
            },
            "module": {
                "__soil_arg_type": "decorated",
                "calls": [
                    {
                        "args": [],
                        "kwargs": {
                            "output_types": {
                                "__soil_arg_type": "lambda"
                            }
                        }
                    },
                    {
                        "args": [
                            {
                                "__soil_arg_type": "function",
                                "module": "unit.soil.test_utils",
                                "name": "my_model"
                            }
                        ],
                        "kwargs": {}
                    }
                ],
                "module": "soil.modulify",
                "name": "modulify"
            }
        }
        assert args == expected

    @patch('soil.data')
    def test_build_arguments_6(self, mock_soil_data):
        mock_soil_data.side_effect = mock_soil_data_fn
        data = soil.data([{'test': 'test'}])
        res, = my_model(test_param=my_model(f1=2, t1=test_fn1))(cross_validation=True)(data)
        plan = res.pipeline.plan
        first = plan[0]
        args = first['args']
        expected = {
            "cross_validation": True,
            "model_params": {
                "test_param": {
                    "__soil_arg_type": "decorated",
                    "calls": [
                        {
                            "args": [
                                {
                                    "__soil_arg_type": "decorated",
                                    "calls": [
                                        {
                                            "args": [],
                                            "kwargs": {
                                                "output_types": {
                                                    "__soil_arg_type": "lambda"
                                                }
                                            }
                                        },
                                        {
                                            "args": [
                                                {
                                                    "__soil_arg_type": "function",
                                                    "module": "unit.soil.test_utils",
                                                    "name": "my_model"
                                                }
                                            ],
                                            "kwargs": {}
                                        }
                                    ],
                                    "module": "soil.modulify",
                                    "name": "modulify"
                                }
                            ],
                            "kwargs": {}
                        },
                        {
                            "args": [],
                            "kwargs": {
                                "f1": 2,
                                "t1": {
                                    "__soil_arg_type": "function",
                                    "module": "unit.soil.test_utils",
                                    "name": "test_fn1"
                                }
                            }
                        }
                    ],
                    "module": "unit.soil.test_utils",
                    "name": "trainable"
                }
            },
            "module": {
                "__soil_arg_type": "decorated",
                "calls": [
                    {
                        "args": [],
                        "kwargs": {
                            "output_types": {
                                "__soil_arg_type": "lambda"
                            }
                        }
                    },
                    {
                        "args": [
                            {
                                "__soil_arg_type": "function",
                                "module": "unit.soil.test_utils",
                                "name": "my_model"
                            }
                        ],
                        "kwargs": {}
                    }
                ],
                "module": "soil.modulify",
                "name": "modulify"
            }
        }
        assert args == expected

    @patch('soil.data')
    def test_build_arguments_7(self, mock_soil_data):
        mock_soil_data.side_effect = mock_soil_data_fn
        data = soil.data([{'test': 'test'}])
        res, = my_model(test_param=my_model(f1=2, t1=test_fn1)(another=1))(cross_validation=True)(data)
        plan = res.pipeline.plan
        first = plan[0]
        args = first['args']
        expected = {
            "cross_validation": True,
            "model_params": {
                "test_param": {
                    "__soil_arg_type": "decorated",
                    "calls": [
                        {
                            "args": [
                                {
                                    "__soil_arg_type": "decorated",
                                    "calls": [
                                        {
                                            "args": [],
                                            "kwargs": {
                                                "output_types": {
                                                    "__soil_arg_type": "lambda"
                                                }
                                            }
                                        },
                                        {
                                            "args": [
                                                {
                                                    "__soil_arg_type": "function",
                                                    "module": "unit.soil.test_utils",
                                                    "name": "my_model"
                                                }
                                            ],
                                            "kwargs": {}
                                        }
                                    ],
                                    "module": "soil.modulify",
                                    "name": "modulify"
                                }
                            ],
                            "kwargs": {}
                        },
                        {
                            "args": [],
                            "kwargs": {
                                "f1": 2,
                                "t1": {
                                    "__soil_arg_type": "function",
                                    "module": "unit.soil.test_utils",
                                    "name": "test_fn1"
                                }
                            }
                        },
                        {
                            "args": [],
                            "kwargs": {
                                "another": 1
                            }
                        }
                    ],
                    "module": "unit.soil.test_utils",
                    "name": "trainable"
                }
            },
            "module": {
                "__soil_arg_type": "decorated",
                "calls": [
                    {
                        "args": [],
                        "kwargs": {
                            "output_types": {
                                "__soil_arg_type": "lambda"
                            }
                        }
                    },
                    {
                        "args": [
                            {
                                "__soil_arg_type": "function",
                                "module": "unit.soil.test_utils",
                                "name": "my_model"
                            }
                        ],
                        "kwargs": {}
                    }
                ],
                "module": "soil.modulify",
                "name": "modulify"
            }
        }
        assert args == expected

    @patch('soil.data')
    def test_build_arguments_8(self, mock_soil_data):
        mock_soil_data.side_effect = mock_soil_data_fn
        data = soil.data([{'test': 'test'}])
        res, = my_model(test_param=trainable(my_model2)(f1=2, t1=test_fn1)(another=1))(cross_validation=True)(data)
        plan = res.pipeline.plan
        first = plan[0]
        args = first['args']
        expected = {
            "cross_validation": True,
            "model_params": {
                "test_param": {
                    "__soil_arg_type": "decorated",
                    "calls": [
                        {
                            "args": [
                                {
                                    "__soil_arg_type": "decorated",
                                    "calls": [
                                        {
                                            "args": [],
                                            "kwargs": {
                                                "output_types": {
                                                    "__soil_arg_type": "lambda"
                                                }
                                            }
                                        },
                                        {
                                            "args": [
                                                {
                                                    "__soil_arg_type": "function",
                                                    "module": "unit.soil.test_utils",
                                                    "name": "my_model2"
                                                }
                                            ],
                                            "kwargs": {}
                                        }
                                    ],
                                    "module": "soil.modulify",
                                    "name": "modulify"
                                }
                            ],
                            "kwargs": {}
                        },
                        {
                            "args": [],
                            "kwargs": {
                                "f1": 2,
                                "t1": {
                                    "__soil_arg_type": "function",
                                    "module": "unit.soil.test_utils",
                                    "name": "test_fn1"
                                }
                            }
                        },
                        {
                            "args": [],
                            "kwargs": {
                                "another": 1
                            }
                        }
                    ],
                    "module": "unit.soil.test_utils",
                    "name": "trainable"
                }
            },
            "module": {
                "__soil_arg_type": "decorated",
                "calls": [
                    {
                        "args": [],
                        "kwargs": {
                            "output_types": {
                                "__soil_arg_type": "lambda"
                            }
                        }
                    },
                    {
                        "args": [
                            {
                                "__soil_arg_type": "function",
                                "module": "unit.soil.test_utils",
                                "name": "my_model"
                            }
                        ],
                        "kwargs": {}
                    }
                ],
                "module": "soil.modulify",
                "name": "modulify"
            }
        }
        assert args == expected

    @patch('soil.data')
    def test_build_arguments_9(self, mock_soil_data):
        mock_soil_data.side_effect = mock_soil_data_fn
        data = soil.data([{'test': 'test'}])
        res, = my_model(test_param=trainable(my_model2)(f1=2, t1=my_model(test=1))
                        (another=1))(cross_validation=True)(data)
        plan = res.pipeline.plan
        first = plan[0]
        args = first['args']
        expected = {
            "cross_validation": True,
            "model_params": {
                "test_param": {
                    "__soil_arg_type": "decorated",
                    "calls": [
                        {
                            "args": [
                                {
                                    "__soil_arg_type": "decorated",
                                    "calls": [
                                        {
                                            "args": [],
                                            "kwargs": {
                                                "output_types": {
                                                    "__soil_arg_type": "lambda"
                                                }
                                            }
                                        },
                                        {
                                            "args": [
                                                {
                                                    "__soil_arg_type": "function",
                                                    "module": "unit.soil.test_utils",
                                                    "name": "my_model2"
                                                }
                                            ],
                                            "kwargs": {}
                                        }
                                    ],
                                    "module": "soil.modulify",
                                    "name": "modulify"
                                }
                            ],
                            "kwargs": {}
                        },
                        {
                            "args": [],
                            "kwargs": {
                                "f1": 2,
                                "t1": {
                                    "__soil_arg_type": "decorated",
                                    "calls": [
                                        {
                                            "args": [
                                                {
                                                    "__soil_arg_type": "decorated",
                                                    "calls": [
                                                        {
                                                            "args": [],
                                                            "kwargs": {
                                                                "output_types": {
                                                                    "__soil_arg_type": "lambda"
                                                                }
                                                            }
                                                        },
                                                        {
                                                            "args": [
                                                                {
                                                                    "__soil_arg_type": "function",
                                                                    "module": "unit.soil.test_utils",
                                                                    "name": "my_model"
                                                                }
                                                            ],
                                                            "kwargs": {}
                                                        }
                                                    ],
                                                    "module": "soil.modulify",
                                                    "name": "modulify"
                                                }
                                            ],
                                            "kwargs": {}
                                        },
                                        {
                                            "args": [],
                                            "kwargs": {
                                                "test": 1
                                            }
                                        }
                                    ],
                                    "module": "unit.soil.test_utils",
                                    "name": "trainable"
                                }
                            }
                        },
                        {
                            "args": [],
                            "kwargs": {
                                "another": 1
                            }
                        }
                    ],
                    "module": "unit.soil.test_utils",
                    "name": "trainable"
                }
            },
            "module": {
                "__soil_arg_type": "decorated",
                "calls": [
                    {
                        "args": [],
                        "kwargs": {
                            "output_types": {
                                "__soil_arg_type": "lambda"
                            }
                        }
                    },
                    {
                        "args": [
                            {
                                "__soil_arg_type": "function",
                                "module": "unit.soil.test_utils",
                                "name": "my_model"
                            }
                        ],
                        "kwargs": {}
                    }
                ],
                "module": "soil.modulify",
                "name": "modulify"
            }
        }
        assert args == expected

    @patch('soil.data')
    def test_build_arguments_10(self, mock_soil_data):
        mock_soil_data.side_effect = mock_soil_data_fn
        data = soil.data([{'test': 'test'}])
        res, = my_model(test_param=my_model(f1=2, t1=my_model(test=1))
                        (another=1))(cross_validation=True)(data)
        plan = res.pipeline.plan
        first = plan[0]
        args = first['args']
        expected = {
            "cross_validation": True,
            "model_params": {
                "test_param": {
                    "__soil_arg_type": "decorated",
                    "calls": [
                        {
                            "args": [
                                {
                                    "__soil_arg_type": "decorated",
                                    "calls": [
                                        {
                                            "args": [],
                                            "kwargs": {
                                                "output_types": {
                                                    "__soil_arg_type": "lambda"
                                                }
                                            }
                                        },
                                        {
                                            "args": [
                                                {
                                                    "__soil_arg_type": "function",
                                                    "module": "unit.soil.test_utils",
                                                    "name": "my_model"
                                                }
                                            ],
                                            "kwargs": {}
                                        }
                                    ],
                                    "module": "soil.modulify",
                                    "name": "modulify"
                                }
                            ],
                            "kwargs": {}
                        },
                        {
                            "args": [],
                            "kwargs": {
                                "f1": 2,
                                "t1": {
                                    "__soil_arg_type": "decorated",
                                    "calls": [
                                        {
                                            "args": [
                                                {
                                                    "__soil_arg_type": "decorated",
                                                    "calls": [
                                                        {
                                                            "args": [],
                                                            "kwargs": {
                                                                "output_types": {
                                                                    "__soil_arg_type": "lambda"
                                                                }
                                                            }
                                                        },
                                                        {
                                                            "args": [
                                                                {
                                                                    "__soil_arg_type": "function",
                                                                    "module": "unit.soil.test_utils",
                                                                    "name": "my_model"
                                                                }
                                                            ],
                                                            "kwargs": {}
                                                        }
                                                    ],
                                                    "module": "soil.modulify",
                                                    "name": "modulify"
                                                }
                                            ],
                                            "kwargs": {}
                                        },
                                        {
                                            "args": [],
                                            "kwargs": {
                                                "test": 1
                                            }
                                        }
                                    ],
                                    "module": "unit.soil.test_utils",
                                    "name": "trainable"
                                }
                            }
                        },
                        {
                            "args": [],
                            "kwargs": {
                                "another": 1
                            }
                        }
                    ],
                    "module": "unit.soil.test_utils",
                    "name": "trainable"
                }
            },
            "module": {
                "__soil_arg_type": "decorated",
                "calls": [
                    {
                        "args": [],
                        "kwargs": {
                            "output_types": {
                                "__soil_arg_type": "lambda"
                            }
                        }
                    },
                    {
                        "args": [
                            {
                                "__soil_arg_type": "function",
                                "module": "unit.soil.test_utils",
                                "name": "my_model"
                            }
                        ],
                        "kwargs": {}
                    }
                ],
                "module": "soil.modulify",
                "name": "modulify"
            }
        }
        assert args == expected
