main script

# jungle.py

# import sys, io, os, inspect
import numpy as np, pandas as pd; _=np
pd.set_option('display.precision', 6)
import matplotlib.pyplot as plt

# from dqsbt.util2 import util as aid # 加载数据
# from util import ttr                # 计算指标
# import candleplotter as cplt

# from qsutil import pkl
# from texttable import Texttable
# from packages import report as rp

# from qsutil import pkl; _cn_dict = pkl.pkl_load('d:/db/amipy/data/codename_dict.pkl')
import momentums_config as cfg
# import utility as ut
from qsutil import utility as ut
from qsutil import bbplot as bbp; _=bbp
import flower as fl; _=fl
p=dict(mfile='D:\\algolab\\study\\momentums\\momentums.py') # __file__
_pn = [n for n in dir(cfg) if not n.startswith('__')]
_pv = [getattr(cfg, n) for n in _pn]
p.update(dict(zip(_pn, _pv)))

from collections import namedtuple
# import pfbroker as pfb
_=namedtuple, 

import jungle as jg

def pre_proc():
    ''' 
    >>> ohlc, ind = pre_proc()
    '''
    ohlc, _codes, _dnames=ut.load(subset=cfg.subset)
    ind = ut.indicator(ohlc, n1=cfg.n1, n2=cfg.n2)
    return ohlc, ind




class RocAlgo(jg.Jungle):
    def __init__(self, 
                 cfg, ind, 
                 weight,
                 nComponent=2, 
                 nAdjustPos=5,
                 dbg=True,
        ):

        super().__init__(cfg, ind, 
                         nComponent=nComponent, nAdjustPos=nAdjustPos, dbg=dbg)
        self.weight = weight
        
    def run(self, ):
        mcdf = self.ind.loc[:, ('close roc'.split())]
        n = len(self.codes)
        for i, row in enumerate(mcdf.itertuples()):
            dt,clos, rocs = row[0], row[1:1+n], row[1+n:]
            dts = self.dt2dts(dt)
            close_dict={ k:v for k,v in  zip(self.codes, clos)}
            _ = dts, close_dict
            
            # 循环时最好依据数据出现的情况的顺序来处理, 
            # 比如首先遇到的是na值, 处理之. 很简单. 
            # 用continue来处理之(继续下一行的循环). 避免太多的elif语句
            if pd.isna(rocs[0]): continue
                
            # 开仓调仓日
            if i%self.nAdjustPos==0: 
                # 开仓
                if self.b.balance is None:
                    # print(i, dts, clos, rocs); _=input('pause')
                    buy_list, sorted_roc = self.sort_rocs(rocs)
                    self.price =[clos[i] for i in self.buy_iloc]
                    # 等权重配置方式
                    # 收盘前5分钟买入
                    code=buy_list
                    bbo = self.o.create_batch_buy_order(i, code,  percent=self.weight)
                    # print('最新订单:\n', self.o); print('\n历史订单:'); self.o.porders() # 最新订单和历史订单
                    self.b._first_batch_buy(bbo)
                    
                    # self.pause()
                    
                    # 收盘后收集收盘价
                    self.check_totcap(i)
    
                    # _=input('建仓之后pause')
                    
                # 调仓
                else:
                    buy_list, sorted_roc = self.sort_rocs(rocs)
                    self.price =[clos[i] for i in self.buy_iloc]
                    if sorted(buy_list)==sorted(self.get_ccc()):
                        # print('调仓日, 动量排序进榜个股与持仓个股名单相同, 无需操作.')
                        pass
                    # 动量选股之后, 有新的进榜个股
                    else:
                        print(f'调仓日: i={i}, buy_list={buy_list}, 头寸个股名单={self.get_ccc()}')
                        remove_gg, addon_gg = self.get_addon_remove(buy_list)
                        # print(f'卖出个股={remove_gg}, 买入个股={addon_gg}')
                        # Todo
                        # .  卖出头寸里的 退榜个股, 及时drop该席位: pos=pos.drop(psn, level=0)
                        # .  买入 新的进榜个股, 及时新建一个单行单列头寸: 设置psn=100+psn
                        #       other_pos=pd.DataFrame(), 
                        #       并联到头寸里: pos=pd.concat([pos, other_pos], axis=1)
                        # .  按照动量顺序重排头寸席位 
                        #       b.sort_position_seat(buylist)
                        
                        
                        if len(remove_gg)==1:
                            code=remove_gg[0]
                            so = self.o.create_sell_order(i, code, 0.)
                            self.b.sell_target_percent(so)
                        
                            if len(addon_gg)==1:
                                code=addon_gg[0]
                                # print(f'准备买入进榜个股={code}')
                                bo = self.o.create_buy_order(i, code, 0.5)
                                self.b.buy_target_percent(bo, buy_list)
                                
                        # 头寸合并,排序
                        # self.b.position = pd.concat([self.b.position, self.b.other_pos], axis=1)
                        self.b.sort_position_seat(buy_list)
                        # print(self.o)
                        # print(self.b.get_position())
                        
                        # self.pause()

                    # 收盘后查资
                    self.check_totcap(i)
    
                    # _=input('建仓之后pause')
                    
            # 非调仓日
            else: 
                # 尚未建仓时
                if self.b.balance==None:
                    pass
                # 有仓位时
                elif self.get_v0()>0:#第一号仓位有的
                    self.check_totcap(i)
                    pass
                # 空仓时, 只登记个时间戳而已
                elif self.get_v0()==0:
                    self.log_dt() # 
                    pass
    
        print(f'nAdjustPos={self.nAdjustPos}, totcap={self.b.totcap:.2f}')
        # self.post_proc()
        # self.draw_c0()
        # self.check_result()
    
