# 导入常用的库
import numpy as np
import pandas as pd 
import datetime as dt
import time
import matplotlib.pyplot as plt
import seaborn as sns
from jqdata import *
import tushare as ts
import seaborn as sns
plt.style.use('fivethirtyeight')

plt.style.use('fivethirtyeight')
# 由于数据量比较多，这里打算从 tushare 获取上证指数所有的价格数据
# tushare 接口，参数为注册时生成的 token
pro = ts.pro_api('xxxxxxxxxxxxxxxx')

# tushare 要求一次最多获取 3000 条数据，所以分两次获取
# 然后将数据合并，按时间排序
df1 = pro.index_daily(ts_code='000001.SH', 
                      start_date='19901219', 
                      end_date='20101231')
df2 = pro.index_daily(ts_code='000001.SH', 
                      start_date='20110101', 
                      end_date='20190630')
df = pd.concat([df2, df1])  # 合并
df = df.sort_values(by=['trade_date'])  # 排序
df['trade_date'] = pd.to_datetime(df['trade_date'])  # 转换为时间类型
df.set_index(['trade_date'], inplace=True)  # 设置索引列
df.index.name = None  # 去掉索引列名

# 将上证指数价格曲线画出来
# 并在对应的牛市年份，画一条竖线来标记
df.close.plot(figsize=(14, 7), title='牛顶间隔展示')
for year in ['1993-02-01', '2001-06-01', 
             '2007-10-01', '2015-06-01']:
    plt.axvline(year,color='r', alpha=0.7)
plt.show()

# 转换成时间格式
best_years = [dt.datetime.strptime(year, '%Y-%m-%d').date() 
              for year in ['1993-02-01', '2001-06-01', 
                           '2007-10-01', '2015-06-01']]

# 计算牛顶时间间隔
gap_year = [days.days for days in np.diff(best_years)]
print('牛顶平均间隔时间：', ['{} days'.format(days) for days in gap_year])

# 计算平均间隔年数，为避免盲目猜测，再计算出均值的偏差值
mean_gap = np.mean(gap_year)
print('平均时间间隔 {} 天，即 {} 年'.format(round(mean_gap, 2), round(mean_gap / 365, 2)))

# 计算平均时间间隔偏差值
std_gap = np.std(gap_year)
print('平均时间间隔偏差 {} 天，即 {} 年'.format(round(std_gap, 2), round(std_gap / 365, 2)))

# 计算下一个牛市出现的合理时间区间
early_year = (best_years[-1] + dt.timedelta(round(mean_gap - std_gap, 0)))
latest_year = (best_years[-1] + dt.timedelta(round(mean_gap + std_gap, 0)))
# 构造 title
early_y = early_year.year
early_m = early_year.month
latest_y = latest_year.year
latest_m = latest_year.month
print('下个牛顶时间范围 {}年{}月 ~ {}年{}月'.format(early_y, early_m, latest_y, latest_m))

title = '评估下个牛顶时间范围 {}年{}月 ~ {}年{}月'.format(early_y, early_m, latest_y, latest_m)
df.close.plot(figsize=(14, 7), title=title)
for year in ['1993-02-01', '2001-06-01', 
             '2007-10-01', '2015-06-01']:
    plt.axvline(year,color='r', alpha=0.6)
    
# 生成 2019-02-11 ~ 2023-12-10 的时间区间
# 如果要填充一个区间，y 就给价格的最大值便好
date_span = pd.date_range(early_year, latest_year)
value_span = [df.close.max() for x in date_span]
plt.fill_between(date_span, value_span, color='orange', alpha=0.8)
plt.show()

# 通过线性回归，刻画整体平均趋势
plt.figure(figsize=(14, 7))
sns.regplot(x=np.arange(0, df.shape[0]), y=df.close.values)
plt.show()

# 为了计算方便，默认都倩收盘价格为准
# 获取上市第一天的收盘价和当前收盘价
start_price = df.close[0]
end_price = df.close[-1]
print('最初收盘价为 {}，当前收盘价为 {}'.format(start_price, end_price))

