### Momentum, Monthly long-only strategy, based on 3-month window

import pandas as pd
import numpy as np
import matplotlib.pyplot as plt
import os
from math import floor


# Collect name of stocks on China A-Share
def collect_stock_name():
    stock_code_list=[]
    for root, dirs, files in os.walk('D:/stock_data'):
        if files:
            for f in files:
                if 'csv' in f:
                    stock_code_list.append(f.split('.csv')[0])
    return stock_code_list


# Collect and merge data of all stocks
def collect_all_data(setting_column_list):

    all_data = pd.DataFrame()
    stock_code_list = collect_stock_name()

    for code in stock_code_list:
        stock_data = pd.read_csv('D:/stock_data/'+code+'.csv')
        stock_data = stock_data[setting_column_list]
        stock_data['date'] = pd.to_datetime(stock_data['date'])
        stock_data = stock_data.set_index('date')
        stock_data = stock_data.resample('M',how='last')
        stock_data = stock_data.reset_index()

        all_data = pd.concat([all_data,stock_data],axis=0)
    all_data.columns = (setting_column_list)

    return all_data



# If one stock is limited up in a specific day, then 'limit_up' is "1"
def limit_up(data):
    data.loc[data['open']>data['close'].shift(1)*1.907,'limit_up'] = 1
    data = data['limit_up'].fillna(0)
    return data


# Get momentum portfolio
def momentum_and_contrarian(index_data,stock_data,start_date, end_date, window=3):
    all_stock = stock_data
    index_data=index_data.sort_index()
    index_data=index_data[start_date:end_date]
    by_month = index_data[['close']].resample('M', how='last')
    by_month= by_month.reset_index()
    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)]
        index_temp=index_data[start_month:end_month]

        grouped=stock_temp.groupby('code')['change'].agg({'return':lambda x: (x+1).prod()-1})
        grouped.sort_values(by='return')
        num=floor(len(grouped)*0.1)
        momentum_code_list=grouped.index[-num:]

        # construting momentum portfolio
        momentum=all_stock[all_stock['code'].isin(momentum_code_list)]
        momentum=all_stock[all_stock['date']>end_month]
        momentum=all_stock[all_stock['date']<=by_month['date'].iloc[i+1]]


    return momentum



# Load stock data
setting_column_list = ['date','code','close']
index_data = pd.read_csv('D:\index_data\index.csv')

all_stock_data['date']=pd.to_datetime(all_stock_data['date'])
all_stock_data['change']=all_stock_data['close']/all_stock_data['close'].shift(1)-1
all_stock_data=all_stock_data.fillna(0)

# Load index data
index_data = pd.read_csv('D:\index_data\index.csv')
index_data['date']=pd.to_datetime(index_data['date'])
index_data=index_data.set_index('date')

# Get momentum portfolio
start_date = pd.to_datetime('20020101')
end_date = pd.to_datetime('20171231')

momentum_portfolio = momentum_and_contrarian(index_data,all_stock_data,start_date, end_date)
momentum_output = momentum_portfolio.groupby('date')['change'].mean()[1:]
momentum_output = momentum_output[start_date:end_date]
momentum_output_capital= (momentum_output+1).cumprod()
momentum_output_capital = momentum_output_capital - momentum_output_capital[0]
momentum_output.to_csv('D:/momentum_return.csv')