class IOptimization(ABC, metaclass=IOptimizationMetaClass):
    """
    The `IOptimization` interface serves as a base for all optimization classes to inherit from.
    It encompasses the following principles:

    1. It provides fundamental functionalities shared by all optimizations.
    2. Non-abstract methods and properties are generally not meant to be overridden.
    3. Methods and properties marked as 'final' should not be overridden.
    4. Strict adherence to data types is expected.
    5. The batch input and output dataclass types are primarily for typing purposes and can accommodate various
       data types. However, it's essential to have the specified required arguments available.
    """

    AXIS_MAPPING: dict[AxisType, int] = {
        "wholesaler": 0,
        "brand": 1,
        "time": 2,
        "vehicle": 3,
    }
    _losses: dict[str, CalculatedMetric] = {}
    _metrics: dict[str, CalculatedMetric] = {}

    @final
    @property
    def _default_metrics(self) -> dict[str, CalculatedMetric]:
        """Default metrics that all optimizers will use."""
        return {
            "ROI": lambda batch: np.sum(self.simulate(batch).vehicle_impacts) / (np.sum(self.vehicle_spends) + EPSILON),
            "Profit": lambda batch: np.sum(self.simulate(batch).vehicle_impacts) - np.sum(self.vehicle_spends),
            "constraint_loss": lambda _: np.mean(np.minimum(self.get_constraints(), 0)),
            "Impact": lambda batch: np.sum(self.simulate(batch).vehicle_impacts),
            "Spend": lambda _: np.sum(self.vehicle_spends),
        }

    def __init__(self, name: str, hyperparameters: HyperparameterConfig):
        self.name = name
        self._hyperparameters = hyperparameters

    def _post_init(self):
        for metric_name, metric in self._default_metrics.items():
            IOptimization.add_metric(self, metric_name, metric)

    # TODO(legendof-selda): get constraints in a user intuitive format. similar to old OptimizationIngestion.
    @abstractproperty
    def constraints(self) -> Constraints:
        """Returns list of constraints used for optimization."""

    def hyperparameters(self) -> HyperparameterConfig:
        """Hyperparameters used."""
        return self._hyperparameters

    @abstractproperty
    def vehicle_spends(self) -> VehicleSpendType:
        """
        The vehicle investments variable which is being optimized.
        Investment amounts for each batch (location*product), time and vehicle.
        """

    @abstractmethod
    def simulate(self, batch: Type[OptimizationInput]) -> Type[OptimizationOutput]:  # noqa: U100
        """For the given `batch` input simulate the impacts received.

        Args:
            batch (OptimizationInput): Optimization input that contains `vehicle_spends: VehicleSpendType` & others.

        Returns:
            Type[OptimizationOutput]: The impacts for the given investment amounts.
        """

    def __call__(self, batch: OptimizationInput) -> Type[OptimizationOutput]:
        return self.simulate(batch)

    @abstractmethod
    def optimize(self, dataset_factory: DatasetFactory, epochs: int, **_kwargs):  # noqa: U100
        """Optimize for the given number of epochs.

        Args:
            dataset_factory (DatasetFactory): A generator that returns `OptimizationInput` in batches.
            epochs (int): Number of epochs to optimize.
        """

    def get_constraints(self) -> np.ndarray:
        """Function to apply and gather all the constraints.

        Returns:
            np.ndarray: Stacked constraints applied on vehicle_spends.
        """
        if len(self.constraints) == 0:
            return np.array([0])
        gathers = []
        for constraint in self.constraints:
            gathered = self.vehicle_spends
            for axis_name, indices in constraint.gathers:
                axis = self.AXIS_MAPPING[axis_name]
                gathered = np.take(gathered, indices, axis=axis)
            constrained = np.sum(gathered) - constraint.max_value
            if not constraint.negate:
                constrained = -constrained
            gathers.append(constrained)
        return np.stack(gathers, axis=0)

    @final
    def create_result(
        self,
        location_type: Literal[LOCATION_TYPES],
        dataset_factory: DatasetFactory,
        encodings: dict[str, Any],
        return_dataframe: bool = True,
    ) -> dict[str, dict[tuple[str, ...], dict[str, float]]] | pd.DataFrame:
        """
        Returns results of current optimized state.

        Args:
            location_type (str): The location type ("wholesaler", "state", "region") the results should be in.
            dataset_factory (DatasetFactory): A generator that returns `OptimizationInput` in batches.
            encodings (dict[str, Any]): Encodings to decode the values in dataset.
            return_dataframe (bool, optional): Return a dataframe instead of dict. Defaults to True.

        Returns:
            dict[str, dict[tuple[str], dict[str, float]]] | pd.DataFrame: The results of the current state of optimizer.
                If `return_dataframe is True`, a DataFrame is returned.
        """
        if location_type == "wholesaler":
            loc_encoding = encodings["wholesaler"]
        else:
            loc_encoding = {
                loc: i
                for i, loc in enumerate(dict.fromkeys(encodings[f"wholesaler_{location_type}_lookup"].values()).keys())
            }
        loc_lookup = np.array(get_lookups(loc_encoding))
        brand_lookup = np.array(get_lookups(encodings, "brand"))
        vehicle_lookup = np.array(get_lookups(encodings, "vehicle"))
        time_lookup = np.array(get_lookups(encodings, "date"))
        veh_dfs = []
        baseline_dfs = []
        df_axes = [location_type, "brand", "date", "signal"]
        batch: Type[OptimizationInput]
        for batch in dataset_factory():
            brands = brand_lookup[np.array(batch.brand_index)]
            locations = loc_lookup[np.array(batch.location_index)]
            vehicles = vehicle_lookup[np.array(batch.vehicle_index)]
            dates = time_lookup[np.array(batch.date_index)]

            output = self.simulate(batch)

            index = tuple(zip(locations, brands))
            index = [(*lb, t) for lb, t in product(index, dates)]
            veh_index = [(*lbt, v) for lbt, v in product(index, vehicles)]

            veh_impacts_df = get_other_rev_components(
                output.vehicle_impacts,
                batch,
                encodings["normalization_factor"],
                pd.MultiIndex.from_tuples(veh_index, names=df_axes),
            )
            total_impacts_df = get_other_rev_components(
                output.yhat,
                batch,
                encodings["normalization_factor"],
                pd.MultiIndex.from_tuples(index, names=df_axes[:-1]),
            )
            baseline_impacts_df = total_impacts_df - veh_impacts_df.groupby(df_axes[:-1], axis=0).sum()
            veh_impacts_df["amount"] = np.reshape(output.vehicle_spends, -1) * encodings["normalization_factor"]
            veh_impacts_df["cores"] = veh_impacts_df["maco"] - veh_impacts_df["amount"]

            veh_dfs.append(veh_impacts_df)
            baseline_dfs.append(baseline_impacts_df)

        vehicle_results_df = pd.concat(veh_dfs, axis=0)
        baseline_results_df = pd.concat(baseline_dfs, axis=0)
        add_col_level(baseline_results_df, "baseline", axis=0, levelname="signal")
        # we agg as location could be repeated since some optimizers may work in lower granularity
        # and location are directly mapped. agg can avoid this issue.
        if vehicle_results_df.index.duplicated().any():
            vehicle_results_df = vehicle_results_df.groupby(vehicle_results_df.index.names).sum()
            baseline_results_df = baseline_results_df.groupby(baseline_results_df.index.names).sum()

        if return_dataframe:
            return pd.concat([baseline_results_df, vehicle_results_df], axis=0).sort_index(axis=0)
        else:
            return {
                "vehicle_results": vehicle_results_df.T.to_dict(),
                "baseline_results": baseline_results_df.T.to_dict(),
            }

    def add_loss(self, name: str, loss: CalculatedMetric):
        """Add the following loss function for tracking.

        Args:
            name (str): Name of the loss function.
            loss (CalculatedMetric): The loss function that will be evaluated.
        """
        self._losses[name] = loss

    def add_metric(self, name: str, metric: CalculatedMetric):
        """Add the following metric function for tracking.

        Args:
            name (str): Name of the metric function.
            metric (CalculatedMetric): The metric function that will be evaluated.
        """
        self._metrics[name] = metric

    def all_losses(self, batch: Type[OptimizationInput]) -> dict[str, float]:
        """Returns a dict of all losses. This property can be overriden if you have a custom dict of losses you track.

        Args:
            batch (OptimizationInput): Optimization input that contains `vehicle_spends: VehicleSpendType` & others.

        Returns:
            dict[str, CalculatedMetric]: Dict of computed losses.
        """
        return {loss: value(batch) for loss, value in self._losses.items()}

    def all_metrics(self, batch: Type[OptimizationInput]) -> dict[str, float]:
        """Returns a dict of all metrics. This property can be overriden if you have a custom dict of metrics you track.

        Args:
            batch (OptimizationInput): Optimization input that contains `vehicle_spends: VehicleSpendType` & others.

        Returns:
            dict[str, CalculatedMetric]: Dict of computed metrics.
        """
        return {metric: value(batch) for metric, value in self._metrics.items()}

