OpenBB/openbb_terminal/custom/custom_controller.py

"""Custom Data Controller Module"""
__docformat__ = "numpy"

import argparse
import logging
from pathlib import Path
from typing import List

import pandas as pd
from prompt_toolkit.completion import NestedCompleter

from openbb_terminal import feature_flags as obbff
from openbb_terminal.custom import custom_model, custom_view
from openbb_terminal.decorators import log_start_end
from openbb_terminal.helper_funcs import (
    EXPORT_ONLY_FIGURES_ALLOWED,
)
from openbb_terminal.menu import session
from openbb_terminal.parent_classes import BaseController
from openbb_terminal.rich_config import console

logger = logging.getLogger(__name__)
# pylint:disable=import-outside-toplevel


class CustomDataController(BaseController):
    """Alternative Controller class"""

    CHOICES_COMMANDS: List[str] = ["load", "plot", "show", "info"]
    CHOICES_MENUS: List[str] = ["qa", "pred"]
    pandas_plot_choices = [
        "line",
        "scatter",
        "bar",
        "barh",
        "hist",
        "box",
        "kde",
        "area",
        "pie",
        "hexbin",
    ]
    PATH = "/custom/"

    def __init__(self, queue: List[str] = None):
        """Constructor"""
        super().__init__(queue)
        self.data = pd.DataFrame()
        self.file = ""
        self.DATA_FILES = [file.name for file in Path("custom_imports").iterdir()]

        if session and obbff.USE_PROMPT_TOOLKIT:
            choices: dict = {c: {} for c in self.controller_choices}
            choices["load"] = {c: None for c in self.DATA_FILES}
            choices["plot"]["-k"] = {c: None for c in self.pandas_plot_choices}
            self.choices = choices

            choices["support"] = self.SUPPORT_CHOICES

            self.completer = NestedCompleter.from_nested_dict(choices)

    def update_runtime_choices(self):
        if session and obbff.USE_PROMPT_TOOLKIT:
            self.choices["plot"] = {c: None for c in self.data.columns}
            self.choices["plot"]["-y"] = {c: None for c in self.data.columns}
            self.choices["plot"]["-x"] = {c: None for c in self.data.columns}
            self.choices["plot"]["--vs"] = {c: None for c in self.data.columns}
            self.choices["show"]["-s"] = {c: None for c in self.data.columns}

        self.completer = NestedCompleter.from_nested_dict(self.choices)

    def print_help(self):
        """Print help"""
        has_data_start = "" if not self.data.empty else "[unvl]"
        has_data_end = "" if not self.data.empty else "[/unvl]"
        help_text = f"""[cmds]
    load            load in custom data set[/cmds]

[param]Current file:[/param]    {self.file or None}[cmds]{has_data_start}

    show            show portion of loaded data
    info            show data info (columns and datatypes)
    plot            plot data from loaded file{has_data_end}[/cmds]
[menu]
>   qa              quantitative analysis,   \t e.g.: decompose, cusum, residuals analysis
>   pred            prediction techniques    \t e.g.: regression, arima, rnn, lstm[/menu]
"""
        console.print(text=help_text, menu="Custom")

    @log_start_end(log=logger)
    def call_load(self, other_args: List[str]):
        """Process load"""
        parser = argparse.ArgumentParser(
            add_help=False,
            formatter_class=argparse.ArgumentDefaultsHelpFormatter,
            prog="load",
            description="Load custom data set into a dataframe",
        )
        parser.add_argument(
            "-f",
            "--file",
            choices=self.DATA_FILES,
            help="File to load in.",
            default="test.csv",
            type=str,
        )

        if other_args and "-" not in other_args[0][0]:
            other_args.insert(0, "-f")
        ns_parser = self.parse_known_args_and_warn(parser, other_args)
        if ns_parser:
            file = Path("custom_imports") / ns_parser.file
            self.data = custom_model.load(file)
            self.data.columns = self.data.columns.map(lambda x: x.lower())
            for col in self.data.columns:
                if col in ["date", "time", "timestamp"]:  # Could be others?
                    self.data[col] = pd.to_datetime(self.data[col])
                    self.data = self.data.set_index(col)
            self.file = ns_parser.file
            self.update_runtime_choices()
        console.print("")

    @log_start_end(log=logger)
    def call_plot(self, other_args: List[str]):
        """Process plot command"""
        parser = argparse.ArgumentParser(
            add_help=False,
            formatter_class=argparse.ArgumentDefaultsHelpFormatter,
            prog="plot",
            description="Plot 2 columns of loaded data",
        )
        parser.add_argument(
            "-y",
            help="Variable to plot along y",
            dest="yvar",
            type=str,
            choices=list(self.data.columns),
        )
        parser.add_argument(
            "-x",
            "--vs",
            help="Variable along x axis",
            dest="xvar",
            type=str,
            choices=list(self.data.columns),
            default="",
        )
        parser.add_argument(
            "-k",
            "--kind",
            default="line",
            help="Type of plot for data",
            choices=self.pandas_plot_choices,
        )
        if other_args and "-" not in other_args[0][0]:
            other_args.insert(0, "-y")
        ns_parser = self.parse_known_args_and_warn(
            parser, other_args, export_allowed=EXPORT_ONLY_FIGURES_ALLOWED
        )
        if ns_parser:
            if self.data.empty:
                logger.error("No data loaded")
                console.print("[red]No data loaded.[/red]\n")
                return
            custom_view.custom_plot(
                self.data,
                ns_parser.yvar,
                ns_parser.xvar,
                kind=ns_parser.kind,
                export=ns_parser.export,
            )
        console.print("")

    @log_start_end(log=logger)
    def call_show(self, other_args: List[str]):
        """Process head command"""
        parser = argparse.ArgumentParser(
            add_help=False,
            formatter_class=argparse.ArgumentDefaultsHelpFormatter,
            prog="show",
            description="Show loaded dataframe",
        )
        parser.add_argument(
            "-s", "--sortcol", nargs="+", type=str, dest="sortcol", default=""
        )
        parser.add_argument(
            "-a", "--ascend", action="store_true", default=False, dest="ascend"
        )
        ns_parser = self.parse_known_args_and_warn(parser, other_args, limit=5)
        if ns_parser:
            if self.data.empty:
                logger.error("No data loaded")
                console.print("[red]No data loaded.[/red]\n")
                return
            if ns_parser.sortcol:
                sort_column = " ".join(ns_parser.sortcol)
                if sort_column not in self.data.columns:
                    logger.warning(
                        "%s is not a valid column. Showing without sorting", sort_column
                    )
                    console.print(
                        f"[red]{sort_column} not a valid column.  Showing without sorting.\n[/red]"
                    )
                else:
                    console.print(
                        self.data.sort_values(
                            by=sort_column, ascending=ns_parser.ascend
                        ).head(ns_parser.limit)
                    )
                    console.print()
                    return

            console.print(self.data.head(ns_parser.limit))
        console.print()

    @log_start_end(log=logger)
    def call_info(self, other_args: List[str]):
        """Process info command"""
        parser = argparse.ArgumentParser(
            add_help=False,
            formatter_class=argparse.ArgumentDefaultsHelpFormatter,
            prog="info",
            description="Show information of custom data.",
        )

        ns_parser = self.parse_known_args_and_warn(parser, other_args)
        if ns_parser:
            if self.data.empty:
                logger.error("No data loaded.")
                console.print("[red]No data loaded.[/red]\n")
                return
            console.print(self.data.info())
        console.print()

    @log_start_end(log=logger)
    def call_qa(self, _):
        """Process qa command"""
        from openbb_terminal.custom.quantitative_analysis import qa_controller

        self.queue = self.load_class(
            qa_controller.QaController,
            custom_df=self.data,
            file=self.file,
            queue=self.queue,
        )

    @log_start_end(log=logger)
    def call_pred(self, _):
        """Process pred command"""
        if obbff.ENABLE_PREDICT:
            from openbb_terminal.custom.prediction_techniques import pred_controller

            self.queue = self.load_class(
                pred_controller.PredictionTechniquesController,
                self.data,
                self.file,
                self.queue,
            )