Source code for aixcalibuha.utils

"""
Package containing utility functions used in different packages.
Contains a statistics analyzer and a visualizer.
"""
import os
from typing import Union, List
from pathlib import Path
from aixcalibuha import CalibrationClass
from ebcpy import load_time_series_data
import pandas as pd


def convert_mat_to_suffix(mat_result_file, variable_names, suffix_files, parquet_engine='pyarrow'):
    """
    Postprocess the mat result files.

    :param str mat_result_file: The path to the MATLAB result file.
    :param List[str] variable_names: The names of the variables to extract.
    :param str suffix_files: The file format for the output files.
        Options: 'csv', 'parquet', 'parquet.snappy', 'parquet.gzip', etc.
    :param str parquet_engine: The engine to use for saving parquet files.
    """
    df = load_time_series_data(mat_result_file, variable_names=variable_names)
    for col in df.columns:
        if isinstance(df[col].dtype, pd.SparseDtype):
            df[col] = df[col].sparse.to_dense()
    df_path = Path(mat_result_file).with_suffix("." + suffix_files)
    df.tsd.save(df_path, engine=parquet_engine)
    os.remove(mat_result_file)
    return df_path


def empty_postprocessing(mat_result, **_kwargs):
    return mat_result


def validate_cal_class_input(
        calibration_classes: Union[CalibrationClass, List[CalibrationClass]]
) -> List[CalibrationClass]:
    """Check if given list contains only CalibrationClass objects or is one
    and return a list in both cases. Else raise an error"""
    if isinstance(calibration_classes, list):
        for cal_class in calibration_classes:
            if not isinstance(cal_class, CalibrationClass):
                raise TypeError(f"calibration_classes is of type {type(cal_class).__name__} "
                                f"but should be CalibrationClass")
    elif isinstance(calibration_classes, CalibrationClass):
        calibration_classes = [calibration_classes]
    else:
        raise TypeError(f"calibration_classes is of type {type(calibration_classes).__name__} "
                        f"but should be CalibrationClass or list")
    return calibration_classes


class MaxIterationsReached(Exception):
    """
    Exception raised for when the calibration
    ends because the maximum number of
    allowed iterations is reached.
    """


class MaxTimeReached(Exception):
    """
    Exception raised for when the calibration
    ends because the maximum calibration time is reached.
    """