# 计算收益率
total_return = (end_price / start_price) - 1
print('持有到目前为止的收益率为 {}%'.format(round(total_return * 100, 2)))

# 计算年平均收益，一年以250个交易日为准
# 平均年化收益率=(投资内收益/本金)×（250/投资天数）× 100%
mean_return = total_return * (250 / df.shape[0])
print('平均每年收益率为 {}%'.format(round(mean_return * 100, 2)))

# 计算年化收益率复利
# 总收益 = 本金 * （年化利率 + 1）的 n 次方，n为交易年数
# 年化利率 = （总收益 / 本金）的 n次开根 - 1
annualized_return = pow(((end_price - start_price) / start_price), 
                        1 / (df.shape[0] / 250)) - 1
print('年化收益率为 {}%'.format(round(annualized_return * 100, 2)))

# 获取每个月的第一个交易日的数据
first_day = []
for i in range(len(df)):
    date = df.index[i]
    if i == 0:
        first_day.append(date)
    else:
        last_date = df.index[i - 1]
        if date.day < last_date.day:
            first_day.append(date)
            
index_df = df.loc[first_day]
index_df.index

# 按照模型进行定投
month_df = index_df.copy()
month_df['pct_change'] = month_df['close'].pct_change()
month_df = month_df[['close', 'pct_change']]  # 按月整合数据

save_money = []
hold_money = []
save_base = 1000
for i in range(len(month_df)):
    if i == 0:
        save_money.append(save_base)
        hold_money.append(save_base)
    else:
        save_money.append(save_money[-1] + save_base)
        hold_money.append(hold_money[-1] * (1 + month_df['pct_change'][i]) + save_base)

month_df['save_money'] = save_money
month_df['hold_money'] = hold_money
month_df.head(10)

# 计算定投、收益曲线
month_df['return_money'] = month_df['hold_money'] - month_df['save_money']
month_df[['save_money', 'hold_money', 'return_money']].plot(figsize=(14, 7))
plt.legend(['累积投入', '累积本息', '累积收入'])
plt.show()

print('累计投入: {}元'.format(month_df['save_money'][-1]))
print('累计收益: {}元'.format(month_df['return_money'][-1]))
print('最终本息累积: {}元'.format(month_df['hold_money'][-1]))
print('绝对收益率为: {}%'.format((month_df['return_money'][-1] / month_df['save_money'][-1]) * 100))

# 从聚宽获取上证指数的信息
index = '000001.XSHG'  # 指数 code
index_info = get_security_info(index)  # 指数信息
start_date = index_info.start_date  # 指数开始时间
end_date = datetime.datetime.now().date()  # 以当天为最后一天
index_name = index_info.display_name  # 指数全称


# 定义一个函数，计算每天的成份股的平均pe/pb
def get_pe_pb(index_code, start_date, end_date):
    def iter_pe_pb():
        # 一个获取PE/PB的生成器
        trade_date = get_trade_days(start_date=start_date, end_date=end_date)   
        for date in trade_date:
            stocks = get_index_stocks(index_code, date)
            q = query(valuation.pe_ratio, 
                      valuation.pb_ratio
                     ).filter(valuation.pe_ratio != None,
                              valuation.pb_ratio != None,
                              valuation.code.in_(stocks))
            df = get_fundamentals(q, date)
            
            # 通过分位值进行过滤异常值
            # 这里并没有采用三倍标准差来去除极值，差异不大
            quantile = df.quantile([0.25, 0.75])
            df_pe = df.pe_ratio[(df.pe_ratio > quantile.pe_ratio.values[0]) &\
                                (df.pe_ratio < quantile.pe_ratio.values[1])]
            df_pb = df.pb_ratio[(df.pb_ratio > quantile.pb_ratio.values[0]) &\
                                (df.pb_ratio < quantile.pb_ratio.values[1])]
            yield date, df_pe.median(), df_pb.median()

    dict_result = [{'date': value[0], 'pe': value[1], 'pb':value[2]} for value in iter_pe_pb()]
    df_result = pd.DataFrame(dict_result)
    df_result.set_index('date', inplace=True)
    return df_result

