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

BID/ASK coverage can have gaps of very different character: an occasional
one- or two-day hole (a quote CRSP could not capture) versus a genuine
multi-year stretch where the field was not collected at all (bid/ask
history for many older NYSE issues does not begin until decades after
their price history does). Bridging the first kind by carrying the
previous quote forward is reasonable; doing that across the second kind
would silently fabricate decades of fake, flat quotes. This script finds
the longest stretch of BID/ASK data that contains only short (fillable)
holes and plots only that stretch.

Usage:
    python plot_bidask.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 10145
MAX_FILLABLE_GAP = 2  # longest run of missing BID/ASK we are willing to bridge


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


def longest_fillable_segment(good, max_gap):
    """Return the (start, end) index range, inclusive, of the longest run of
    quotes in which every hole is at most max_gap days long. Returns None if
    there is no usable data at all."""
    n = len(good)
    segments = []
    seg_start = None
    i = 0
    while i < n:
        if good[i]:
            if seg_start is None:
                seg_start = i
            i += 1
            continue
        j = i
        while j < n and not good[j]:
            j += 1
        gap_len = j - i
        if seg_start is not None and gap_len <= max_gap and j < n:
            i = j  # bridge a short internal gap, segment continues
        else:
            if seg_start is not None:
                segments.append((seg_start, i - 1))
            seg_start = None
            i = j
    if seg_start is not None:
        segments.append((seg_start, n - 1))
    if not segments:
        return None
    return max(segments, key=lambda seg: seg[1] - seg[0])


def fill_short_gaps(stock, start, end):
    for i in range(start + 1, end + 1):
        if np.isnan(stock.BID[i]) or np.isnan(stock.ASK[i]):
            retx = stock.RETX[i]
            is_number = retx.replace('.', '', 1).replace('-', '', 1).isdigit()
            payoff = 1 + abs(float(retx)) if is_number else 1.0
            stock.at[i, 'BID'] = stock.BID[i - 1] * payoff
            stock.at[i, 'ASK'] = stock.ASK[i - 1] * payoff


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)
    good = (stock.BID.notna() & stock.ASK.notna()).to_numpy()

    segment = longest_fillable_segment(good, MAX_FILLABLE_GAP)
    if segment is None:
        print(f'PERMNO {permno} has no BID/ASK coverage in this file; nothing to plot.')
        sys.exit(0)

    start, end = segment
    fill_short_gaps(stock, start, end)

    stock = stock.iloc[start:end + 1].reset_index(drop=True)
    stock['midquote'] = 0.5 * (stock.BID + stock.ASK) / 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}_midquote.png'

    plot_gradient_series(xdata, stock.midquote.to_numpy(), '$',
                          legend_label, outfile)
    print(f'Wrote {outfile}')
    print(f'Used {end - start + 1} rows, {stock.date.iloc[0]} to {stock.date.iloc[-1]} '
          f'(longest stretch with no BID/ASK hole longer than {MAX_FILLABLE_GAP} days)')