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)

[((136, 198), 'sklearn.ensemble.RandomForestRegressor', 'ensemble.RandomForestRegressor', ([], {'n_estimators': '(2)', 'random_state': '(1)'}), '(n_estimators=2, random_state=1)\n', (166, 198), False, 'from sklearn import ensemble\n'), ((255, 303), 'm2cgen.assemblers.RandomForestModelAssembler', 'assemblers.RandomForestModelAssembler', (['estimator'], {}), '(estimator)\n', (292, 303), False, 'from m2cgen import assemblers, ast\n'), ((763, 796), 'tests.utils.cmp_exprs', 'utils.cmp_exprs', (['actual', 'expected'], {}), '(actual, expected)\n', (778, 796), False, 'from tests import utils\n'), ((842, 905), 'sklearn.ensemble.RandomForestRegressor', 'ensemble.RandomForestRegressor', ([], {'n_estimators': '(2)', 'random_state': '(13)'}), '(n_estimators=2, random_state=13)\n', (872, 905), False, 'from sklearn import ensemble\n'), ((970, 1018), 'm2cgen.assemblers.RandomForestModelAssembler', 'assemblers.RandomForestModelAssembler', (['estimator'], {}), '(estimator)\n', (1007, 1018), False, 'from m2cgen import assemblers, ast\n'), ((1687, 1720), 'tests.utils.cmp_exprs', 'utils.cmp_exprs', (['actual', 'expected'], {}), '(actual, expected)\n', (1702, 1720), False, 'from tests import utils\n'), ((1763, 1827), 'sklearn.ensemble.RandomForestClassifier', 'ensemble.RandomForestClassifier', ([], {'n_estimators': '(2)', 'random_state': '(13)'}), '(n_estimators=2, random_state=13)\n', (1794, 1827), False, 'from sklearn import ensemble\n'), ((1902, 1950), 'm2cgen.assemblers.RandomForestModelAssembler', 'assemblers.RandomForestModelAssembler', (['estimator'], {}), '(estimator)\n', (1939, 1950), False, 'from m2cgen import assemblers, ast\n'), ((2928, 2961), 'tests.utils.cmp_exprs', 'utils.cmp_exprs', (['actual', 'expected'], {}), '(actual, expected)\n', (2943, 2961), False, 'from tests import utils\n'), ((704, 719), 'm2cgen.ast.NumVal', 'ast.NumVal', (['(0.5)'], {}), '(0.5)\n', (714, 719), False, 'from m2cgen import assemblers, ast\n'), ((1628, 1643), 'm2cgen.ast.NumVal', 'ast.NumVal', (['(0.5)'], {}), '(0.5)\n', (1638, 1643), False, 'from m2cgen import assemblers, ast\n'), ((2869, 2884), 'm2cgen.ast.NumVal', 'ast.NumVal', (['(0.5)'], {}), '(0.5)\n', (2879, 2884), False, 'from m2cgen import assemblers, ast\n'), ((406, 421), 'm2cgen.ast.NumVal', 'ast.NumVal', (['(1.0)'], {}), '(1.0)\n', (416, 421), False, 'from m2cgen import assemblers, ast\n'), ((610, 625), 'm2cgen.ast.NumVal', 'ast.NumVal', (['(1.0)'], {}), '(1.0)\n', (620, 625), False, 'from m2cgen import assemblers, ast\n'), ((643, 658), 'm2cgen.ast.NumVal', 'ast.NumVal', (['(2.0)'], {}), '(2.0)\n', (653, 658), False, 'from m2cgen import assemblers, ast\n'), ((1296, 1311), 'm2cgen.ast.NumVal', 'ast.NumVal', (['(1.0)'], {}), '(1.0)\n', (1306, 1311), False, 'from m2cgen import assemblers, ast\n'), ((1329, 1344), 'm2cgen.ast.NumVal', 'ast.NumVal', (['(2.0)'], {}), '(2.0)\n', (1339, 1344), False, 'from m2cgen import assemblers, ast\n'), ((1534, 1549), 'm2cgen.ast.NumVal', 'ast.NumVal', (['(2.0)'], {}), '(2.0)\n', (1544, 1549), False, 'from m2cgen import assemblers, ast\n'), ((1567, 1582), 'm2cgen.ast.NumVal', 'ast.NumVal', (['(3.0)'], {}), '(3.0)\n', (1577, 1582), False, 'from m2cgen import assemblers, ast\n'), ((497, 514), 'm2cgen.ast.FeatureRef', 'ast.FeatureRef', (['(0)'], {}), '(0)\n', (511, 514), False, 'from m2cgen import assemblers, ast\n'), ((536, 551), 'm2cgen.ast.NumVal', 'ast.NumVal', (['(1.5)'], {}), '(1.5)\n', (546, 551), False, 'from m2cgen import assemblers, ast\n'), ((1183, 1200), 'm2cgen.ast.FeatureRef', 'ast.FeatureRef', (['(0)'], {}), '(0)\n', (1197, 1200), False, 'from m2cgen import assemblers, ast\n'), ((1222, 1237), 'm2cgen.ast.NumVal', 'ast.NumVal', (['(1.5)'], {}), '(1.5)\n', (1232, 1237), False, 'from m2cgen import assemblers, ast\n'), ((1421, 1438), 'm2cgen.ast.FeatureRef', 'ast.FeatureRef', (['(0)'], {}), '(0)\n', (1435, 1438), False, 'from m2cgen import assemblers, ast\n'), ((1460, 1475), 'm2cgen.ast.NumVal', 'ast.NumVal', (['(2.5)'], {}), '(2.5)\n', (1470, 1475), False, 'from m2cgen import assemblers, ast\n'), ((2124, 2141), 'm2cgen.ast.FeatureRef', 'ast.FeatureRef', (['(0)'], {}), '(0)\n', (2138, 2141), False, 'from m2cgen import assemblers, ast\n'), ((2163, 2178), 'm2cgen.ast.NumVal', 'ast.NumVal', (['(1.5)'], {}), '(1.5)\n', (2173, 2178), False, 'from m2cgen import assemblers, ast\n'), ((2512, 2529), 'm2cgen.ast.FeatureRef', 'ast.FeatureRef', (['(0)'], {}), '(0)\n', (2526, 2529), False, 'from m2cgen import assemblers, ast\n'), ((2551, 2566), 'm2cgen.ast.NumVal', 'ast.NumVal', (['(2.5)'], {}), '(2.5)\n', (2561, 2566), False, 'from m2cgen import assemblers, ast\n'), ((2273, 2288), 'm2cgen.ast.NumVal', 'ast.NumVal', (['(0.0)'], {}), '(0.0)\n', (2283, 2288), False, 'from m2cgen import assemblers, ast\n'), ((2310, 2325), 'm2cgen.ast.NumVal', 'ast.NumVal', (['(1.0)'], {}), '(1.0)\n', (2320, 2325), False, 'from m2cgen import assemblers, ast\n'), ((2381, 2396), 'm2cgen.ast.NumVal', 'ast.NumVal', (['(1.0)'], {}), '(1.0)\n', (2391, 2396), False, 'from m2cgen import assemblers, ast\n'), ((2418, 2433), 'm2cgen.ast.NumVal', 'ast.NumVal', (['(0.0)'], {}), '(0.0)\n', (2428, 2433), False, 'from m2cgen import assemblers, ast\n'), ((2661, 2676), 'm2cgen.ast.NumVal', 'ast.NumVal', (['(1.0)'], {}), '(1.0)\n', (2671, 2676), False, 'from m2cgen import assemblers, ast\n'), ((2698, 2713), 'm2cgen.ast.NumVal', 'ast.NumVal', (['(0.0)'], {}), '(0.0)\n', (2708, 2713), False, 'from m2cgen import assemblers, ast\n'), ((2769, 2784), 'm2cgen.ast.NumVal', 'ast.NumVal', (['(0.0)'], {}), '(0.0)\n', (2779, 2784), False, 'from m2cgen import assemblers, ast\n'), ((2806, 2821), 'm2cgen.ast.NumVal', 'ast.NumVal', (['(1.0)'], {}), '(1.0)\n', (2816, 2821), False, 'from m2cgen import assemblers, ast\n')]