# -*- coding: utf-8 -*-
"""
Plot daily returns and cumulative growth of $1 for one PERMNO from the
CRSP daily USStocks files, benchmarked against the CRSP value-weighted
market return (vwretd).

Usage:
    python plot_returns.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', 'TICKER', 'RET', 'vwretd'],
                converters={'RET': str})
    stock.drop_duplicates(subset='date', keep='first', ignore_index=True,
                           inplace=True)
    return stock


def to_numeric_return(ret_series):
    # A handful of days carry a one-letter CRSP code (e.g. 'C' for "no prior
    # price to compute a return from") instead of a number. We can't compute
    # a return on those days, so they are treated as 0 for the cumulative
    # growth chart and reported separately.
    numeric = pd.to_numeric(ret_series, errors='coerce')
    n_missing = numeric.isna().sum()
    return numeric.fillna(0.0), n_missing


if __name__ == '__main__':
    stock = load_stock(permno)
    stock['ret'], n_missing = to_numeric_return(stock.RET)
    stock['vwretd'] = stock.vwretd.fillna(0.0)

    xdata = np.array(pd.to_datetime(stock.date, format='%Y%m%d'))
    cum_stock = (1 + stock.ret).cumprod()
    cum_mkt = (1 + stock.vwretd).cumprod()

    fig, (ax1, ax2) = plt.subplots(2, 1, figsize=(12, 8), sharex=True)

    ax1.plot(xdata, stock.ret * 100, color='steelblue', linewidth=0.6)
    ax1.set_ylabel('Daily return (%)')
    ax1.grid()

    ax2.plot(xdata, cum_stock, 'r-', linewidth=1.2,
             label=f'{permno} {stock.TICKER.iloc[-1]}')
    ax2.plot(xdata, cum_mkt, 'k--', linewidth=1, label='CRSP value-weighted market (vwretd)')
    ax2.set_ylabel('Growth of $1')
    ax2.xaxis.set_major_locator(mdates.AutoDateLocator())
    ax2.grid()
    ax2.legend(loc='upper left')

    fig.tight_layout()
    outfile = f'{permno}_returns.png'
    fig.savefig(outfile, dpi=300)
    plt.show()
    print(f'Wrote {outfile}')
    print(f'{n_missing} day(s) had a non-numeric RET code and were treated as 0 for cumulative growth')