# -*- coding: utf-8 -*-
"""
Plot the split/dividend-adjusted daily closing price for one PERMNO
from the CRSP daily USStocks files.

Usage:
    python plot_price.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
import matplotlib.dates as mdates

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', 'COMNAM', 'PERMCO',
                'TICKER', 'PRIMEXCH', 'SHRCLS', 'SHROUT',
                'OPENPRC', 'BIDLO', 'ASKHI', 'PRC', 'VOL',
                'CFACPR', 'CFACSHR', 'RETX', 'RET',
                'BID', 'ASK', 'vwretx'],
                converters={'RETX': str, 'RET': str})
    stock.drop_duplicates(subset='date', keep='first', ignore_index=True,
                           inplace=True)
    return stock


def clean_stock(stock):
    # RETX/RET are read as text because CRSP marks a handful of days (e.g. the
    # first trade after a listing gap) with a one-letter code such as 'C'
    # instead of a number. Those days are left alone; a blank RETX means a
    # no-trade/halt day, which we treat as flat price and zero volume.
    n = len(stock)
    for i in range(n):
        if stock.RETX[i].isalpha():
            continue
        if stock.RETX[i] == '':
            stock.at[i, 'RETX'] = '0'
            stock.at[i, 'RET'] = '0'
            stock.at[i, 'VOL'] = 0
            stock.at[i, 'BIDLO'] = stock.BIDLO[i - 1]
            stock.at[i, 'ASKHI'] = stock.ASKHI[i - 1]
            stock.at[i, 'BID'] = stock.BID[i - 1]
            stock.at[i, 'ASK'] = stock.ASK[i - 1]
            stock.at[i, 'PRC'] = stock.PRC[i - 1]
            stock.at[i, 'OPENPRC'] = stock.OPENPRC[i - 1]
    return stock


def plot_gradient_series(xdata, ydata, ylabel, legend_label, figfilepath,
                          cmap='Oranges'):
    xlims = mdates.date2num([xdata[0], xdata[-1]])

    xv, zv = np.meshgrid(np.linspace(0, 5000, 100), np.linspace(0, 5000, 100))

    nmax = len(str(int(ydata.max())))
    factor = 10 ** (nmax - 1)
    ymax = np.ceil(np.ceil(ydata.max()) / factor) * factor

    fig, ax = plt.subplots(figsize=(12, 8))
    ax.imshow(zv, cmap=cmap, origin='lower',
              extent=[xlims[0], xlims[1], 0, ymax], aspect='auto', alpha=0.7)
    ax.set_ylim(0, 1.05 * ydata.max())
    ax.fill_between(xdata, ydata, ymax, color='white')
    ax.plot(xdata, ydata, 'r-', linewidth=1)

    plt.ylabel(ylabel)
    plt.grid()
    plt.legend([legend_label], loc="upper left")
    fig.savefig(figfilepath, dpi=300)
    plt.show()


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

    # PRC is negative on days CRSP substitutes a bid/ask average for a
    # missing close, so take the absolute value before adjusting for splits.
    stock['adj_price'] = stock.PRC.abs() / stock.CFACPR

    xdata = np.array(pd.to_datetime(stock.date, format='%Y%m%d'))
    legend_label = f'{permno} {stock.TICKER.iloc[-1]}'
    outfile = f'{permno}_price.png'

    plot_gradient_series(xdata, stock.adj_price.to_numpy(), '$',
                          legend_label, outfile)
    print(f'Wrote {outfile}')