df_pe_pb = get_pe_pb(index, start_date, end_date)
df_pe_pb.head(10)

# 可视化PE/PB曲线图
df_pe_pb.plot(figsize=(14, 7), subplots=True)
plt.show()

# 将PE/PB趋势与指数趋势一起展示，以作观察
_, axs = plt.subplots(ncols=2, figsize=(14, 5))
close = get_price(index, start_date=start_date, end_date=end_date).close
_df = pd.DataFrame()
_df['close'] = close
_df['pe'] = df_pe_pb.pe
_df['pb'] = df_pe_pb.pb
_df[['close', 'pe']].plot(secondary_y=['pe'], ax=axs[0], alpha=.8)
_df[['close', 'pb']].plot(secondary_y=['pb'], ax=axs[1], alpha=.8)
plt.show()

# 分析PE/PB数据分布情况
_, axs = plt.subplots(nrows=2, ncols=2, figsize=(14, 7))
sns.distplot(df_pe_pb.pe, ax=axs[0][0])
sns.boxplot(df_pe_pb.pe, ax=axs[0][1])
sns.distplot(df_pe_pb.pb, ax=axs[1][0])
sns.boxplot(df_pe_pb.pb, ax=axs[1][1])
plt.show()

# 观察PE/PB之间的关系
sns.jointplot(x='pb',y='pe', data=df_pe_pb, height=7)
plt.show()

# 将PE分成十个分位，查看各分位PE数量
pe_array = df_pe_pb.pe.values
value_counts = pd.cut(pe_array, 10).value_counts()
print(value_counts)

plt.figure(figsize=(14, 4))
sns.barplot(x=np.arange(0, len(value_counts)), 
            y=value_counts.values)
plt.show()

# 刻画PE整体趋势的中等分位区间（40%~60%）
def show_quantile():
    _df = pd.DataFrame()
    df = df_pe_pb.copy()
    df.index.name = None

    _df['pe'] = df.pe
    _df = _df
    p_high = [_df.pe.quantile(i / 10.0) for i in [4, 5, 6]] 
    for p_h, i in zip(p_high, [4, 5, 6]):
        _df[str(i / 10 * 100)+'%'] = p_h

    low_p = _df[_df.pe < _df.pe.iloc[-1]]
    quantile_now = low_p.shape[0] / _df.shape[0]  # 当前百分位值
    last_p = _df.pe[-1]

    _df.plot(figsize=(14, 7))
show_quantile()

# 计算比值
low = value_counts[0:4].sum()
medin = value_counts[4:6].sum()
high = value_counts[6:10].sum()
print('比值({}：{}：{})'.format(low, medin, high))

# 获取每个月的第一个交易日
first_day = []
for i in range(len(df_pe_pb)):
    date = df_pe_pb.index[i]
    if i == 0:
        first_day.append(date)
    else:
        last_date = df_pe_pb.index[i - 1]
        if date.day < last_date.day:
            first_day.append(date)

# 按月计算价格与涨跌幅度
close = get_price(index, start_date=df_pe_pb.index[0], end_date=df_pe_pb.index[-1])['close']
df = df_pe_pb.copy()
df['close'] = close
df = df.loc[first_day]
df['pct_change'] = df.close.pct_change()
df.head(10)

miden_estimation = (38.492, 49.856)  # 中等估值的pe区间
save_money = []  # 每月定存
back_money = []  # 回收资金
hold_money = []  # 持仓资金
base_money = 1000  # 定投基准

