option volume chart

theme_colors = {
    'darkest_black': '#0d1117',
    'darker_black': '#161b22',
    'dark_black': '#21262d',
    'darkest_gray': '#89929b',
    'darker_gray': '#c6cdd5',
    'dark_gray': '#ecf2f8',
    'white': '#ffffff',
    'red': '#fa7970',
    'orange': '#faa356',
    'green': '#7ce38b',
    'sky_blue': '#a2d2fb',
    'blue': '#77bdfb',
    'purple': '#cea5fb'
}

def theme_light(legend_position: str = 'bottom') -> ggplot:
    darkest_black = theme_colors['darkest_black']
    darker_black = theme_colors['darker_black']
    dark_black = theme_colors['dark_black']
    white = theme_colors['white']
    return theme_none() + theme(
        line=element_line(color=darkest_black, size=1),
        rect=element_rect(color=darkest_black, fill=white, size=1),
        text=element_text(color=darkest_black),
        axis_ontop=True,
        axis_ticks=element_line(color=darker_black),
        legend_background=element_rect(size=1),
        legend_position=legend_position,
        panel_grid_major=element_line(color=dark_black, size=0.5, linetype='dashed'),
        panel_grid_minor=element_blank(),
        plot_title=element_text(hjust=0.5, color=darkest_black),
        tooltip=element_rect(color=dark_black),
        axis_tooltip=element_rect(color=darker_black)
    )

def create_uw_option_df(ticker: str,
                        dates_to_exclude: list[str] = [],
                        dates_to_include: list[str] = []) -> pl.DataFrame:
    """
    API documentation:
    https://api.unusualwhales.com/docs#/operations/PublicApi.OptionContractController.option_contracts

    Function:
    Pass in a ticker and optional expiration dates to exclude or include
    then return a Polars DataFrame containing option volume and transaction
    data, by strike, for the entire montage.
    """
    # Collect json data from the Unusual Whales API
    base_url = ('https://api.unusualwhales.com/api/stock/'
                '{}/option-contracts')
    full_url = base_url.format(ticker)
    r = requests.get(full_url, headers=headers)

    # Create DataFrame then format columns andcalculate strike price
    # before returning to DataFrame the caller
    df = pl.DataFrame(r.json()['data'])
    df = (
        df
        .with_columns(
            [
                pl.col('high_price').cast(pl.Float64),
                pl.col('implied_volatility').cast(pl.Float64),
                pl.col('last_price').cast(pl.Float64),
                pl.col('low_price').cast(pl.Float64),
                pl.col('nbbo_ask').cast(pl.Float64),
                pl.col('nbbo_bid').cast(pl.Float64),
                pl.col('total_premium').cast(pl.Float64)
            ]
        )
        .with_columns(
            [
                (
                    pl.col('option_symbol').str.slice(-9, 1).alias('option_type')
                ),
                (
                    pl.col('option_symbol').str.slice(-8).cast(pl.Float64) / 1000
                ).alias('strike_price'),
            ]
        )
        .with_columns(
            pl.col('strike_price').cast(pl.Utf8).alias('strike_price_str')
        )
        .with_columns(
            pl.concat_str(
                [pl.col('strike_price').cast(pl.Utf8), pl.col('option_type')], separator=''
            ).alias('strike')
        )
        .with_columns(
            pl.concat_str(
                [
                    (pl.lit('20') + pl.col('option_symbol').str.slice(len(ticker), 2)),
                    pl.col('option_symbol').str.slice(len(ticker) + 2, 2),
                    pl.col('option_symbol').str.slice(len(ticker) + 4, 2)
                ],
                separator='-'
            ).alias('expiration_date_str')
        )
    )

    # Filter out any dates that are in the dates_to_ignore list
    # (if not empty) and return the final results
    if dates_to_exclude:
        df = df.filter(~pl.col('expiration_date_str').is_in(dates_to_exclude))
    elif dates_to_include:
        df = df.filter(pl.col('expiration_date_str').is_in(dates_to_include))
    return df

