from werkzeug.datastructures import FileStorage

from backend.classification import fit_classification_model
from backend.classification.model import ClassificationResult
from backend.dataset import Dataset
from backend.dataset.model import DatasetParams, SplittedDataset
from backend.regression import fit_regression_model
from backend.regression.model import RegressionResult
from backend.tree.model import DecisionTreeParams


def run_regression(
    path: str | None,
    file: FileStorage,
    dataset_params: DatasetParams,
    tree_params: DecisionTreeParams,
) -> RegressionResult:
    try:
        dataset: Dataset = Dataset(path=path, file=file)
        data = dataset.read(dataset_params)
        splitted_dataset: SplittedDataset = dataset.split_regression(
            data=data,
            params=dataset_params,
            random_state=tree_params.random_state,
        )
        result = fit_regression_model(
            data=splitted_dataset,
            params=tree_params,
        )
    finally:
        dataset.remove()
    return result


def run_classification(
    path: str | None,
    file: FileStorage,
    dataset_params: DatasetParams,
    tree_params: DecisionTreeParams,
) -> ClassificationResult:
    try:
        dataset: Dataset = Dataset(path=path, file=file)
        data = dataset.read(dataset_params)
        splitted_dataset: SplittedDataset = dataset.split_classification(
            data=data,
            params=dataset_params,
            random_state=tree_params.random_state,
        )
        result = fit_classification_model(
            data=splitted_dataset,
            params=tree_params,
        )
    finally:
        dataset.remove()
    return result