def trade():
    for i in range(len(df)):
        pe = df['pe'][i]  # 估值位

        if i == 0:  # 初始买入
            # 1.计算买入金额
            save_money.append(base_money)
            # 2. 计算回收金额
            back_money.append(0)
            # 3. 计算持仓变化
            hold_money.append(base_money)
            continue

        if pe <= miden_estimation[0]:  # 执行买入计算
            # 1.计算买入金额
            save_money.append(base_money)
            # 2. 计算回收金额
            back_money.append(0)
            # 3. 计算持仓变化
            hold_money.append(hold_money[-1] * (1 + df['pct_change'][i]) + base_money)
        elif pe >= miden_estimation[-1]:  # 执行卖出计算 
            # 1. 计算买入金额
            save_money.append(0)
            # 2. 计算回收金额
            back_money.append(base_money)
            # 3. 计算持仓变化
            hold_money.append(hold_money[-1] * (1 + df['pct_change'][i]) - base_money)
        else:
            # 1.计算买入金额
            save_money.append(0)
            # 2. 计算回收金额
            back_money.append(0)
            # 3. 计算持仓变化
            hold_money.append(hold_money[-1] * (1 + df['pct_change'][i]))
trade()

df['save_money'] = save_money  # 定投金额
df['save_money_cumsum'] = df['save_money'].cumsum()  # 定投累计金额
df['hold_money'] = hold_money  # 持仓金额
df['back_money'] = back_money  # 回收金额
df['back_money_cumsum'] = df['back_money'].cumsum()  # 累计回收金额
df['total_money'] = df['hold_money'] + df['back_money_cumsum']  # 总资金
df['return_money'] = df['total_money'] - df['save_money_cumsum']  # 持续收益
df['return_rate'] = (df['total_money'] / df['save_money_cumsum']) - 1  # 持续收益率
df[['save_money_cumsum', 'total_money', 'back_money_cumsum', 'return_money']].plot(figsize=(14, 7))
plt.legend(['累积定投', '累计本息', '回收资金', '收益曲线'])
plt.show()

print('累计投入: {}元'.format(df['save_money_cumsum'][-1]))
print('累计收益: {}元'.format(df['return_money'][-1]))
print('最终本息累积: {}元'.format(df['total_money'][-1]))
print('绝对收益率为: {}%'.format((df['return_money'][-1] / df['save_money_cumsum'][-1]) * 100))

# 展示各年投入金额
money_year = {}
for date in df.index:
    year = date.year
    if year in money_year.keys():
        money_year[year] = money_year[year] + df.loc[date, 'save_money']
    else:
        money_year[year] = df.loc[date, 'save_money']
        
money_mean = mean(list(money_year.values()))
years_count = len(money_year) - 1
money_year = {key: [value] for key, value in money_year.items()}

df_money_year = pd.DataFrame(money_year, index=[''])
df_money_year = df_money_year.T
df_money_year.plot(figsize=(14, 4), kind='bar')
plt.hlines(money_mean, 0, years_count, color='orange')
plt.legend(['年均投入', '定投年金'])
plt.show()

# 展示各年的收益
return_year = {}
for date in df.index:
    year = date.year
    return_year[year] = df.loc[date, 'return_rate']
return_year = {key: [value] for key, value in return_year.items()}
return_df = pd.DataFrame(return_year, index=['return']).T
return_df['diff'] = return_df['return'].diff()
return_df['diff'].fillna(return_df['return'], inplace=True)
return_df[['diff']].plot(figsize=(14, 4), kind='bar')
plt.legend(['各年收益率'])
plt.show()

# 获取每个月的第一个交易日
first_day = []
for i in range(len(df_pe_pb)):
    date = df_pe_pb.index[i]
    if i == 0:
        first_day.append(date)
    else:
        last_date = df_pe_pb.index[i - 1]
        if date.day < last_date.day:
            first_day.append(date)

# 按月计算价格与涨跌幅度
close = get_price(index, start_date=df_pe_pb.index[0], end_date=df_pe_pb.index[-1])['close']
df = df_pe_pb.copy()
df['close'] = close
df = df.loc[first_day]
df['pct_change'] = df.close.pct_change()
df.head(10)

def get_how_value(pe):
    how_value = [15.708,21.447,27.129,32.811,38.492,44.174,
                49.856,55.538,61.219,66.901,72.583]
    for i, value in zip(range(0, len(how_value)) , how_value):
        # zip 包装了整数倍的分位值与对应的pe值区间
        if how_value[i] <= pe < how_value[i + 1]:
            location = i + 1
            _how_value = 5 - location  # 以5为中等值
            return _how_value  # 返回基于中位的买入或卖出倍数

