#import pandas as pd
#import mpl_finance

class Stock(object):
    pass

def get_data(
code='000300.XSHG',
start_date='2015-11-01',
end_date='2015-11-15',
subset = None, # ('2015-11-10', '2015-11-11')
freqs = '1d 1m 5m 30m 60m'.split(),
):
    
    '''
    >>> stk = get_data()
    '''
    stk = Stock()
    for freq in freqs:
        history = get_price(code, start_date=start_date, end_date=end_date, frequency=freq,fields= ['open',  'high', 'low','close', 'volume'])
        stk.__setattr__('ohlc_'+freq, history)

    # 修改日线的日期索引到: @ 15:00:00
    midnight = stk.ohlc_1d.index
    noon = stk.ohlc_1d.index + datetime.timedelta(hours=12)
    deltat = datetime.timedelta(hours=15)
    stk.ohlc_1d.index +=deltat

    df = get_all_securities(['index',])
    stk.dname = df.loc[code]['display_name']
    stk.name = df.loc[code]['name']
    #print(name,dname)
    
    return stk

def plot_ohlc_multiple_timeframe(stk, subset=None):
    '''
    >>> plot_ohlc_multiple_timeframe(stk, subset=('2015-11-10','2015-11-11'))
    >>> plot_ohlc_multiple_timeframe(stk, None)
    '''
    #subset=None
    fig,ax=plt.subplots()
    for attr in map(lambda x: 'ohlc_'+x, freqs[:2]):
        if subset is not None: print('subset={}, {}'.format(*subset))
        if subset is not None:
            ohlc = stk.__getattribute__(attr).loc[subset[0]:subset[1],]
            if attr=='ohlc_1d':
                midnight = ohlc.index + datetime.timedelta(hours=-15)
                noon = ohlc.index + datetime.timedelta(hours=-3)
        else:
            ohlc = stk.__getattribute__(attr)        
            if attr=='ohlc_1d':
                midnight = ohlc.index + datetime.timedelta(hours=-15)
                noon = ohlc.index + datetime.timedelta(hours=-3)

        x = mpl.dates.date2num(ohlc.index)
        ax.plot(x, ohlc.close, label=attr[5:])
        #print(stk.__getattribute__(attr).close[:2])


    x = mpl.dates.date2num(midnight)
    y1,y2= ax.get_ylim() 
    ax.vlines(x,y1,y2, color='k', lw=0.1, label='子夜')
    x = mpl.dates.date2num(noon)
    ax.vlines(x,y1,y2, color='c', lw=0.3, label='中午')



    #'xtls = ax.get_xmajorticklabels()
    #_new_t = mpl.dates.num2date(t._x) #.strftime('%Y\n%m%d')
    xticks = ax.get_xticks()
    if subset is not None:
        new_t= [mpl.dates.num2date(xtick).strftime('%Y\n%m%d\n%H:%M')  for xtick in xticks]
    else:
        new_t= [mpl.dates.num2date(xtick).strftime('%Y\n%m%d')  for xtick in xticks]

    ax.set_xticklabels(new_t)


    ax.legend()
    title = '日线分钟线欢聚一堂\n'
    title += '证券名称: {}, 时间段: {}, {}'.format(stk.dname, ohlc.index[0], ohlc.index[-1])
    fig.suptitle(title, fontsize=18)
    fig.set_size_inches(18,8)

if __name__=='__main__':
    stk = get_data()
    plot_ohlc_multiple_timeframe(stk, None)
    plot_ohlc_multiple_timeframe(stk, subset=('2015-11-10','2015-11-11'))