class IOptimizationMetaClass(ABCMeta):
    """Meta class that will call _post_init after the instance is initialized."""

    def __call__(cls, *args, **kwargs):
        instance = ABCMeta.__call__(cls, *args, **kwargs)
        instance._post_init()
        return instance

@dataclass
class OptimizationInput(ABC):
    """
    This is an abstract OptimizationInput class that is mainly used for typing.
    It is not mandatory that your Input must be a dataclass. It can be a NamedTuple or any other class that resembles a
    dataclass.
    Only ensure that the following attributes below exists.
    """

    vehicle_spends: VehicleSpendType
    price: Annotated[TensorLike, TensorMeta((Batch, Time), np.float32)]
    price_normalization: Annotated[TensorLike, TensorMeta((Batch,), np.float32)]
    maco_cost: Annotated[TensorLike, TensorMeta((Batch, Time), np.float32)]
    vehicle_index: Annotated[TensorLike, TensorMeta((Vehicle,), np.int32)]
    brand_index: Annotated[TensorLike, TensorMeta((Batch,), np.int32)]
    location_index: Annotated[TensorLike, TensorMeta((Batch,), np.int32)]
    date_index: Annotated[TensorLike, TensorMeta((Time,), np.int32)]
    ...

@dataclass
class OptimizationOutput(ABC):
    """
    This is an abstract OptimizationOutput class that is mainly used for typing.
    It is not mandatory that your Output must be a dataclass. It can be a NamedTuple or any other class that resembles a
    dataclass.
    Only ensure that the following attributes below exists.
    """

    vehicle_spends: VehicleSpendType
    vehicle_impacts: ImpactsType
    yhat: YHatType
    ...