from apiflask import Schema, fields
from apiflask.validators import OneOf

from backend.metric.dto import MetrciDto
from backend.tree.dto import DecisionTreeParamsDto, RuleDto, TreeNodeDto


class RegressionTreeDto(DecisionTreeParamsDto):
    criterion = fields.String(
        load_default="squared_error",
        validate=OneOf(["squared_error", "friedman_mse", "absolute_error", "poisson"]),
    )


class RegressionResultDto(Schema):
    rules = fields.List(fields.Nested(RuleDto()))
    tree = fields.List(fields.Nested(TreeNodeDto()))
    mse = fields.Nested(MetrciDto())
    mae = fields.Nested(MetrciDto())
    rmse = fields.Nested(MetrciDto())
    rmae = fields.Nested(MetrciDto())
    r2 = fields.Nested(MetrciDto())