# -*- coding: utf-8 -*-
"""
Plot split-adjusted daily trading volume for one PERMNO from the CRSP
daily USStocks files.

Usage:
    python plot_volume.py [PERMNO]

If no PERMNO is given, the value set below is used.
"""
import sys
import numpy as np
import pandas as pd
import matplotlib.pyplot as plt

datadir = 'C:/CRSP_daily/USStocks/'
permno = int(sys.argv[1]) if len(sys.argv) > 1 else 10866


def load_stock(permno):
    filename = datadir + str(permno) + '.csv'
    stock = pd.read_csv(filename,
                usecols=['date', 'TICKER', 'VOL', 'CFACSHR', 'RETX'],
                converters={'RETX': str})
    stock.drop_duplicates(subset='date', keep='first', ignore_index=True,
                           inplace=True)
    stock.VOL = stock.VOL.fillna(0)
    return stock


if __name__ == '__main__':
    stock = load_stock(permno)

    # VOL is in raw (pre-split) shares traded; multiplying by CFACSHR restates
    # every day's volume in today's share-count terms so a split does not
    # show up as an artificial jump in trading activity.
    stock['adj_vol'] = stock.VOL * stock.CFACSHR

    xdata = np.array(pd.to_datetime(stock.date, format='%Y%m%d'))

    fig, ax = plt.subplots(figsize=(12, 5))
    ax.bar(xdata, stock.adj_vol / 1000, width=1.5, color='steelblue')
    plt.ylabel('Adjusted volume (thousands of shares)')
    plt.grid(axis='y')

    legend_label = f'{permno} {stock.TICKER.iloc[-1]}'
    plt.legend([legend_label], loc="upper left")

    outfile = f'{permno}_volume.png'
    fig.savefig(outfile, dpi=300)
    plt.show()
    print(f'Wrote {outfile}')
