path: root/openbb_platform/providers/intrinio/openbb_intrinio/models/options_snapshots.py

"""Intrinio Options Snapshots Model."""

# pylint: disable=unused-argument

import gzip
from datetime import (
    date as dateType,
    datetime,
)
from io import BytesIO
from typing import Any, Dict, List, Optional, Union

import numpy as np
from openbb_core.app.model.abstract.error import OpenBBError
from openbb_core.provider.abstract.fetcher import Fetcher
from openbb_core.provider.standard_models.options_snapshots import (
    OptionsSnapshotsData,
    OptionsSnapshotsQueryParams,
)
from openbb_core.provider.utils.helpers import amake_request
from pandas import DataFrame, NaT, Series, read_csv, to_datetime
from pydantic import Field
from pytz import timezone


class IntrinioOptionsSnapshotsQueryParams(OptionsSnapshotsQueryParams):
    """Intrinio Options Snapshots Query.

    Source: https://docs.intrinio.com/documentation/web_api/get_options_snapshots_v2
    """

    date: Optional[Union[dateType, datetime, str]] = Field(
        default=None,
        description="The date of the data. Can be a datetime or an ISO datetime string."
        + " Data appears to go back to around 2022-06-01"
        + " Example: '2024-03-08T12:15:00+0400'",
    )
    only_traded: bool = Field(
        default=True,
        description="Only include options that have been traded during the session, default is True."
        + " Setting to false will dramatically increase the size of the response - use with caution.",
    )


class IntrinioOptionsSnapshotsData(OptionsSnapshotsData):
    """Intrinio Options Snapshots Data. Warning: This is a large file."""

    bid: List[Union[float, None]] = Field(
        default_factory=list,
        description="The last bid price at the time.",
        json_schema_extra={"x-unit_measurement": "currency"},
    )
    bid_size: List[Union[int, None]] = Field(
        default_factory=list,
        description="The size of the last bid price.",
    )
    bid_timestamp: List[Union[datetime, None]] = Field(
        default_factory=list,
        description="The timestamp of the last bid price.",
    )
    ask: List[Union[float, None]] = Field(
        default_factory=list,
        description="The last ask price at the time.",
        json_schema_extra={"x-unit_measurement": "currency"},
    )
    ask_size: List[Union[int, None]] = Field(
        default_factory=list,
        description="The size of the last ask price.",
    )
    ask_timestamp: List[Union[datetime, None]] = Field(
        default_factory=list,
        description="The timestamp of the last ask price.",
    )
    total_bid_volume: List[Union[int, None]] = Field(
        default_factory=list,
        description="Total volume of bids.",
    )
    bid_high: List[Union[float, None]] = Field(
        default_factory=list,
        description="The highest bid price.",
        json_schema_extra={"x-unit_measurement": "currency"},
    )
    bid_low: List[Union[float, None]] = Field(
        default_factory=list,
        description="The lowest bid price.",
        json_schema_extra={"x-unit_measurement": "currency"},
    )
    total_ask_volume: List[Union[int, None]] = Field(
        default_factory=list,
        description="Total volume of asks.",
    )
    ask_high: List[Union[float, None]] = Field(
        default_factory=list,
        description="The highest ask price.",
        json_schema_extra={"x-unit_measurement": "currency"},
    )
    ask_low: List[Union[float, None]] = Field(
        default_factory=list,
        description="The lowest ask price.",
        json_schema_extra={"x-unit_measurement": "currency"},
    )


