import pandas as pd
import datetime
import data.stock as st
import data.strategy.base as stbase
import matplotlib.pyplot as plt

# 获取三只股票的数据代码
names = ['亿纬锂能', '比亚迪', '赣锋锂业']
codes = []
for name in names:
    codes.append(st.get_single_stock_code(name))

sharp_d = []
sharp_y = []
# 获取三只股票的price数据，并计算三只股票的夏普比率，将其与股票代码成对存放在一个series中
# 获取price数据的同时，画出close的折线图，并标明图例
for code in codes:
    data = st.get_single_price(code, '2019-01-01', datetime.date.today(), 'daily')
    plt.plot(data.index, data['close'], label=code)  # 这句用data['close'].plot(label=code)效果相同
    plt.xticks(rotation=30)
    daily_sp, yearly_sp = stbase.calculate_sharpe(data)
    sharp_d.append([code, daily_sp])
    sharp_y.append([code, yearly_sp])
plt.legend()

# 画出夏普比率柱状图
sharp_d = pd.DataFrame(sharp_d, columns=['code', 'sharp_d'])  # 将日夏普比率保存为df
sharp_y = pd.DataFrame(sharp_y, columns=['code', 'sharp_y'])  # 将年夏普比率保存为df
# 将code字段设置成日夏普比率和年夏普比率的index
sharp_d = sharp_d.set_index('code')
sharp_y = sharp_y.set_index('code')
# 将日夏普比率和年夏普比率数据合并到一个dataframe变量sharp中
sharp = sharp_d
sharp['sharp_y'] = sharp_y['sharp_y']
# 设置图形名称并打印
sharp.plot.bar(title='sharp')
plt.xticks(rotation=0)
plt.show()