%matplotlib notebook
import numpy as np
import pandas as pd
import tushare as ts
import matplotlib.pyplot as plt
import matplotlib.dates as mdates
import matplotlib.ticker as mtick


def filter_dr(dr, dtype):
    """
    从利率的 dataframe 中获取指定类型的利率
    @dr: 利率数据
    @dtype: 利率类型，例如：'定期存款整存整取(一年)'
    """

    # 过滤出指定类型的利率数据
    result = dr.loc[ dr['deposit_type'] == dtype ]
    # 过滤无效数据
    result = result.loc[ result['rate'] != '--' ]
    # 丢弃类型数据,只保留时间和利率值
    result = result.drop(['deposit_type'], axis=1)
    
    # 时间格式
    result.date = pd.to_datetime(result.date)
    result.rate = pd.to_numeric(result.rate, errors='coerce')
    
    # result.sort_values('date', inplace=True)
    # result.set_index(['date'], inplace = True) 
    return result


def plot_dr(dr, title):
    """
    绘制利率图
    @dr : 利率数据
    @title : 图的标题
    """

    # drtypes = pd.unique(dr['deposit_type'])
    drtypes = ['定期存款整存整取(一年)', '定期存款整存整取(三年)', '定期存款整存整取(五年)']

    fig = plt.figure(figsize=(10,6))
    ax = fig.add_subplot(1,1,1)
    
    for drtype in drtypes:
        result = filter_dr(dr, drtype)
        ax.plot(result['date'], result['rate'], label=drtype)

    fmt='%.2f%%'
    ax.yaxis.set_major_formatter(mtick.FormatStrFormatter(fmt))
    ax.xaxis.set_major_formatter(mdates.DateFormatter('%Y-%m-%d'))
    ymin, ymax = ax.get_ylim()
    ax.set_yticks(np.round(np.linspace(ymin, ymax, 20), 2))
    plt.xticks(rotation=60)
    plt.xlabel('日期')
    plt.ylabel('利率(%)')
    plt.title(title)
    plt.subplots_adjust(bottom=0.2, right=0.9, left=0.1)
    plt.legend()
    plt.show()

# 使用 tushare 获取利率数据
deposit_rate = ts.get_deposit_rate()

# 画利率走势图
plot_dr(deposit_rate, '利率走势图')

# 查看一年期利率数据
filter_dr(deposit_rate, '定期存款整存整取(一年)')