from dataclasses import dataclass
from typing import List

from backend.metric.model import MetricValue
from backend.tree.model import Rule, TreeNode


@dataclass
class RegressionTreeParams:
    criterion: str
    splitter: str
    max_depth: int
    min_samples_split: int
    min_samples_leaf: int
    min_weight_fraction_leaf: float
    max_features: str
    random_state: int
    max_leaf_nodes: int
    min_impurity_decrease: float
    ccp_alpha: float


@dataclass
class RegressionResult:
    tree: List[TreeNode]
    rules: List[Rule]
    mse: MetricValue
    mae: MetricValue
    rmse: MetricValue
    rmae: MetricValue
    r2: MetricValue