from enum import StrEnum
import numpy as np
import pandas as pd
import loguru


class DatasetSplit(StrEnum):
    VALIDATION = "validation"
    TEST = "test"


class OptimizationInputType:
    def __init__(self, wide_format_optimization_input: pd.DataFrame):
        self.logger = loguru.logger
        self._unpack_optimization_input(wide_format_optimization_input)

    def _unpack_optimization_input(
        self, wide_format_optimization_input: pd.DataFrame
    ) -> None:
        self.id = wide_format_optimization_input["id"]
        self.logger.debug(f"Size: {self.id.shape[0]}.")

        cost_cols = [
            col
            for col in wide_format_optimization_input.columns
            if col.endswith("_cost")
        ]
        self.cost = wide_format_optimization_input[cost_cols].to_numpy()
        self.cost_producer = [col.removesuffix("_cost") for col in cost_cols]
        self.logger.debug(f"Producer in cost array: {self.cost_producer}.")

        quality_cols = [
            col
            for col in wide_format_optimization_input.columns
            if col.endswith("_quality")
        ]
        self.quality = wide_format_optimization_input[quality_cols].to_numpy()
        self.quality_producer = [col.removesuffix("_quality") for col in quality_cols]
        self.logger.debug(f"Producer in quality array: {self.quality_producer}.")

        reference_cols = [
            col
            for col in wide_format_optimization_input.columns
            if col.endswith("_reference")
        ]
        self.reference = wide_format_optimization_input[reference_cols].to_numpy()
        self.reference_producer = [
            col.removesuffix("_reference") for col in reference_cols
        ]
        self.logger.debug(f"Producer in reference array: {self.reference_producer}.")
