# coding=utf-8
# ======================================
# File:     test_user_defined.py
# Author:   Jackie PENG
# Contact:  jackie.pengzhao@gmail.com
# Created:  2020-02-12
# Desc:
#   Test file for user-defined
# strategies.
# ======================================
import unittest

from numba.core.utils import benchmark

import qteasy as qt
import numpy as np

from qteasy.parameter import Parameter
from qteasy.datatypes import StgData


def market_value_weighted(stock_return, mv, mv_cat, bp_cat, mv_target, bp_target):
    """ 根据mv_target和bp_target计算市值加权收益率

    """
    sel = (mv_cat == mv_target) & (bp_cat == bp_target)
    mv_total = np.nansum(mv[sel])
    mv_weight = mv / mv_total
    return_total = np.nansum(stock_return[sel] * mv_weight[sel])
    return return_total


class MultiFactors(qt.FactorSorter):

    def __init__(self, par_values: tuple = (0.5, 0.3, 0.7)):
        super().__init__(
                pars=[Parameter((0.01, 0.99), name='size_gate_percentile', par_type='float'),
                      Parameter((0.01, 0.49), name='bp_small_percentile', par_type='float'),
                      Parameter((0.50, 0.99), name='bp_large_percentile', par_type='float')],
                name='MultiFactor',
                description='根据Fama-French三因子回归模型估算HS300成分股的alpha值选股',
                data_types=[StgData('pb', freq='d', asset_type='E'),
                            StgData('total_mv', freq='d', asset_type='E'),
                            StgData('close', freq='d', asset_type='E'),
                            StgData('close-000300.SH', freq='d', asset_type='IDX')],  # 参考数据
                window_length=20,
                use_latest_data_cycle=True,
                max_sel_count=10,  # 最多选出10支股票
                sort_ascending=True,  # 选择因子最小的股票
                condition='less',  # 仅选择因子小于某个值的股票
                lbound=0,  # 仅选择因子小于0的股票
                ubound=0,  # 仅选择因子小于0的股票
        )
        if par_values:
            self.update_par_values(*par_values)

    def realize(self):

        size_gate_percentile, bp_small_percentile, bp_large_percentile = self.get_pars('size_gate_percentile',
                                                                                       'bp_small_percentile',
                                                                                       'bp_large_percentile')
        pb_data, mv_data, close_data = self.get_data('pb_E_d', 'total_mv_E_d', 'close_E_d')  # 获取策略数据

        # 读取投资组合的数据PB和total_MV的最新值
        pb = pb_data[-1]  # 当前所有股票的PB值
        mv = mv_data[-1]  # 当前所有股票的市值
        pre_close = close_data[-2]  # 当前所有股票的前收盘价
        close = close_data[-1]  # 当前所有股票的最新收盘价

        # 读取参考数据(r)
        r = self.get_data('close-000300.SH_IDX_d')  # 获取参考数据
        market_pre_close = r[-2]  # HS300的昨收价
        market_close = r[-1]  # HS300的收盘价

        # 计算账面市值比，为pb的倒数
        bp = pb ** -1
        # 计算市值的50%的分位点,用于后面的分类
        size_gate = np.nanquantile(mv, size_gate_percentile)
        # 计算账面市值比的30%和70%分位点,用于后面的分类
        bm_30_gate = np.nanquantile(bp, bp_small_percentile)
        bm_70_gate = np.nanquantile(bp, bp_large_percentile)
        # 计算每只股票的当日收益率
        stock_return = pre_close / close - 1

        # 根据每只股票的账面市值比和市值，给它们分配bp分类和mv分类
        # 市值小于size_gate的cat为1，否则为2
        mv_cat = np.ones_like(mv)
        mv_cat += (mv > size_gate).astype('float')
        # bp小于30%的cat为1，30%～70%之间为2，大于70%为3
        bp_cat = np.ones_like(bp)
        bp_cat += (bp > bm_30_gate).astype('float')
        bp_cat += (bp > bm_70_gate).astype('float')

        # 获取小市值组合的市值加权组合收益率
        smb_s = (market_value_weighted(stock_return, mv, mv_cat, bp_cat, 1, 1) +
                 market_value_weighted(stock_return, mv, mv_cat, bp_cat, 1, 2) +
                 market_value_weighted(stock_return, mv, mv_cat, bp_cat, 1, 3)) / 3
        # 获取大市值组合的市值加权组合收益率
        smb_b = (market_value_weighted(stock_return, mv, mv_cat, bp_cat, 2, 1) +
                 market_value_weighted(stock_return, mv, mv_cat, bp_cat, 2, 2) +
                 market_value_weighted(stock_return, mv, mv_cat, bp_cat, 2, 3)) / 3
        smb = smb_s - smb_b
        # 获取大账面市值比组合的市值加权组合收益率
        hml_b = (market_value_weighted(stock_return, mv, mv_cat, bp_cat, 1, 3) +
                 market_value_weighted(stock_return, mv, mv_cat, bp_cat, 2, 3)) / 2
        # 获取小账面市值比组合的市值加权组合收益率
        hml_s = (market_value_weighted(stock_return, mv, mv_cat, bp_cat, 1, 1) +
                 market_value_weighted(stock_return, mv, mv_cat, bp_cat, 2, 1)) / 2
        hml = hml_b - hml_s

        # 计算市场收益率
        market_return = market_pre_close / market_close - 1

        coff_pool = []
        # 对每只股票进行回归获取其alpha值
        for rtn in stock_return:
            x = np.array([[market_return, smb, hml, 1.0]])
            y = np.array([[rtn]])
            # OLS估计系数
            coff = np.linalg.lstsq(x, y)[0][3][0]
            coff_pool.append(coff)

        # 以alpha值为股票组合的选股因子执行选股
        factors = np.array(coff_pool)

        return factors


