# -*- coding: utf-8 -*-
"""Generate probabilities of forecastings crossing limit thresholds.
Created on Thu Jan 27 13:36:18 2022
"""
import random
import numpy as np
import pandas as pd
from autots.evaluator.auto_model import model_forecast, back_forecast
from autots.evaluator.auto_ts import AutoTS
from autots.models.model_list import all_result_path, diff_window_motif_list
from autots.tools.plotting import (
    plot_forecast_with_intervals,
    plot_risk_score_bar,
    plot_simulation_paths,
    grays,
)


def extract_result_windows(forecasts, model_name=None):
    """Standardize result windows from different models into a 3d numpy array."""
    # First check if result_windows is directly on the forecasts object (for ensembles)
    result_windows = getattr(forecasts, "result_windows", None)
    if result_windows is None:
        # Otherwise, check the model attribute (for regular models)
        model = getattr(forecasts, "model", None)
        if model is None or getattr(model, "result_windows", None) is None:
            raise ValueError(
                "Model did not expose result windows. Ensure return_result_windows is enabled."
            )
        result_windows = model.result_windows

    if model_name is None:
        model_name = forecasts.model_name

    if isinstance(result_windows, dict) or model_name in diff_window_motif_list:
        result_windows = np.moveaxis(np.array(list(result_windows.values())), 0, -1)
    else:
        result_windows = np.asarray(result_windows)

    if result_windows.ndim == 4:
        result_windows = result_windows[0]

    if result_windows.ndim != 3:
        raise ValueError(
            "Result windows must be three dimensional, received shape {}".format(
                result_windows.shape
            )
        )

    transformer = getattr(forecasts, "transformer", None)
    model = getattr(forecasts, "model", None)
    combination_transformer = (
        getattr(model, "combination_transformer", None) if model is not None else None
    )

    need_inverse_transformer = False
    if transformer is not None and hasattr(transformer, "inverse_transform"):
        transformations = getattr(transformer, "transformations", None)
        if transformations is None:
            need_inverse_transformer = True
        else:
            need_inverse_transformer = bool(transformations)

    need_inverse_combination = combination_transformer is not None and hasattr(
        combination_transformer, "inverse_transform"
    )

    if not need_inverse_transformer and not need_inverse_combination:
        return np.asarray(result_windows)

    index = forecasts.forecast.index
    columns = forecasts.forecast.columns
    transformed_array = []
    # bring these back to the original feature space, as they aren't already
    for samp in result_windows:
        inv_df = pd.DataFrame(samp, index=index, columns=columns)
        if need_inverse_combination:
            inv_df = combination_transformer.inverse_transform(inv_df)
        if need_inverse_transformer:
            inv_df = transformer.inverse_transform(inv_df)
        if isinstance(inv_df, pd.DataFrame):
            transformed_array.append(inv_df.to_numpy(copy=False))
        else:
            transformed_array.append(np.asarray(inv_df))
    return np.stack(transformed_array, axis=0)


def extract_window_index(forecasts):
    windows = getattr(getattr(forecasts, "model", None), "windows", None)
    if windows is None:
        raise ValueError("window indexes not supported by this model yet")
    return windows


def set_limit_forecast(
    df_train,
    forecast_length,
    model_name="SeasonalNaive",
    model_param_dict={
        'method': "median",
        "lag_1": 28,
        "lag_2": None,
    },
    model_transform_dict={
        'fillna': 'nearest',
        'transformations': {},
        'transformation_params': {},
    },
    prediction_interval=0.9,
    frequency='infer',
    model_forecast_kwargs={
        "verbose": 1,
        "n_jobs": "auto",
        "random_seed": 321,
    },
    future_regressor_train=None,
    future_regressor_forecast=None,
):
    """Helper function for forecast limits set by forecast algorithms."""
    forecasts = model_forecast(
        model_name=model_name,
        model_param_dict=model_param_dict,
        model_transform_dict=model_transform_dict,
        df_train=df_train,
        forecast_length=forecast_length,
        frequency=frequency,
        prediction_interval=prediction_interval,
        future_regressor_train=future_regressor_train,
        future_regressor_forecast=future_regressor_forecast,
        return_model=True,
        **model_forecast_kwargs,
    )
    return forecasts.upper_forecast.values, forecasts.lower_forecast.values


def set_limit_forecast_historic(
    df_train,
    forecast_length,
    model_name="SeasonalNaive",
    model_param_dict={
        'method': "median",
        "lag_1": 28,
        "lag_2": None,
    },
    model_transform_dict={
        'fillna': 'nearest',
        'transformations': {},
        'transformation_params': {},
    },
    prediction_interval=0.9,
    frequency='infer',
    model_forecast_kwargs={
        "verbose": 2,
        "n_jobs": "auto",
        "random_seed": 321,
    },
    future_regressor_train=None,
    future_regressor_forecast=None,
    eval_periods=None,
):
    """Helper function for forecast limits set by forecast algorithms."""
    forecasts = back_forecast(
        df=df_train,
        n_splits="auto",
        model_name=model_name,
        model_param_dict=model_param_dict,
        model_transform_dict=model_transform_dict,
        forecast_length=forecast_length,
        frequency=frequency,
        prediction_interval=prediction_interval,
        future_regressor_train=future_regressor_train,
        eval_periods=eval_periods,
        **model_forecast_kwargs,
    )
    return forecasts.upper_forecast.values, forecasts.lower_forecast.values


