#!/usr/bin/env python
"""Perform a functional test for client helper functions."""
import os

import pytest

import orion.core.cli
from orion.core.utils.exceptions import InvalidResult
from orion.core.worker.consumer import Consumer
from orion.testing import OrionState


def test_interrupt(monkeypatch, capsys):
    """Test interruption from within user script."""
    with OrionState() as cfg:
        storage = cfg.storage

        monkeypatch.chdir(os.path.dirname(os.path.abspath(__file__)))

        user_args = ["-x~uniform(-50, 50, precision=5)"]

        error_code = orion.core.cli.main(
            [
                "hunt",
                "--config",
                "./orion_config.yaml",
                "--exp-max-trials",
                "1",
                "--worker-trials",
                "1",
                "python",
                "black_box.py",
                "interrupt_trial",
            ]
            + user_args
        )

        assert error_code == 130

        captured = capsys.readouterr()
        assert captured.out == "Orion is interrupted.\n"
        assert captured.err == ""

        exp = list(storage.fetch_experiments({"name": "voila_voici"}))
        exp = exp[0]
        exp_id = exp["_id"]
        trials = list(storage.fetch_trials(uid=exp_id))
        assert len(trials) == 1
        assert trials[0].status == "interrupted"


def test_interrupt_diff_code(monkeypatch, capsys, storage):
    """Test interruption from within user script with custom int code"""
    monkeypatch.chdir(os.path.dirname(os.path.abspath(__file__)))

    user_args = ["-x~uniform(-50, 50, precision=5)"]

    # Set local to 200
    orion.core.config.worker.interrupt_signal_code = 200

    # But child won't be passed ORION_INTERRUPT_CODE and therefore will send default code 130
    def empty_env(self, trial, results_file=None):
        return os.environ

    with monkeypatch.context() as m:
        m.setattr(Consumer, "get_execution_environment", empty_env)

        # Interrupt won't be interpreted properly and trials will be marked as broken
        error_code = orion.core.cli.main(
            [
                "hunt",
                "--config",
                "./orion_config.yaml",
                "--worker-max-broken",
                "2",
                "--exp-max-trials",
                "2",
                "--worker-trials",
                "2",
                "python",
                "black_box.py",
                "interrupt_trial",
            ]
            + user_args
        )

        assert error_code == 0

        exp = list(storage.fetch_experiments({"name": "voila_voici"}))
        exp = exp[0]
        exp_id = exp["_id"]
        trials = list(storage.fetch_trials(uid=exp_id))
        assert len(trials) == 2
        assert trials[0].status == "broken"

    # This time we use true `get_execution_environment which pass properly int code to child.
    error_code = orion.core.cli.main(
        [
            "hunt",
            "--config",
            "./orion_config.yaml",
            "--exp-max-trial",
            "1",
            "--worker-trials",
            "1",
            "python",
            "black_box.py",
            "interrupt_trial",
        ]
        + user_args
    )

    assert error_code == 130

    captured = capsys.readouterr()
    assert "Orion is interrupted.\n" in captured.out
    assert captured.err == ""

    exp = list(storage.fetch_experiments({"name": "voila_voici"}))
    exp = exp[0]
    exp_id = exp["_id"]
    trials = list(storage.fetch_trials(uid=exp_id))
    assert len(trials) == 3
    assert trials[-1].status == "interrupted"


# TODO:

# test no call to any report

# Add all this in DOC.


@pytest.mark.parametrize("fct", ["report_bad_trial", "report_objective"])
def test_report_no_name(monkeypatch, fct):
    """Test report helper functions with default names"""

    with OrionState() as cfg:
        monkeypatch.chdir(os.path.dirname(os.path.abspath(__file__)))

        user_args = ["-x~uniform(-50, 50, precision=5)"]

        orion.core.cli.main(
            [
                "hunt",
                "--config",
                "./orion_config.yaml",
                "--exp-max-trials",
                "2",
                "--worker-trials",
                "2",
                "python",
                "black_box.py",
                fct,
                "--objective",
                "1.0",
            ]
            + user_args
        )

        exp = list(cfg.storage.fetch_experiments({"name": "voila_voici"}))
        exp = exp[0]
        exp_id = exp["_id"]
        trials = list(cfg.storage.fetch_trials(uid=exp_id))
        assert len(trials) == 2
        assert trials[0].status == "completed"
        assert trials[0].results[0].name == "objective"
        assert trials[0].results[0].type == "objective"
        assert trials[0].results[0].value == 1.0