class IndexEnhancement(qt.GeneralStg):

    def __init__(self, par_values: tuple = (0.35, 0.8, 5)):
        super().__init__(
                # 参数1:沪深300指数权重阈值，低于它的股票不被选中，参数2: 初始权重，参数3: 连续涨跌天数，作为强弱势判断阈值
                pars=[Parameter((0.01, 0.99), name='weight_threshold', par_type='float'),
                      Parameter((0.51, 0.99), name='initial_weight', par_type='float'),
                      Parameter((2, 20), name='price_days', par_type='int')],
                name='IndexEnhancement',
                description='跟踪HS300指数选股，并根据连续上涨/下跌趋势判断强弱势以增强权重',
                # 利用HS300权重设定选股权重, 根据收盘价判断强弱势
                data_types=[StgData('wt_idx|000300.SH', freq='m', asset_type='ANY', window_length=1),
                            StgData('close', freq='d', asset_type='ANY', window_length=20, )],
                use_latest_data_cycle=True,
        )
        if par_values:
            self.update_par_values(*par_values)

    def realize(self):
        weight_threshold, init_weight, price_days = self.get_pars('weight_threshold', 'initial_weight', 'price_days')
        index_weight, prices = self.get_data('wt_idx|000300.SH_ANY_m', 'close_ANY_d')
        # 读取投资组合的权重wt和最近price_days天的收盘价
        wt = index_weight[-1]  # 当前所有股票的权重值
        pre_close = prices[-price_days - 1:-1]
        close = prices[-price_days:]  # 当前所有股票的最新连续收盘价
        # 计算连续price_days天的收益
        stock_returns = pre_close - close  # 连续p天的收益

        # 设置初始选股权重为0.8
        weights = init_weight * np.ones_like(wt)

        # 剔除所有没有权重（NaN）的股票
        weights[np.where(np.isnan(wt))] = 0

        # 剔除掉权重小于weight_threshold的股票
        weights[wt < weight_threshold] = 0

        # 找出强势股，将其权重设为1, 找出弱势股，将其权重设置为 init_weight - (1 - init_weight)
        up_trends = np.all(stock_returns > 0, axis=0)
        weights[up_trends] = 1.0
        down_trend_weight = init_weight - (1 - init_weight)
        down_trends = np.all(stock_returns < 0, axis=0)
        weights[down_trends] = down_trend_weight

        # 实际选股权重为weights * HS300权重
        weights *= wt
        # print(f'select weight is: \n{weights}')

        return weights


