class BFNet(EconomicNetwork):
    result_type: ClassVar[type] = EconomicIntermediaries
    num_mask_steps: ClassVar[int] = 1

    def get_baseline(self, impacts: ImpactsIntermediaries, training=False) -> tuple[tf.Tensor, dict]:  # noqa: U100
        """Getting baseline from the previous week actual sales by removing impacts

        Args:
            impacts (ImpactsIntermediaries): Intermediaries object which contains additive and multiplicative impacts.
            training (bool, optional):  Flag to set the trainable params. Defaults to False.

        Returns:
            tuple[tf.Tensor, dict[str, tf.Tensor]] : Returning the baseline by removing the additive and
             multiplicative effect of previous week
        """
        additive_shifted_impacts = [impact[:, :-1] for impact in impacts.additive_impacts]
        multiplicative_shifted_impacts = [impact[:, :-1] for impact in impacts.multiplicative_impacts]
        baseline = apply_inverse_impacts(
            self.yph[:, :-1],
            additive_impacts=additive_shifted_impacts,
            multiplicative_impacts=multiplicative_shifted_impacts,
        )
        return tf.concat([tf.zeros_like(baseline[:, :1]), baseline], axis=1), {}

class EconomicNetwork(TrainableModule):
    def __init__(
        self,
        encodings: Mapping[str, int | Mapping[str, Any] | Iterable | str],
        name: str | None = None,
        hyperparameters: Hyperparams | None = None,
        enable_bdl_metrics: bool = False,
    ):
        super().__init__(
            name=name,
            hyperparameters=hyperparameters,
        )
        self.data_encodings = encodings
        self.enable_bdl_metrics = enable_bdl_metrics

    def create_encodings(self, encodings, input_shapes):
        # TODO: this list needs to be removed
        cont_encodings_default = {
            "tot_pct_black_african_american",
            "tot_pct_hispanic",
            "pct_2021_pop_age_20_29",
            "pct_2021_no_high_school_diploma",
            "pct_2021_college_or_advanced_degree",
            "household_income_median",
            "pct_2021_income_upto_24999",
            "pct_2021_income_over_200000",
            "tot_pop_median_age",
            "GOP",
            "DEM",
            "density",
            "brand_product_market_share",
            "avg_snow",
            "avg_prcp",
            "avg_temp",
            "marital_status_2021_pop_15_years_and_over",
            "pct_2021_income_75000_124999",
            "pct_2021_married",
            "pct_2021_pop_age_40_49",
            "pct_2021_pop_age_65_over",
            "tot_pct_asian",
            "tot_pct_native_hawaiian_other_pacific_islander",
            "tot_pct_other_race",
            "tot_pct_white",
            "wholesaler_pop",
            "lan",
            "lon",
            "lon_from_center",
        }

        cont_encodings_used = {
            feature: self.hyperparameters.get_bool(
                feature, default=feature in cont_encodings_default, help=f"Use {feature} for hierarchical embedding."
            )
            for feature in input_shapes.continuous_hier_params.keys()
        }
        self.continous_encodings = {
            feature: encodings[feature] for feature, used in cont_encodings_used.items() if used
        }
        base_encodings = (
            {
                k: encodings[k]
                for k in (
                    "region",
                    "state",
                    "wholesaler",
                    "brand",
                    "date",
                )
                if getattr(input_shapes, f"{k}_index", ...) is not ...
            }
            | self.continous_encodings
            | {"max_real_date": encodings.get("max_real_date", max(encodings["date"].keys()))}
        )
        self.encodings = {
            k: encodings[k]
            for k in (
                "global",
                "weather",
                "temperature",
                "price_dev",
                "price_ratio",
                "holiday",
                "distribution",
                "pre_investment",
                "covid",
                "national_trend",
                "regional_trend",
                "date",
                "granularity",
            )
            if getattr(input_shapes, f"{k}_index", ...) is not ...
        } | base_encodings

        def get_encodings(indices, keys, specific_maps={}):
            return (
                {key: encodings[specific_maps.get(key, key)] for key in keys} | base_encodings
                if all(getattr(input_shapes, f"{idx}_index", ...) is not ... for idx in indices)
                else {}
            )

        self.vehicle_encodings = get_encodings(["vehicle", "parent_vehicle"], ["vehicle", "parent_vehicle"])

        self.global_encodings = get_encodings(
            ["global", "global_parent"],
            ["global", "signal", "parent_signal"],
            specific_maps={"signal": "global", "parent_signal": "global_parent"},
        )

        self.weather_encodings = get_encodings(
            ["weather", "weather_parent"],
            ["granularity", "weather", "signal", "parent_signal"],
            specific_maps={"signal": "weather", "parent_signal": "weather_parent"},
        )

        self.temperature_encodings = get_encodings(
            ["temperature", "temperature_parent"],
            ["granularity", "temperature", "signal", "parent_signal"],
            specific_maps={"signal": "temperature", "parent_signal": "temperature_parent"},
        )

    @cached_property
    def result_type(self):
        return EconomicIntermediaries

    @cached_property
    def hlayer_lookup(self):
        assert self._built
        return {
            "/".join(layer.weights.name.split("/")[:-1]): layer
            for layer in self.submodules
            if isinstance(layer, HierchicalEmbedding)
        }

    @cached_property
    def annotated_var_lookup(self):
        assert self._built
        return {v.name: v for v in self.all_variables if v._annotated_shape is not None}

    def get_closest_match(self, name, lookup):
        if name in lookup:
            return lookup[name], name
        logger.debug(f"{name} is not an exact match")
        match, score, _ = rapidfuzz.process.extract(name, lookup.keys(), limit=1)[0]
        logger.debug(f"{match} is the closest match with a score of {score}/100.")
        return lookup[match], match

    def get_hier_layer(self, name):
        # TODO: generalize for all layers by updating scoped layer names to be unique
        #       currently generalized only for hierarchical because can get scoped name
        #       via hlayer.weights.name
        return self.get_closest_match(name, self.hlayer_lookup)

    def get_annotated_variable(self, name):
        return self.get_closest_match(name, self.annotated_var_lookup)

    def build(self, input_shapes):  # noqa: U100
        self.create_encodings(self.data_encodings, input_shapes)
        self.window_radius_revenue = self.hyperparameters.get_int(
            "window_radius_revenue", default=4, min=0, max=6, help="The window size for smoothing revenue."
        )
        self.window_radius_signal = self.hyperparameters.get_int(
            "window_radius_signal", default=4, min=0, max=6, help="The window size for smoothing signals."
        )
        self.WINDOW_SIZE_REVENUE = 1 + self.window_radius_revenue * 2
        self.WINDOW_SIZE_SIGNAL = 1 + self.window_radius_signal * 2
        self.window_std_revenue = self.hyperparameters.get_float(
            "window_std_revenue",
            default=self.window_radius_revenue / 2.0,
            min=0.0,
            max=4.0,
            help="The ratio of the window size to use for the stddev of the gaussian smoothing for revenue.",
        )
        self.window_std_signal = self.hyperparameters.get_float(
            "window_std_signal",
            default=self.window_radius_signal / 2.0,
            min=0.0,
            max=4.0,
            help="The ratio of the window size to use for the stddev of the gaussian smoothing for signals.",
        )
        self.hier_lambda = self.hyperparameters.get_float(
            "hier_lambda", default=1e-4, min=1e-06, max=1e-02, logdist=True, help="The weightage for hier loss."
        )
        self.aux_lambda = self.hyperparameters.get_float(
            "aux_lambda", default=1.0, min=1e-03, max=1.0, logdist=True, help="The weightage for aux loss."
        )
        self.impacts_layer = self.hyperparameters.get_submodule(
            "impacts",
            module_type=Impacts,
            kwargs=dict(
                continous_encodings=self.continous_encodings,
                encodings=self.encodings,
                vehicle_encodings=self.vehicle_encodings,
                global_encodings=self.global_encodings,
                weather_encodings=self.weather_encodings,
                temperature_encodings=self.temperature_encodings,
                wibble_means=np.array(self.data_encodings["wholesaler_means"], dtype=np.float32)[:, :, 0],
            ),
            help="The layer for calculating the impacts to apply on top of the baseline.",
        )
        self.pre_2021_mse_weight = self.hyperparameters.get_float(
            "pre_2021_mse_weight",
            default=0.2,
            min=0.0,
            max=1.0,
            help="The weight to use for predictions before 2021 relative to the weight of points in 2021 forwards.",
        )
        self.post_2021_mse_scale = self.hyperparameters.get_float(
            "post_2021_mse_scale",
            default=1.0,
            min=1.0,
            max=10.0,
            help="After post 2021 mse weight jumps to 1, this is the incremental increase till end of time axis",
        )

    def calculate_impacts(self, training=False, debug=False, skip_metrics=False) -> ImpactsIntermediaries:
        return self.impacts_layer(
            ImpactsInput(
                global_effect=self.global_effect_ph,
                weather_effect=self.weather_effect_ph,
                temperature_effect=self.temperature_effect_ph,
                holiday_effect=self.holiday_ph,
                national_trend=self.national_trend_ph,
                regional_trend=self.regional_trend_ph,
                price_ratio=self.price_ratio_ph,
                distribution=self.distribution_ph,
                prices=self.price_ph,
                invest_spends=self.invest_spend_ph,
                dates_since_start=self.dates_since_start_ph,
                investment_axis_scale=self.investment_axis_scale,
                hierarchies=dict(
                    hierarchy=dict(self.hierarchical_placeholders.items()),
                    vehicle_hierarchy=dict(self.vehicle_hierarchical_placeholders.items()),
                    global_hierarchy=dict(self.global_hierarchical_placeholders.items()),
                    weather_hierarchy=dict(self.weather_hierarchical_placeholders.items()),
                    temperature_hierarchy=dict(self.temperature_hierarchical_placeholders.items()),
                ),
                mask=tf.expand_dims(self.yhat_mask_ph, -1) if self.yhat_mask_ph is not None else None,
                yearly_week_number=self.yearly_week_number_ph,
                brand_size=self.brand_size_ph,
                preinvestment_slope=self.preinvestment_slope_ph,
                preinvestment_intercept=self.preinvestment_intercept_ph,
                date_index=self.date_index_ph,
                brand_index=self.brand_index_ph,
            ),
            training=training,
            skip_metrics=skip_metrics,
            debug=debug,
        )

    num_mask_steps: ClassVar[int] = 0

    def process_baseline(
        self,
        result_type: type,
        baseline: tf.Tensor,
        impacts: ImpactsIntermediaries,
        inputs: EconomicModelInput,
        training: bool = False,
        debug: bool = False,
        skip_metrics: bool = False,
        prefix: str = "",
        **kwargs,
    ):
        yhat = apply_impacts(
            baseline,
            additive_impacts=impacts.additive_impacts,
            multiplicative_impacts=impacts.multiplicative_impacts,
        )
        y_mask = self.yhat_mask_ph
        if training and y_mask is not None:
            for i in range(1, self.num_mask_steps + 1):
                y_mask = y_mask * tf.concat([tf.zeros_like(y_mask[:, :i]), y_mask[:, :-i]], axis=1, name=f"mask_{i}")
        # Might be None, but we only pass as a mask to things which will handle None fine
        weighted_batch_mask = self.instability_loss_mult
        assert isinstance(self.data_encodings["date"], Mapping), "data_encoding['date'] should be a mapping"
        if self.yph is not None and not skip_metrics:
            start_2021_mask = self.hierarchical_placeholders["date"] > max(
                i for k, i in self.data_encodings["date"].items() if k < "2021"
            )
            start_2021_cumsum = tf.math.cumsum(tf.cast(start_2021_mask, tf.float32)) - tf.constant(
                1.0, dtype=tf.float32
            )
            start_2021_weights = (
                tf.constant(1.0, dtype=tf.float32)
                + start_2021_cumsum / tf.reduce_max(start_2021_cumsum) * self.post_2021_mse_scale
            )
            if y_mask is None:
                mse_weights = tf.where(
                    start_2021_mask,
                    start_2021_weights,
                    tf.constant(self.pre_2021_mse_weight, dtype=tf.float32),
                )
            else:
                mse_weights = y_mask * tf.where(
                    start_2021_mask,
                    start_2021_weights,
                    tf.constant(self.pre_2021_mse_weight, dtype=tf.float32),
                )
            mse_per_signal = metrics.weighted_mean(tf.math.squared_difference(self.yph, yhat), mask=mse_weights, axis=1)
            mse = tf.math.reduce_sum(mse_per_signal, name="mean_mse")
            mean_sales = metrics.weighted_mean(
                self.yph,
                mask=y_mask,
                name=f"{prefix}mean_sales",
                axis=1,
                keepdims=True,
            )
            self.add_loss(f"{prefix}mse", mse, f"{prefix}mse")
            if y_mask is None:
                data_point_mask = tf.cast(start_2021_mask, dtype=yhat.dtype)
            else:
                data_point_mask = tf.cast(start_2021_mask, dtype=yhat.dtype) * y_mask
            r_squared = metrics.r_squared(self.yph, yhat, data_point_mask, axis=1)
            self.add_metric(f"{prefix}r_squared", r_squared, sample_weights=weighted_batch_mask)
            self.add_metric(f"{prefix}pos_r_squared", tf.maximum(0.0, r_squared), sample_weights=weighted_batch_mask)
            if self.enable_bdl_metrics and impacts.bud_light_effect is not None:
                mult_diff, mult_mask = bdl_mult_diff(
                    self.yph,
                    inputs,
                    impacts.bud_light_effect.impact,
                    impacts.bud_light_effect.mf_mask,
                    self.data_encodings,
                )
                self.add_metric(f"{prefix}pos_bdl_mult_diff", tf.math.maximum(mult_diff, 0.0), sample_weights=mult_mask)
                self.add_metric(f"{prefix}neg_bdl_mult_diff", tf.math.minimum(mult_diff, 0.0), sample_weights=mult_mask)
                self.add_metric(f"{prefix}bdl_mult_diff", mult_diff, sample_weights=mult_mask)
            all_losses = self.get_all_losses()
            total_loss = tf.math.add_n(all_losses.values())
        else:
            mse = None
            r_squared = None
            mean_sales = None
            all_losses = None
            total_loss = None
        return result_type(
            impacts=impacts,
            baseline=baseline,
            yhat=yhat,
            y_true=self.yph,
            y_smooth=self.y_smooth,
            all_losses=all_losses,
            total_loss=total_loss,
            mask=y_mask,
            mse=mse,
            r_squared=r_squared if debug else None,
            train_r2=(
                metrics.weighted_mean(r_squared, mask=weighted_batch_mask, axis=0, name=f"{prefix}r_squared")
                if r_squared is not None
                else None
            ),
            step=self._step_var,
            # Would be nice to make this None when not debugging
            inputs=inputs,
            mean_sales=mean_sales if debug else None,
            **kwargs,
        )

    def __call__(
        self, batch: EconomicModelInput, training: bool = False, debug: bool = False, skip_metrics: bool = False
    ):
        self.create_network_phs(batch, training=training)
        impacts = self.calculate_impacts(training=training, debug=debug, skip_metrics=skip_metrics)
        baseline, kwargs = self.get_baseline(impacts, training=training, skip_metrics=skip_metrics, debug=debug)
        result = self.process_baseline(
            self.result_type,
            baseline,
            impacts,
            batch,
            training=training,
            skip_metrics=skip_metrics,
            debug=debug,
            **kwargs,
        )
        return result

    def get_baseline(
        self,
        impacts: ImpactsIntermediaries,  # noqa: U100
        training: bool = False,  # noqa: U100
        debug: bool = False,  # noqa: U100
        skip_metrics: bool = False,  # noqa: U100
    ) -> tuple[tf.Tensor, dict]:
        """
        Returns baseline and any kwargs for the Intermediaries result type.
        """
        return tf.constant(0.0, dtype=tf.float32), {}

    def create_network_phs(self, batch: EconomicModelInput, training: bool = False):
        self.date_index_ph = batch.date_index
        self.brand_index_ph = batch.brand_index
        smoothing_weights_revenue = get_smoothing_weights(
            self.WINDOW_SIZE_REVENUE, "gaussian", std=self.window_std_revenue
        )
        smoothing_weights_signal = get_smoothing_weights(
            self.WINDOW_SIZE_SIGNAL, "gaussian", std=self.window_std_signal
        )

        def get_rolling_mean(
            signals, window_size=None, center=True, smoothing_weights=smoothing_weights_signal, **kwargs
        ):
            if window_size is None:
                window_size = self.WINDOW_SIZE_SIGNAL
            return (
                rolling_mean(
                    signals,
                    window_size=window_size,
                    center=center,
                    smoothing_weights=smoothing_weights,
                    **kwargs,
                )
                if signals is not None
                else None
            )

        self.global_effect_ph = get_rolling_mean(batch.global_signals)
        self.weather_effect_ph = get_rolling_mean(batch.weather_signals)
        self.temperature_effect_ph = get_rolling_mean(batch.temperature_signals)
        self.holiday_ph = get_rolling_mean(batch.holiday_signals)

        self.yearly_week_number_ph = batch.yearly_week_number
        self.brand_size_ph = batch.brand_size
        self.national_trend_ph = batch.national_trend
        self.regional_trend_ph = batch.regional_trend
        self.before_2021_mask_ph = batch.before_2021_mask
        self.yhat_mask_ph = batch.no_prediction_mask
        self.y_true_mask_ph = self.yhat_mask_ph
        if batch.feature_masks is not None:
            reduced = tf.math.reduce_all(batch.feature_masks, axis=2)
            self.y_true_mask_ph = reduced if self.y_true_mask_ph is None else reduced & self.y_true_mask_ph
            if training:
                self.yhat_mask_ph = reduced if self.yhat_mask_ph is None else reduced & self.yhat_mask_ph
        if batch.no_train_mask is not None and training:
            self.yhat_mask_ph = (
                batch.no_train_mask
                if self.yhat_mask_ph is None
                else tf.logical_and(batch.no_train_mask, self.yhat_mask_ph)
            )
        self.yhat_mask_ph = tf.cast(self.yhat_mask_ph, dtype=tf.float32) if self.yhat_mask_ph is not None else None
        self.y_true_mask_ph = (
            tf.cast(self.y_true_mask_ph, dtype=tf.float32) if self.y_true_mask_ph is not None else None
        )
        expanded_mask = self.yhat_mask_ph[:, :, None] if self.yhat_mask_ph is not None else None
        self.price_ph = get_rolling_mean(batch.price_devs, mask=expanded_mask)
        self.price_ratio_ph = get_rolling_mean(batch.price_ratios, mask=expanded_mask)
        # even when training is False, rolling mean should be built on the all the masks.
        self.y_smooth = get_rolling_mean(
            batch.true_sales, mask=self.y_true_mask_ph, smoothing_weights=smoothing_weights_revenue
        )
        self.yph = self.y_smooth if training else batch.true_sales
        self.distribution_ph = batch.distributions
        self.instability_loss_mult = batch.instability_loss_mult
        self.invest_spend_ph = batch.vehicle_spends
        self.dates_since_start_ph = batch.weeks_since_restart
        self.preinvestment_slope_ph = batch.preinvestment_slope
        self.preinvestment_intercept_ph = batch.preinvestment_intercept
        self.investment_axis_scale = batch.investment_axis_scale

        # TODO(@legendof-selda): automate indexes encodings as well
        # We append all the continous values in batch using self.continous_encodings
        self.hierarchical_placeholders = {
            k: getattr(batch, f"{k}_index")
            for k in self.encodings.keys()
            if getattr(batch, f"{k}_index", None) is not None
        } | {feature: batch.continuous_hier_params[feature] for feature in self.continous_encodings.keys()}
        (
            self.vehicle_hierarchical_placeholders,
            self.distribution_hierarchical_placeholders,
            self.global_hierarchical_placeholders,
            self.weather_hierarchical_placeholders,
            self.temperature_hierarchical_placeholders,
        ) = get_hiers(batch, self.hierarchical_placeholders.keys()).values()

    def _get_loss_lambda(self, loss_category: str) -> float:
        if hasattr(self, f"{loss_category}_lambda"):
            return getattr(self, f"{loss_category}_lambda")
        else:
            logger.warning(f"Loss category {loss_category} was not recognized and was given a lambda of 1.")
            return 1.0

    def get_kwargs_to_save(self, **kwargs):
        return super().get_kwargs_to_save(**kwargs) | {"encodings": self.data_encodings}

    def add_global_me_tracker(self, batch, radius, desired_num_points):
        if batch.global_signals is None:
            logger.warning("Tried to add a global_me tracker when there is no global_signals data.")
        else:
            assert self.impacts_layer.global_me is not None, "global_me not None if tracker added"
            self.tracked["global_me"] = make_me_tracker(
                self,
                batch.global_signals,
                batch,
                self.impacts_layer.global_me,
                radius=radius,
                desired_num_points=desired_num_points,
            )

    def add_weather_me_tracker(self, batch, radius, desired_num_points):
        if batch.weather_signals is None:
            logger.warning("Tried to add a weather_me tracker when there is no weather_signals data.")
        else:
            assert self.impacts_layer.weather_me is not None, "weather_me not None if tracker added"
            self.tracked["weather_me"] = make_me_tracker(
                self,
                batch.weather_signals,
                batch,
                self.impacts_layer.weather_me,
                radius=radius,
                desired_num_points=desired_num_points,
            )

    def add_temperature_me_tracker(self, batch, radius, desired_num_points):
        if batch.temperature_signals is None:
            logger.warning("Tried to add a temperature_me tracker when there is no temperature_signals data.")
        else:
            assert self.impacts_layer.temperature_me is not None, "temperature_me not None if tracker added"
            self.tracked["temperature_me"] = make_me_tracker(
                self,
                batch.temperature_signals,
                batch,
                self.impacts_layer.temperature_me,
                radius=radius,
                desired_num_points=desired_num_points,
            )

    def add_national_trend_tracker(self, batch, radius, desired_num_points):
        if batch.national_trend is None:
            logger.warning("Tried to add national_trend tracker when there is not national_trend data.")
        else:
            assert (
                self.impacts_layer.industry_national_trend_me is not None
            ), "industry_national_trend_me not None if tracker added"
            self.tracked["national_trend"] = make_me_tracker(
                self,
                batch.national_trend,
                batch,
                self.impacts_layer.industry_national_trend_me,
                radius=radius,
                desired_num_points=desired_num_points,
            )

    def add_regional_trend_tracker(self, batch, radius, desired_num_points):
        if batch.regional_trend is None:
            logger.warning("Tried to add regional_trend tracker when there is not regional_trend data.")
        else:
            assert (
                self.impacts_layer.industry_regional_trend_me is not None
            ), "industry_regional_trend_me not None if tracker added"
            self.tracked["regional_trend"] = make_me_tracker(
                self,
                batch.regional_trend,
                batch,
                self.impacts_layer.industry_regional_trend_me,
                radius=radius,
                desired_num_points=desired_num_points,
            )

    def add_price_ratio_tracker(self, batch, price_radius, desired_num_points):
        if batch.price_ratios is None:
            logger.warning("Tried to add a price_ratio tracker when there is no price_ratios data.")
        else:
            assert self.impacts_layer.price_ratio is not None, "price_ratio not None if tracker added"
            self.tracked["price_ratio"] = make_me_tracker(
                self,
                batch.price_ratios,
                batch,
                self.impacts_layer.price_ratio,
                radius=price_radius,
                desired_num_points=desired_num_points,
            )

    def add_holiday_me_tracker(self, batch, desired_num_points):
        if batch.holiday_signals is None:
            logger.warning("Tried to add a holiday_me tracker when there is no holiday_signals data.")
        else:
            assert self.impacts_layer.holiday_me is not None, "holiday_me not None if tracker added"
            center = (
                to_numpy(get_smoothing_weights(self.WINDOW_SIZE_SIGNAL, "gaussian", std=self.window_std_signal)).max()
                / 2
            )
            self.tracked["holiday_me"] = make_me_tracker(
                self,
                batch.holiday_signals,
                batch,
                self.impacts_layer.holiday_me,
                radius=center,
                center=center,
                desired_num_points=desired_num_points,
            )

    def add_price_elasticity_tracker(self, batch, price_range, desired_num_points):
        if batch.price_devs is None:
            logger.warning("Tried to add a price_elasticity tracker when there is no price_devs data.")
        else:
            self.tracked["price_elasticity"] = make_price_tracker(
                self, batch, price_range=price_range, desired_num_points=desired_num_points
            )

    def add_distribution_tracker(self, batch, desired_num_points):
        if batch.distributions is None:
            logger.warning("Tried to add a distribution tracker when there is no distributions data.")
        else:
            assert self.impacts_layer.distribution_layer is not None, "distribution_layer not None if tracker added"
            self.tracked["distribution"] = make_me_tracker(
                self,
                batch.distributions,
                batch,
                self.impacts_layer.distribution_layer,
                radius=0.99,
                center=1.0,
                spacing=AxisSpacing.QUADRATIC,
                desired_num_points=desired_num_points,
            )

    def add_price_me_tracker(self, batch, price_radius, desired_num_points):
        if batch.price_devs is None:
            logger.warning("Tried to add a price_me tracker when there is no price_devs data.")
        else:
            assert self.impacts_layer.pricing_lead_lag_me is not None, "pricing_lead_lag_me cannot be none"
            self.tracked["price_me"] = make_me_tracker(
                self,
                batch.price_devs,
                batch,
                self.impacts_layer.pricing_lead_lag_me,
                radius=price_radius,
                desired_num_points=desired_num_points,
            )

    def add_periodic_me_tracker(self, batch):
        if batch.date_index is None:
            logger.warning("Tried to add a periodic_me tracker when there is no date data in the hierarchy.")
        else:
            self.tracked["periodic_me"] = make_periodic_tracker(self, batch)

    def add_roicurve_tracker(
        self, batch, max_val, desired_num_points, spacing, dynamic_range, investment_axis_scale_range
    ):
        if batch.vehicle_spends is None or self.impacts_layer.roicurve is None:
            logger.warning("Tried to add a roicurve tracker when there is no vehicle_spends data.")
        else:
            self.tracked["roicurve"] = make_roi_tracker(
                self,
                batch,
                max_val=max_val,
                desired_num_points=desired_num_points,
                spacing=spacing,
                dynamic_range=dynamic_range,
                investment_axis_scale_range=investment_axis_scale_range,
            )
            self.curve_trackers = ["roicurve"]

    # TODO(@ruler501)  out a way to simplify this
    def add_curve_trackers(  # noqa: C901 Haven't d out how to simplify yet
        self,
        batch: EconomicModelInput,
        max_val: float = 10.0,
        max_roi_val: float = 10.0,
        price_range=1.0,
        radius=3.0,
        price_radius=0.25,
        curves: list[str] | Literal["all"] = "all",
        desired_num_points: int = 256,
        dynamic_range: bool | float = False,
        investment_axis_scale_range: bool = False,
        spacing: AxisSpacing = AxisSpacing.QUADRATIC,
    ):
        if is_option_included(curves, "global_me"):
            self.add_global_me_tracker(batch, radius, desired_num_points)

        if is_option_included(curves, "weather_me"):
            self.add_weather_me_tracker(batch, radius, desired_num_points)

        if is_option_included(curves, "temperature_me"):
            self.add_temperature_me_tracker(batch, radius, desired_num_points)

        if is_option_included(curves, "national_trend"):
            self.add_national_trend_tracker(batch, radius, desired_num_points)

        if is_option_included(curves, "regional_trend"):
            self.add_regional_trend_tracker(batch, radius, desired_num_points)

        if is_option_included(curves, "price_ratio"):
            self.add_price_ratio_tracker(batch, price_radius, desired_num_points)

        if is_option_included(curves, "holiday_me"):
            self.add_holiday_me_tracker(batch, desired_num_points)

        if is_option_included(curves, "price_elasticity"):
            self.add_price_elasticity_tracker(batch, price_range, desired_num_points)

        if is_option_included(curves, "distribution"):
            self.add_distribution_tracker(batch, desired_num_points)

        if is_option_included(curves, "price_me"):
            self.add_price_me_tracker(batch, price_radius, desired_num_points)

        if is_option_included(curves, "periodic_me"):
            self.add_periodic_me_tracker(batch)

        if self.impacts_layer.roicurve is not None:
            self.add_roicurve_tracker(
                batch, max_roi_val, desired_num_points, spacing, dynamic_range, investment_axis_scale_range
            )

    def viz_impacts(
        self,
        epoch: EpochSpec = None,
        intermediaries: EconomicIntermediaries | list[EconomicIntermediaries] | None = None,
        num_viz: int = 2,
        separate_decay: bool = False,
        collapse_level: CollapseLevel | dict[str, CollapseLevel] = 0,
        collapse_lead_lag: bool | None = None,
        use_grans: list[str] | None = None,
        group_signals: bool = True,
        linearization_method: LinearizationMethod = "cumprod",
        **kwargs,
    ):
        if intermediaries is None:
            intermediaries = self.get_tracker_at("intermediaries", epoch=epoch)

        impacts = OutputImpact(
            encodings=self.data_encodings,
            intermediaries=intermediaries,
            separate_decay=separate_decay,
            collapse_level=collapse_level,
            collapse_lead_lag=collapse_lead_lag,
            with_groups=group_signals,
            linearization_method=linearization_method,
            combine_granularities_flag=True,
            is_animation_call=(epoch == "all"),
        )
        all_impacts_df = impacts.df
        if isinstance(all_impacts_df, list):
            all_impacts_df = pd.concat(all_impacts_df, axis=0 if epoch == "all" else 1)
        return OutputImpact(final_df=all_impacts_df, is_animation_call=(epoch == "all")).visualize(
            wibbles=use_grans,
            **kwargs,
        )

    def viz_impacts_change(
        self,
        epoch: Sequence[int] | Literal["all"] = "all",
        intermediaries: list[EconomicIntermediaries] | None = None,
        separate_decay: bool = False,
        collapse_level: CollapseLevel | dict[str, CollapseLevel] = 0,
        **kwargs,
    ):
        if intermediaries is None:
            intermediaries = self.get_tracker_at("intermediaries", epoch=epoch)
        impacts = OutputImpact(
            encodings=self.data_encodings,
            intermediaries=intermediaries,
            separate_decay=separate_decay,
            collapse_level=collapse_level,
            combine_granularities_flag=True,
            is_animation_call=(epoch == "all"),
        )
        return impacts.visualize_impacts_change(**kwargs)

    def viz_price_elasticity(
        self,
        epoch: EpochSpec = None,
        intermediaries: EconomicIntermediaries | list[EconomicIntermediaries] | None = None,
        num_viz: int = 2,
        return_df_only: bool = False,
        **kwargs,
    ):
        if intermediaries is None:
            intermediaries = self.get_tracker_at("price_elasticity", epoch=epoch)
            if intermediaries is None:
                return None
        price_elasticity = OutputPriceElasticity(
            price_values=intermediaries,
            encodings=self.data_encodings,
            combine_granularities_flag=True,
            is_animation_call=(epoch == "all"),
        )

        if return_df_only:
            return price_elasticity.df

        first_df = price_elasticity.df[0] if isinstance(price_elasticity.df, list) else price_elasticity.df
        return price_elasticity.visualize(
            price_devs={v: i for i, v in enumerate(first_df.columns.unique("signal"))},
            wibbles={v: i for i, v in enumerate(first_df.columns.unique("granularity")[:num_viz])},
            **kwargs,
        )

    def viz_mixed_effects(
        self,
        name: str,
        intermediaries: CurveTrackers | list[CurveTrackers] | None = None,
        epoch: EpochSpec = None,
        num_viz: int = 2,
        group_by_granularity: bool = True,
        return_df_only: bool = False,
        **kwargs,
    ):
        if intermediaries is None:
            intermediaries = self.get_tracker_at(name, epoch=epoch)
            if intermediaries is None:
                return None
        mixed_effects = OutputMixedEffect(
            curve_values=intermediaries,
            encodings=self.data_encodings,
            combine_granularities_flag=True,
            is_animation_call=(epoch == "all"),
        )
        me_df = mixed_effects.df

        if return_df_only:
            return me_df

        first_df = me_df[0] if isinstance(me_df, list) else me_df
        signals = first_df.columns.unique("signal")
        granularities = first_df.columns.unique("granularity")
        if group_by_granularity:
            granularities = granularities[:num_viz]
        else:
            signals = signals[:num_viz]

        signals = first_df.columns.unique("signal")
        granularities = first_df.columns.unique("granularity")
        if group_by_granularity:
            granularities = granularities[:num_viz]
        else:
            signals = signals[:num_viz]
        return mixed_effects.visualize(
            signal_encodings={v: i for i, v in enumerate(signals)},
            wibble_encodings={v: i for i, v in enumerate(granularities)},
            group_by_granularity=group_by_granularity,
            **kwargs,
        )

    def viz_curve(
        self,
        name: str,
        intermediaries: CurveTrackers | list[CurveTrackers] | None = None,
        epoch: EpochSpec = None,
        num_viz: int = 2,
        plot_slope: bool = False,
        dynamic_range: bool | float = False,
        use_granularities: Sequence[str] | None = None,
        use_vehicles: Sequence[str] | None = None,
        **kwargs,
    ) -> dict[str, Figure]:
        if intermediaries is None:
            intermediaries = self.get_tracker_at(name, epoch=epoch)
            if intermediaries is None:
                return {}
        curves = OutputCurve(
            curve_values=intermediaries,
            encodings=self.data_encodings,
            dynamic_range=dynamic_range,
            combine_granularities_flag=True,
            is_animation_call=(epoch == "all"),
        )
        return curves.visualize(
            vehicle_names=use_vehicles,
            wibbles=use_granularities,
            plot_slope=plot_slope,
            num_viz=num_viz,
            **kwargs,
        )

    def viz_all_curves(self, **kwargs):
        if self.impacts_layer.roicurve is None:
            return {}
        assert hasattr(self, "curve_trackers"), "You need to add the curve trackers first."
        all_plots = {}
        for name in self.curve_trackers:
            logger.info(f"Plotting {name}")
            all_plots |= {f"{name} {key}": v for key, v in self.viz_curve(name, **kwargs).items()}
        return all_plots

    def viz_decay(
        self,
        epoch: EpochSpec = None,
        intermediaries: EconomicIntermediaries | list[EconomicIntermediaries] | None = None,
        max_time: int = 20,
        num_viz: int = 2,
        desired_num_points: int = 256,
        use_granularities: Sequence[str] | None = None,
        use_vehicles: Sequence[str] | None = None,
        **kwargs,
    ):
        if intermediaries is None:
            intermediaries = self.get_tracker_at("intermediaries", epoch=epoch)
            if intermediaries is None:
                return None
            intermediaries = [
                RecursiveNamespace.parse(inter) if isinstance(inter, Mapping) else inter
                for inter in (intermediaries if isinstance(intermediaries, list) else [intermediaries])
            ]
        if isinstance(intermediaries, list):
            intermediaries = [inter for inter in intermediaries if inter.impacts.roicurve is not None]
            if len(intermediaries) == 0:
                return None
        elif intermediaries.impacts.roicurve is None:
            return None
        decay_df = OutputDecay(
            encodings=self.data_encodings,
            intermediaries=intermediaries,
            combine_granularities_flag=True,
            max_time=max_time,
            desired_num_points=desired_num_points,
            is_animation_call=(epoch == "all"),
        ).df
        return OutputCurve(final_df=decay_df, is_animation_call=(epoch == "all")).visualize(
            vehicle_names=use_vehicles,
            wibbles=use_granularities,
            plot_slope=False,
            num_viz=num_viz,
            **kwargs,
        )

    def viz_cumulative_decay(
        self,
        epoch: EpochSpec = None,
        intermediaries: EconomicIntermediaries | list[EconomicIntermediaries] | None = None,
        max_time: int = 20,
        num_viz: int = 2,
        desired_num_points: int = 256,
        use_granularities: Sequence[str] | None = None,
        use_vehicles: Sequence[str] | None = None,
        **kwargs,
    ):
        if intermediaries is None:
            intermediaries = self.get_tracker_at("intermediaries", epoch=epoch)
            if intermediaries is None:
                return None
            intermediaries = [
                RecursiveNamespace.parse(inter) if isinstance(inter, Mapping) else inter
                for inter in (intermediaries if isinstance(intermediaries, list) else [intermediaries])
            ]
        if isinstance(intermediaries, list):
            intermediaries = [inter for inter in intermediaries if inter.impacts.roicurve is not None]
            if len(intermediaries) == 0:
                return None
        elif intermediaries.impacts.roicurve is None:
            return None
        decay_df = OutputDecay(
            encodings=self.data_encodings,
            intermediaries=intermediaries,
            combine_granularities_flag=True,
            max_time=max_time,
            desired_num_points=desired_num_points,
            is_animation_call=(epoch == "all"),
        ).cumulative_df
        return OutputCurve(final_df=decay_df, is_animation_call=(epoch == "all")).visualize(
            vehicle_names=use_vehicles,
            wibbles=use_granularities,
            plot_slope=False,
            num_viz=num_viz,
            **kwargs,
        )

    DATASET_FILE_NAME = "dataset.pkl"