class EventRiskForecast(object):
    """Generate a risk score (0 to 1, but usually close to 0) for a future event exceeding user specified upper or lower bounds.

    Event risk was created for the use case of predicting the risk of an out of stock event (lower limit ~= 0), or a very high demand day, and similar issues.
    It can also be used to identify "regressions", that is, cases where the real world data (run as a holdout) was unlikely, to raise an alert that something is amiss.
    This is a capability somewhat unique to motif models, as most probabilistic models "fan out" rather than creating a mix of quality outcomes.

    Upper and lower limits can be one of four types, and may each be different.
    1. None (no risk score calculated for this direction)
    2. Float in range [0, 1] historic quantile of series (which is historic min and max at edges) is chosen as limit.
    3. A dictionary of {"model_name": x,  "model_param_dict": y, "model_transform_dict": z, "prediction_interval": 0.9} to generate a forecast as the limits
        Primarily intended for simple forecasts like SeasonalNaive, but can be used with any AutoTS model
    4. a custom input numpy array or pandas DataFrame of shape (forecast_length, num_series)

    This can be used to find the "middle" limit too, flip so upper=lower and lower=upper, then abs(U - (1 - L)).
    In some cases it may help to drop the results from the first forecast timestep or two.

    This functions by generating multiple outcome forecast possiblities in two ways.
    If a 'Motif' type model is passed, it uses all the k neighbors motifs as outcome paths (recommended)
    All other AutoTS models will generate the possible outcomes by utilizing multiple prediction_intervals (more intervals = slower but more resolution).
    The risk score is then the % of outcome forecasts which cross the limit.
    (less than or equal for lower, greater than or equal for upper)

    Only accepts `wide` style dataframe input.
    Methods are class_methods and can be used standalone. They default to __init__ inputs, but can be overriden.
    Results are usually a numpy array of shape (forecast_length, num_series)

    Args:
        df_train (pd.DataFrame): `wide style data, pd.DatetimeIndex for index and one series per column
        forecast_length (int): number of forecast steps to make
        frequency (str): frequency of timesteps
        prediction_interval (float): float or list of floats for probabilistic forecasting
            if a list, the first item in the list is the one used for .fit default
        model_forecast_kwargs (dict): AutoTS kwargs to pass to generaet_result_windows, .fit_forecast, and forecast-style limits
        model_name, model_param_dict, model_transform_dict: for model_forecast in generate_result_windows
        future_regressor_train, future_regressor_forecast: regressor arrays if used

    Methods:
        fit
        predict
        predict_historic
        generate_result_windows
        generate_risk_array
        generate_historic_risk_array
        set_limit
        plot

    Attributes:
        result_windows, forecast_df, up_forecast_df, low_forecast_df
        lower_limit_2d, upper_limit_2d, upper_risk_array, lower_risk_array
        window_index, prediction_object
    """

    def __init__(
        self,
        df_train,
        forecast_length,
        frequency: str = "infer",
        prediction_interval=0.9,
        lower_limit=0.05,
        upper_limit=0.95,
        model_name="BallTreeRegressionMotif",
        model_param_dict={
            "window": 3,
            "point_method": "midhinge",
            "distance_metric": "euclidean",
            "k": 100,
            "sample_fraction": 5000000,
            "comparison_transformation": {
                "fillna": "cubic",
                "transformations": {"0": "AlignLastDiff"},
                "transformation_params": {
                    "0": {
                        "rows": 364,
                        "displacement_rows": 1,
                        "quantile": 1.0,
                        "decay_span": None,
                    }
                },
            },
            "combination_transformation": {
                "fillna": "time",
                "transformations": {"0": "AlignLastDiff"},
                "transformation_params": {
                    "0": {
                        "rows": 7,
                        "displacement_rows": 1,
                        "quantile": 1.0,
                        "decay_span": 2,
                    }
                },
            },
            "extend_df": True,
            "mean_rolling_periods": 12,
            "macd_periods": 74,
            "std_rolling_periods": 30,
            "max_rolling_periods": 12,
            "min_rolling_periods": None,
            "quantile90_rolling_periods": 10,
            "quantile10_rolling_periods": 10,
            "ewm_alpha": None,
            "ewm_var_alpha": None,
            "additional_lag_periods": None,
            "abs_energy": False,
            "rolling_autocorr_periods": None,
            "nonzero_last_n": None,
            "datepart_method": None,
            "polynomial_degree": None,
            "regression_type": None,
            "holiday": False,
            "scale_full_X": False,
            "series_hash": True,
            "frac_slice": None,
        },
        model_transform_dict={
            "fillna": "akima",
            "transformations": {
                "0": "Log",
                "1": "SinTrend",
                "2": "ChangepointDetrend",
                "3": "RobustScaler",
            },
            "transformation_params": {
                "0": {},
                "1": {},
                "2": {
                    "model": "Linear",
                    "changepoint_spacing": 5040,
                    "changepoint_distance_end": 520,
                    "datepart_method": "common_fourier",
                },
                "3": {},
            },
        },
        model_forecast_kwargs={
            "max_generations": 30,
            "verbose": 1,
            "n_jobs": "auto",
            "random_seed": 321,
            "transformer_list": "no_expanding",
        },
        future_regressor_train=None,
        future_regressor_forecast=None,
    ):
        self.name = "EventRiskForecast"
        self.df_train = df_train
        self.forecast_length = forecast_length
        self.frequency = frequency
        self.prediction_interval = prediction_interval
        self.lower_limit = lower_limit
        self.upper_limit = upper_limit
        self.model_forecast_kwargs = model_forecast_kwargs
        self.model_name = model_name
        self.model_param_dict = model_param_dict
        self.model_transform_dict = model_transform_dict
        self.future_regressor_train = future_regressor_train
        self.future_regressor_forecast = future_regressor_forecast
        self.outcome_shape = (forecast_length, df_train.shape[1])
        self.outcome_columns = df_train.columns
        self.outcome_index = None
        self.result_windows = None
        self.window_index = None
        self.prediction_object = None

    def __repr__(self):
        """Print."""
        return f"{self.name} object with window generating model: {self.model_name, self.model_param_dict, self.model_transform_dict}"

    def fit(
        self,
        df_train=None,
        forecast_length=None,
        prediction_interval=None,
        models_mode="event_risk",
        model_list=[
            "UnivariateMotif",
            "MultivariateMotif",
            "SectionalMotif",
            "ARCH",
            'MetricMotif',
            'SeasonalityMotif',
        ],
        ensemble=None,
        autots_kwargs=None,
        future_regressor_train=None,
    ):
        """Shortcut for generating model params.

        args specified are those suggested for an otherwise normal AutoTS run

        Args:
            df_train (pd.DataFrame): wide style only
            model_method (str): event_risk here is used by motif models
            model_list (list): suggesting the use of motif models
            ensemble (list): must be None or empty list to get motif result windows
            autots_kwargs (dict): all other args passed in as kwargs
                if None, defaults to class model_forecast_kwargs, for blank pass empty dict
        """
        autots_kwargs = (
            self.model_forecast_kwargs if autots_kwargs is None else autots_kwargs
        )
        future_regressor_train = (
            self.future_regressor_train
            if future_regressor_train is None
            else future_regressor_train
        )
        prediction_interval = (
            self.prediction_interval
            if prediction_interval is None
            else prediction_interval
        )
        if isinstance(prediction_interval, list):
            prediction_interval = prediction_interval[0]
        model = AutoTS(
            forecast_length=(
                self.forecast_length if forecast_length is None else forecast_length
            ),
            prediction_interval=prediction_interval,
            models_mode=models_mode,
            model_list=model_list,
            ensemble=ensemble,
            **autots_kwargs,
        ).fit(
            self.df_train if df_train is None else df_train,
            future_regressor=future_regressor_train,
        )
        self.model_name = model.best_model_name
        self.model_param_dict = model.best_model_params
        self.model_transform_dict = model.best_model_transformation_params
        return self

    def generate_result_windows(
        self,
        df_train=None,
        forecast_length=None,
        frequency=None,
        prediction_interval=None,
        model_name=None,
        model_param_dict=None,
        model_transform_dict=None,
        model_forecast_kwargs=None,
        future_regressor_train=None,
        future_regressor_forecast=None,
    ):
        """For event risk forecasting. Params default to class init but can be overridden here.

        Returns:
            result_windows (numpy.array): (num_samples/k, forecast_length, num_series/columns)
        """
        df_train = self.df_train if df_train is None else df_train
        forecast_length = (
            self.forecast_length if forecast_length is None else forecast_length
        )
        prediction_interval = (
            self.prediction_interval
            if prediction_interval is None
            else prediction_interval
        )
        frequency = self.frequency if frequency is None else frequency
        model_name = self.model_name if model_name is None else model_name
        model_param_dict = (
            self.model_param_dict if model_param_dict is None else model_param_dict
        )
        model_transform_dict = (
            self.model_transform_dict
            if model_transform_dict is None
            else model_transform_dict
        )
        model_forecast_kwargs = (
            self.model_forecast_kwargs
            if model_forecast_kwargs is None
            else model_forecast_kwargs
        )
        future_regressor_train = (
            self.future_regressor_train
            if future_regressor_train is None
            else future_regressor_train
        )
        future_regressor_forecast = (
            self.future_regressor_forecast
            if future_regressor_forecast is None
            else future_regressor_forecast
        )

        if model_name in all_result_path:
            if isinstance(prediction_interval, list):
                prediction_interval = prediction_interval[0]
            if model_name in diff_window_motif_list:
                model_param_dict = {
                    **model_param_dict,
                    **{"return_result_windows": True},
                }
            forecasts = model_forecast(
                model_name=model_name,
                model_param_dict=model_param_dict,
                model_transform_dict=model_transform_dict,
                df_train=df_train,
                forecast_length=forecast_length,
                frequency=frequency,
                prediction_interval=prediction_interval,
                future_regressor_train=future_regressor_train,
                future_regressor_forecast=future_regressor_forecast,
                return_model=True,
                **model_forecast_kwargs,
            )
            result_windows = extract_result_windows(forecasts, model_name=model_name)
            lower_forecast = forecasts.lower_forecast
            upper_forecast = forecasts.upper_forecast
        else:
            if isinstance(prediction_interval, float):
                prediction_interval = list(set([prediction_interval, 0.95, 0.8, 0.5]))
            result_windows_list = []
            point_included = False
            for interval in prediction_interval:
                forecasts = model_forecast(
                    model_name=model_name,
                    model_param_dict=model_param_dict,
                    model_transform_dict=model_transform_dict,
                    df_train=df_train,
                    forecast_length=forecast_length,
                    frequency=frequency,
                    prediction_interval=interval,
                    future_regressor_train=future_regressor_train,
                    future_regressor_forecast=future_regressor_forecast,
                    return_model=True,
                    **model_forecast_kwargs,
                )
                if not point_included:
                    result_windows_list.append(forecasts.forecast)
                    lower_forecast = forecasts.lower_forecast
                    upper_forecast = forecasts.upper_forecast
                    point_included = True
                result_windows_list.append(forecasts.upper_forecast)
                result_windows_list.append(forecasts.lower_forecast)
            result_windows = np.array(result_windows_list)
        try:
            self.window_index = extract_window_index(forecasts)
        except ValueError:
            self.window_index = None
        self.prediction_object = forecasts
        return result_windows, forecasts.forecast, upper_forecast, lower_forecast

    @staticmethod
    def set_limit(
        limit,
        target_shape,
        df_train,
        direction="upper",
        period="forecast",
        forecast_length=None,
        eval_periods=None,
    ):
        """Handles all limit input styles and returns numpy array.

        Args:
            limit: see class overview for input options, also accepts pandas DataFrame
            target_shape (tuple): of (forecast_length, num_series)
            df_train (pd.DataFrame): training data
            direction (str): whether it is the "upper" or "lower" limit
            period (str): "forecast" or "historic" only used for limits defined by forecast algorithm params
            forecast_length (int): needed only for historic of forecast algorithm defined limit
            eval_periods (int): only for historic forecast limit, only runs on the tail n (this) of data
        """
        # handle a predefined array or dataframe
        if isinstance(limit, pd.DataFrame):
            limit_columns = getattr(limit, "columns", None)
            if limit_columns is not None and list(limit_columns) != list(
                df_train.columns
            ):
                try:
                    limit = limit.reindex(columns=df_train.columns)
                except Exception as exc:
                    raise ValueError(
                        f"Unable to align {direction}_limit DataFrame columns to training data"
                    ) from exc
            limit_array = np.asarray(limit)
        elif isinstance(limit, np.ndarray):
            limit_array = limit
        else:
            limit_array = None

        if limit_array is not None:
            assert (
                limit_array.ndim == 2
            ), f"{direction}_limit, if array, must be 2d np array of shape forecast_length, n_series"
            assert (
                limit_array.shape == target_shape
            ), f"{direction}_limit, if array, must be 2d np array of shape forecast_length, n_series"
            return limit_array
        # handle a limit as a quantile defined by float
        elif isinstance(limit, float) or isinstance(limit, int):
            assert (
                limit >= 0 and limit <= 1
            ), f"{direction}_limit if float must be in the range [0, 1], received {limit}"
            return np.repeat(
                np.nanquantile(df_train, limit, axis=0).reshape(1, -1),
                target_shape[0],
                axis=0,
            )
        # handle None
        elif limit is None:
            return None
        # handle a limit defined by a forecast algorithm
        elif isinstance(limit, dict):
            if period == "historic":
                upper, lower = set_limit_forecast_historic(
                    df_train=df_train,
                    forecast_length=forecast_length,
                    model_name=limit.get("model_name", "SeasonalNaive"),
                    model_param_dict=limit.get("model_param_dict", {}),
                    model_transform_dict=limit.get("model_transform_dict", {}),
                    prediction_interval=limit.get("prediction_interval", 0.9),
                    frequency=limit.get("frequency", 'infer'),
                    model_forecast_kwargs=limit.get(
                        "model_forecast_kwargs",
                        {
                            "verbose": 1,
                            "n_jobs": "auto",
                            "random_seed": 321,
                        },
                    ),
                    future_regressor_train=limit.get("future_regressor_train", None),
                    future_regressor_forecast=limit.get(
                        "future_regressor_forecast", None
                    ),
                    eval_periods=eval_periods,
                )
            else:
                upper, lower = set_limit_forecast(
                    df_train=df_train,
                    forecast_length=target_shape[0],
                    model_name=limit.get("model_name", "SeasonalNaive"),
                    model_param_dict=limit.get("model_param_dict", {}),
                    model_transform_dict=limit.get("model_transform_dict", {}),
                    prediction_interval=limit.get("prediction_interval", 0.9),
                    frequency=limit.get("frequency", 'infer'),
                    model_forecast_kwargs=limit.get(
                        "model_forecast_kwargs",
                        {
                            "verbose": 1,
                            "n_jobs": "auto",
                            "random_seed": 321,
                        },
                    ),
                    future_regressor_train=limit.get("future_regressor_train", None),
                    future_regressor_forecast=limit.get(
                        "future_regressor_forecast", None
                    ),
                )
            if direction == "upper":
                return upper
            elif direction == "lower":
                return lower
            elif direction == "both":
                return upper, lower
        else:
            raise ValueError(
                f"{direction}_limit was not recognized dtype, input was {limit}"
            )

    @staticmethod
    def generate_risk_array(result_windows, limit, direction="upper"):
        """Given a df and a limit, returns a 0/1 array of whether limit was equaled or exceeded."""
        if direction == "upper":
            return (result_windows >= limit).astype(int).sum(
                axis=0
            ) / result_windows.shape[0]
        elif direction == "lower":
            return (result_windows <= limit).astype(int).sum(
                axis=0
            ) / result_windows.shape[0]
        else:
            raise ValueError(
                f"arg `direction`: {direction} not recognized in generate_risk_array"
            )

    @staticmethod
    def generate_historic_risk_array(df, limit, direction="upper"):
        """Given a df and a limit, returns a 0/1 array of whether limit was equaled or exceeded."""
        if direction == "upper":
            return (df >= limit).astype(int)
        elif direction == "lower":
            return (df <= limit).astype(int)
        else:
            raise ValueError(
                f"arg `direction`: {direction} not recognized in generate_risk_array"
            )

    def predict(self):
        """Returns forecast upper, lower risk probability arrays for input limits."""
        self.upper_limit_2d = self.set_limit(
            self.upper_limit, self.outcome_shape, self.df_train, direction="upper"
        )
        self.lower_limit_2d = self.set_limit(
            self.lower_limit, self.outcome_shape, self.df_train, direction="lower"
        )
        if self.upper_limit_2d is None and self.lower_limit_2d is None:
            raise ValueError(
                "both upper and lower limits are None, at least one must be specified"
            )

        (
            self.result_windows,
            self.forecast_df,
            self.up_forecast_df,
            self.low_forecast_df,
        ) = self.generate_result_windows()
        self.outcome_index = self.forecast_df.index

        self.upper_risk_array = None
        if self.upper_limit_2d is not None:
            self.upper_risk_array = self.generate_risk_array(
                self.result_windows, self.upper_limit_2d, direction="upper"
            )

        self.lower_risk_array = None
        if self.lower_limit_2d is not None:
            self.lower_risk_array = self.generate_risk_array(
                self.result_windows, self.lower_limit_2d, direction="lower"
            )
        return pd.DataFrame(
            self.upper_risk_array,
            columns=self.outcome_columns,
            index=self.outcome_index,
        ), pd.DataFrame(
            self.lower_risk_array,
            columns=self.outcome_columns,
            index=self.outcome_index,
        )

    def predict_historic(self, upper_limit=None, lower_limit=None, eval_periods=None):
        """Returns upper, lower risk probability arrays for input limits for the historic data.
        If manual numpy array limits are used, the limits will need to be appropriate shape (for df_train and eval_periods if used)

        Args:
            upper_limit: if different than the version passed to init
            lower_limit: if different than the version passed to init
            eval_periods (int): only assess the n most recent periods of history
        """
        upper_limit = self.upper_limit if upper_limit is None else upper_limit
        lower_limit = self.lower_limit if lower_limit is None else lower_limit
        if eval_periods is not None:
            target_shape = (eval_periods, self.df_train.shape[1])
            train_df = self.df_train.tail(eval_periods)
        else:
            target_shape = self.df_train.shape
            train_df = self.df_train
        self.historic_upper_limit_2d = self.set_limit(
            upper_limit,
            target_shape,
            self.df_train,
            direction="upper",
            period="historic",
            forecast_length=self.forecast_length,
            eval_periods=eval_periods,
        )
        self.historic_lower_limit_2d = self.set_limit(
            lower_limit,
            target_shape,
            self.df_train,
            direction="lower",
            period="historic",
            forecast_length=self.forecast_length,
            eval_periods=eval_periods,
        )
        if (
            self.historic_upper_limit_2d is None
            and self.historic_lower_limit_2d is None
        ):
            raise ValueError(
                "both upper and lower limits are None, at least one must be specified"
            )

        self.historic_upper_risk_array = None
        if self.historic_upper_limit_2d is not None:
            self.historic_upper_risk_array = self.generate_historic_risk_array(
                train_df, self.historic_upper_limit_2d, direction="upper"
            )

        self.historic_lower_risk_array = None
        if self.historic_lower_limit_2d is not None:
            self.historic_lower_risk_array = self.generate_historic_risk_array(
                train_df, self.historic_lower_limit_2d, direction="lower"
            )

        return pd.DataFrame(
            self.historic_upper_risk_array,
            columns=self.outcome_columns,
            index=train_df.index,
        ), pd.DataFrame(
            self.historic_lower_risk_array,
            columns=self.outcome_columns,
            index=train_df.index,
        )

    def query_risk(
        self,
        dates=None,
        series=None,
        include_limits=True,
        include_forecast=False,
        historic=False,
        return_json=False,
    ):
        """Query a specific slice of risk forecast results with minimal token usage.

        Designed for LLM-friendly output with compact representation.

        Args:
            dates (str, datetime, list, slice): Date(s) to query.
                - Single date: "2024-01-15" or datetime object
                - Date range: slice("2024-01-01", "2024-01-31")
                - List of dates: ["2024-01-15", "2024-01-20"]
                - None: all dates
            series (str, list): Series name(s) to query.
                - Single series: "sales"
                - Multiple series: ["sales", "revenue"]
                - None: all series
            include_limits (bool): Include the limit thresholds used
            include_forecast (bool): Include the underlying forecast values
            historic (bool): Query historic risk (from predict_historic) instead of forecast risk
            return_json (bool): Return JSON string instead of dict

        Returns:
            dict or str: Compact risk data

        Examples:
            >>> # Single series, single date
            >>> risk.query_risk(dates="2024-01-15", series="sales")
            {'upper_risk': {'sales': {'2024-01-15': 0.05}}, 'lower_risk': {...}}

            >>> # Include limits and forecast
            >>> risk.query_risk(
            ...     dates=slice("2024-01-01", "2024-01-07"),
            ...     series=["sales"],
            ...     include_limits=True,
            ...     include_forecast=True
            ... )
        """
        # Determine which arrays to use
        if historic:
            upper_risk_df = (
                pd.DataFrame(
                    self.historic_upper_risk_array,
                    columns=self.outcome_columns,
                    index=self.df_train.tail(
                        self.historic_upper_risk_array.shape[0]
                        if self.historic_upper_risk_array is not None
                        else 0
                    ).index
                    if self.historic_upper_risk_array is not None
                    else None,
                )
                if hasattr(self, 'historic_upper_risk_array')
                and self.historic_upper_risk_array is not None
                else None
            )

            lower_risk_df = (
                pd.DataFrame(
                    self.historic_lower_risk_array,
                    columns=self.outcome_columns,
                    index=self.df_train.tail(
                        self.historic_lower_risk_array.shape[0]
                        if self.historic_lower_risk_array is not None
                        else 0
                    ).index
                    if self.historic_lower_risk_array is not None
                    else None,
                )
                if hasattr(self, 'historic_lower_risk_array')
                and self.historic_lower_risk_array is not None
                else None
            )

            upper_limit_df = (
                pd.DataFrame(
                    self.historic_upper_limit_2d,
                    columns=self.outcome_columns,
                    index=upper_risk_df.index,
                )
                if hasattr(self, 'historic_upper_limit_2d')
                and self.historic_upper_limit_2d is not None
                and upper_risk_df is not None
                else None
            )

            lower_limit_df = (
                pd.DataFrame(
                    self.historic_lower_limit_2d,
                    columns=self.outcome_columns,
                    index=lower_risk_df.index,
                )
                if hasattr(self, 'historic_lower_limit_2d')
                and self.historic_lower_limit_2d is not None
                and lower_risk_df is not None
                else None
            )
        else:
            upper_risk_df = (
                pd.DataFrame(
                    self.upper_risk_array,
                    columns=self.outcome_columns,
                    index=self.outcome_index,
                )
                if hasattr(self, 'upper_risk_array')
                and self.upper_risk_array is not None
                else None
            )

            lower_risk_df = (
                pd.DataFrame(
                    self.lower_risk_array,
                    columns=self.outcome_columns,
                    index=self.outcome_index,
                )
                if hasattr(self, 'lower_risk_array')
                and self.lower_risk_array is not None
                else None
            )

            upper_limit_df = (
                pd.DataFrame(
                    self.upper_limit_2d,
                    columns=self.outcome_columns,
                    index=self.outcome_index,
                )
                if hasattr(self, 'upper_limit_2d') and self.upper_limit_2d is not None
                else None
            )

            lower_limit_df = (
                pd.DataFrame(
                    self.lower_limit_2d,
                    columns=self.outcome_columns,
                    index=self.outcome_index,
                )
                if hasattr(self, 'lower_limit_2d') and self.lower_limit_2d is not None
                else None
            )

        # Validate we have some data
        if upper_risk_df is None and lower_risk_df is None:
            raise ValueError(
                "No risk data available. Run predict() or predict_historic() first."
            )

        # Use first available for index reference
        ref_df = upper_risk_df if upper_risk_df is not None else lower_risk_df

        # Handle series selection
        if series is None:
            selected_series = self.outcome_columns.tolist()
        elif isinstance(series, str):
            selected_series = [series]
        else:
            selected_series = list(series)

        # Validate series exist
        missing = set(selected_series) - set(self.outcome_columns)
        if missing:
            raise ValueError(f"Series not found: {missing}")

        # Handle date selection
        if dates is None:
            date_slice = ref_df.index
        elif isinstance(dates, slice):
            date_slice = ref_df.loc[dates.start : dates.stop].index
        elif isinstance(dates, (list, pd.Index)):
            date_slice = pd.DatetimeIndex(dates)
        else:
            # Single date
            date_slice = pd.DatetimeIndex([pd.to_datetime(dates)])

        # Build result dictionary
        result = {
            'model': self.model_name,
            'type': 'historic_risk' if historic else 'forecast_risk',
        }

        # Extract upper risk values
        if upper_risk_df is not None:
            upper_risk_data = {}
            for col in selected_series:
                col_data = {}
                for dt in date_slice:
                    if dt in upper_risk_df.index:
                        val = upper_risk_df.loc[dt, col]
                        col_data[dt.isoformat()] = float(val) if pd.notna(val) else None
                upper_risk_data[col] = col_data
            result['upper_risk'] = upper_risk_data

        # Extract lower risk values
        if lower_risk_df is not None:
            lower_risk_data = {}
            for col in selected_series:
                col_data = {}
                for dt in date_slice:
                    if dt in lower_risk_df.index:
                        val = lower_risk_df.loc[dt, col]
                        col_data[dt.isoformat()] = float(val) if pd.notna(val) else None
                lower_risk_data[col] = col_data
            result['lower_risk'] = lower_risk_data

        # Add limits if requested
        if include_limits:
            if upper_limit_df is not None:
                upper_limit_data = {}
                for col in selected_series:
                    col_data = {}
                    for dt in date_slice:
                        if dt in upper_limit_df.index:
                            val = upper_limit_df.loc[dt, col]
                            col_data[dt.isoformat()] = (
                                float(val) if pd.notna(val) else None
                            )
                    upper_limit_data[col] = col_data
                result['upper_limit'] = upper_limit_data

            if lower_limit_df is not None:
                lower_limit_data = {}
                for col in selected_series:
                    col_data = {}
                    for dt in date_slice:
                        if dt in lower_limit_df.index:
                            val = lower_limit_df.loc[dt, col]
                            col_data[dt.isoformat()] = (
                                float(val) if pd.notna(val) else None
                            )
                    lower_limit_data[col] = col_data
                result['lower_limit'] = lower_limit_data

        # Add forecast if requested
        if (
            include_forecast
            and hasattr(self, 'forecast_df')
            and self.forecast_df is not None
        ):
            forecast_data = {}
            for col in selected_series:
                col_data = {}
                for dt in date_slice:
                    if dt in self.forecast_df.index:
                        val = self.forecast_df.loc[dt, col]
                        col_data[dt.isoformat()] = float(val) if pd.notna(val) else None
                forecast_data[col] = col_data
            result['forecast'] = forecast_data

        if return_json:
            import json

            return json.dumps(result, indent=2)

        return result

    def plot(
        self,
        column_idx=0,
        column=None,
        grays=grays,
        up_low_color=["#ff4500", "#ff5349"],
        bar_color="#6495ED",
        bar_ylim=[0.0, 0.5],
        figsize=(14, 8),
        result_windows=None,
        lower_limit_2d=None,
        upper_limit_2d=None,
        upper_risk_array=None,
        lower_risk_array=None,
    ):
        """Plot a sample of the risk forecast outcomes.

        Args:
            column_idx (int): positional index of series to sample for plot
            column (str): optional column name to select instead of column_idx
            grays (list of str): list of hex codes for colors for the potential forecasts
            up_low_colors (list of str): two hex code colors for upper and lower limits
            bar_color (str): hex color for bar graph
            bar_ylim (list): passed to ylim of plot, sets scale of axis of barplot
            figsize (tuple): passed to figsize of output figure
        """
        import matplotlib.pyplot as plt

        result_windows = (
            self.result_windows if result_windows is None else result_windows
        )
        lower_limit_2d = (
            self.lower_limit_2d if lower_limit_2d is None else lower_limit_2d
        )
        upper_limit_2d = (
            self.upper_limit_2d if upper_limit_2d is None else upper_limit_2d
        )
        upper_risk_array = (
            self.upper_risk_array if upper_risk_array is None else upper_risk_array
        )
        lower_risk_array = (
            self.lower_risk_array if lower_risk_array is None else lower_risk_array
        )

        if column is not None:
            if hasattr(self.outcome_columns, "get_loc"):
                try:
                    column_idx = self.outcome_columns.get_loc(column)
                except KeyError as exc:
                    raise ValueError(
                        f"Column '{column}' is not present in outcome columns"
                    ) from exc
            else:
                columns_list = list(self.outcome_columns)
                if column not in columns_list:
                    raise ValueError(
                        f"Column '{column}' is not present in outcome columns"
                    )
                column_idx = columns_list.index(column)

        column_name = self.outcome_columns[column_idx]
        simulations = result_windows[:, :, column_idx]
        horizon_length = simulations.shape[1]
        horizon_index = self.outcome_index
        if horizon_index is None:
            horizon_index = np.arange(horizon_length)
        fig, (ax1, ax2) = plt.subplots(
            nrows=2, ncols=1, gridspec_kw={'height_ratios': [2, 1]}, figsize=figsize
        )
        fig.suptitle(f"{column_name} Event Risk Forecasting")

        # Plot simulation paths using utility function
        colors = random.choices(grays, k=simulations.shape[0])
        plot_simulation_paths(
            simulations,
            index=horizon_index,
            colors=colors,
            alpha=0.9,
            linewidth=1.2,
            ax=ax1,
        )

        upper_color = up_low_color[0] if len(up_low_color) >= 1 else "#ff4500"
        lower_color = up_low_color[1] if len(up_low_color) >= 2 else up_low_color[0]
        if upper_limit_2d is not None:
            ax1.plot(
                horizon_index,
                upper_limit_2d[:, column_idx],
                color=upper_color,
                linewidth=2,
                linestyle="--",
                label="Upper Limit",
            )
        if lower_limit_2d is not None:
            ax1.plot(
                horizon_index,
                lower_limit_2d[:, column_idx],
                color=lower_color,
                linewidth=2,
                linestyle="--",
                label="Lower Limit",
            )

        ax1.set_ylabel("Value")
        ax1.set_facecolor("#f7f7f7")
        ax1.grid(axis="y", linestyle="--", alpha=0.25)
        if upper_limit_2d is not None or lower_limit_2d is not None:
            ax1.legend(loc="upper right", frameon=False)

        # handle one being None
        try:
            up_risk = np.asarray(upper_risk_array[:, column_idx])
        except Exception:
            up_risk = np.zeros(horizon_length)
        try:
            low_risk = np.asarray(lower_risk_array[:, column_idx])
        except Exception:
            low_risk = np.zeros(horizon_length)
        risk_total = up_risk + low_risk

        # Plot risk score using utility function
        plot_risk_score_bar(
            risk_total,
            index=horizon_index,
            bar_color=bar_color,
            bar_ylim=bar_ylim,
            title="Combined Risk Score",
            ylabel="Combined Risk",
            xlabel="Forecast Horizon",
            ax=ax2,
        )

        fig.tight_layout(rect=[0, 0, 1, 0.96])

    def plot_eval(
        self,
        df_test,
        column_idx=0,
        actuals_color=["#00BFFF"],
        up_low_color=["#ff4500", "#ff5349"],
        bar_color="#6495ED",
        bar_ylim=[0.0, 0.5],
        figsize=(14, 8),
        lower_limit_2d=None,
        upper_limit_2d=None,
        upper_risk_array=None,
        lower_risk_array=None,
    ):
        """Plot a sample of the risk forecast with known value vs risk score.

        Args:
            df_test (pd.DataFrame): dataframe of known values (dt index, series)
            column_idx (int): positional index of series to sample for plot
            actuals_color (list of str): list of one hex code for line of known actuals
            up_low_colors (list of str): two hex code colors for lower and upper
            bar_color (str): hex color for bar graph
            bar_ylim (list): passed to ylim of plot, sets scale of axis of barplot
            figsize (tuple): passed to figsize of output figure
        """
        import matplotlib.pyplot as plt

        lower_limit_2d = (
            self.lower_limit_2d if lower_limit_2d is None else lower_limit_2d
        )
        upper_limit_2d = (
            self.upper_limit_2d if upper_limit_2d is None else upper_limit_2d
        )
        upper_risk_array = (
            self.upper_risk_array if upper_risk_array is None else upper_risk_array
        )
        lower_risk_array = (
            self.lower_risk_array if lower_risk_array is None else lower_risk_array
        )
        col = self.outcome_columns[column_idx]

        fig, (ax1, ax2) = plt.subplots(
            nrows=2, ncols=1, gridspec_kw={'height_ratios': [2, 1]}, figsize=figsize
        )
        fig.suptitle(f'{col} Event Risk Forecasting Evaluation')

        # Prepare data for plotting
        plot_df = df_test[col].to_frame()
        if lower_limit_2d is not None:
            plot_df['lower_limit'] = lower_limit_2d[:, column_idx]
        if upper_limit_2d is not None:
            plot_df['upper_limit'] = upper_limit_2d[:, column_idx]

        # Use plot_forecast_with_intervals for the top panel
        # Map colors: actuals get first color, limits get up_low_color
        color_map = {col: actuals_color[0] if actuals_color else '#00BFFF'}
        lower_color = up_low_color[1] if len(up_low_color) >= 2 else up_low_color[0]
        upper_color = up_low_color[0] if len(up_low_color) >= 1 else "#ff4500"

        if 'lower_limit' in plot_df.columns:
            color_map['lower_limit'] = lower_color
        if 'upper_limit' in plot_df.columns:
            color_map['upper_limit'] = upper_color

        plot_forecast_with_intervals(
            plot_df,
            actual_col=col,
            forecast_col=None,  # No forecast, just actuals
            lower_col='lower_limit' if 'lower_limit' in plot_df.columns else None,
            upper_col='upper_limit' if 'upper_limit' in plot_df.columns else None,
            title=None,  # Title set on fig
            colors=color_map,
            include_bounds=False,  # We'll plot limits as lines, not bands
            ax=ax1,
        )

        # handle one being None
        try:
            up_risk = upper_risk_array[:, column_idx]
        except Exception:
            up_risk = 0
        try:
            low_risk = lower_risk_array[:, column_idx]
        except Exception:
            low_risk = 0

        risk_total = up_risk + low_risk
        plot_risk_score_bar(
            risk_total,
            index=plot_df.index,
            bar_color=bar_color,
            bar_ylim=bar_ylim,
            title="Risk Score",
            ylabel="upper & lower risk",
            xlabel="",
            ax=ax2,
        )
        # Remove x-tick labels for consistency with original
        ax2.set_xticks([])