miden_estimation = (38.492, 49.856)  # 中等估值的pe区间
save_money = []  # 每月定存
back_money = []  # 回收资金
hold_money = []  # 持仓资金
base_money = 1000  # 定投基准

def trade():
    for i in range(len(df)):
        pe = df['pe'][i]  # 估值位
        how_value = get_how_value(pe)
        
        if i == 0:  # 初始买入
            # 1.计算买入金额
            save_money.append(base_money)
            # 2. 计算回收金额
            back_money.append(0)
            # 3. 计算持仓变化
            hold_money.append(base_money)
            continue

        if how_value > 0:  # 执行买入计算
            # 1.计算买入金额
            save_money.append(base_money * how_value)
            # 2. 计算回收金额
            back_money.append(0)
            # 3. 计算持仓变化
            hold_money.append(hold_money[-1] * (1 + df['pct_change'][i]) + base_money * how_value)
        else:  # 执行卖出计算 
            # 1. 计算买入金额
            save_money.append(0)
            # 2. 计算回收金额
            back_money.append(base_money * -how_value * 11)
            # 3. 计算持仓变化
            hold_money.append(hold_money[-1] * (1 + df['pct_change'][i]) - base_money * -how_value * 11)
trade()

df['save_money'] = save_money  # 定投金额
df['save_money_cumsum'] = df['save_money'].cumsum()  # 定投累计金额
df['hold_money'] = hold_money  # 持仓金额
df['back_money'] = back_money  # 回收金额
df['back_money_cumsum'] = df['back_money'].cumsum()  # 累计回收金额
df['total_money'] = df['hold_money'] + df['back_money_cumsum']  # 总资金
df['return_money'] = df['total_money'] - df['save_money_cumsum']  # 持续收益
df['return_rate'] = (df['total_money'] / df['save_money_cumsum']) - 1  # 持续收益率
df[['save_money_cumsum', 'total_money', 'back_money_cumsum', 'return_money']].plot(figsize=(14, 7))
plt.legend(['累积定投', '累计本息', '回收资金', '收益曲线'])
plt.show()

print('累计投入: {}元'.format(df['save_money_cumsum'][-1]))
print('累计收益: {}元'.format(df['return_money'][-1]))
print('最终本息累积: {}元'.format(df['total_money'][-1]))
print('绝对收益率为: {}%'.format((df['return_money'][-1] / df['save_money_cumsum'][-1]) * 100))

# 展示各年投入金额
money_year = {}
for date in df.index:
    year = date.year
    if year in money_year.keys():
        money_year[year] = money_year[year] + df.loc[date, 'save_money']
    else:
        money_year[year] = df.loc[date, 'save_money']
        
money_mean = mean(list(money_year.values()))
years_count = len(money_year) - 1
money_year = {key: [value] for key, value in money_year.items()}

df_money_year = pd.DataFrame(money_year, index=[''])
df_money_year = df_money_year.T
df_money_year.plot(figsize=(14, 4), kind='bar')
plt.hlines(money_mean, 0, years_count, color='orange')
plt.legend(['年均投入', '定投年金'])
plt.show()

# 展示各年的收益
return_year = {}
for date in df.index:
    year = date.year
    return_year[year] = df.loc[date, 'return_rate']
return_year = {key: [value] for key, value in return_year.items()}
return_df = pd.DataFrame(return_year, index=['return']).T
return_df['diff'] = return_df['return'].diff()
return_df['diff'].fillna(return_df['return'], inplace=True)
return_df[['diff']].plot(figsize=(14, 4), kind='bar')
plt.legend(['各年收益率'])
plt.show()

