OpenBB/openbb_terminal/forecast/forecast_view.py

"""Forecast View"""
__docformat__ = "numpy"

import logging
import os
from typing import Dict, List, Optional

import matplotlib.pyplot as plt
import pandas as pd
import seaborn as sns
from darts.utils.statistics import plot_acf
from pandas.plotting import register_matplotlib_converters

from openbb_terminal.config_plot import PLOT_DPI
from openbb_terminal.config_terminal import theme
from openbb_terminal.decorators import log_start_end
from openbb_terminal.forecast import forecast_model, helpers
from openbb_terminal.helper_funcs import export_data, plot_autoscale, print_rich_table
from openbb_terminal.rich_config import console

logger = logging.getLogger(__name__)

register_matplotlib_converters()


@log_start_end(log=logger)
def show_options(
    datasets: Dict[str, pd.DataFrame],
    dataset_name: str = None,
    export: str = "",
    sheet_name: str = None,
):
    """Plot custom data

    Parameters
    ----------
    datasets: dict
        The loaded in datasets
    dataset_name: str
        The name of the dataset you wish to show options for
    sheet_name: str
        Optionally specify the name of the sheet the data is exported to.
    export: str
        Format to export image
    """
    if not datasets:
        console.print(
            "Please load in a dataset by using the 'load' command before using this feature."
        )
    else:
        option_tables = forecast_model.get_options(datasets, dataset_name)

        for dataset, data_values in option_tables.items():
            print_rich_table(
                data_values,
                headers=list(data_values.columns),
                show_index=False,
                title=f"Options for dataset: '{dataset}'",
            )

            export_data(
                export,
                os.path.dirname(os.path.abspath(__file__)),
                f"{dataset}_options",
                data_values.set_index("column"),
                sheet_name,
            )


@log_start_end(log=logger)
def display_plot(
    data: pd.DataFrame,
    columns: List[str],
    export: str = "",
    sheet_name: str = None,
    external_axes: Optional[List[plt.axes]] = None,
):
    """Plot data from a dataset
    Parameters
    ----------
    data: pd.DataFrame
        The dataframe to plot
    columns: List[str]
        The columns to show
    sheet_name: str
        Optionally specify the name of the sheet the data is exported to.
    export: str
        Format to export image
    external_axes:Optional[List[plt.axes]]
        External axes to plot on
    """

    # Check that there's at least a valid dataframe
    if external_axes is None:
        _, ax = plt.subplots(figsize=plot_autoscale(), dpi=PLOT_DPI)
    else:
        ax = external_axes[0]

    # Only do if data is not plotted, otherwise an error will occur
    if "date" in data.columns and "date" not in columns:
        data = data.set_index("date")

    for column in columns:
        ax.plot(data[column], label=column)

    theme.style_primary_axis(ax)
    if external_axes is None:
        theme.visualize_output()

    ax.legend()

    export_data(
        export,
        os.path.dirname(os.path.abspath(__file__)),
        "plot",
        sheet_name,
    )


@log_start_end(log=logger)
def display_seasonality(
    data: pd.DataFrame,
    column: str = "close",
    export: str = "",
    sheet_name: str = None,
    m: Optional[int] = None,
    max_lag: int = 24,
    alpha: float = 0.05,
    external_axes: Optional[List[plt.axes]] = None,
):
    """Plot seasonality from a dataset

    Parameters
    ----------
    data: pd.DataFrame
        The dataframe to plot
    column: str
        The column of the dataframe to analyze
    sheet_name: str
        Optionally specify the name of the sheet the data is exported to.
    export: str
        Format to export image
    m: Optional[int]
        Optionally, a time lag to highlight on the plot. Default is none.
    max_lag: int
        The maximal lag order to consider. Default is 24.
    alpha: float
        The confidence interval to display. Default is 0.05.
    external_axes:Optional[List[plt.axes]]
        External axes to plot on
    """

    if not data.empty:
        _, series = helpers.get_series(data, column)
        if external_axes is None:
            _, ax = plt.subplots(figsize=plot_autoscale(), dpi=PLOT_DPI)
        else:
            ax = external_axes[0]

        # TODO: Add darts check_seasonality here
        plot_acf(
            series, m=m, max_lag=max_lag, alpha=alpha, axis=ax, default_formatting=False
        )

        theme.style_primary_axis(ax)

        if external_axes is None:
            theme.visualize_output()

    export_data(
        export,
        os.path.dirname(os.path.abspath(__file__)),
        "plot",
        sheet_name,
    )