class GridTrading(qt.GeneralStg):

    def __init__(self, par_values: tuple = (2.0, 3.0, 0.3, 0.5, 300)):
        super().__init__(
                # 仓位配置的阈值：参数1:低仓位阈值，参数2: 高仓位阈值，参数3：低仓位比例，参数4:高仓位比例，参数5:计算天数
                pars=[Parameter((0.5, 3.0), name='low_threshold', par_type='float'),
                      Parameter((2.0, 10.0), name='high_threshold', par_type='float'),
                      Parameter((0.01, 0.5), name='low_position', par_type='float'),
                      Parameter((0.5, 0.99), name='high_position', par_type='float'),
                      Parameter((10, 300), name='days', par_type='int')],
                name='GridTrading',
                description='根据过去300份钟的股价均值和标准差，改变投资金额的仓位',
                data_types=[StgData('close', freq='1min', asset_type='ANY')],  # 使用分钟收盘价调整
                window_length=300,
                use_latest_data_cycle=False,  # 高频数据不需要使用当前数据区间
        )
        if par_values:
            self.update_par_values(*par_values)

    def realize(self):
        """策略输出PT信号，即仓位目标信号"""

        low_threshold, high_threshold, low_pos, hi_pos, days = self.get_pars('low_threshold',
                                                                             'high_threshold',
                                                                             'low_position',
                                                                             'high_position',
                                                                             'days')
        h = self.get_data('close_ANY_1min')

        # 读取最近N天的收盘价
        close = h[- days:]  # 最新连续收盘价
        current_close = h[-1]  # 当天的收盘价

        # 计算N天的平均价和标准差，并计算仓位阈值
        close_mean = np.nanmean(close)
        close_std = np.nanstd(close)
        hi_positive = close_mean + high_threshold * close_std
        low_positive = close_mean + low_threshold * close_std
        low_negative = close_mean - low_threshold * close_std
        hi_negative = close_mean - high_threshold * close_std

        # 根据当前的实际价格确定目标仓位，并将目标仓位作为信号输出
        pos = np.zeros_like(close_mean)
        pos = np.where(current_close > hi_positive, hi_pos, pos)
        pos = np.where((hi_positive >= current_close > low_positive), low_pos, pos)
        pos = np.where((low_positive >= current_close > low_negative), 0, pos)
        pos = np.where((low_negative >= current_close > hi_negative), - low_pos, pos)
        pos = np.where(current_close >= hi_negative, - hi_pos, pos)

        return pos


class AlphaPT(qt.GeneralStg):

    def realize(self):
        # 从历史数据编码中读取四种历史数据的最新数值
        total_mv = self.get_data('total_mv_E_d')[-1]  # 总市值
        total_liab = self.get_data('total_liab_E_q')[-1]  # 总负债
        cash_equ = self.get_data('c_cash_equ_end_period_E_q')[-1]  # 现金及现金等价物总额
        ebitda = self.get_data('ebitda_E_q')[-1]  # ebitda，息税折旧摊销前利润

        # 选股因子为EV/EBIDTA，使用下面公式计算
        factors = (total_mv + total_liab - cash_equ) / ebitda
        # 处理交易信号，将所有小于0的因子变为NaN
        factors = np.where(factors < 0, np.nan, factors)
        # 选出数值最小的30个股票的序号
        arg_partitioned = factors.argpartition(30)
        selected = arg_partitioned[:30]  # 被选中的30个股票的序号，此时股票可能有NaN被选中的情况，需要去掉
        not_selected = arg_partitioned[30:]  # 未被选中的其他股票的序号（包括因子为NaN的股票）

        # 如果选出的股票中有因子为NaN的，则剔除掉
        selected = selected[~np.isnan(selected)]
        sel_count = len(selected)

        # 开始生成PT交易信号
        signal = np.zeros_like(factors)
        # 所有被选中的股票的持仓目标被设置为0.03，表示持有3.3%
        signal[selected] = 1 / sel_count
        # 其余未选中的所有股票持仓目标在PT信号模式下被设置为0，代表目标仓位为0
        signal[not_selected] = 0

        return signal