# 查看动态pe形态的分位趋势图
def get_quantile(index_code, p, n, data):
    """指数百分位展示。
    Args:
        index_code: 指数 code。
        p: 可以是 pe，也可以是 pb。
        n: 指用于计算指数估值百分位的区间，如果是5指近5年数据。
        data: 包含有 pe/pb 的 DataFrame。
    Returns:
        计算后的DataFrame。
    """
    # 这里的计算按一年244个交易日计算
    windows = int(n * 244)  # 将时间取整数
    _df = data.copy()
    _df.index.name = None
    price = get_price(index_code, start_date=_df.index[0], end_date=_df.index[-1])
    _df['close'] = price.close
    _df['quantile'] = _df[p].rolling(windows).apply(lambda x: pd.Series(x).rank().iloc[-1] / 
                                                  pd.Series(x).shape[0], raw=True)
    _df.dropna(inplace=True)
    _df['quantile'].plot(figsize=(14, 7))
    # 画出适中估值区间
    plt.fill_between(_df.index, y1=0.4, y2=0.6, color='orange', alpha=0.7)
    plt.annotate('适中估值区', (_df.index[-1], 0.5))
    return _df
# 展示指数百分位趋势图
df_quantile = get_quantile(index, 'pe', 7.45, df_pe_pb)

# pe动态分位图与指数高低位的比较
def show_quantile(index_code, p, n, data):
    """指数百分位展示。
    Args:
        index_code: 指数 code。
        p: 可以是 pe，也可以是 pb。
        n: 指用于计算指数估值百分位的区间，如果是5指近5年数据。
        data: 包含有 pe/pb 的 DataFrame。
    Returns:
        None.
    """
    # 这里的计算按一年244个交易日计算
    windows = int(n * 244)  # 将时间取整数
    _df = data.copy()
    _df.index.name = None
    price = get_price(index_code, start_date=_df.index[0], end_date=_df.index[-1])
    _df['close'] = price.close
    _df['quantile'] = _df[p].rolling(windows).apply(lambda x: pd.Series(x).rank().iloc[-1] / 
                                                  pd.Series(x).shape[0], raw=True)
    _df.dropna(inplace=True)
    _df[['quantile', 'close']].plot(figsize=(14, 10), subplots=True)
    # 画出适中估值区间
    
# 展示指数百分位趋势图
show_quantile(index, 'pe', 7.5, df_pe_pb)

# 获取每个月的第一个交易日
first_day = []
for i in range(len(df_quantile)):
    date = df_quantile.index[i]
    if i == 0:
        first_day.append(date)
    else:
        last_date = df_quantile.index[i - 1]
        if date.day < last_date.day:
            first_day.append(date)

# 按月计算价格与涨跌幅度
close = get_price(index, start_date=df_quantile.index[0], end_date=df_quantile.index[-1])['close']
df = df_quantile.copy()
df['close'] = close
df = df.loc[first_day]
df['pct_change'] = df.close.pct_change()
df.head(10)

save_money = []  # 每月定存
back_money = []  # 回收资金
hold_money = []  # 持仓资金
base_money = 1000  # 定投基准

def trade():
    for i in range(len(df)):
        quantile = df['quantile'][i]  # 估值位
        multiple = int((0.5 - quantile) * 10)  # 定投倍数计算

        if i == 0:  # 初始买入
            # 1.计算买入金额
            _save_money = base_money * multiple
            save_money.append(_save_money)
            # 2. 计算回收金额
            _back_money = 0
            back_money.append(_back_money)
            # 3. 计算持仓变化
            _hold_money = _save_money
            hold_money.append(_hold_money)
            continue

        if multiple >=0:  # 执行买入计算
            # 1.计算买入金额
            _save_money = base_money * multiple
            save_money.append(_save_money)
            # 2. 计算回收金额
            _back_money = 0
            back_money.append(_back_money)
            # 3. 计算持仓变化
            _hold_money = hold_money[-1] * (1 + df['pct_change'][i]) + _save_money
            hold_money.append(_hold_money)
        else:  # 执行卖出计算 
            # 1. 计算买入金额
            _save_money = 0
            save_money.append(_save_money)
            # 2. 计算回收金额
            _back_money = base_money * (2 ** -multiple)  # 按2的指数倍卖出
            _hold_money = hold_money[-1] * (1 + df['pct_change'][i])
            if _back_money > _hold_money:
                _back_money = _hold_money
            if quantile >= 1.0:
                _back_money = _hold_money  # 如果达到100%分位，清仓
            back_money.append(_back_money)
            # 3. 计算持仓变化
            _hold_money = _hold_money - _back_money
            hold_money.append(_hold_money)