@pytest.mark.parametrize("fct", ["report_bad_trial", "report_objective"])
def test_report_with_name(monkeypatch, fct):
    """Test report helper functions with custom names"""
    with OrionState() as cfg:
        storage = cfg.storage
        monkeypatch.chdir(os.path.dirname(os.path.abspath(__file__)))

        user_args = ["-x~uniform(-50, 50, precision=5)"]

        orion.core.cli.main(
            [
                "hunt",
                "--config",
                "./orion_config.yaml",
                "--exp-max-trials",
                "2",
                "--worker-trials",
                "2",
                "python",
                "black_box.py",
                fct,
                "--objective",
                "1.0",
                "--name",
                "metric",
            ]
            + user_args
        )

        exp = list(storage.fetch_experiments({"name": "voila_voici"}))
        exp = exp[0]
        exp_id = exp["_id"]
        trials = list(storage.fetch_trials(uid=exp_id))
        assert len(trials) == 2
        assert trials[0].status == "completed"
        assert trials[0].results[0].name == "metric"
        assert trials[0].results[0].type == "objective"
        assert trials[0].results[0].value == 1.0


@pytest.mark.parametrize("fct", ["report_bad_trial", "report_objective"])
def test_report_with_bad_objective(monkeypatch, fct):
    """Test report helper functions with bad objective types"""
    with OrionState() as cfg:
        monkeypatch.chdir(os.path.dirname(os.path.abspath(__file__)))

        user_args = ["-x~uniform(-50, 50, precision=5)"]

        with pytest.raises(InvalidResult) as exc:
            orion.core.cli.main(
                [
                    "hunt",
                    "--config",
                    "./orion_config.yaml",
                    "--exp-max-trials",
                    "2",
                    "--worker-trials",
                    "2",
                    "python",
                    "black_box.py",
                    fct,
                    "--objective",
                    "oh oh",
                ]
                + user_args
            )

        assert "must contain a type `objective` with type float/int" in str(exc.value)


def test_report_with_bad_trial_no_objective(monkeypatch):
    """Test bad trial report helper function with default objective."""
    with OrionState() as cfg:
        storage = cfg.storage

        monkeypatch.chdir(os.path.dirname(os.path.abspath(__file__)))

        user_args = ["-x~uniform(-50, 50, precision=5)"]

        orion.core.cli.main(
            [
                "hunt",
                "--config",
                "./orion_config.yaml",
                "--exp-max-trials",
                "2",
                "--worker-trials",
                "2",
                "python",
                "black_box.py",
                "report_bad_trial",
            ]
            + user_args
        )

        exp = list(storage.fetch_experiments({"name": "voila_voici"}))
        exp = exp[0]
        exp_id = exp["_id"]
        trials = list(storage.fetch_trials(uid=exp_id))
        assert len(trials) == 2
        assert trials[0].status == "completed"
        assert trials[0].results[0].name == "objective"
        assert trials[0].results[0].type == "objective"
        assert trials[0].results[0].value == 1e10


def test_report_with_bad_trial_with_data(monkeypatch):
    """Test bad trial report helper function with additional data."""
    with OrionState() as cfg:
        storage = cfg.storage
        monkeypatch.chdir(os.path.dirname(os.path.abspath(__file__)))

        user_args = ["-x~uniform(-50, 50, precision=5)"]

        orion.core.cli.main(
            [
                "hunt",
                "--config",
                "./orion_config.yaml",
                "--exp-max-trials",
                "2",
                "--worker-trials",
                "2",
                "python",
                "black_box.py",
                "report_bad_trial",
                "--data",
                "another",
            ]
            + user_args
        )

        exp = list(storage.fetch_experiments({"name": "voila_voici"}))
        exp = exp[0]
        exp_id = exp["_id"]
        trials = list(storage.fetch_trials(uid=exp_id))
        assert len(trials) == 2
        assert trials[0].status == "completed"
        assert trials[0].results[0].name == "objective"
        assert trials[0].results[0].type == "objective"
        assert trials[0].results[0].value == 1e10

        assert trials[0].results[1].name == "another"
        assert trials[0].results[1].type == "constraint"
        assert trials[0].results[1].value == 1.0


def test_no_report(monkeypatch, capsys):
    """Test script call without any results reported."""
    with OrionState() as cfg:
        storage = cfg.storage
        monkeypatch.chdir(os.path.dirname(os.path.abspath(__file__)))

        user_args = ["-x~uniform(-50, 50, precision=5)"]

        errorcode = orion.core.cli.main(
            [
                "hunt",
                "--config",
                "./orion_config.yaml",
                "--exp-max-trials",
                "2",
                "--worker-trials",
                "2",
                "python",
                "black_box.py",
                "no_report",
            ]
            + user_args
        )

        assert errorcode == 1

        captured = capsys.readouterr()
        assert captured.out == ""
        assert "Cannot parse result file" in captured.err
