class OutputCurve(ModelOutputItem):
    """Class to visualize roi curve"""

    def __init__(
        self,
        curve_values: CurveTrackers | list[CurveTrackers] | None = None,
        encodings: Encodings | None = None,
        dynamic_range: bool | float = False,
        combine_granularities_flag: bool = False,
        final_df: pd.DataFrame | list[pd.DataFrame] | None = None,
        is_animation_call: bool = False,
    ):
        super().__init__(final_df, [curve_values, encodings])
        self.final_df = final_df
        self.dynamic_range = dynamic_range
        self.is_animation_call = is_animation_call
        if self.final_df is None:
            assert encodings is not None, "Encodings is required"
            self.curve_values = curve_values if isinstance(curve_values, list) else [curve_values]
            self.encodings = encodings
            self.combine_granularities_flag = combine_granularities_flag

    def get_curve_df(
        self,
        curve_values: CurveTrackers,
    ) -> pd.DataFrame | tuple[pd.DataFrame, pd.Series]:
        if isinstance(curve_values, Mapping):
            curve_values = RecursiveNamespace.parse(curve_values)
        index, curve_np = get_curve_matrix(
            to_numpy(curve_values.spends),
            to_numpy(curve_values.impact),
            to_numpy(curve_values.slope),
            self.encodings["normalization_factor"],
            self.dynamic_range,
        )

        if self.dynamic_range:
            index = pd.Index(index, name="percentage_spend")
        else:
            index = pd.Index(index, name="spend")
        granularity_names, level_names = get_granularity_names(self.encodings, curve_values.input)
        vehicle_names = to_signal_names(curve_values.signal_names)
        columns = pd.MultiIndex.from_tuples(
            [
                (*names, signal, vehicle)
                for names in granularity_names
                for signal in ["impact", "slope"]
                for vehicle in vehicle_names
            ],
            names=(*level_names, "value", "signal"),
        )
        curve_df = pd.DataFrame(
            curve_np.reshape(index.shape[0], -1),
            index=index,
            columns=columns,
        )
        step_cols = pd.MultiIndex.from_tuples(
            [(*names, "step", vehicle) for names in granularity_names for vehicle in vehicle_names],
            names=(*level_names, "value", "signal"),
        )
        df = pd.concat(
            [
                curve_df,
                pd.DataFrame(
                    np.full((len(curve_df.index), len(step_cols)), to_numpy(curve_values.step), dtype=np.int32),
                    index=curve_df.index,
                    columns=step_cols,
                ),
            ],
            axis=1,
        )
        if self.dynamic_range:
            range_index = pd.MultiIndex.from_tuples(
                [(*names, vehicle) for names in granularity_names for vehicle in vehicle_names],
                names=(*level_names, "signal"),
            )
            max_vals = tf.math.reduce_max(curve_values.spends, axis=1, keepdims=True)
            range_series = pd.Series(
                self.encodings["normalization_factor"] * to_numpy(max_vals[:, 0]).reshape(-1),
                index=range_index,
                name="max_val",
            )
            return df, range_series
        else:
            return df

    def get_plot_range(self, curve_df: pd.DataFrame, cols: list[str]) -> tuple[float, float]:
        values = curve_df[cols].values
        return (0.0, float(1.1 * np.where(np.isnan(values), 0.0, values).max()))

    @cached_property
    def df(
        self,
    ) -> pd.DataFrame | list[pd.DataFrame] | tuple[pd.DataFrame | list[pd.DataFrame], pd.Series | list[pd.Series]]:
        """Get the curve dataframe"""
        if self.final_df is not None:
            return self.final_df

        curve_dfs: list[pd.DataFrame] = []
        range_series: list[pd.Series] = []
        for step_curve_values in self.curve_values:
            step_curve_result = self.get_curve_df(step_curve_values)
            if self.dynamic_range:
                curve_df, range_ser = step_curve_result
                if TYPE_CHECKING:
                    assert isinstance(curve_df, pd.DataFrame)
                    assert isinstance(range_ser, pd.Series)

                curve_dfs.append(combine_granularities(curve_df) if self.combine_granularities_flag else curve_df)
                range_series.append(range_ser)
            else:
                if TYPE_CHECKING:
                    assert isinstance(step_curve_result, pd.DataFrame)

                curve_dfs.append(
                    combine_granularities(step_curve_result) if self.combine_granularities_flag else step_curve_result
                )

        final_curve_dfs = curve_dfs[0] if len(self.curve_values) == 1 else curve_dfs
        if self.dynamic_range:
            final_range_series = range_series[0] if len(self.curve_values) == 1 else range_series
            return final_curve_dfs, final_range_series

        return final_curve_dfs

    @property
    def visualization_df(self):
        if not self.dynamic_range:
            return self.df

        return self.df[0]

    def visualize(
        self,
        wibbles: Sequence[str] | None = None,
        vehicle_names: Sequence[str] | None = None,
        wibble_encodings: dict[str, int] | None = None,
        vehicle_encodings: dict[str, int] | None = None,
        group_by: Literal["granularity", "signal"] | None = None,
        group_by_granularity: bool = True,
        show_both=False,
        **kwargs,
    ) -> dict[str, go.Figure]:
        """Visualize curve

        Args:
            wibbles (Sequence[str] | None, optional): Granularity names. Defaults to None.
            vehicle_names (Sequence[str] | None, optional): Vehicle names. Defaults to None.
            wibble_encodings (dict[str, int] | None, optional): Granularity encodings. Defaults to None.
            vehicle_encodings (dict[str, int] | None, optional): Vehcile encodings. Defaults to None.
            group_by (Literal["granularity", "signal"] | None, optional): Level name by which to group and
            view the plots. Defaults to None.
            group_by_granularity (bool, optional): Granularity names by which to group and view the plots.
            Defaults to True.
            show_both (bool, optional): Flag to indicate whether to group plots both by granularity level and
            signal level. Defaults to False.
            show_plots (bool, optional): Flag to indicate whether to show the plots. Defaults to True.

        Returns:
            dict[str, go.Figure]: Prepared line or scatter plots of curve dataframe.
        """
        all_plots = {}
        if show_both:
            all_plots |= self.visualize(
                granularities=wibble_encodings,
                vehicles=vehicle_encodings,
                granularity_names=wibbles,
                vehicle_names=vehicle_names,
                group_by="granularity",
                show_both=False,
                **kwargs,
            )
            all_plots |= self.visualize(
                granularities=wibble_encodings,
                vehicles=vehicle_encodings,
                granularity_names=wibbles,
                vehicle_names=vehicle_names,
                group_by="signal",
                show_both=False,
                **kwargs,
            )
            return all_plots

        # Prepare animation frame and curve_df
        animation_frame = "step" if self.is_animation_call else None
        curve_df = self.visualization_df
        if isinstance(curve_df, list):
            if self.is_animation_call:
                curve_df = pd.concat(curve_df, axis=0)
            else:
                curve_df = pd.concat(curve_df, axis=1)
        if TYPE_CHECKING:
            assert isinstance(curve_df, pd.DataFrame)
        if wibbles is None:
            if wibble_encodings:
                wibbles = list(wibble_encodings.keys())
            else:
                wibbles = curve_df.columns.unique("granularity").tolist()

        level_name = curve_df.columns.names[-1] if "signal" not in curve_df.columns.names else "signal"
        if vehicle_names is None:
            if vehicle_encodings is None:
                vehicle_names = list(curve_df.columns.unique(level_name))
            else:
                vehicle_names = list(vehicle_encodings.keys())
        if group_by is None:
            group_by = "granularity" if group_by_granularity else level_name

        all_plots = make_plots(
            curve_df,
            wibbles if group_by == "granularity" else vehicle_names,
            group_by,
            vehicle_names if group_by == "granularity" else wibbles,
            animation_frame,
            get_range=self.get_plot_range,
            **kwargs,
        )
        return all_plots