#?coding=utf-8
from?__future__?import?division
from?backtest?import?*
import?os
import?pandas?as?pd
from?math?import?floor


#?計算復權價格
def?cal_right_price（input_stock_data?type=‘前復權‘）:
????“““
????:param?input_stock_data:?標準股票數據，需要‘收盤價‘?‘漲跌幅‘
????:param?type:?確定是前復權還是后復權，分別為‘后復權‘，‘前復權‘
????:return:?新增一列‘后復權價‘/‘前復權價‘的stock_data
????“““
????#?計算收盤復權價
????stock_data?=?input_stock_data.copy（）
????num?=?{‘后復權‘:?0?‘前復權‘:?-1}
????price1?=?stock_data[‘close‘].iloc[num[type]]
????stock_data[‘復權價_temp‘]?=?（stock_data[‘change‘]?+?1.0）.cumprod（）
????price2?=?stock_data[‘復權價_temp‘].iloc[num[type]]
????stock_data[‘復權價‘]?=?stock_data[‘復權價_temp‘]?*?（price1?/?price2）
????stock_data.pop（‘復權價_temp‘）

????#?計算開盤復權價
????stock_data[‘復權價_開盤‘]?=?stock_data[‘復權價‘]?/?（stock_data[‘close‘]?/?stock_data[‘open‘]）
????return?stock_data[[‘復權價_開盤‘?‘復權價‘]]


#?獲取股票數據
def?get_stock_data（）:
????#?遍歷數據文件夾中所有股票文件的文件名，得到股票代碼列表
????stock_code_list?=?[]
????#?此處為股票數據文件的本地路徑，請自行修改
????for?root?dirs?files?in?os.walk（‘e:/data/stock?data‘）:
????????if?files:
????????????for?f?in?files:
????????????????if?‘.csv‘?in?f:
????????????????????stock_code_list.append（f.split（‘.csv‘）[0]）

????all_stock?=?pd.Dataframe（）

????for?code?in?stock_code_list:
????????#?此處為股票數據文件的本地路徑，請自行修改
????????stock_data?=?pd.read_csv（‘e:/data/stock?data/‘?+?code?+?‘.csv‘?parse_dates=[‘date‘]）
????????stock_data?=?stock_data[[‘code‘?‘date‘?‘open‘?‘close‘?‘change‘]].sort_values（by=‘date‘）
????????stock_data.reset_index（drop=True?inplace=True）
????????#?計算復權價
????????stock_data[[‘open‘?‘close‘]]?=?cal_right_price（stock_data?type=‘后復權‘）
????????#?判斷每天開盤是否漲停
????????stock_data.ix[stock_data[‘open‘]?>?stock_data[‘close‘].shift（1）?*?1.097?‘limit_up‘]?=?1
????????stock_data[‘limit_up‘].fillna（0?inplace=True）

????????all_stock?=?all_stock.append（stock_data?ignore_index=True）

????return?all_stock[[‘code‘?‘date‘?‘change‘?‘limit_up‘]]


def?momentum_and_contrarian（all_stock?start_date?end_date?window=3）:
????“““
????:param?all_stock:?所有股票的數據集
????:param?start_date:?起始日期（包含排名期）
????:param?end_date:?結束日期
????:param?window:?排名期的月份數，默認為3個月
????:return:?返回動量策略和反轉策略的收益率和資金曲線
????“““
????#?取出指數數據作為交易天數的參考標準?此處為指數數據文件的本地路徑，請自行修改
????index_data?=?pd.read_csv（‘e:/data/index?data/sh000001.csv‘?parse_dates=True?index_col=‘date‘）
????index_data.sort_index（inplace=True）
????index_data?=?index_data[start_date:end_date]
????#?轉換成月度數據
????by_month?=?index_data[[‘close‘]].resample（‘M‘?how=‘last‘）
????by_month.reset_index（inplace=True）

????momentum_portfolio_all?=?pd.Dataframe（）
????contrarian_portfolio_all?=?pd.Dataframe（）

????for?i?in?range（window?len（by_month）?-?1）:
????????start_month?=?by_month[‘date‘].iloc[i?-?window]??#?排名期第一個月
????????end_month?=?by_month[‘date‘].iloc[i]??#?排名期最后一個月

????????#?取出在排名期內的數據
????????stock_temp?=?all_stock[（all_stock[‘date‘]?>?start_month）?&?（all_stock[‘date‘]?<=?end_month）]

????????#?將指數在這段時間的數據取出