def aggregate_uw_option_df_by_strike_price(df: pl.DataFrame) -> pl.DataFrame:
    """
    Aggregate volume and open interest by strike price across
    the option montage included in the passed-in DataFrame.
    """
    new_df = (
        df
        .group_by('strike')
        .agg(
            [
                pl.sum('ask_volume').alias('total_ask_volume'),
                pl.sum('bid_volume').alias('total_bid_volume'),
                pl.sum('mid_volume').alias('total_mid_volume'),
                pl.sum('no_side_volume').alias('total_no_side_volume'),
                pl.sum('open_interest').alias('total_open_interest'),
                pl.sum('volume').alias('total_volume'),
                pl.min('strike_price'),
                pl.min('strike_price_str'),
                pl.min('option_type')
            ]
        )
    )
    return new_df.sort('strike_price', descending=False)

def configure_agg_df_for_plotting(df: pl.DataFrame) -> pl.DataFrame:
    """
    Create a "vertical" DataFrame that plots much more easily
    with Lets-Plot than the "horizontal" condensed DataFrame
    """
    # Create an ask-specific DataFrame
    ask_side_columns = [
        'strike',
        'strike_price',
        'strike_price_str',
        'option_type',
        'total_ask_volume',
    ]
    ask_side_df = df.select(ask_side_columns)
    ask_side_df = ask_side_df.rename({'total_ask_volume': 'volume'})
    ask_side_df = (
        ask_side_df
        .with_columns(
            pl.lit('ask').alias('side')
        )
        .with_columns(
            pl.when(pl.col('option_type') == 'P')
            .then(pl.col('volume') * -1)
            .otherwise(pl.col('volume'))
        )
    )

    # Create a bid-specific DataFrame
    bid_side_columns = [
        'strike',
        'strike_price',
        'strike_price_str',
        'option_type',
        'total_bid_volume',
    ]
    bid_side_df = df.select(bid_side_columns)
    bid_side_df = bid_side_df.rename({'total_bid_volume': 'volume'})
    bid_side_df = (
        bid_side_df
        .with_columns(
            pl.lit('bid').alias('side')
        )
        .with_columns(
            pl.when(pl.col('option_type') == 'P')
            .then(pl.col('volume') * -1)
            .otherwise(pl.col('volume'))
        )
    )

    # Create a mid-specific DataFrame
    mid_side_columns = [
        'strike',
        'strike_price',
        'strike_price_str',
        'option_type',
        'total_mid_volume',
    ]
    mid_side_df = df.select(mid_side_columns)
    mid_side_df = mid_side_df.rename({'total_mid_volume': 'volume'})
    mid_side_df = (
        mid_side_df
        .with_columns(
            pl.lit('mid').alias('side')
        )
        .with_columns(
            pl.when(pl.col('option_type') == 'P')
            .then(pl.col('volume') * -1)
            .otherwise(pl.col('volume'))
        )
    )

    # Create a no-side-specific DataFrame
    no_side_columns = [
        'strike',
        'strike_price',
        'strike_price_str',
        'option_type',
        'total_no_side_volume',
    ]
    no_side_df = df.select(no_side_columns)
    no_side_df = no_side_df.rename({'total_no_side_volume': 'volume'})
    no_side_df = (
        no_side_df
        .with_columns(
            pl.lit('no_side').alias('side')
        )
        .with_columns(
            pl.when(pl.col('option_type') == 'P')
            .then(pl.col('volume') * -1)
            .otherwise(pl.col('volume'))
        )
    )
    return (
        ask_side_df
        .vstack(bid_side_df)
        .vstack(mid_side_df)
        .vstack(no_side_df)
    )


def plot_option_volume_by_strike_price(ticker: str, df: pl.DataFrame) -> ggplot:
    """
    Plot the volume of option contracts by strike price
    """
    tooltip_content = layer_tooltips().line('@volume')
    color_values = [
        theme_colors['green'],
        theme_colors['red'],
        theme_colors['blue'],
        theme_colors['darker_gray']
    ]
    return ggplot(df, aes(x='volume', y='strike_price_str')) + \
           geom_bar(aes(fill='side'), stat='identity', color='black', size=0.5, tooltips=tooltip_content) + \
           ggsize(800, 700) + \
           ggtitle(f'{ticker} Option Volume') + \
           theme_light() + \
           scale_fill_manual(values=color_values)