import numpy as np
import matplotlib.pyplot as plt
import cvxopt as opt
from cvxopt import blas, solvers
import pandas as pd
from cvxopt import solvers, matrix

df=pd.read_csv('stocks.csv')

stock_list=[row['stock'] for index,row in df.iterrows()]

#选股池
instruments=stock_list

#回测时间
start_date='2019-01-03'
end_date='2021-01-22'

def initialize(context):
    
    context.set_commission(PerOrder(buy_cost=0.0003,sell_cost=0.0003,min_cost=5))

##每天执行一次
def handle_data(context,data):
    
    ##技术指标所需要的最大天数
    
    if context.trading_day_index<100:
        return
    
    if (context.trading_day_index)%20!=0:
        return 
    
    print(context.trading_day_index)
    
    price_df=pd.DataFrame()
    
    stock_num=120
    count=0
    
    for stock in instruments:
        count+=1
        sid=context.symbol(stock)
        prices = data.history(sid,'price',5,'1d')
        price_df=pd.concat([price_df,prices],axis=1)
        price_df.fillna(0,inplace=True)
        if(count==stock_num):
            break
       
    ## 120,5
    price_np=price_df.to_numpy().T
   
    
    ##cov
    ##(120,120)
    cov_df= pd.DataFrame(np.cov(price_np))
    P=matrix(cov_df.to_numpy())
    
    q=-matrix(np.mean(price_np,axis=1),(stock_num,1))
    
    
    ##Ax=b
    A=matrix(1.,(1,stock_num))
    b=matrix(1.)
    
    
    ##GX<=h
    G = -matrix(np.eye(stock_num)) 
    h = opt.matrix(0.0, (stock_num,1))
  
    
    ##solver:
    solvers.options['show_progress'] = False
    sol = solvers.qp(P,q,G,h,A,b)
    weight=list(sol['x'])

    stock_selected_list=[]
    sum_weight=0.0
    cur=0
    for stock in stock_list:      
        sid=context.symbol(stock)
        if data.can_trade(sid):
            stock_selected_list.append(stock)
            sum_weight+=weight[cur]
        else:
            weight[cur]=0
        cur+=1
    
    cur=0
    for stock in stock_list:
        if weight[cur]!=0:
            sid=context.symbol(stock)
            price=data.current(sid,'price')
            context.order_target_percent(sid,weight[cur]/sum_weight)
        cur+=1
        
    
            
        
        
    
    
    

        
        
  
    
        

m=M.trade.v2(
    instruments=instruments,
    start_date=start_date,
    end_date=end_date,
    initialize=initialize,
    handle_data=handle_data,
    order_price_field_buy='open',
    order_price_field_sell='open',
    capital_base=1000000
)