trade()

df['save_money'] = save_money  # 定投金额
df['save_money_cumsum'] = df['save_money'].cumsum()  # 定投累计金额
df['hold_money'] = hold_money  # 持仓金额
df['back_money'] = back_money  # 回收金额
df['back_money_cumsum'] = df['back_money'].cumsum()  # 累计回收金额
df['total_money'] = df['hold_money'] + df['back_money_cumsum']  # 总资金
df['return_money'] = df['total_money'] - df['save_money_cumsum']  # 持续收益
df['return_rate'] = (df['total_money'] / df['save_money_cumsum']) - 1  # 持续收益率
df[['save_money_cumsum', 'total_money', 'back_money_cumsum', 'return_money']].plot(figsize=(14, 7))
plt.legend(['累积定投', '累计本息', '回收资金', '收益曲线'])
plt.show()

print('累计投入: {}元'.format(df['save_money_cumsum'][-1]))
print('累计收益: {}元'.format(df['return_money'][-1]))
print('最终本息累积: {}元'.format(df['total_money'][-1]))
print('绝对收益率为: {}%'.format((df['return_money'][-1] / df['save_money_cumsum'][-1]) * 100))

# 展示各年投入金额
money_year = {}
for date in df.index:
    year = date.year
    if year in money_year.keys():
        money_year[year] = money_year[year] + df.loc[date, 'save_money']
    else:
        money_year[year] = df.loc[date, 'save_money']
        
money_mean = mean(list(money_year.values()))
years_count = len(money_year) - 1
money_year = {key: [value] for key, value in money_year.items()}

df_money_year = pd.DataFrame(money_year, index=[''])
df_money_year = df_money_year.T
df_money_year.plot(figsize=(14, 4), kind='bar')
plt.hlines(money_mean, 0, years_count, color='orange')
plt.legend(['年均投入', '定投年金'])
plt.show()

# 展示各年的收益
return_year = {}
for date in df.index:
    year = date.year
    return_year[year] = df.loc[date, 'return_rate']
return_year = {key: [value] for key, value in return_year.items()}
return_df = pd.DataFrame(return_year, index=['return']).T
return_df['diff'] = return_df['return'].diff()
return_df['diff'].fillna(return_df['return'], inplace=True)
return_df[['diff']].plot(figsize=(14, 4), kind='bar')
plt.legend(['各年收益率'])
plt.show()

index = '000300.XSHG'  # 指数 code
index_info = get_security_info(index)  # 指数信息
start_date = index_info.start_date  # 指数开始时间
end_date = datetime.datetime.now().date()  # 以当天为最后一天
index_name = index_info.display_name  # 指数全称

df_pe_pb = get_pe_pb(index, start_date, end_date)
df_pe_pb.head(10)

# 展示指数百分位趋势图
df_quantile = get_quantile(index, 'pe', 7.45, df_pe_pb)

# 展示指数百分位趋势图
show_quantile(index, 'pe', 7.45, df_pe_pb)

# 获取每个月的第一个交易日
first_day = []
for i in range(len(df_quantile)):
    date = df_quantile.index[i]
    if i == 0:
        first_day.append(date)
    else:
        last_date = df_quantile.index[i - 1]
        if date.day < last_date.day:
            first_day.append(date)

# 按月计算价格与涨跌幅度
close = get_price(index, start_date=df_quantile.index[0], end_date=df_quantile.index[-1])['close']
df = df_quantile.copy()
df['close'] = close
df = df.loc[first_day]
df['pct_change'] = df.close.pct_change()
df.head(10)

