class BetaGammaDecay(Module):
    """Class to learn decayed impacts for the ensuing time periods after a spend in a media vehicle"""

    def __init__(
        self,
        encodings: dict[str, Any],
        hierarchy_categories: list[str | list[str]] | None = None,
        hyperparameters: Hyperparams | None = None,
        name: str | None = None,
    ):
        """Creates a betagammadecay object to learn decayed impacts using beta and gamma parameters.

        Args:
            hierarchy (pd.DataFrame): The hierarchy that the impact learns on.
            hyperparameters (Hyperparams | None, optional): Dictionary of hyperparameters for buidling this layer.
                                                            Defaults to None.
            name (str | None, optional): Name of the layer. Defaults to None.
        """
        super().__init__(hyperparameters=hyperparameters, name=name)
        self.encodings = encodings
        self.hierarchy_categories = hierarchy_categories

    def build(self, input_shapes):  # noqa: U100
        """Build the layer parameters needed for calculating decays.

        Args:
            input_shapes (InputShapes): The effect and hierarchy shapes.
        """
        self.gamma_min = self.hyperparameters.get_float(
            "gamma_min",
            default=0.01,
            min=0.00,
            max=1.0,
            help="The minimum possible value to learn for the exponential decay factor.",
        )
        self.gamma_max = self.hyperparameters.get_float(
            "gamma_max",
            default=1.0,
            min=self.gamma_min,
            max=1.0,
            help="The maximum possible value to learn for the exponential decay factor.",
        )
        self.beta_min = self.hyperparameters.get_float(
            "beta_min",
            default=0.01,
            min=0.00,
            max=1.0,
            help="The minimum possible value to learn for the first step of the decay.",
        )
        self.beta_max = self.hyperparameters.get_float(
            "beta_max",
            default=1.0,
            min=self.beta_min,
            max=1.0,
            help="The maximum possible value to learn for the first step of the decay.",
        )
        self.betagamma_emb_layer = self.hyperparameters.get_submodule(
            name="betagamma_hier",
            module_type=HierchicalEmbedding,
            kwargs=dict(
                encodings=self.encodings,
                columns=self.hierarchy_categories,
                dropped_columns=[],
                shape=[2],
                feature_names=["beta", "gamma"],
            ),
            help="The embedding layer for the decay parameters.",
        )

    @property
    def gamma_range(self) -> float:
        return self.gamma_max - self.gamma_min

    @property
    def beta_range(self) -> float:
        return self.beta_max - self.beta_min

    def __call__(
        self, batch: BetaGammaDecayInput, training: bool = False, debug: bool = False, skip_metrics: bool = False
    ) -> BetaGammaDecayIntermediaries:
        """Calculate decays, total impacts using the learned beta gamma parameters

        Args:
            impact_by_signal_instant (TensorLike): Instant impacts
            hierarchy (dict[str, TensorLike]): Hierarchy placeholder for Hierarchial embedding variable.
            training (bool, optional): Whether this is a training or inference run. Defaults to False.

        Returns:
            BetaGammaDecayIntermediaries: Intermediate calculations for beta gamma decay - beta, gamma, impacts etc.
        """
        # batch x vehicles x 2
        betagamma_emb = self.betagamma_emb_layer(
            batch.hierarchy, training=training, skip_metrics=skip_metrics, debug=debug
        )
        beta_emb, gamma_emb = tf.unstack(betagamma_emb, num=2, axis=2)
        if self.gamma_max < 1:
            gamma = 1.0 - 1.0 / (
                transform_softbounded(
                    gamma_emb,
                    max_val=3.0,
                    min_val=-3,
                    name="gamma",
                    add_loss=self.add_loss,
                    mult=0.1,
                    enabled=not skip_metrics,
                )
                * (1 / (1 - self.gamma_max) - 1 / (1 - self.gamma_min))
                + 1 / (1 - self.gamma_min)
            )
        else:
            gamma = 1.0 - 1.0 / (softplus(gamma_emb) + 1 / (1 - self.gamma_min))
        # batch x vehicles
        beta = (
            transform_softbounded(
                beta_emb,
                max_val=3.0,
                min_val=-3,
                name="beta",
                mult=0.1,
                add_loss=self.add_loss,
                enabled=not skip_metrics,
            )
            * self.beta_range
            + self.beta_min
        )
        # batch x 1 x vehicle
        decayed_impact_mult = tf.expand_dims(1 + beta * (1 - tf.math.pow(gamma, DECAY_LENGTH)) / (1 - gamma), 1)
        impact_by_signal_total = batch.impact_by_signal_instant * decayed_impact_mult
        impact_by_signal_decayed = exp_moving_avg(
            batch.impact_by_signal_instant,
            beta,
            gamma,
            name="impact_by_signal_decayed",
            decay_length=DECAY_LENGTH,
        )
        impact = tf.math.reduce_sum(impact_by_signal_decayed, axis=2)
        return BetaGammaDecayIntermediaries(
            beta_emb=beta_emb if debug else None,
            gamma_emb=gamma_emb if debug else None,
            beta=beta,
            gamma=gamma,
            decayed_impact_mult=decayed_impact_mult,
            impact_by_signal_total=impact_by_signal_total,
            impact_by_signal=impact_by_signal_decayed,
            impact=impact,
            signal_names=tf.gather(
                tf.convert_to_tensor(tuple(f"{vehicle}_decayed" for vehicle in get_lookups(self.encodings["vehicle"]))),
                batch.hierarchy["vehicle"][0],
            ),
        )