class IntrinioOptionsSnapshotsFetcher(
    Fetcher[
        IntrinioOptionsSnapshotsQueryParams,
        List[IntrinioOptionsSnapshotsData],
    ]
):
    """Intrinio Options Snapshots Fetcher."""

    @staticmethod
    def transform_query(params: Dict[str, Any]) -> IntrinioOptionsSnapshotsQueryParams:
        """Transform the query params."""
        transformed_params = params.copy()
        if "date" in transformed_params:
            if isinstance(transformed_params["date"], datetime):
                dt = transformed_params["date"]
                dt = dt.astimezone(tz=timezone("America/New_York"))
            if isinstance(transformed_params["date"], dateType):
                dt = transformed_params["date"]  # type: ignore
                if isinstance(dt, dateType):
                    dt = datetime(
                        dt.year,
                        dt.month,
                        dt.day,
                        16,
                        15,
                        0,
                        0,
                        tzinfo=timezone("America/New_York"),
                    )
            if isinstance(transformed_params["date"], str):
                dt = datetime.fromisoformat(transformed_params["date"])
            else:
                try:
                    dt = datetime.fromisoformat(str(transformed_params["date"]))  # type: ignore
                except ValueError as exc:
                    raise OpenBBError(
                        "Invalid date format. Please use '2024-03-08T12:15-0400'."
                    ) from exc

            transformed_params["date"] = (
                dt.strftime("%Y-%m-%dT%H:%M:%S.%f%z")
                .replace("+", "-")
                .replace("T00:", "T20:")
                if isinstance(dt, datetime)
                else dt
            )
        return IntrinioOptionsSnapshotsQueryParams(**transformed_params)

    @staticmethod
    async def aextract_data(
        query: IntrinioOptionsSnapshotsQueryParams,
        credentials: Optional[Dict[str, str]],
        **kwargs: Any,
    ) -> DataFrame:
        """Return the raw data from the Intrinio endpoint."""
        api_key = credentials.get("intrinio_api_key") if credentials else ""

        # This gets the URL to the actual file.
        url = f"https://api-v2.intrinio.com/options/snapshots?api_key={api_key}"
        if query.date:
            url += f"&at_datetime={query.date}"

        try:
            response = await amake_request(url, **kwargs)
        except Exception as exc:
            raise OpenBBError("Could not fetch data from Intrinio.") from exc

        if isinstance(response, dict) and "error" in response:
            raise OpenBBError(
                f"{response.get('error')}. Message: {response.get('message')}"
            )
        urls = []
        # Get the URL to the CSV file.
        if response.get("snapshots"):  # type: ignore
            for d in response["snapshots"]:  # type: ignore
                if d.get("files"):
                    for f in d["files"]:
                        if f.get("url"):
                            urls.append(f.get("url"))
        if not urls:
            raise OpenBBError("No snapshots found.")

        async def response_callback(response, _):
            """Response Callback."""
            return await response.read()

        async def get_csv(url) -> DataFrame:
            """Return the CSV data."""
            try:
                response = await amake_request(
                    url, response_callback=response_callback, **kwargs
                )
                df = DataFrame()
                if isinstance(response, bytes):
                    file = gzip.decompress(response)
                    df = read_csv(BytesIO(file))

                return df

            except Exception as exc:
                try:
                    df = read_csv(response)
                    return df
                except Exception:
                    raise OpenBBError("Could not read file from URL.") from exc

        # There should only be one URL with this bulk data.
        return await get_csv(urls[0])

    @staticmethod
    def transform_data(
        query: IntrinioOptionsSnapshotsQueryParams,
        data: DataFrame,
        **kwargs: Any,
    ) -> List[IntrinioOptionsSnapshotsData]:
        """Return the transformed data."""
        df = data
        if df.empty:
            raise OpenBBError("Empty CSV file")
        COL_MAP = {
            "CONTRACT ID": "contract_symbol",
            "OPEN INTEREST": "open_interest",
            "TRADE PRICE": "last_price",
            "TRADE SIZE": "last_size",
            "TOTAL TRADE VOLUME": "volume",
            "LAST TRADE TIMESTAMP": "last_timestamp",
            "TRADE HIGH PRICE": "high",
            "TRADE LOW PRICE": "low",
            "ASK PRICE": "ask",
            "ASK SIZE&qu