save_money = []  # 每月定存
back_money = []  # 回收资金
hold_money = []  # 持仓资金
base_money = 1000  # 定投基准

def trade():
    for i in range(len(df)):
        quantile = df['quantile'][i]  # 估值位
        multiple = int((0.5 - quantile) * 10)  # 定投倍数计算

        if i == 0:  # 初始买入
            # 1.计算买入金额
            _save_money = base_money * multiple
            save_money.append(_save_money)
            # 2. 计算回收金额
            _back_money = 0
            back_money.append(_back_money)
            # 3. 计算持仓变化
            _hold_money = _save_money
            hold_money.append(_hold_money)
            continue

        if multiple >=0:  # 执行买入计算
            # 1.计算买入金额
            _save_money = base_money * multiple
            save_money.append(_save_money)
            # 2. 计算回收金额
            _back_money = 0
            back_money.append(_back_money)
            # 3. 计算持仓变化
            _hold_money = hold_money[-1] * (1 + df['pct_change'][i]) + _save_money
            hold_money.append(_hold_money)
        else:  # 执行卖出计算 
            # 1. 计算买入金额
            _save_money = 0
            save_money.append(_save_money)
            # 2. 计算回收金额
            _back_money = base_money * (2 ** -multiple)  # 按2的指数倍卖出
            _hold_money = hold_money[-1] * (1 + df['pct_change'][i])
            if _back_money > _hold_money:
                _back_money = _hold_money
            if quantile >= 1.0:
                _back_money = _hold_money  # 如果达到100%分位，清仓
            back_money.append(_back_money)
            # 3. 计算持仓变化
            _hold_money = _hold_money - _back_money
            hold_money.append(_hold_money)
trade()

df['save_money'] = save_money  # 定投金额
df['save_money_cumsum'] = df['save_money'].cumsum()  # 定投累计金额
df['hold_money'] = hold_money  # 持仓金额
df['back_money'] = back_money  # 回收金额
df['back_money_cumsum'] = df['back_money'].cumsum()  # 累计回收金额
df['total_money'] = df['hold_money'] + df['back_money_cumsum']  # 总资金
df['return_money'] = df['total_money'] - df['save_money_cumsum']  # 持续收益
df['return_rate'] = (df['total_money'] / df['save_money_cumsum']) - 1  # 持续收益率
df[['save_money_cumsum', 'total_money', 'back_money_cumsum', 'return_money']].plot(figsize=(14, 7))
plt.legend(['累积定投', '累计本息', '回收资金', '收益曲线'])
plt.show()

print('累计投入: {}元'.format(df['save_money_cumsum'][-1]))
print('累计收益: {}元'.format(df['return_money'][-1]))
print('最终本息累积: {}元'.format(df['total_money'][-1]))
print('绝对收益率为: {}%'.format((df['return_money'][-1] / df['save_money_cumsum'][-1]) * 100))

# 展示各年投入金额
money_year = {}
for date in df.index:
    year = date.year
    if year in money_year.keys():
        money_year[year] = money_year[year] + df.loc[date, 'save_money']
    else:
        money_year[year] = df.loc[date, 'save_money']
        
money_mean = mean(list(money_year.values()))
years_count = len(money_year) - 1
money_year = {key: [value] for key, value in money_year.items()}

df_money_year = pd.DataFrame(money_year, index=[''])
df_money_year = df_money_year.T
df_money_year.plot(figsize=(14, 4), kind='bar')
plt.hlines(money_mean, 0, years_count, color='orange')
plt.legend(['年均投入', '定投年金'])
plt.show()

# 展示各年的收益
return_year = {}
for date in df.index:
    year = date.year
    return_year[year] = df.loc[date, 'return_rate']
return_year = {key: [value] for key, value in return_year.items()}
return_df = pd.DataFrame(return_year, index=['return']).T
return_df['diff'] = return_df['return'].diff()
return_df['diff'].fillna(return_df['return'], inplace=True)
return_df[['diff']].plot(figsize=(14, 4), kind='bar')
plt.legend(['各年收益率'])
plt.show()