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', False, 'from sklearn import ensemble\n'), (12, 'm2cgen.assemblers.RandomForestModelAssembler', 'assemblers.RandomForestModelAssembler', False, 'from m2cgen import assemblers, ast\n'), (29, 'tests.utils.cmp_exprs', 'utils.cmp_exprs', False, 'from tests import utils\n'), (33, 'sklearn.ensemble.RandomForestRegressor', 'ensemble.RandomForestRegressor', False, 'from sklearn import ensemble\n'), (37, 'm2cgen.assemblers.RandomForestModelAssembler', 'assemblers.RandomForestModelAssembler', False, 'from m2cgen import assemblers, ast\n'), (60, 'tests.utils.cmp_exprs', 'utils.cmp_exprs', False, 'from tests import utils\n'), (64, 'sklearn.ensemble.RandomForestClassifier', 'ensemble.RandomForestClassifier', False, 'from sklearn import ensemble\n'), (69, 'm2cgen.assemblers.RandomForestModelAssembler', 'assemblers.RandomForestModelAssembler', False, 'from m2cgen import assemblers, ast\n'), (100, 'tests.utils.cmp_exprs', 'utils.cmp_exprs', False, 'from tests import utils\n'), (26, 'm2cgen.ast.NumVal', 'ast.NumVal', False, 'from m2cgen import assemblers, ast\n'), (57, 'm2cgen.ast.NumVal', 'ast.NumVal', False, 'from m2cgen import assemblers, ast\n'), (97, 'm2cgen.ast.NumVal', 'ast.NumVal', False, 'from m2cgen import assemblers, ast\n'), (17, 'm2cgen.ast.NumVal', 'ast.NumVal', False, 'from m2cgen import assemblers, ast\n'), (23, 'm2cgen.ast.NumVal', 'ast.NumVal', False, 'from m2cgen import assemblers, ast\n'), (24, 'm2cgen.ast.NumVal', 'ast.NumVal', False, 'from m2cgen import assemblers, ast\n'), (47, 'm2cgen.ast.NumVal', 'ast.NumVal', False, 'from m2cgen import assemblers, ast\n'), (48, 'm2cgen.ast.NumVal', 'ast.NumVal', False, 'from m2cgen import assemblers, ast\n'), (54, 'm2cgen.ast.NumVal', 'ast.NumVal', False, 'from m2cgen import assemblers, ast\n'), (55, 'm2cgen.ast.NumVal', 'ast.NumVal', False, 'from m2cgen import assemblers, ast\n'), (20, 'm2cgen.ast.FeatureRef', 'ast.FeatureRef', False, 'from m2cgen import assemblers, ast\n'), (21, 'm2cgen.ast.NumVal', 'ast.NumVal', False, 'from m2cgen import assemblers, ast\n'), (44, 'm2cgen.ast.FeatureRef', 'ast.FeatureRef', False, 'from m2cgen import assemblers, ast\n'), (45, 'm2cgen.ast.NumVal', 'ast.NumVal', False, 'from m2cgen import assemblers, ast\n'), (51, 'm2cgen.ast.FeatureRef', 'ast.FeatureRef', False, 'from m2cgen import assemblers, ast\n'), (52, 'm2cgen.ast.NumVal', 'ast.NumVal', False, 'from m2cgen import assemblers, ast\n'), (76, 'm2cgen.ast.FeatureRef', 'ast.FeatureRef', False, 'from m2cgen import assemblers, ast\n'), (77, 'm2cgen.ast.NumVal', 'ast.NumVal', False, 'from m2cgen import assemblers, ast\n'), (87, 'm2cgen.ast.FeatureRef', 'ast.FeatureRef', False, 'from m2cgen import assemblers, ast\n'), (88, 'm2cgen.ast.NumVal', 'ast.NumVal', False, 'from m2cgen import assemblers, ast\n'), (80, 'm2cgen.ast.NumVal', 'ast.NumVal', False, 'from m2cgen import assemblers, ast\n'), (81, 'm2cgen.ast.NumVal', 'ast.NumVal', False, 'from m2cgen import assemblers, ast\n'), (83, 'm2cgen.ast.NumVal', 'ast.NumVal', False, 'from m2cgen import assemblers, ast\n'), (84, 'm2cgen.ast.NumVal', 'ast.NumVal', False, 'from m2cgen import assemblers, ast\n'), (91, 'm2cgen.ast.NumVal', 'ast.NumVal', False, 'from m2cgen import assemblers, ast\n'), (92, 'm2cgen.ast.NumVal', 'ast.NumVal', False, 'from m2cgen import assemblers, ast\n'), (94, 'm2cgen.ast.NumVal', 'ast.NumVal', False, 'from m2cgen import assemblers, ast\n'), (95, 'm2cgen.ast.NumVal', 'ast.NumVal', False, 'from m2cgen import assemblers, ast\n')]