OpenBB/openbb_terminal/stocks/discovery/ark_model.py

""" ARK Model """
__docformat__ = "numpy"

import json
import logging
from datetime import timedelta

import numpy as np
import pandas as pd
import requests
import yfinance as yf
from bs4 import BeautifulSoup
from pandas.core.frame import DataFrame

from openbb_terminal.decorators import log_start_end
from openbb_terminal.helper_funcs import get_user_agent
from openbb_terminal.rich_config import console

logger = logging.getLogger(__name__)


@log_start_end(log=logger)
def get_ark_orders(
    buys_only: bool = False,
    sells_only: bool = False,
    fund: str = "",
) -> DataFrame:

    """Returns ARK orders in a Dataframe

    Parameters
    ----------
    buys_only: bool
        Flag to filter on buys only
    sells_only: bool
        Flag to sort on sells only
    fund: str
        Optional filter by fund

    Returns
    -------
    DataFrame
        ARK orders data frame with the following columns -
        (ticker, date, shares, weight, fund, direction)
    """
    url_orders = "https://cathiesark.com/ark-funds-combined/trades"

    raw_page = requests.get(url_orders, headers={"User-Agent": get_user_agent()}).text

    parsed_script = BeautifulSoup(raw_page, "lxml").find(
        "script", {"id": "__NEXT_DATA__"}
    )

    parsed_json = json.loads(parsed_script.string)

    df_orders = pd.json_normalize(parsed_json["props"]["pageProps"]["arkTrades"])

    if df_orders.empty:
        return pd.DataFrame()

    df_orders.drop(
        [
            "hidden",
            "images.thumbnail",
            "cusip",
            "estimated_price",
            "updated_at",
            "created_at",
            "region",
            "country",
            "isADR",
            "companyName",
            "clinicalTrialsSearchHandle",
            "wasSPACBuy",
            "currencyMultiplier",
            "useRapidAPI",
            "description",
            "quandlTicker",
            "customThumbnail",
            "custom_thumbnail",
            "id",
        ],
        axis=1,
        inplace=True,
    )

    df_orders["date"] = pd.to_datetime(df_orders["date"], format="%Y-%m-%d").dt.date

    if df_orders.empty:
        console.print("The ARK orders aren't available at the moment.\n")
        return pd.DataFrame()
    if fund:
        df_orders = df_orders[df_orders.fund == fund]
    if buys_only:
        df_orders = df_orders[df_orders.direction == "Buy"]
    if sells_only:
        df_orders = df_orders[df_orders.direction == "Sell"]

    return df_orders


@log_start_end(log=logger)
def add_order_total(data: DataFrame) -> DataFrame:
    """Takes an ARK orders dataframe and pulls data from Yahoo Finance to add
    volume, open, close, high, low, and total columns

    Parameters
    ----------
    data: DataFrame
        ARK orders data frame with the following columns -
        (ticker, date, shares, weight, fund, direction)

    Returns
    -------
    DataFrame
        ARK orders data frame with the following columns -
        (ticker, date, shares, volume, open, close, high, low, total, weight, fund, direction)
    """
    start_date = data["date"].iloc[-1] - timedelta(days=1)

    tickers = " ".join(data["ticker"].unique())

    prices = yf.download(tickers, start=start_date, progress=False)

    if prices.empty:
        return pd.DataFrame()
    for i, candle in enumerate(["Volume", "Open", "Close", "High", "Low", "Total"]):
        data.insert(i + 3, candle.lower(), 0)

    pd.options.mode.chained_assignment = None
    for i, _ in data.iterrows():
        if np.isnan(
            prices["Open"][data.loc[i, "ticker"]][
                data.loc[i, "date"].strftime("%Y-%m-%d")
            ]
        ):
            for candle in ["Volume", "Open", "Close", "High", "Low", "Total"]:
                data.loc[i, candle.lower()] = 0
            continue

        for candle in ["Volume", "Open", "Close", "High", "Low"]:
            data.loc[i, candle.lower()] = prices[candle][data.loc[i, "ticker"]][
                data.loc[i, "date"].strftime("%Y-%m-%d")
            ]

        data.loc[i, "total"] = data.loc[i, "close"] * data.loc[i, "shares"]

    pd.options.mode.chained_assignment = "warn"

    return data