from logitbias_parser import parse
from logitbias import (
    OpAdd,
    OpMultiply,
    OpDivide,
    OpMod,
    OpAnd,
    OpOr,
    OpCompare,
    OpIn,
    OpVariableAccess,
    OpArrayAccess,
    escape_to_zero,
    eval,
    logitbias_simplify,
)
from logitbias_json import json2logitbias, logitbias_eval_all
import pytest

parser_tests = [
    ("1", 1),
    ("1 + 2", OpAdd(1, 2)),
    ("!(1 + 2)", OpCompare(False, OpAdd(1, 2), "==")),
    ("1 - 2", OpAdd(1, 2, rhs_flip_sign=True)),
    ("1 + 2 * 3", OpAdd(1, OpMultiply(2, 3))),
    ("length > 0.25", OpCompare(OpVariableAccess("length"), 0.25, ">")),
    (
        "pitch <= pitches[-1]",
        OpCompare(OpVariableAccess("pitch"), OpArrayAccess("pitches", -1), "<="),
    ),
    ("index === 1", OpCompare(OpVariableAccess("index"), 1, "==")),
    ("index == 1", OpCompare(OpVariableAccess("index"), 1, "==")),
    ("index != 1", OpCompare(OpVariableAccess("index"), 1, "!=")),
    (
        "1.5 <= beat < 1.75",
        OpAnd(
            OpCompare(1.5, OpVariableAccess("beat"), "<="),
            OpCompare(OpVariableAccess("beat"), 1.75, "<"),
        ),
    ),
    ("pitch % 12 in [1,2,3]", OpIn(OpMod(OpVariableAccess("pitch"), 12), [1, 2, 3])),
    ("0.33 in [1/3, 0.3]", OpIn(0.33, [OpDivide(1, 3), 0.3])),
    ("true", True),
    ("false", False),
]


@pytest.mark.parametrize("input,expected", parser_tests)
def test_parser(input, expected):
    assert parse(input) == expected


def raiser():
    raise Exception("whoops")


def test_eval_raises():
    with pytest.raises(Exception):
        eval(parse("pitch > 0"), {"pitch": raiser})


eval_tests = [
    ("1+2+3-4*5/2", {}, 1 + 2 + 3 - 4 * 5 / 2),
    ("!(1 == 2)", {}, True),
    ("pitch % 12 in [1,2,3,5,6,8,10,11]", {"pitch": 61}, True),
    ("pitch % 12 in [1,2,3,5,6,8,10,11]", {"pitch": 57}, False),
    ("index === 2", {"index": 2}, True),
    ("pitch <= pitches[-1]", {"pitch": 61, "pitches": [60, 61, 62]}, True),
    ("pitch > pitches[-1]", {"pitch": 61, "pitches": [60, 61, 62]}, False),
    ("pitch > 0", {"pitch": lambda: 1}, True),
    ("1 < 2 || pitch > 0", {"pitch": raiser}, True),
    ("0.333333 in [1/3]", {}, True),
    ("0.3 in [1/3, 0.4, 5+6]", {}, False),
]


@pytest.mark.parametrize("input,env,expected", eval_tests)
def test_eval(input, env, expected):
    assert eval(parse(input), env) == expected


json_tests = [
    ({"true": 1}, [(True, 1)]),
    (
        {
            "true": {
                "length > 0.25": -100,
                "length > 0.50": -1000,
            },
            "pitch >= pitches[-1]": {
                "index === 1": -101,
            },
            "pitch <= pitches[-1]": {
                "index === 2": -102,
            },
        },
        [
            (OpAnd(True, OpCompare(OpVariableAccess("length"), 0.25, ">")), -100),
            (OpAnd(True, OpCompare(OpVariableAccess("length"), 0.5, ">")), -1000),
            (
                OpAnd(
                    OpCompare(
                        OpVariableAccess("pitch"), OpArrayAccess("pitches", -1), ">="
                    ),
                    OpCompare(OpVariableAccess("index"), 1, "=="),
                ),
                -101,
            ),
            (
                OpAnd(
                    OpCompare(
                        OpVariableAccess("pitch"), OpArrayAccess("pitches", -1), "<="
                    ),
                    OpCompare(OpVariableAccess("index"), 2, "=="),
                ),
                -102,
            ),
        ],
    ),
]


@pytest.mark.parametrize("input,expected", json_tests)
def test_json(input, expected):
    assert json2logitbias(input) == expected


eval_all_inp1 = {
    "true": {
        "length > 0.25": -100,
        "length > 0.50": -1000,
    },
    "pitch >= pitches[-1]": {
        "index === 1": -101,
    },
}
eval_all_tests = [
    (
        eval_all_inp1,
        {"length": 0.3, "pitch": 61, "pitches": [60, 61, 62], "index": 1},
        -100,
    ),
    (
        eval_all_inp1,
        {"length": 0.6, "pitch": 63, "pitches": [60, 61, 62], "index": 1},
        -1201,
    ),
]


@pytest.mark.parametrize("input,env,expected", eval_all_tests)
def test_eval_all(input, env, expected):
    assert logitbias_eval_all(json2logitbias(input), env) == expected


simplify_inp1 = {
    "true": {
        "length > 0.25": -100,
        "length > 0.50": -1000,
    },
    "pitch >= pitches[-1]": {
        "index === 1": -101,
    },
    "pitch == 0 && nonvar == 1": -10000,
    "pitch == 50 && escape": -1,
}
simplify_tests = [
    (
        simplify_inp1,
        {"length": 0.3, "pitch": 61, "pitches": [60, 61, 62], "index": 1, "nonvar": 0},
        "foo",
        [None],
        [-100],
    ),
    (
        simplify_inp1,
        {"pitch": 70, "pitches": [60, 61, 62], "index": 1},
        "length",
        [0.3, 0.7],
        [-201, -1201],
    ),
    (
        simplify_inp1,
        {
            "length": 0,
            "index": 0,
            "pitch": 50,
            "pitches": [60, 61, 62],
            "escape": escape_to_zero,
        },
        "nonvar",
        [None],
        [0],
    ),
]


@pytest.mark.parametrize("input,env,variable,vals,expecteds", simplify_tests)
def test_simplify(input, env, variable, vals, expecteds):
    f = logitbias_simplify(json2logitbias(input), env, variable)
    for val, expected in zip(vals, expecteds):
        if val is None:
            assert f == expected
        else:
            assert f(val) == expected