#%%


def utest2(cfg, ind, nAdjustPos=5, nComponent=2, weight=(0.5, 0.5)):
    '''
    >>> ohlc, ind = pre_proc()
    >>> jgl = utest2(cfg, ind, 8, 2, (0.5, 0.5)) # 14.4w
    >>> jgl = utest2(cfg, ind, 4, 2, (0.9, 0.1)) # 15.5w
    >>> jgl = utest2(cfg, ind, 4, 3, (0.6, 0.2, 0.2)) # 15.5w
    '''
    jgl = RocAlgo(cfg, ind, 
                  weight=weight, 
                  nComponent=nComponent,
                  nAdjustPos=nAdjustPos, 
    ) 
    jgl.run()
    return jgl



def save_etf(start_date='1900-01-01', end_date='2020-12-31'):
    ''' 下载主要etf数据, 并且保存到db. 
    >>> op=DbOperation()
    >>> op.save_etf()
    df_dict = save_etf()
    '''
    import util.tdxhq as hq, time
    start_date='2020-01-01'
    major_etf_codes = '512880 512660 512800 512170 159928'.split()
    print('Downloading 主要etf的日线数据 from tdxhq......\n')
    print(f"start date = {start_date:s}, end date={end_date:s}")
    df_dict={}
    for i, code in enumerate(major_etf_codes):
        print(f'{i:2d}', 'downloading.......:' , code)
        df=hq.test_get_k_data(code, sdate=start_date, edate=end_date)
        df.index=pd.DatetimeIndex(df.index)
        
        df = df.assign(code=code[:6])
        df = df.rename(columns={'vol':'volume'})
        # df = df.drop(columns='date')
        # cols = list(df)
        cols = 'open high low close volume amount code'.split()
        df = df.loc[:, cols]
        # df.to_sql('table_etf', conn, if_exists='append', dtype={'code':'CHAR'})
        time.sleep(2)
        df_dict.update({code:df})
    return df_dict

def _draw(df_dict):
    fig, ax = plt.subplots(len(df_dict), 1)    
    for i, (k,v) in enumerate(df_dict.items()):
        (v.close/v.close[0]).plot(ax=ax[i], label=k, legend=True)
    
    fig, ax = plt.subplots(1, 1)    
    for i, (k,v) in enumerate(df_dict.items()):
        (v.close/v.close[0]).plot(ax=ax, label=k, legend=True)