# 本代码由可视化策略环境自动生成 2022年4月30日 13:28
# 本代码单元只能在可视化模式下编辑。您也可以拷贝代码，粘贴到新建的代码单元或者策略，然后修改。


m1 = M.instruments.v2(
    start_date='2018-01-01',
    end_date='2022-04-20',
    market='CN_STOCK_A',
    instrument_list='',
    max_count=0
)

m2 = M.input_features.v1(
    features="""close_0
open_0
price_limit_status_0
return_0
"""
)

m3 = M.general_feature_extractor.v7(
    instruments=m1.data,
    features=m2.data,
    start_date='',
    end_date='',
    before_start_days=0
)

m4 = M.derived_feature_extractor.v3(
    input_data=m3.data,
    features=m2.data,
    date_col='date',
    instrument_col='instrument',
    drop_na=False,
    remove_extra_columns=False,
    user_functions={}
)

#计算下一个交易日的收益
df_total = m4.data.read()

#获取指数数据
df_index_000905 = DataSource('bar1d_index_CN_STOCK_A').read(instruments=['000905.HIX'], start_date='2018-01-01')
df_index_000905['returns'] = (df_index_000905.close / df_index_000905.close.shift(1) -1) + 1
df_index_000905['中证500'] = df_index_000905['returns'].cumprod()
df_index_000905['中证500'] = df_index_000905['中证500'] .fillna(1)

#涨停和跌停数
def cal_limit(x):
    x['high_limit'] = len(x[x.price_limit_status_0==3])
    x['low_limit'] = len(x[x.price_limit_status_0==1])
    x['上涨数'] = len(x[x.return_0>1])
    return x
df_limit = df_total.groupby('date').apply(cal_limit)
df_limit = df_limit.drop_duplicates(subset=['date'],keep='last').sort_values('date').reset_index()
df_limit = pd.merge(left=df_limit,right=df_index_000905[['date','中证500']],left_on='date',right_on='date')

#与指数合并
options = {'chart': {'type': 'spline', 'height': 600},}
T.plot(df_limit[['date','high_limit','low_limit','中证500']].set_index('date'),title='每日涨停跌停数',panes=[['high_limit', 'low_limit', '60%'], ['中证500', '40%']],options=options)

#上涨股票数
T.plot(df_limit[['上涨数','date','中证500']].set_index('date'),title='每日上涨股票数',panes=[['上涨数', '60%'], ['中证500','40%']],options=options)
def cal(x):
    x = x.sort_values('date')
    x['next_return'] = x.return_0.shift(-1)
    x['pre_limit_status_0'] = x.price_limit_status_0.shift(1)
    return x
df_total = df_total.groupby('instrument').apply(cal)

import talib as ta
#根据收益率计算指数净值
df_yj = df_total.copy().reset_index(drop=True)
df_yj = df_yj[df_yj.price_limit_status_0==3]
#获取交易日里为了填充空值
df_trade = DataSource("trading_days").read(start_date="2018-01-01")
df_trade = df_trade[df_trade.country_code=='CN']
df_trade.sort_values('date')

def cal_return(x):
    return x.next_return.mean()
df_yj = df_yj.groupby('date').apply(cal_return)
df_yj = pd.DataFrame(df_yj).rename(columns={0:'return'})
df_yj = pd.merge(left=df_trade,right=df_yj,left_on='date',right_index=True,how='left')
df_yj = df_yj.set_index('date')
df_yj['return'] = df_yj['return'].fillna(1)
df_yj['return'] = df_yj['return'].shift(1)
df_yj['return'] = df_yj['return'].fillna(1)
df_yj['净值'] = df_yj['return'].cumprod()

N = 5
df_yj["rsi"] = ta.RSI(df_yj['return'].values, timeperiod=N)
#合并指数数据
df_yj = pd.merge(left=df_yj,right=df_index_000905[['date','中证500']],left_on='date',right_on='date')
df_yj = df_yj.set_index('date')
T.plot(df_yj[['净值']],title='涨停溢价')
options = {'chart': {'type': 'spline', 'height': 600},}
T.plot(df_yj[['rsi','中证500']],panes=[['rsi', '60%'], ['中证500','40%']],options=options,title='涨停溢价RSI')

#昨日连扳指数
df_lb = df_total.copy().reset_index(drop=True)
df_lb = df_lb[(df_lb.price_limit_status_0==3)&(df_lb.pre_limit_status_0==3)]

def cal_return(x):
    return x.next_return.mean()
df_lb = df_lb.groupby('date').apply(cal_return)
df_lb = pd.DataFrame(df_lb).rename(columns={0:'return'})
df_lb['return'] = df_lb['return'].shift(1)
df_lb['return'] = df_lb['return'].fillna(1)
df_lb['净值'] = df_lb['return'].cumprod()
T.plot(df_lb['净值'])