@log_start_end(log=logger)
def display_corr(
    dataset: pd.DataFrame,
    export: str = "",
    sheet_name: str = None,
    external_axes: Optional[List[plt.axes]] = None,
):
    """Plot correlation coefficients for dataset features

    Parameters
    ----------
    dataset : pd.DataFrame
        The dataset fore calculating correlation coefficients
    sheet_name: str
        Optionally specify the name of the sheet the data is exported to.
    export: str
        Format to export image
    external_axes:Optional[List[plt.axes]]
        External axes to plot on
    """

    if external_axes is None:
        _, ax = plt.subplots(figsize=plot_autoscale(), dpi=PLOT_DPI)
    else:
        ax = external_axes[0]

    # correlation
    correlation = forecast_model.corr_df(dataset)
    sns.heatmap(
        correlation,
        vmin=correlation.values.min(),
        vmax=1,
        square=True,
        linewidths=0.1,
        annot=True,
        annot_kws={"size": 8},
        cbar_kws=dict(use_gridspec=True, location="left"),
    )

    ax.set_title("Correlation Matrix")

    for t in ax.get_yticklabels():
        t.set_fontsize(7)
        t.set_fontweight("bold")
        t.set_horizontalalignment("left")

    for t in ax.get_xticklabels():
        t.set_fontsize(7)
        t.set_fontweight("bold")
        t.set_rotation(60)
        t.set_horizontalalignment("right")

    theme.style_primary_axis(ax)

    if external_axes is None:
        theme.visualize_output()

    export_data(
        export,
        os.path.dirname(os.path.abspath(__file__)),
        "plot",
        sheet_name,
    )


@log_start_end(log=logger)
def show_df(
    data: pd.DataFrame,
    limit: int = 15,
    limit_col: int = 10,
    name: str = "",
    export: str = "",
    sheet_name: str = None,
):
    console.print(
        f"[green]{name} dataset has shape (row, column): {data.shape}\n[/green]"
    )
    if len(data.columns) > limit_col:
        console.print(
            f"[red]Dataframe has more than {limit_col} columns."
            " If you have extra screen space, consider increasing the `limit_col`,"
            " else export to see all of the data.[/red]\n"
        )
        data = data.iloc[:, :limit_col]
    print_rich_table(
        data.head(limit),
        headers=list(data.columns),
        show_index=True,
        title=f"Dataset {name} | Showing {limit} of {len(data)} rows",
    )

    export_data(
        export,
        os.path.dirname(os.path.abspath(__file__)),
        f"{name}_show",
        data,
        sheet_name,
    )


@log_start_end(log=logger)
def describe_df(
    data: pd.DataFrame, name: str = "", export: str = "", sheet_name: str = None
):
    new_df = forecast_model.describe_df(data)
    print_rich_table(
        new_df,
        headers=list(data.describe().columns),
        show_index=True,
        title=f"Showing Descriptive Statistics for Dataset {name}",
    )

    export_data(
        export,
        os.path.dirname(os.path.abspath(__file__)),
        f"{name}_show",
        sheet_name,
    )


@log_start_end(log=logger)
def export_df(
    data: pd.DataFrame, export: str, name: str = "", sheet_name: str = None
) -> None:
    export_data(
        export,
        os.path.dirname(os.path.abspath(__file__)),
        f"{name}_show",
        data,
        sheet_name,
    )