class TemporalNet(EconomicNetwork):
    result_type: ClassVar[type[tf.experimental.ExtensionType]] = TemporalNetIntermediaries
    num_mask_steps: ClassVar[int] = 0

    def build(self, input_shapes):
        """Build the layers needed for temporal net.

        Args:
            input_shapes (Tuple[tf.Tensor, ...]): Tuple of tensor shapes of positional arguments passed to `__call__()`.
        """
        super().build(input_shapes)
        self.baseline_layer = self.hyperparameters.get_submodule(
            "baseline",
            module_type=LinearBaseline,
            kwargs=dict(
                starting_sales=np.array(self.data_encodings["wholesaler_means"], dtype=np.float32)[None, :, :, 0],
                num_starts=self.data_encodings["total_restarts"],
                encodings=self.encodings,
            ),
            help="A piecewise linear curve to serve as the baseline for long term predictions.",
        )

    def get_baseline(
        self, impacts: ImpactsIntermediaries, training=False, debug=False, skip_metrics=False  # noqa: U100
    ) -> tuple[tf.Tensor, dict[str, LinearBaselineIntermediaries]]:
        """Calling the linear baseline layer to compute the baseline

        Args:
            impacts (ImpactsIntermediaries): Additive or multiplicative impacts on baseline
            training (bool, optional): Whether training the model parameters or not. Defaults to False.

        Returns:
            tf.Tensor: calculated baseline
            dict[str,LinearBaselineIntermediaries]: Dictionary of intermediate calculations for baseline,
                                                    like slope, intercept, etc.
        """
        baseline_layer = self.baseline_layer(
            LinearBaselineInput(
                dates_since_start=self.dates_since_start_ph,
                sales_num_restarts=self.sales_num_restarts_ph,
                hierarchy=dict(self.hierarchical_placeholders.items()),
                mask=self.yhat_mask_ph,
            ),
            training=training,
            skip_metrics=skip_metrics,
            debug=debug,
        )
        return baseline_layer.baseline, {"baseline_layer": baseline_layer}

    def create_network_phs(self, batch: EconomicModelInput, training=False):
        """
        Method to convert dataset object data into tensors for network inputs.
        Also handles data objects passed in kwargs externally.
        """
        super().create_network_phs(batch, training=training)
        self.sales_num_restarts_ph = batch.num_restarts

    def python_train_step(
        self, batch, optimizer: tf.optimizers.Optimizer, return_grads: bool = False
    ) -> tuple[dict[str, tf.Tensor], dict[str, tf.Tensor], tf.Tensor]:
        from wt_ml.layers.layer_utils import to_dense

        self.clear()
        for child in self.submodules:
            if isinstance(child, Module):
                child.clear()
        with tf.GradientTape() as tape:
            intermediaries = self(batch, training=True)
            total_loss = self.get_total_loss()
        trn_vars = self.trn_vars
        gradients = tape.gradient(total_loss, trn_vars)
        optimizer.apply_gradients(zip(gradients, trn_vars))
        if self.baseline_layer.use_perfect_adjustment:
            if isinstance(intermediaries, Mapping):
                intermediaries = intermediaries[self.NetTypes[0].__name__.lower()]
            self.baseline_layer.do_perfect_adjustment(batch, intermediaries)
        step = self._step_var.assign_add(1)
        if return_grads:
            gradients_tracker = {
                variable.name: to_dense(grad) for grad, variable in zip(gradients, trn_vars) if grad is not None
            }
            return (self.targets(), self.get_all_losses() | {"loss": total_loss}, step, gradients_tracker)
        else:
            return (self.targets(), self.get_all_losses() | {"loss": total_loss}, step)