from sklearn import ensemble

from m2cgen import assemblers, ast
from tests import utils


def test_single_condition():
    estimator = ensemble.RandomForestRegressor(n_estimators=2, random_state=1)

    estimator.fit([[1], [2]], [1, 2])

    assembler = assemblers.RandomForestModelAssembler(estimator)
    actual = assembler.assemble()

    expected = ast.BinNumExpr(
        ast.BinNumExpr(
            ast.NumVal(1.0),
            ast.IfExpr(
                ast.CompExpr(
                    ast.FeatureRef(0),
                    ast.NumVal(1.5),
                    ast.CompOpType.LTE),
                ast.NumVal(1.0),
                ast.NumVal(2.0)),
            ast.BinNumOpType.ADD),
        ast.NumVal(0.5),
        ast.BinNumOpType.MUL)

    assert utils.cmp_exprs(actual, expected)


def test_two_conditions():
    estimator = ensemble.RandomForestRegressor(n_estimators=2, random_state=13)

    estimator.fit([[1], [2], [3]], [1, 2, 3])

    assembler = assemblers.RandomForestModelAssembler(estimator)
    actual = assembler.assemble()

    expected = ast.BinNumExpr(
        ast.BinNumExpr(
            ast.IfExpr(
                ast.CompExpr(
                    ast.FeatureRef(0),
                    ast.NumVal(1.5),
                    ast.CompOpType.LTE),
                ast.NumVal(1.0),
                ast.NumVal(2.0)),
            ast.IfExpr(
                ast.CompExpr(
                    ast.FeatureRef(0),
                    ast.NumVal(2.5),
                    ast.CompOpType.LTE),
                ast.NumVal(2.0),
                ast.NumVal(3.0)),
            ast.BinNumOpType.ADD),
        ast.NumVal(0.5),
        ast.BinNumOpType.MUL)

    assert utils.cmp_exprs(actual, expected)


def test_multi_class():
    estimator = ensemble.RandomForestClassifier(
        n_estimators=2, random_state=13)

    estimator.fit([[1], [2], [3]], [1, -1, 1])

    assembler = assemblers.RandomForestModelAssembler(estimator)
    actual = assembler.assemble()

    expected = ast.BinVectorNumExpr(
        ast.BinVectorExpr(
            ast.IfExpr(
                ast.CompExpr(
                    ast.FeatureRef(0),
                    ast.NumVal(1.5),
                    ast.CompOpType.LTE),
                ast.VectorVal([
                    ast.NumVal(0.0),
                    ast.NumVal(1.0)]),
                ast.VectorVal([
                    ast.NumVal(1.0),
                    ast.NumVal(0.0)])),
            ast.IfExpr(
                ast.CompExpr(
                    ast.FeatureRef(0),
                    ast.NumVal(2.5),
                    ast.CompOpType.LTE),
                ast.VectorVal([
                    ast.NumVal(1.0),
                    ast.NumVal(0.0)]),
                ast.VectorVal([
                    ast.NumVal(0.0),
                    ast.NumVal(1.0)])),
            ast.BinNumOpType.ADD),
        ast.NumVal(0.5),
        ast.BinNumOpType.MUL)

    assert utils.cmp_exprs(actual, expected)

