"""
Testcases for variables:
- Input
- Output
- ActiveLearning
- VariableGroup
"""
import numpy as np
from profit.util.variable import Variable


def test_input():
    from profit.util.variable import InputVariable

    NAME = "v"
    KIND = "uniform"
    CONSTRAINTS = (2, 5)
    SIZE = (3, 1)
    DTYPE = np.float64
    VARIABLE_DICT = {
        "name": NAME,
        "kind": KIND,
        "constraints": CONSTRAINTS,
        "size": SIZE,
        "value": None,
        "dtype": DTYPE,
    }

    v1 = Variable.create(NAME, KIND, SIZE, constraints=CONSTRAINTS)
    constraints_str = ", ".join(str(c) for c in CONSTRAINTS)
    v2 = Variable.create_from_str(
        NAME,
        SIZE,
        "{kind}({constraints})".format(kind=KIND, constraints=constraints_str),
    )
    v3 = InputVariable(NAME, KIND, SIZE, CONSTRAINTS)
    variables = (v1, v2, v3)
    assert all(v.name == NAME for v in variables)
    assert all(v.size == SIZE for v in variables)
    assert all(v.dtype == DTYPE for v in variables)
    for v in variables:
        v.generate_values()
    VARIABLE_DICT["value"] = variables[0].value
    assert all(
        CONSTRAINTS[0] <= v.value.min() and v.value.max() <= CONSTRAINTS[1]
        for v in variables
    )
    assert np.all(
        v.named_value == np.array(v.value, dtype=[(v.name, v.dtype)]) for v in variables
    )
    assert np.all(v.as_dict() == VARIABLE_DICT for v in variables)
    assert np.all(
        v[key] == value for key, value in VARIABLE_DICT.items() for v in variables
    )

    CONST = 1.4
    v4 = Variable.create_from_str(NAME, SIZE, f"{CONST}")
    assert v4.kind.lower() == "constant"
    v4.generate_values()
    assert np.all(v4.value == np.ones(SIZE) * CONST)


def test_output():
    from profit.util.variable import OutputVariable

    NAME = "f"
    KIND = "Output"
    DEPENDENT = "t"
    SIZE = (3, 1)
    DTYPE = np.float64
    VARIABLE_DICT = {
        "name": NAME,
        "kind": KIND,
        "dependent": DEPENDENT,
        "size": SIZE,
        "value": None,
        "dtype": DTYPE,
    }

    DEPENDENT_SIZE = (200, 1)
    DEPENDENT_VARIABLE = Variable.create_from_str(
        "t", DEPENDENT_SIZE, "Independent(0, 100, 200)"
    )
    v1 = Variable.create(NAME, KIND, SIZE, dependent=DEPENDENT_VARIABLE)
    v2 = Variable.create_from_str(NAME, SIZE, f"{KIND}({DEPENDENT})")
    v3 = OutputVariable(NAME, KIND, SIZE, DEPENDENT)
    variables = (v1, v2, v3)
    for v in variables:
        v.resolve_dependent(DEPENDENT_VARIABLE)
    assert all(v.name == NAME for v in variables)
    assert all(v.size == (SIZE[0], DEPENDENT_SIZE[0]) for v in variables)
    assert all(v.dtype == DTYPE for v in variables)
    VARIABLE_DICT["value"] = variables[0].value
    assert all(
        v.dependent[0].as_dict() == DEPENDENT_VARIABLE.as_dict() for v in variables
    )
    assert np.all(
        v.named_value == np.array(v.value, dtype=[(v.name, v.dtype)]) for v in variables
    )
    assert np.all(v.as_dict() == VARIABLE_DICT for v in variables)
    assert np.all(
        v[key] == value for key, value in VARIABLE_DICT.items() for v in variables
    )


def test_active_learning():
    from profit.util.variable import ActiveLearningVariable

    NAME = "v"
    KIND = "ActiveLearning"
    DISTR = "uniform"
    CONSTRAINTS = (0, 2)
    SIZE = (10, 1)
    DTYPE = np.float64
    VARIABLE_DICT = {
        "name": NAME,
        "kind": KIND,
        "distr": DISTR,
        "constraints": CONSTRAINTS,
        "size": SIZE,
        "value": None,
        "dtype": DTYPE,
    }

    v1 = Variable.create(NAME, KIND, SIZE, distr=DISTR, constraints=CONSTRAINTS)
    constraints_str = ", ".join(str(c) for c in CONSTRAINTS)
    v2 = Variable.create_from_str(
        NAME,
        SIZE,
        "{kind}({constraints})".format(kind=KIND, constraints=constraints_str),
    )
    v3 = ActiveLearningVariable(NAME, KIND, SIZE, DISTR, CONSTRAINTS)
    variables = (v1, v2, v3)

    assert all(v.name == NAME for v in variables)
    assert all(v.size == SIZE for v in variables)
    assert all(v.dtype == DTYPE for v in variables)

    VARIABLE_DICT["value"] = variables[0].value
    assert all(v.distr == DISTR for v in variables)
    assert np.all(
        v.named_value == np.array(v.value, dtype=[(v.name, v.dtype)]) for v in variables
    )
    assert np.all(v.as_dict() == VARIABLE_DICT for v in variables)
    assert np.all(
        v[key] == value for key, value in VARIABLE_DICT.items() for v in variables
    )

    # Distribute prediction points for AL according to distribution
    Xpred = v1.create_Xpred(SIZE)
    assert np.all(Xpred == np.linspace(*CONSTRAINTS, SIZE[0]).reshape(SIZE))

    DISTR = "log"
    CONSTRAINTS = (1e-2, 1e3)
    v4 = ActiveLearningVariable(NAME, KIND, SIZE, DISTR, CONSTRAINTS)
    Xpred = v4.create_Xpred(SIZE)
    assert np.all(
        Xpred
        == CONSTRAINTS[0]
        * np.exp(
            (np.log(CONSTRAINTS[1]) - np.log(CONSTRAINTS[0]))
            * np.linspace(0, 1, SIZE[0])
        ).reshape(SIZE)
    )


def test_variablegroup():
    from profit.util.variable import VariableGroup

    SIZE = (3, 1)
    var = {
        "v": "Uniform(0, 1)",
        "u": "ActiveLearning(3, 4)",
        "t": "Independent(1, 5, 100)",
        "f": "Output(t)",
    }

    variables = [
        Variable.create_from_str(name, SIZE, v_str) for name, v_str in var.items()
    ]
    group = VariableGroup(samples=SIZE[0])
    group.add(variables)

    assert all(group[v.name] is v for v in variables)
    # assert group.all  # Not working yet
    variables[1].value = np.random.random(SIZE)
    variables[-1].value = np.random.random(variables[-1].size)

    assert all(
        np.all(group.input[:, i] == v.value.flat) for i, v in enumerate(variables[:2])
    )
    assert all(np.all(group.named_input[v.name] == v.value) for v in variables[:2])
    assert np.all(group.output == variables[-1].value)
    assert np.all(group.named_output["f"] == variables[-1].value)
    assert group.formatted_output.dtype.names == tuple(
        f"f({t})" for t in variables[2].value.flat
    )
    assert all(group.as_dict[v.name] == v.as_dict() for v in variables)
    group.delete_variable(0)
    assert all(
        np.all(group.input[:, i] == v.value.flat) for i, v in enumerate(variables[1:2])
    )
    group.delete_sample(0)
    assert (
        group.input.shape[0]
        == group.output.shape[0]
        == variables[1].value.shape[0]
        == variables[-1].value.shape[0]
        == SIZE[0] - 1
    )
