OpenBB/openbb_terminal/portfolio/portfolio_optimization/po_engine.py

from typing import Dict, List, Tuple
import pandas as pd
from openbb_terminal.portfolio.portfolio_optimization import (
    excel_model,
    optimizer_model,
    optimizer_helper,
)
from openbb_terminal.portfolio.portfolio_optimization.parameters import params_view
from openbb_terminal.rich_config import console


class PoEngine:
    """Portfolio Optimization Engine"""

    def __init__(
        self,
        symbols_categories: Dict[str, Dict[str, str]] = None,
        symbols_file_path: str = None,
        parameters_file_path: str = None,
    ):
        """Initialize the engine

        Parameters
        ----------
        symbols_categories : Dict[str, float], optional
            Categories, by default None
        symbols_file_path : str, optional
            Symbols file path, by default None
        parameters_file_path : str, optional
            Parameters file path, by default None
        """

        self._categories: Dict[str, Dict[str, str]] = {}
        self._symbols: List[str] = []
        self._weights: Dict[str, float] = {}
        self._returns: pd.DataFrame = None
        self._params: Dict[str, float] = {}
        self._current_model: str

        if symbols_categories is not None:
            self._symbols, self._categories = PoEngine.__parse_dictionary(
                symbols_categories
            )

        elif symbols_file_path is not None:
            self._symbols, self._categories = excel_model.load_allocation(
                symbols_file_path
            )
        else:
            raise ValueError("symbols or file_path must be provided")

        if parameters_file_path is not None:
            self._params, self._current_model = params_view.load_file(
                parameters_file_path
            )

    @staticmethod
    def __parse_dictionary(
        symbols_categories: Dict[str, Dict[str, str]]
    ) -> Tuple[List[str], Dict[str, Dict[str, str]]]:
        """Parse the categories dictionary

        Parameters
        ----------
        symbols_categories : Dict[str, Dict[str, str]]
            Categories

        Returns
        -------
        Tuple[List[str], Dict[str, Dict[str, str]]]
            Symbols and categories
        """
        if isinstance(symbols_categories, dict):
            symbols = PoEngine.__get_symbols_from_categories(symbols_categories)
        else:
            raise TypeError("'symbols_categories' must be a dictionary.")

        for symbol in symbols:
            symbols_categories["CURRENCY"].setdefault(symbol, "USD")
            symbols_categories["CURRENT_INVESTED_AMOUNT"].setdefault(symbol, "0")

        return symbols, symbols_categories

    @staticmethod
    def __get_symbols_from_categories(
        symbols_categories: Dict[str, Dict[str, str]]
    ) -> List[str]:
        """Get the symbols from the categories dictionary

        Parameters
        ----------
        symbols_categories : Dict[str, Dict[str, str]], optional
            Categories

        Returns
        -------
        List[str]
            List of symbols
        """

        try:
            symbols = []
            for item in symbols_categories.items():
                _, values = item
                for v in values.keys():
                    symbols.append(v)

            return list(set(symbols))

        except Exception:
            console.print(
                "Unsupported dictionary format. See `portfolio.po.load` examples for correct format."
            )
            return []

    def get_symbols(self):
        return self._symbols

    def get_available_categories(self) -> List[str]:
        """Get the available categories

        Returns
        -------
        List[str]
            Available categories
        """
        available_categories = list(self._categories.keys())

        if "CURRENT_INVESTED_AMOUNT" in available_categories:
            available_categories.remove("CURRENT_INVESTED_AMOUNT")

        return available_categories

    def get_category(self, category: str = None) -> Dict[str, str]:
        """Get the category

        Parameters
        ----------
        category : str, optional
            Category, by default None

        Returns
        -------
        Dict[str, str]
            Category
        """
        if category is None:
            console.print("No category provided. Please provide a category.")
            return {}

        d = self.get_categories_dict()
        if category in d:
            return d[category]
        return {}

    def get_categories_dict(self) -> Dict[str, Dict[str, str]]:
        """Get the categories

        Returns
        -------
        Dict[str, Dict[str, str]]
            Categories
        """
        if not self._categories:
            console.print("No categories found. Use 'load' to load a file.")
            return {}
        return self._categories

    def get_category_df(self, category: str = None) -> pd.DataFrame:
        """Get the category df

        Returns
        -------
        pd.DataFrame
            Category DataFrame
        """
        if category is None:
            console.print("No category provided. Please provide a category.")
            return pd.DataFrame()
        if not self._categories:
            console.print("No categories found. Use 'load' to load a file.")
            return pd.DataFrame()
        return optimizer_model.get_categories(
            weights=self._weights, categories=self._categories, column=category
        )

    def set_weights(self, weights: Dict[str, float]):
        """Set the weights

        Parameters
        ----------
        weights : Dict[str, float]
            Weights
        """
        self._weights = weights

    def get_weights(self, warning=True) -> Dict[str, float]:
        """Get the weights

        Parameters
        ----------
        warning : bool, optional
            Display warning, by default True

        Returns
        -------
        Dict[str, float]
            Weights
        """
        if not self._weights:
            if warning:
                console.print("No weights found. Please perform some optimization.")
            return {}
        return self._weights

    def get_weights_df(self, warning=True) -> pd.DataFrame:
        """Get the weights

        Parameters
        ----------
        warning : bool, optional
            Display warning, by default True

        Returns
        -------
        pd.DataFrame
            Weights
        """
        if not self._weights:
            if warning:
                console.print("No weights found. Please perform some optimization.")
            return pd.DataFrame()
        return optimizer_helper.dict_to_df(self._weights)

    def set_params(self, params: Dict[str, float], update=False):
        """Set the parameters

        Parameters
        ----------
        params : Dict[str, float]
            Parameters
        update : bool, optional
            Update the current model, by default False
        """
        if update:
            self._params.update(params)
        else:
            self._params = params

    def get_params(self) -> Dict:
        """Get the parameters

        Returns
        -------
        Dict
            Parameters
        """
        return self._params

    def set_params_from_file(self, file_path: str):
        """Set the parameters from a file

        Parameters
        ----------
        file_path : str
            File path
        """
        self._params, self._current_model = params_view.load_file(file_path)

    def set_current_model(self, model: str):
        """Set the current model

        Parameters
        ----------
        model : str
            Model
        """
        self._current_model = model

    def get_current_model(self) -> str:
        """Get the current model

        Returns
        -------
        str
            Current model
        """
        return self._current_model

    def set_returns(self, returns: pd.DataFrame):
        """Set the stock returns

        Parameters
        ----------
        returns : pd.DataFrame
            Stock returns
        """
        self._returns = returns

    def get_returns(self) -> pd.DataFrame:
        """Get the stock returns

        Returns
        -------
        pd.DataFrame
            Stock returns
        """
        if self._returns.empty:
            console.print("No returns found. Please perform some optimization.")
            return pd.DataFrame()
        return self._returns