[((8, 16, 8, 78), 'sklearn.ensemble.RandomForestRegressor', 'ensemble.RandomForestRegressor', (), '', False, 'from sklearn import ensemble\n'), ((12, 16, 12, 64), 'm2cgen.assemblers.RandomForestModelAssembler', 'assemblers.RandomForestModelAssembler', ({(12, 54, 12, 63): 'estimator'}, {}), '(estimator)', False, 'from m2cgen import assemblers, ast\n'), ((29, 11, 29, 44), 'tests.utils.cmp_exprs', 'utils.cmp_exprs', ({(29, 27, 29, 33): 'actual', (29, 35, 29, 43): 'expected'}, {}), '(actual, expected)', False, 'from tests import utils\n'), ((33, 16, 33, 79), 'sklearn.ensemble.RandomForestRegressor', 'ensemble.RandomForestRegressor', (), '', False, 'from sklearn import ensemble\n'), ((37, 16, 37, 64), 'm2cgen.assemblers.RandomForestModelAssembler', 'assemblers.RandomForestModelAssembler', ({(37, 54, 37, 63): 'estimator'}, {}), '(estimator)', False, 'from m2cgen import assemblers, ast\n'), ((60, 11, 60, 44), 'tests.utils.cmp_exprs', 'utils.cmp_exprs', ({(60, 27, 60, 33): 'actual', (60, 35, 60, 43): 'expected'}, {}), '(actual, expected)', False, 'from tests import utils\n'), ((64, 16, 65, 40), 'sklearn.ensemble.RandomForestClassifier', 'ensemble.RandomForestClassifier', (), '', False, 'from sklearn import ensemble\n'), ((69, 16, 69, 64), 'm2cgen.assemblers.RandomForestModelAssembler', 'assemblers.RandomForestModelAssembler', ({(69, 54, 69, 63): 'estimator'}, {}), '(estimator)', False, 'from m2cgen import assemblers, ast\n'), ((100, 11, 100, 44), 'tests.utils.cmp_exprs', 'utils.cmp_exprs', ({(100, 27, 100, 33): 'actual', (100, 35, 100, 43): 'expected'}, {}), '(actual, expected)', False, 'from tests import utils\n'), ((26, 8, 26, 23), 'm2cgen.ast.NumVal', 'ast.NumVal', ({(26, 19, 26, 22): '0.5'}, {}), '(0.5)', False, 'from m2cgen import assemblers, ast\n'), ((57, 8, 57, 23), 'm2cgen.ast.NumVal', 'ast.NumVal', ({(57, 19, 57, 22): '0.5'}, {}), '(0.5)', False, 'from m2cgen import assemblers, ast\n'), ((97, 8, 97, 23), 'm2cgen.ast.NumVal', 'ast.NumVal', ({(97, 19, 97, 22): '0.5'}, {}), '(0.5)', False, 'from m2cgen import assemblers, ast\n'), ((17, 12, 17, 27), 'm2cgen.ast.NumVal', 'ast.NumVal', ({(17, 23, 17, 26): '1.0'}, {}), '(1.0)', False, 'from m2cgen import assemblers, ast\n'), ((23, 16, 23, 31), 'm2cgen.ast.NumVal', 'ast.NumVal', ({(23, 27, 23, 30): '1.0'}, {}), '(1.0)', False, 'from m2cgen import assemblers, ast\n'), ((24, 16, 24, 31), 'm2cgen.ast.NumVal', 'ast.NumVal', ({(24, 27, 24, 30): '2.0'}, {}), '(2.0)', False, 'from m2cgen import assemblers, ast\n'), ((47, 16, 47, 31), 'm2cgen.ast.NumVal', 'ast.NumVal', ({(47, 27, 47, 30): '1.0'}, {}), '(1.0)', False, 'from m2cgen import assemblers, ast\n'), ((48, 16, 48, 31), 'm2cgen.ast.NumVal', 'ast.NumVal', ({(48, 27, 48, 30): '2.0'}, {}), '(2.0)', False, 'from m2cgen import assemblers, ast\n'), ((54, 16, 54, 31), 'm2cgen.ast.NumVal', 'ast.NumVal', ({(54, 27, 54, 30): '2.0'}, {}), '(2.0)', False, 'from m2cgen import assemblers, ast\n'), ((55, 16, 55, 31), 'm2cgen.ast.NumVal', 'ast.NumVal', ({(55, 27, 55, 30): '3.0'}, {}), '(3.0)', False, 'from m2cgen import assemblers, ast\n'), ((20, 20, 20, 37), 'm2cgen.ast.FeatureRef', 'ast.FeatureRef', ({(20, 35, 20, 36): '0'}, {}), '(0)', False, 'from m2cgen import assemblers, ast\n'), ((21, 20, 21, 35), 'm2cgen.ast.NumVal', 'ast.NumVal', ({(21, 31, 21, 34): '1.5'}, {}), '(1.5)', False, 'from m2cgen import assemblers, ast\n'), ((44, 20, 44, 37), 'm2cgen.ast.FeatureRef', 'ast.FeatureRef', ({(44, 35, 44, 36): '0'}, {}), '(0)', False, 'from m2cgen import assemblers, ast\n'), ((45, 20, 45, 35), 'm2cgen.ast.NumVal', 'ast.NumVal', ({(45, 31, 45, 34): '1.5'}, {}), '(1.5)', False, 'from m2cgen import assemblers, ast\n'), ((51, 20, 51, 37), 'm2cgen.ast.FeatureRef', 'ast.FeatureRef', ({(51, 35, 51, 36): '0'}, {}), '(0)', False, 'from m2cgen import assemblers, ast\n'), ((52, 20, 52, 35), 'm2cgen.ast.NumVal', 'ast.NumVal', ({(52, 31, 52, 34): '2.5'}, {}), '(2.5)', False, 'from m2cgen import assemblers, ast\n'), ((76, 20, 76, 37), 'm2cgen.ast.FeatureRef', 'ast.FeatureRef', ({(76, 35, 76, 36): '0'}, {}), '(0)', False, 'from m2cgen import assemblers, ast\n'), ((77, 20, 77, 35), 'm2cgen.ast.NumVal', 'ast.NumVal', ({(77, 31, 77, 34): '1.5'}, {}), '(1.5)', False, 'from m2cgen import assemblers, ast\n'), ((87, 20, 87, 37), 'm2cgen.ast.FeatureRef', 'ast.FeatureRef', ({(87, 35, 87, 36): '0'}, {}), '(0)', False, 'from m2cgen import assemblers, ast\n'), ((88, 20, 88, 35), 'm2cgen.ast.NumVal', 'ast.NumVal', ({(88, 31, 88, 34): '2.5'}, {}), '(2.5)', False, 'from m2cgen import assemblers, ast\n'), ((80, 20, 80, 35), 'm2cgen.ast.NumVal', 'ast.NumVal', ({(80, 31, 80, 34): '0.0'}, {}), '(0.0)', False, 'from m2cgen import assemblers, ast\n'), ((81, 20, 81, 35), 'm2cgen.ast.NumVal', 'ast.NumVal', ({(81, 31, 81, 34): '1.0'}, {}), '(1.0)', False, 'from m2cgen import assemblers, ast\n'), ((83, 20, 83, 35), 'm2cgen.ast.NumVal', 'ast.NumVal', ({(83, 31, 83, 34): '1.0'}, {}), '(1.0)', False, 'from m2cgen import assemblers, ast\n'), ((84, 20, 84, 35), 'm2cgen.ast.NumVal', 'ast.NumVal', ({(84, 31, 84, 34): '0.0'}, {}), '(0.0)', False, 'from m2cgen import assemblers, ast\n'), ((91, 20, 91, 35), 'm2cgen.ast.NumVal', 'ast.NumVal', ({(91, 31, 91, 34): '1.0'}, {}), '(1.0)', False, 'from m2cgen import assemblers, ast\n'), ((92, 20, 92, 35), 'm2cgen.ast.NumVal', 'ast.NumVal', ({(92, 31, 92, 34): '0.0'}, {}), '(0.0)', False, 'from m2cgen import assemblers, ast\n'), ((94, 20, 94, 35), 'm2cgen.ast.NumVal', 'ast.NumVal', ({(94, 31, 94, 34): '0.0'}, {}), '(0.0)', False, 'from m2cgen import assemblers, ast\n'), ((95, 20, 95, 35), 'm2cgen.ast.NumVal', 'ast.NumVal', ({(95, 31, 95, 34): '1.0'}, {}), '(1.0)', False, 'from m2cgen import assemblers, ast\n')]