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, 'sklearn.ensemble.RandomForestRegressor', 'ensemble.RandomForestRegressor', ([], {'n_estimators': '(2)', 'random_state': '(1)'}), False, 'from sklearn import ensemble\n'), (12, 'm2cgen.assemblers.RandomForestModelAssembler', 'assemblers.RandomForestModelAssembler', (['estimator'], {}), False, 'from m2cgen import assemblers, ast\n'), (29, 'tests.utils.cmp_exprs', 'utils.cmp_exprs', (['actual', 'expected'], {}), False, 'from tests import utils\n'), (33, 'sklearn.ensemble.RandomForestRegressor', 'ensemble.RandomForestRegressor', ([], {'n_estimators': '(2)', 'random_state': '(13)'}), False, 'from sklearn import ensemble\n'), (37, 'm2cgen.assemblers.RandomForestModelAssembler', 'assemblers.RandomForestModelAssembler', (['estimator'], {}), False, 'from m2cgen import assemblers, ast\n'), (60, 'tests.utils.cmp_exprs', 'utils.cmp_exprs', (['actual', 'expected'], {}), False, 'from tests import utils\n'), (64, 'sklearn.ensemble.RandomForestClassifier', 'ensemble.RandomForestClassifier', ([], {'n_estimators': '(2)', 'random_state': '(13)'}), False, 'from sklearn import ensemble\n'), (69, 'm2cgen.assemblers.RandomForestModelAssembler', 'assemblers.RandomForestModelAssembler', (['estimator'], {}), False, 'from m2cgen import assemblers, ast\n'), (100, 'tests.utils.cmp_exprs', 'utils.cmp_exprs', (['actual', 'expected'], {}), False, 'from tests import utils\n'), (26, 'm2cgen.ast.NumVal', 'ast.NumVal', (['(0.5)'], {}), False, 'from m2cgen import assemblers, ast\n'), (57, 'm2cgen.ast.NumVal', 'ast.NumVal', (['(0.5)'], {}), False, 'from m2cgen import assemblers, ast\n'), (97, 'm2cgen.ast.NumVal', 'ast.NumVal', (['(0.5)'], {}), False, 'from m2cgen import assemblers, ast\n'), (17, 'm2cgen.ast.NumVal', 'ast.NumVal', (['(1.0)'], {}), False, 'from m2cgen import assemblers, ast\n'), (23, 'm2cgen.ast.NumVal', 'ast.NumVal', (['(1.0)'], {}), False, 'from m2cgen import assemblers, ast\n'), (24, 'm2cgen.ast.NumVal', 'ast.NumVal', (['(2.0)'], {}), False, 'from m2cgen import assemblers, ast\n'), (47, 'm2cgen.ast.NumVal', 'ast.NumVal', (['(1.0)'], {}), False, 'from m2cgen import assemblers, ast\n'), (48, 'm2cgen.ast.NumVal', 'ast.NumVal', (['(2.0)'], {}), False, 'from m2cgen import assemblers, ast\n'), (54, 'm2cgen.ast.NumVal', 'ast.NumVal', (['(2.0)'], {}), False, 'from m2cgen import assemblers, ast\n'), (55, 'm2cgen.ast.NumVal', 'ast.NumVal', (['(3.0)'], {}), False, 'from m2cgen import assemblers, ast\n'), (20, 'm2cgen.ast.FeatureRef', 'ast.FeatureRef', (['(0)'], {}), False, 'from m2cgen import assemblers, ast\n'), (21, 'm2cgen.ast.NumVal', 'ast.NumVal', (['(1.5)'], {}), False, 'from m2cgen import assemblers, ast\n'), (44, 'm2cgen.ast.FeatureRef', 'ast.FeatureRef', (['(0)'], {}), False, 'from m2cgen import assemblers, ast\n'), (45, 'm2cgen.ast.NumVal', 'ast.NumVal', (['(1.5)'], {}), False, 'from m2cgen import assemblers, ast\n'), (51, 'm2cgen.ast.FeatureRef', 'ast.FeatureRef', (['(0)'], {}), False, 'from m2cgen import assemblers, ast\n'), (52, 'm2cgen.ast.NumVal', 'ast.NumVal', (['(2.5)'], {}), False, 'from m2cgen import assemblers, ast\n'), (76, 'm2cgen.ast.FeatureRef', 'ast.FeatureRef', (['(0)'], {}), False, 'from m2cgen import assemblers, ast\n'), (77, 'm2cgen.ast.NumVal', 'ast.NumVal', (['(1.5)'], {}), False, 'from m2cgen import assemblers, ast\n'), (87, 'm2cgen.ast.FeatureRef', 'ast.FeatureRef', (['(0)'], {}), False, 'from m2cgen import assemblers, ast\n'), (88, 'm2cgen.ast.NumVal', 'ast.NumVal', (['(2.5)'], {}), False, 'from m2cgen import assemblers, ast\n'), (80, 'm2cgen.ast.NumVal', 'ast.NumVal', (['(0.0)'], {}), False, 'from m2cgen import assemblers, ast\n'), (81, 'm2cgen.ast.NumVal', 'ast.NumVal', (['(1.0)'], {}), False, 'from m2cgen import assemblers, ast\n'), (83, 'm2cgen.ast.NumVal', 'ast.NumVal', (['(1.0)'], {}), False, 'from m2cgen import assemblers, ast\n'), (84, 'm2cgen.ast.NumVal', 'ast.NumVal', (['(0.0)'], {}), False, 'from m2cgen import assemblers, ast\n'), (91, 'm2cgen.ast.NumVal', 'ast.NumVal', (['(1.0)'], {}), False, 'from m2cgen import assemblers, ast\n'), (92, 'm2cgen.ast.NumVal', 'ast.NumVal', (['(0.0)'], {}), False, 'from m2cgen import assemblers, ast\n'), (94, 'm2cgen.ast.NumVal', 'ast.NumVal', (['(0.0)'], {}), False, 'from m2cgen import assemblers, ast\n'), (95, 'm2cgen.ast.NumVal', 'ast.NumVal', (['(1.0)'], {}), False, 'from m2cgen import assemblers, ast\n')]