class FastExperiments(unittest.TestCase):
    """This test case is created to have experiments done that can be quickly called from Command line"""

    def setUp(self):
        from qteasy import DataSource, QT_CONFIG
        self.test_ds = DataSource(
                source_type='db',
                host=QT_CONFIG['test_db_host'],
                port=QT_CONFIG['test_db_port'],
                user=QT_CONFIG['test_db_user'],
                password=QT_CONFIG['test_db_password'],
                db_name=QT_CONFIG['test_db_name']
        )

    def test_multi_factors(self):
        """ test strategy MultiFactors"""
        self.assertEqual(1, 1)
        shares = qt.filter_stock_codes(index='000300.SH', date='20210101')
        print(len(shares), shares[:10])

        alpha = MultiFactors()
        op = qt.Operator(alpha,
                         signal_type='PT',
                         run_timing='close',  # 在周期结束（收盘）时运行
                         run_freq='M',  # 每月执一次选股（每周或每天都可以）
                         )

        op.set_blender("0.8*s0", 'Group_1')
        qt.run(op=op,
               mode=1,
               invest_start='20210101',
               invest_end='20220501',
               asset_type='E',
               invest_cash_amounts=[1000000],
               asset_pool=shares,
               trade_batch_size=100,
               sell_batch_size=1,
               trade_log=True)
        self.assertTrue(True)

    def test_index_enhancement(self):
        """ tests for self-defined strategies"""
        self.assertEqual(1, 1)

        shares = qt.filter_stock_codes(index='000300.SH', date='20210101')
        print(len(shares), shares[:10])

        alpha = IndexEnhancement()
        op = qt.Operator(alpha,
                         signal_type='PT',
                         run_timing='close',  # 在周期结束（收盘）时运行
                         run_freq='d',  # 每天执行一次选股
                         )

        op.op_type = 'stepwise'
        op.set_blender("0.8*s0", 'Group_1')
        qt.run(op=op,
               mode=1,
               invest_start='20210101',
               invest_end='20220501',
               asset_type='E',
               invest_cash_amounts=[1000000],
               asset_pool=shares,
               trade_batch_size=100,
               sell_batch_size=1,
               trade_log=True)
        self.assertTrue(True)

    def test_grid_trading(self):
        """ test for strategy GridTrading"""
        self.assertEqual(1, 1)

        alpha = GridTrading()
        op = qt.Operator(alpha, signal_type='PT', run_timing='close', run_freq='1min')

        op.set_blender("1.0*s0", 'Group_1')
        qt.run(
                op=op,
                mode=1,
                invest_start='20220401',
                invest_end='20220731',
                invest_cash_amounts=[1000000],
                asset_type='IDX',
                asset_pool=['000300.SH'],
                trade_batch_size=1.,
                sell_batch_size=1.,
                trade_log=True,
                allow_sell_short=True,
        )
        self.assertTrue(True)

    def test_get_history_data(self):
        """
        """
        # gen_value normal data
        data = qt.get_history_data(htypes='open, high, low, close',
                                   start='20211230',
                                   end='20220110',
                                   shares='000001.SZ',
                                   freq='d',
                                   adj='n',
                                   asset_type='E')
        print(f'not adjusted data: \n{data}')
        # gen_value back adjusted price for one asset
        data = qt.get_history_data(htypes='open, high, low, close',
                                   start='20211230',
                                   end='20220110',
                                   shares='000001.SZ',
                                   freq='d',
                                   adj='b',
                                   asset_type='E')
        print(f'adjusted data: \n{data}')
        # gen_value back adjusted prices for both assets mixed
        data = qt.get_history_data(htypes='open, high, low, close',
                                   start='20211230',
                                   end='20220110',
                                   shares='000001.SZ, 512000.SH',
                                   freq='d',
                                   adj='b',
                                   asset_type='E, FD')
        print(f'adjusted data: \n{data}')
        self.assertTrue(True)

    def test_fast_experiments(self):
        """ test for strategy GridTrading"""
        shares = qt.filter_stock_codes(index='000300.SH', date='20190501')  # 选择股票池，包括2019年5月以来所有沪深300指数成分股
        # 设置回测的运行参数
        # shares = ['000001.SZ', '000002.SZ', '000063.SZ', '000069.SZ', '000100.SZ', '000157.SZ', '000166.SZ',
        #           '000333.SZ', '000338.SZ', '000413.SZ', '000415.SZ', '000423.SZ', '000425.SZ', '000538.SZ',
        #           '000568.SZ', '000596.SZ', '000625.SZ', '000627.SZ', '000629.SZ', '000630.SZ', '000651.SZ',
        #           '000656.SZ', '000661.SZ', '000671.SZ', '000703.SZ', '000709.SZ', '000723.SZ', '000725.SZ',
        #           '000728.SZ', '000768.SZ', '000776.SZ', '000783.SZ', '000786.SZ', '000858.SZ', '000876.SZ',
        #           '000895.SZ', '000898.SZ', '000938.SZ', '000961.SZ', '000963.SZ', '001979.SZ', '002001.SZ',
        #           '002007.SZ', '002008.SZ', '002010.SZ', '002024.SZ', '002027.SZ', '002032.SZ', '002044.SZ',
        #           '002050.SZ']
        qt.configure(mode=1,  # mode=1表示回测模式
                     invest_start='20160405',  # 回测开始日期
                     invest_end='20210201',  # 回测结束日期
                     asset_type='E',  # 投资品种为股票
                     backtest_price_adj='b',
                     asset_pool=shares,  # shares包含同期沪深300指数的成份股
                     trade_batch_size=100,  # 买入批量为100股
                     sell_batch_size=1,  # 卖出批量为整数股
                     trade_log=True,  # 生成交易记录
                     )

        # alpha = MultiFactors()  # 生成一个交易策略的实例，名为alpha
        op = qt.Operator('dma', signal_type='PT',
                         run_timing='close',  # 在周期结束（收盘）时运行
                         run_freq='M',  # 每月执行一次选股（每周或每天都可以）
                         )  # 生成交易员对象，操作alpha策略，交易信号的类型为‘PT'，意思是生成的信号代表持仓比例，例如1代表100%持有股票，0.35表示持有股票占资产的35%
        # op.set_blender('1.0*s0')  # 交易策略混合方式，只有一个策略，不需要混合
        qt.run(op=op)  # 开始运行


if __name__ == '__main__':
    unittest.main()