text:
# coding=utf-8
# operator.py

# ======================================
# This file contains Operator class, that
# merges and applies investment strategies
# to generate operation signals with
# given history data.
# ======================================

import pandas as pd
import numpy as np
from .finance import CashPlan
from .history import HistoryPanel
from .utilfuncs import str_to_list
from .strategy import RollingTiming
from .strategy import SimpleSelecting
from .strategy import SimpleTiming
from .built_in import AVAILABLE_STRATEGIES, BUILT_IN_STRATEGY_DICT

from .utilfuncs import unify, mask_to_signal

# TODO：
# TODO：作为完整的交易信号，为了实现更加贴近实际的交易信号，交易信号应该包括交易方向和头寸位置两个主要参数（对于股票来说
# TODO：只有多头头寸）
# TODO：position > 0时，表示多头头寸
# TODO：position < 0时，表示空头头寸
# TODO：两种不同的头寸位置配合开仓（signal>0）或平仓（signal<0）才能完整地表示所有的交易方式
# TODO：另外，还需要加入更多交易相关信息，如限价单、市价单、交易数量等等，总之，之前仅用singal表示交易信号的方式太过于简单了

# TODO: A reference data type should affiliated to strategies,
# TODO: which is useful when a reference data is needed. for
# TODO: example, a relative change rate is based on the difference
# TODO: between the stock prices and reference prices.
class Operator:
    """交易操作生成类，通过简单工厂模式创建择时属性类和选股属性类，并根据这两个属性类的结果生成交易清单

    根据输入的参数生成Operator对象，在对象中创建相应的策略类型:

    input:
            :param selecting_types: 一个包含多个字符串的列表，表示不同选股策略，后续可以考虑把字符串列表改为逗号分隔的纯字符串输入
            :param timing_types: 字符串列表，表示不同择时策略，后续可以考虑把字符串列表改为逗号分隔的纯字符串输入
            :param ricon_types: 字符串列表，表示不同风控策略，后续可以考虑把字符串列表改为逗号分隔的纯字符串输入

        Operator对象其实就是若干个不同类型的操作策略的容器对象，
        在一个Operator对象中，可以包含任意多个"策略对象"，而运行Operator生成交易信号的过程，就是调用这些不同的交易策略，并通过
        不同的方法对这些交易策略的结果进行组合的过程

        目前在Operator对象中支持五种策略生成器，而每种类型的策略生成器用不同的方式生成策略相关信号，可以在Operator对象使用若干不同
        生成类型的策略。同时，Operator对象在生成最终交易信号时，成员策略会被用于三种用途，分别用于创建多空蒙板、选股蒙板以及交易信号矩阵
        最后将两种蒙板和一种信号矩阵组合生成最终的交易信号清单。

        在同一个Operator对象中，不同生成类型的策略可以被用于不同的用途，具体如下：

         usage \ generator | RollingTiming | SimpleSelecting | Simple_Timing | FactoralSelecting | ReferenceTiming |
         ==================|===============|=================|===============|===================|=================|
         long/short mask   |       Yes     |        Yes      |       Yes     |        Yes        |       Yes       |
         ------------------|---------------|-----------------|---------------|-------------------|-----------------|
         Portfolio mask    |       No      |        Yes      |       No      |        Yes        |       Yes       |
         ------------------|---------------|-----------------|---------------|-------------------|-----------------|
         signal matrix     |       Yes     |        No       |       Yes     |        No         |       Yes       |

        ==五种策略信号生成器==

        目前Operator支持五种不同生成类型的策略，它们并不仅局限于某一种用途，不同生成器之间的区别在于策略利用历史数据并生成最终结果的
        方法不一样。五种生成类型的策略分别如下：

            1,  RollingTiming 逐品种滚动时序信号生成器，用于生成择时信号的策略

                这类策略的共同特征是对投资组合中的所有投资产品逐个考察其历史数据，根据其历史数据，在历史数据的粒度上生成整个时间段上的
                时间序列信号。时间序列信号可以为多空信号，即用>0的数字表示多头头寸，<0的数字代表空头头寸，0代表中性头寸。也可以表示交
                易信号，即>0的数字表示建多仓或平空仓，<0的数字表示见空仓或平多仓。

                这种策略生成器将投资组合中的每一个投资产品逐个处理，每个投资产品中的NA值可以单独处理，与其他投资品种不相关、互不影响，
                同时，每个投资产品可以应用不同的参数生成信号，是最为灵活的择时信号生成器。

                另外，为了避免前视偏差，滚动择时策略仅利用一小段有限的历史数据（被称为时间窗口）来生成每一个时间点上的信号，同时确保
                时间窗口总是处于需要计算多空位置那一点的过去。这种技术称为"时间窗口滚动"。这样的滚动择时信号生成方法适用于未来数据会
                对当前的信号产生影响的情况下。采用滚动择时策略生成方法，可以确保每个时间点的交易信号只跟过去一段时间有关，从而彻底排除
                前视偏差可能给策略带来的影响。

                不过，由于时间窗口滚动的计算方式需要反复提取一段时间窗口内的数据，并反复计算，因此计算复杂度与数据总量M与时间窗口长度N
                的乘积M*N成正比，效率显著低于简单时序信号生成策略，因此，在可能的情况下（例如，简单移动平均值相关策略不受未来价格影响）
                应该尽量使用简单时序信号生成策略，以提升执行速度。

            2,  SimpleSelecting 简单投资组合分配器，用于周期性地调整投资组合中每个个股的权重比例

                这类策略的共同特征是周期性运行，且运行的周期与其历史数据的粒度不同。在每次运行时，根据其历史数据，为潜在投资组合中的每
                一个投资产品分配一个权重，并最终确保所有的权重值归一化。权重为0时表示该投资产品被从组合中剔除，而权重的大小则代表投资
                过程中分配投资资金的比例。

                这种方式生成的策略可以用于生成周期性选股蒙板，也可以用于生成周期性的多空信号模板。

                这种生成方式的策略是针对历史数据区间运行的，是运算复杂度最低的一类生成方式，对于数量超大的投资组合，可以采用这种方式生
                成投资策略。但仅仅局限于部分周期性运行的策略。

            3,  SimpleTiming 逐品种简单时序信号生成器，用于生成择时信号的策略

                这类策略的共同特征是对投资组合中的所有投资产品逐个考察其历史数据，并在历史数据的粒度上生成整个时间段上的时间序列信号。
                这种策略生成方法与逐品种滚动时序信号生成策略的信号产生方法类似，只是缺少了"滚动"的操作，时序信号是一次性在整个历史区间
                上生成的，并不考虑未来数据对当前信号的影响。这类方法生成的信号既可以代表多空信号，也可以代表交易信号。

                同时，简单时序信号生成器也保留了滚动时序信号生成器的灵活性优点：每个投资产品独立处理，不同数据的NA值互不关联，互不影响，
                同时每个不同的投资产品可以应用完全不同的策略参数。最大的区别就在于信号不是滚动生成的。

                正因为没有采用滚动计算的方式，因此简单时序信号生成器的计算复杂度只有O(M),与历史数据数量M成正比，远小于滚动生成器。
                不过，其风险在于策略本身是否受未来数据影响，如果策略本身不受未来数据的影响，则采用简单时序生成器是一种较优的选择，例如，
                基于移动平均线相交的相交线策略、基于过去N日股价变动的股价变动策略本身具备不受未来信息影响的特点，使用滚动时序生成器和
                简单时序生成器能得到相同的结果，而简单时序生成器的计算耗时大大低于滚动时序生成器，因此应该采用简单滚动生成器。又例如，
                基于指数平滑均线或加权平均线的策略，或基于波形降噪分析的策略，其输出信号受未来信息的影响，如果使用简单滚动生成器将会
                导致未来价格信息对回测信号产生影响，因此不应该使用简单时序信号生成器。

            4,  FactoralSelecting 因子选股投资组合分配器，用于周期性地调整投资组合中每个个股的权重比例

                这类策略的共同特征是周期性运行，且运行的周期与其历史数据的粒度不同。在每次运行时，根据其历史数据，为每一个股票计算一个
                选股因子，这个选股因子可以根据任意选定的数据根据任意可能的逻辑生成。生成选股因子后，可以通过对选股因子的条件筛选和
                排序执行选股操作。用户可以在策略属性层面定义筛选条件和排序方法，同时可以选择不同的选股权重分配方式

                这种方式生成的策略可以用于生成周期性选股蒙板，也可以用于生成周期性的多空信号模板。

                这种生成方式的策略是针对历史数据区间运行的，是运算复杂度最低的一类生成方式，对于数量超大的投资组合，可以采用这种方式生
                成投资策略。但仅仅局限于部分周期性运行的策略。

            5,  ReferenceTiming 参考数据信号生成器

                这类策略并不需要所选择股票本身的数据计算策略输出，而是利用参考数据例如大盘、宏观经济数据或其他数据来生成统一的股票多空
                或选股信号模版。其计算的基本方法与Timing类型生成器基本一致，但是同时针对所有的投资组合进行计算，因此信号可以用于多空
                蒙板和选股信号蒙本，计算的基础为参考信号

                这种方式生成的策略可以用于生成周期性选股蒙板，也可以用于生成周期性的多空信号模板，同时也可以用于直接生成交易信号。

        ==策略的三种用途==

        在Operator对象中，包含的策略可以有无限多个，但是Operator会将策略分别用于三种不同的用途，并把它们的输出混合成为最终的策略输出。

        Operator对象可以同时将多个策略用于同一个用途，不过，为了确保输出唯一，多个策略的输出将被以某种方式混合，混合的方式是Operator
        对象的属性，定义了同样用途的不同策略输出结果的混合方式，以下是三种用途及其混合方式的介绍：

            用途1,  生成多空蒙板：
                多空蒙板定义了整个历史时期内所有股票的多空仓位时间序列，如前所述，1代表多头满仓，0代表空仓，-1代表空头满仓。数值为1到
                -1之间的小数时，代表一定的仓位百分比。当在Operator对象中定义了多个策略用于生成多空蒙板时，Operator会分别使用策略生
                成数个蒙板，再将它们混合起来。

                多空蒙板的混合方式由多空混合字符串来确定，字符串的格式为"[chg|pos]-0/9|cumulative"(此处应该使用正则表达式)

                'str-T': T为浮点数，当多个策略多空蒙板的总体信号强度达到阈值T时，总体输出为1(或者-1)，否则为0
                'pos-N': N为正整数，取值区间为1到len(timing)的值，表示在N个策略为多时状态为多，否则为空
                    这种类型有一个变体：
                    'pos-N-T': T为信号强度阈值，忽略信号强度达不到该阈值的多空蒙板信号，将剩余的多空蒙板进行计数，信号数量达到或
                    超过N时，输出为1（或者-1），否则为0
                'avg': 平均信号强度，所有多空蒙板的信号强度的平均值
                'comboo': 在每个策略发生反转时都会产生交易信号，信号的强度不经过衰减，但是通常第一个信号产生后，后续信号就再无意义

            用途2,  生成选股蒙板：
                选股蒙板定义了每一个时刻整个投资组合中每一个投资产品被分配到的权重。同样，如果定义了多个策略，也必须将它们的输出结果混
                合成一个

                选股蒙板的混合方式由一个逻辑表达式来确定，例如'0 and (1 or 2 and (3 or 4))'
                上面的表达式表示了如何根据五个选股蒙板来确定一个个股是否被选中而且权重多大。在目前的系统中，qteasy只是简单地将and运算
                处理为乘法，or运算处理为加法。在只有0和1的情况下这样做是没有问题的，但是在普遍蒙板中存在大量介于0和1之间的浮点数的
                时候，就需要注意了，如果蒙板0中某个股票的权重为0.5,在蒙板1中的权重为0.5，那么0 and 1 与0 or 1的结果分别应该是什么？

                目前这个问题的解决方式是：
                    0.5 and 0.5 = 0.5 * 0.5 = 0.25,
                    0.5 or 0.5 = 0.5 + 0.5 = 1
                完成上述计算后重新unify整个蒙板

                想到还有另一种解决方式：
                    0.5 and 0.5 = 0.5 * 0.5 = 0.25,
                    0.5 or 0.5 = 1 - (1 - 0.5) * (1 - 0.5) = 0.75
                同样在完成上述计算后unify整个蒙板

                孰优孰劣，还需要观察和试验，但是现在先把后一种方式写入代码中，后续再进行验证

            用途3:  生成交易信号矩阵：
                交易信号矩阵是由策略直接生成的交易信号组成的，如前所述，1代表开多仓或平空仓，-1代表平多仓或开空仓。与其他用途一样，如果
                多个策略被用于同样的用途，应该把多个策略的输出混合成一个最终输出。

                交易信号矩阵的混合方式跟多空蒙板的混合方式相似，以混合字符串定义。混合字符串的格式为"[chg|pos]-0/9|cumulative"

                'chg-N': N为正整数，取值区间为1到len(timing)的值，表示多空状态在第N次信号反转时反转
                'pos-N': N为正整数，取值区间为1到len(timing)的值，表示在N个策略为多时状态为多，否则为空
                'cumulative': 在每个策略发生反转时都会产生交易信号，但是信号强度为1/len(timing)

        ==策略的组合==

        以上三类策略通过不同的方式混合后，可以任意组合一种复杂的策略，因此，在qteasy系统中，复杂的策略是可以由很多个简单的策略组合而来
        的。
        在一个Operator对象中，作为容器可以容纳任意多个任意类型的策略，所有的策略以用途分成三组，所有的策略可以引用不同的历史数据，生成
        同样大小尺度的结果（也就是说，生成的结果有相同的历史区间，相同的时间粒度），最后这些结果被通过某种方法"混合"起来，形成每个用途
        的最终的结果，即多空模版、选股蒙板以及交易信号矩阵。
        三种用途的结果又再次被组合起来，变成整个Operator对象的输出。

        目前采用的组合方式是：

        mask_to_signal(多空模版 * 选股蒙板) + 交易信号
        其中mask_to_signal()函数的作用是将蒙板转化为交易信号，这样输出的就是交易信号

        未来将第三类策略升级为单品种信号生成策略后，信号的组合方式就可以变为：

        mask_to_signal(多空蒙板 * 选股蒙板）+ (交易信号 * 选股蒙板)
        这样同样可以输出一组交易信号

        但这样做还会有问题，预先生成交易信号在交易过程中存在MOQ时可能会发生这样的情况，在试图分多次建仓买入股票时，由于股票价值较高，导致
        分批建仓的信号每次都无法买入，解决的思路有两个，第一是在回测时不仅接受交易信号，还接受目标仓位，那么如果第一次建仓不够买入一手
        股票，到后续的目标仓位时总能有足够资金建仓。第二种是修改回测程序，在每次操作后记录理论操作数量和实际操作数量的差值，等下次有同方向
        操作的时候补齐差额。孰优孰劣？目前还没有想清楚。

    """

    # 对象初始化时需要给定对象中包含的选股、择时、风控组件的类型列表

    AVAILABLE_LS_BLENDER_TYPES = ['avg', 'avg_pos', 'pos', 'str', 'combo', 'none']

    def __init__(self, selecting_types=None,
                 timing_types=None,
                 ricon_types=None):
        """根据生成具体的对象

        input:
            :param selecting_types: 一个包含多个字符串的列表，表示不同选股策略，后续可以考虑把字符串列表改为逗号分隔的纯字符串输入
            :param timing_types: 字符串列表，表示不同择时策略，后续可以考虑把字符串列表改为逗号分隔的纯字符串输入
            :param ricon_types: 字符串列表，表示不同风控策略，后续可以考虑把字符串列表改为逗号分隔的纯字符串输入
        """
        # 对象属性：
        # 交易信号通用属性：

        # 如果对象的种类未在参数中给出，则直接指定最简单的策略种类
        if selecting_types is None:
            selecting_types = ['all']
        if isinstance(selecting_types, str):
            selecting_types = str_to_list(selecting_types)
        if timing_types is None:
            timing_types = ['long']
        if isinstance(timing_types, str):
            timing_types = str_to_list(timing_types)
        if ricon_types is None:
            ricon_types = ['ricon_none']
        if isinstance(ricon_types, str):
            ricon_types = str_to_list(ricon_types)
        # 在Operator对象中，对每一种类型的策略，需要三个列表对象加上一个字符串作为基本数据结构，存储相关信息：
        # 对于每一类型的策略，第一个列表是_stg_types， 按照顺序保存所有相关策略对象的种类字符串，如['MACD', 'DMA', 'MACD']
        # 第二个列表是_stg，按照顺序存储所有的策略对象，如[Timing(MACD), Timing(timing_DMA), Timing(MACD)]
        # 第三个列表是_stg_history_data, 列表中按照顺序保存用于不同策略的历史数据切片，格式均为np.ndarray，维度为三维
        # 字符串则是"混合"字符串，代表最终将所有的同类策略混合到一起的方式，对于不同类型的策略，混合方式定义不同
        # 以上的数据结构对于所有类型的策略都基本相同
        self._timing_types = []
        self._timing = []
        self._timing_history_data = []
        self._ls_blender = 'pos-1'  # 默认的择时策略混合方式
        for timing_type in timing_types:
            # 通过字符串比较确认timing_type的输入参数来生成不同的具体择时策略对象，使用.lower()转化为全小写字母
            if isinstance(timing_type, str):
                if timing_type.lower() not in AVAILABLE_STRATEGIES:
                    raise KeyError(f'built-in timing strategy \'{timing_type}\' not found!')
                self._timing_types.append(timing_type)
                self._timing.append(BUILT_IN_STRATEGY_DICT[timing_type]())
            # 当传入的对象是一个strategy时，直接
            elif isinstance(timing_type, (RollingTiming, SimpleTiming)):
                self._timing_types.append(timing_type.stg_type)
                self._timing.append(timing_type)
            else:
                raise TypeError(f'The timing strategy type \'{type(timing_type)}\' is not supported!')
        # 根据输入参数创建不同的具体选股策略对象。selecting_types及selectings属性与择时策略对象属性相似
        # 都是列表，包含若干相互独立的选股策略（至少一个）
        self._selecting_type = []
        self._selecting = []
        self._selecting_history_data = []
        # 生成选股蒙板生成策略清单
        cur_type = 0
        str_list = []

        for selecting_type in selecting_types:
            if cur_type == 0:
                str_list.append(str(cur_type))
            else:
                str_list.append(f' or {str(cur_type)}')
            cur_type += 1
            if isinstance(selecting_type, str):
                if selecting_type.lower() not in AVAILABLE_STRATEGIES:
                    raise KeyError(f'KeyError: built-in selecting type \'{selecting_type}\' not found!')
                self._selecting_type.append(selecting_type)
                self._selecting.append(BUILT_IN_STRATEGY_DICT[selecting_type]())
            elif isinstance(selecting_type, (SimpleSelecting, SimpleTiming)):
                self._selecting_type.append(selecting_type.stg_type)
                self._selecting.append(selecting_type)
            else:
                raise TypeError(f'Type Error, the type of object {type(selecting_type)} is not supported!')
        self._selecting_blender_string = ''.join(str_list)
        # create selecting blender by selecting blender string
        self._selecting_blender = self._exp_to_blender

        # 根据输入参数生成不同的风险控制策略对象
        self._ricon_type = []
        self._ricon = []
        self._ricon_history_data = []
        self._ricon_blender = 'add'
        for ricon_type in ricon_types:
            if isinstance(ricon_type, str):
                if ricon_type.lower() not in AVAILABLE_STRATEGIES:
                    raise KeyError(f'ricon type {ricon_type} not available!')
                self._ricon_type.append(ricon_type)
                self._ricon.append(BUILT_IN_STRATEGY_DICT[ricon_type]())
            elif isinstance(ricon_type, (RollingTiming, SimpleTiming)):
                self._ricon_type.append(ricon_type.stg_type)
                self._ricon.append(ricon_type)
            else:
                raise TypeError(f'Type Error, the type of passed object {type(ricon_type)} is not supported!')

    @property
    def timing(self):
        """返回operator对象的所有timing对象"""
        return self._timing

    @property
    def timing_count(self):
        """返回operator对象中的所有timing对象的数量"""
        return len(self.timing)

    @property
    def ls_blender(self):
        """返回operator对象中的多空蒙板混合器"""
        return self._ls_blender

    @property
    def selecting(self):
        """返回operator对象的所有selecting对象"""
        return self._selecting

    @property
    def selecting_count(self):
        """返回operator对象的所有selecting对象的数量"""
        return len(self.selecting)

    @property
    def selecting_blender(self):
        """返回operator对象的所有选股策略的选股结果混合器"""
        return self._selecting_blender_string

    @property
    def selecting_blender_expr(self):
        """返回operator对象的所有选股策略的选股结果的混合器表达式"""
        return self._selecting_blender

    @property
    def ricon(self):
        """返回operator对象的所有ricon对象"""
        return self._ricon

    @property
    def ricon_count(self):
        """返回operator对象的所有ricon对象的数量"""
        return len(self.ricon)

    @property
    def ricon_blender(self):
        """返回operator对象的所有ricon对象的混合器"""
        return self._ricon_blender

    @property
    def strategies(self):
        """返回operator对象的所有策略子对象"""
        stg = [item for item in self.timing + self.selecting + self.ricon]
        return stg

    @property
    def strategy_count(self):
        """返回operator对象的所有策略的数量"""
        return len(self.strategies)

    @property
    def strategy_benders(self):
        return [self.ls_blender, self.selecting_blender, self.ricon_blender]

    @property
    def op_data_types(self):
        """返回operator对象所有策略子对象所需数据类型的集合"""

        d_types = [typ for item in self.strategies for typ in item.data_types]
        d_types = list(set(d_types))
        d_types.sort()
        return d_types

    @property
    def op_data_freq(self):
        """返回operator对象所有策略子对象所需数据的采样频率"""
        d_freq = [stg.data_freq for stg in self.strategies]
        d_freq = list(set(d_freq))
        assert len(d_freq) == 1, f'ValueError, there are multiple history data frequency required by strategies'
        return d_freq[0]

    @property
    def opt_space_par(self):
        """一次性返回operator对象中所有参加优化（opt_tag != 0）的子策略的参数空间Space信息
            改属性的返回值是一个元组，包含ranges, types两个列表，这两个列表正好可以直接用作Space对象的创建参数，用于创建一个合适的
            Space对象

            一个完整的投资策略由三类多个不同的子策略组成。每个子策略都有自己特定的参数空间，它们的参数空间信息存储在stg.par_boes以及
            stg.par_types等属性中。通常，我们在优化参数的过程中，希望仅对部分策略的参数空间进行搜索，而其他的策略保持参数不变。为了实现
            这样的控制，我们可以给每一个子策略一个属性opt_tag优化标签，通过设置这个标签的值，可以控制这个子策略是否参与优化：
            {0: 'No optimization, 不参与优化，这个子策略在整个优化过程中将始终使用同一组参数',
             1: 'Normal optimization, 普通优化，这个子策略在优化过程中使用不同的参数，这些参数都是从它的参数空间中按照规律取出的',
             2: 'enumerate optimization, 枚举优化，在优化过程中使用不同的参数，但并非取自其参数空间，而是来自一个预设的枚举对象'}

         这个函数遍历operator对象中所有子策略，根据其优化类型选择它的参数空间信息，组合成一个多维向量用于创建可以用于生成所有相关
         策略的参数的高维空间

         return: ranges, types
        """
        ranges = []
        types = []
        for stg in self.strategies:
            if stg.opt_tag == 0:
                pass  # 策略参数不参与优化
            elif stg.opt_tag == 1:
                # 所有的策略参数全部参与优化，且策略的每一个参数作为一个个体参与优化
                ranges.extend(stg.par_boes)
                types.extend(stg.par_types)
            elif stg.opt_tag == 2:
                # 所有的策略参数全部参与优化，但策略的所有参数组合作为枚举同时参与优化
                ranges.append(stg.par_boes)
                types.extend(['enum'])
        return ranges, types

    @property
    def opt_types(self):
        """返回所有策略的优化类型标签"""
        return [stg.opt_tag for stg in self.strategies]

    @property
    def max_window_length(self):
        """ 计算并返回operator对象所有子策略中最长的策略形成期。在准备回测或优化历史数据时，以此确保有足够的历史数据供策略形成

        :return: int
        """
        return max(stg.window_length for stg in self.strategies)

    @property
    def _exp_to_blender(self):
        """选股策略混合表达式解析程序，将通常的中缀表达式解析为前缀运算队列，从而便于混合程序直接调用

        系统接受的合法表达式为包含 '*' 与 '+' 的中缀表达式，符合人类的思维习惯，使用括号来实现强制
        优先计算，如 '0 + (1 + 2) * 3'; 表达式中的数字0～3代表选股策略列表中的不同策略的索引号
        上述表达式虽然便于人类理解，但是不利于快速计算，因此需要转化为前缀表达式，其优势是没有括号
        按照顺序读取并直接计算，便于程序的运行。为了节省系统运行开销，在给出混合表达式的时候直接将它
        转化为前缀表达式的形式并直接存储在blender列表中，在混合时直接调用并计算即可
        input： =====
            no input parameter
        return：===== s2: 前缀表达式
            :rtype: list: 前缀表达式
        """
        # TODO: extract expression with re module
        prio = {'or' : 0,
                'and': 1}
        # 定义两个队列作为操作堆栈
        s1 = []  # 运算符栈
        s2 = []  # 结果栈
        # 读取字符串并读取字符串中的各个元素（操作数和操作符），当前使用str对象的split()方法进行，要
        # 求字符串中个元素之间用空格或其他符号分割，应该考虑写一个self.__split()方法，不依赖空格对
        # 字符串进行分割
        exp_list = self._selecting_blender_string.split()
        # 使用list()的问题是，必须确保表达式中不存在超过一位数字的数，如12等，同时需要去除字符串中的特殊字符如空格等
        # exp_list = list(self._selecting_blender_string.
        #                 replace(' ', '').
        #                 replace('_', '').
        #                 replace('.', '').
        #                 replace('-', ''))
        # 开始循环读取所有操作元素
        while exp_list:
            s = exp_list.pop()
            # 从右至左逐个读取表达式中的元素（数字或操作符）
            # 并按照以下算法处理
            if s.isdigit():
                # 1，如果元素是数字则压入结果栈
                s2.append(s)

            elif s == ')':
                # 2，如果元素是反括号则压入运算符栈
                s1.append(s)
            elif s == '(':
                # 3，扫描到（时，依次弹出所有运算符直到遇到），并把该）弹出
                while s1[-1] != ')':
                    s2.append(s1.pop())
                s1.pop()
            elif s in prio.keys():
                # 4，扫描到运算符时：
                if s1 == [] or s1[-1] == ')' or prio[s] >= prio[s1[-1]]:
                    # 如果这三种情况则直接入栈
                    s1.append(s)
                else:
                    # 否则就弹出s1中的符号压入s2，并将元素放回队列
                    s2.append(s1.pop())
                    exp_list.append(s)
            else:
                raise ValueError(f'unidentified characters found in blender string: \'{s}\'')
        while s1:
            s2.append(s1.pop())
        s2.reverse()  # 表达式解析完成，生成前缀表达式
        return s2

    @property
    def ready(self):
        """ assess if the operator is ready to generate

        :return:
        """
        raise NotImplementedError

    def add_strategy(self, stg, usage):
        """add strategy"""
        raise NotImplementedError

    def remove_strategy(self, stg):
        """remove strategy"""
        raise NotImplementedError

    def clear(self):
        """clear all strategies

        :return:
        """
        raise NotImplementedError

    def set_opt_par(self, opt_par):
        """optimizer接口函数，将输入的opt参数切片后传入stg的参数中

        :param opt_par:
            :type opt_par:Tuple
            一组参数，可能包含多个策略的参数，在这里被分配到不同的策略中

        :return
            None

        本函数与set_parameter()不同，在优化过程中非常有用，可以同时将参数设置到几个不同的策略中去，只要这个策略的opt_tag不为零
        在一个包含多个Strategy的Operator中，可能同时有几个不同的strategy需要寻优。这时，为了寻找最优解，需要建立一个Space，包含需要寻优的
        几个strategy的所有参数空间。通过这个space生成参数空间后，空间中的每一个向量实际上包含了不同的策略的参数，因此需要将他们原样分配到不
        同的策略中。

        举例如下：

            一个Operator对象有三个strategy，分别需要2， 3， 3个参数，而其中第一和第三个策略需要参数寻优，这个operator的所有策略参数可以写
            成一个2+3+2维向量，其中下划线的几个是需要寻优的策略的参数：
                     stg1:   stg2:       stg3:
                     tag=1   tag=0       tag=1
                    [p0, p1, p2, p3, p4, p5, p6, p7]
                     ==  ==              ==  ==  ==
            为了寻优方便，可以建立一个五维参数空间，用于生成五维参数向量：
                    [v0, v1, v2, v3, v4]
            set_opt_par函数遍历Operator对象中的所有strategy函数，检查它的opt_tag值，只要发现opt_tag不为0，则将相应的参数赋值给strategy：
                     stg1:   stg2:       stg3:
                     tag=1   tag=0       tag=1
                    [p0, p1, p2, p3, p4, p5, p6, p7]
                     ==  ==              ==  ==  ==
                    [v0, v1]            [v2, v3, v4]

            在另一种情况下，一个策略的参数本身就以一个tuple的形式给出，一系列的合法参数组以enum的形式形成一个完整的参数空间，这时，opt_tag被
            设置为2，此时参数空间中每个向量的一个分量就包含了完整的参数信息，例子如下：

            一个Operator对象有三个strategy，分别需要2， 3， 3个参数，而其中第一和第三个策略需要参数寻优，这个operator的所有策略参数可以写
            成一个2+3+2维向量，其中下划线的几个是需要寻优的策略的参数，注意其中stg3的opt_tag被设置为2：
                     stg1:   stg2:       stg3:
                     tag=1   tag=0       tag=2
                    [p0, p1, p2, p3, p4, p5, p6, p7]
                     ==  ==              ==  ==  ==
            为了寻优建立一个3维参数空间，用于生成五维参数向量：
                    [v0, v1, v2]，其中v2 = (i0, i1, i2)
            set_opt_par函数遍历Operator对象中的所有strategy函数，检查它的opt_tag值，对于opt_tag==2的策略，则分配参数给这个策略
                     stg1:   stg2:       stg3:
                     tag=1   tag=0       tag=2
                    [p0, p1, p2, p3, p4, p5, p6, p7]
                     ==  ==              ==  ==  ==
                    [v0, v1]         v2=[i0, i1, i2]
        """
        s = 0
        k = 0
        # 依次遍历operator对象中的所有策略：
        for stg in self.strategies:
            # 优化标记为0：该策略的所有参数在优化中不发生变化
            if stg.opt_tag == 0:
                pass
            # 优化标记为1：该策略参与优化，用于优化的参数组的类型为上下界
            elif stg.opt_tag == 1:
                k += stg.par_count
                stg.set_pars(opt_par[s:k])
                s = k
            # 优化标记为2：该策略参与优化，用于优化的参数组的类型为枚举
            elif stg.opt_tag == 2:
                # 在这种情况下，只需要取出参数向量中的一个分量，赋值给策略作为参数即可。因为这一个分量就包含了完整的策略参数tuple
                k += 1
                stg.set_pars(opt_par[s])
                s = k

    # TODO: 完善本函数的Docstring，添加详细的使用介绍和示例
    def set_blender(self, blender_type, *args, **kwargs):
        """ 统一的blender混合器属性设置入口

        :param blender_type:
            :type blender_type: str, 一个字符串，用于指定被设置属性的混合器的类型，接收到可识别
            的输入后，调用不同的混合器设置函数，设置混合器

        :return
            None

        """
        if isinstance(blender_type, str):
            if blender_type.lower() in ['selecting', 'sel', 'select']:
                self._set_selecting_blender(*args, **kwargs)
            elif blender_type.lower() in ['ls', 'longshort', 'long_short']:
                self._set_ls_blender(*args, **kwargs)
            elif blender_type.lower() in ['ricon']:
                self._set_ricon_blender(*args, **kwargs)
            else:
                raise ValueError(f'wrong input! \'{blender_type}\' is not a valid input, '
                                 f'choose from [\'selecting\', \'sel\', \'ls\', \'ricon\']')
        else:
            raise TypeError(f'blender_type should be a string, got {type(blender_type)} instead')
        pass

    def set_parameter(self,
                      stg_id: str,
                      pars: [tuple, dict] = None,
                      opt_tag: int = None,
                      par_boes: [tuple, list] = None,
                      par_types: [list, str] = None,
                      sample_freq: str = None,
                      window_length: int = None,
                      data_types: [str, list] = None,
                      **kwargs):
        """ 统一的策略参数设置入口，stg_id标识接受参数的具体成员策略
            stg_id的格式为'x-n'，其中x为's/t/r'中的一个字母，n为一个整数

            这里应该有本函数的详细介绍

            :param stg_id:
                :type stg_id: str, 策略ID字符串，格式为x-N，表示第N个x类策略，x的取值范围为{'s', 't', 'r'},分别表示选股、择时和风控策略

            :param pars:
                :type pars: tuple or dict , 需要设置的策略参数，格式为tuple

            :param opt_tag:
                :type opt_tag: int, 优化类型，0: 不参加优化，1: 参加优化, 2: 以枚举类型参加优化

            :param par_boes:
                :type par_boes: tuple or list, 策略取值范围列表,一个包含若干tuple的列表,代表参数中一个元素的取值范围，如
                [(0, 1), (0, 100), (0, 100)]

            :param par_types:
                :type par_types: str or list, 策略参数类型列表，与par_boes配合确定策略参数取值范围类型，详情参见Space类的介绍

            :param sample_freq:
                :type sample_freq: str, 采样频率，策略运行时的采样频率

            :param window_length:
                :type window_length: int, 窗口长度：策略计算的前视窗口长度

            :param data_types:
                :type data_types: str or list, 策略计算所需历史数据的数据类型

            :return:
        """
        assert isinstance(stg_id, str), f'stg_id should be a string like \'t-0\', got {stg_id} instead'
        l = stg_id.split('-')
        assert len(l) == 2 and l[1].isdigit(), f'stg_id should be a string like \'t-0\', got {stg_id} instead'
        if l[0].lower() == 's':
            assert int(l[1]) < self.selecting_count, \
                f'ValueError: trying to set parameter for the {int(l[1]) + 1}-th selecting strategy but there\'s only' \
                f' {self.selecting_count} selecting strategy(s)'
            strategy = self.selecting[int(l[1])]
        elif l[0].lower() == 't':
            assert int(l[1]) < self.timing_count, \
                f'ValueError: trying to set parameter for the {int(l[1]) + 1}-th timing strategy but there\'s only' \
                f' {self.timing_count} timing strategies'
            strategy = self.timing[int(l[1])]
        elif l[0].lower() == 'r':
            assert int(l[1]) < self.ricon_count, \
                f'ValueError: trying to set parameter for the {int(l[1]) + 1}-th ricon strategy but there\'s only ' \
                f'{self.ricon_count} ricon strategies'
            strategy = self.ricon[int(l[1])]
        else:
            raise ValueError(f'The identifier of strategy is not recognized, should be like \'t-0\', got {stg_id}')
        if pars is not None:
            if strategy.set_pars(pars):
                pass
            else:
                raise ValueError(f'parameter setting error')
        if opt_tag is not None:
            strategy.set_opt_tag(opt_tag)
        if par_boes is not None:
            strategy.set_par_boes(par_boes)
        if par_types is not None:
            strategy.par_types = par_types
        has_sf = sample_freq is not None
        has_wl = window_length is not None
        has_dt = data_types is not None
        if has_sf or has_wl or has_dt:
            strategy.set_hist_pars(sample_freq=sample_freq,
                                   window_length=window_length,
                                   data_types=data_types)
        # set up additional properties of the class if they do exist:
        strategy.set_custom_pars(**kwargs)

    # =================================================
    # 下面是Operation模块的公有方法：
    def info(self, verbose = False):
        """ 打印出当前交易操作对象的信息，包括选股、择时策略的类型，策略混合方法、风险控制策略类型等等信息
            如果策略包含更多的信息，还会打印出策略的一些具体信息，如选股策略的信息等
            在这里调用了私有属性对象的私有属性，不应该直接调用，应该通过私有属性的公有方法打印相关信息
            首先打印Operation木块本身的信息
            :type verbose: bool

        """
        print('OPERATION MODULE INFO:')
        print('=' * 25)
        print('Information of the Module')
        print('=' * 25)
        # 打印各个子模块的信息：
        # 首先打印Selecting模块的信息
        print('Total count of SimpleSelecting strategies:', len(self._selecting))
        print('the blend type of selecting strategies is', self._selecting_blender_string)
        print('Parameters of SimpleSelecting Strategies:')
        for sel in self.selecting:
            sel.info()
        print('=' * 25)

        # 接着打印 timing模块的信息
        print('Total count of timing strategies:', len(self._timing))
        print('The blend type of timing strategies is', self._ls_blender)
        print('Parameters of timing Strategies:')
        for tmg in self.timing:
            tmg.info()
        print('=' * 25)

        # 最后打印Ricon模块的信息
        print('Total count of Risk Control strategies:', len(self._ricon))
        print('The blend type of Risk Control strategies is', self._ricon_blender)
        for ric in self.ricon:
            ric.info()
        print('=' * 25)

    # TODO 临时性使用cashplan作为参数之一，理想中应该只用一个"start_date"即可，这个Start_date可以在core.run()中具体指定，因为
    # TODO 在不同的运行模式下，start_date可能来源是不同的：
    def prepare_data(self, hist_data: HistoryPanel, cash_plan: CashPlan):
        """ 在create_signal之前准备好相关历史数据，检查历史数据是否符合所有策略的要求：

            检查hist_data历史数据的类型正确；
            检查cash_plan投资计划的类型正确；
            检查hist_data是否为空（要求不为空）；
            在hist_data中找到cash_plan投资计划中投资时间点的具体位置
            检查cash_plan投资计划中的每个投资时间点均有价格数据，也就是说，投资时间点都在交易日内
            检查cash_plan投资计划中第一次投资时间点前有足够的数据量，用于滚动回测
            检查cash_plan投资计划中最后一次投资时间点在历史数据的范围内
            从hist_data中根据各个量化策略的参数选取正确的切片放入各个策略数据仓库中

            然后，根据operator对象中的不同策略所需的数据类型，将hist_data数据仓库中的相应历史数据
            切片后保存到operator的各个策略历史数据属性中，供operator调用生成交易清单。

        :param hist_data:
            :type hist_data: HistoryPanel
            历史数据,一个HistoryPanel对象，应该包含operator对象中的所有策略运行所需的历史数据，包含所有
            个股所有类型的数据，例如，operator对象中存在两个交易策略，分别需要的数据类型如下：
                策略        所需数据类型
                ------------------------------
                策略A:   close, open, high
                策略B:   close, eps

            hist_data中就应该包含close、open、high、eps四种类型的数据
            数据覆盖的时间段和时间频率也必须符合上述要求

        :param cash_plan:
            :type cash_plan: CashPlan
            一个投资计划，临时性加入，在这里仅检查CashPlan与历史数据区间是否吻合，是否会产生数据量不够的问题

        :return:
            None
        """
        # 确保输入的历史数据是HistoryPanel类型
        if not isinstance(hist_data, HistoryPanel):
            raise TypeError(f'Historical data should be HistoryPanel, got {type(hist_data)}')
        # TODO: 临时性处理方式
        # 确保cash_plan的数据类型正确
        if not isinstance(cash_plan, CashPlan):
            raise TypeError(f'cash plan should be CashPlan object, got {type(cash_plan)}')
        # 确保输入的历史数据不为空
        if hist_data.is_empty:
            raise ValueError(f'history data can not be empty!')
        # 默认截取部分历史数据，截取的起点是cash_plan的第一个投资日，在历史数据序列中找到正确的对应位置
        first_cash_pos = np.searchsorted(hist_data.hdates, cash_plan.first_day)
        last_cash_pos = np.searchsorted(hist_data.hdates, cash_plan.last_day)
        # 确保回测操作的起点前面有足够的数据用于满足回测窗口的要求
        # TODO: 这里应该提高容错度，设置更好的回测历史区间设置方法，尽量使用户通过较少的参数设置就能完成基
        # TODO: 本的运行，不用过分强求参数之间的关系完美无缺，如果用户输入的参数之间有冲突，根据优先级调整
        # TODO: 相关参数即可，毋须责备求全。
        # TODO: 当运行模式为0时，不需要判断cash_pos与max_window_length的关系
        assert first_cash_pos >= self.max_window_length, \
            f'InputError, History data starts on {hist_data.hdates[0]} does not have enough data to cover' \
            f' first cash date {cash_plan.first_day}, ' \
            f'expect {self.max_window_length} cycles, got {first_cash_pos} records only'
        # 确保最后一个投资日也在输入的历史数据范围内
        # TODO: 这里应该提高容错度，如果某个投资日超出了历史数据范围，可以丢弃该笔投资，仅输出警告信息即可
        # TODO: 没必要过度要求用户的输入完美无缺。
        assert last_cash_pos < len(hist_data.hdates), \
            f'InputError, Not enough history data record to cover complete investment plan, history data ends ' \
            f'on {hist_data.hdates[-1]}, last investment on {cash_plan.last_day}'
        # 确认cash_plan的所有投资时间点都在价格清单中能找到（代表每个投资时间点都是交易日）
        invest_dates_in_hist = [invest_date in hist_data.hdates for invest_date in cash_plan.dates]
        if not all(invest_dates_in_hist):
            np_dates_in_hist = np.array(invest_dates_in_hist)
            where_not_in = [cash_plan.dates[i] for i in list(np.where(np_dates_in_hist == False)[0])]
            raise ValueError(f'Cash investment should be on trading days, '
                             f'following dates are not valid!\n{where_not_in}')
        # 确保op的策略都设置了参数
        assert all(stg.has_pars for stg in self.strategies),\
            f'One or more strategies has no parameter set properly!'
        # 确保op的策略都设置了混合方式
        assert self.selecting_blender != ''
        assert self.ls_blender != ''
        assert self.ricon_blender != ''
        # 使用循环方式，将相应的数据切片与不同的交易策略关联起来
        self._selecting_history_data = [hist_data[stg.data_types, :, (first_cash_pos - stg.window_length):]
                                        for stg in self.selecting]
        # 用于择时仓位策略的数据需要包含足够的数据窗口用于滚动计算
        self._timing_history_data = [hist_data[stg.data_types, :, (first_cash_pos - stg.window_length):]
                                     for stg in self.timing]
        self._ricon_history_data = [hist_data[stg.data_types, :, (first_cash_pos - stg.window_length):]
                                    for stg in self.ricon]

    # TODO: 需要改进：
    # TODO: 供回测或实盘交易的交易信号应该转化为交易订单，并支持期货交易，因此生成的交易订单应该包含四类：
    # TODO: 1，Buy-开多仓，2，sell-平多仓，3，sel_short-开空仓，4，buy_to_cover-平空仓
    # TODO: 应该创建标准的交易订单模式，并且通过一个函数把交易信号转化为交易订单，以供回测或实盘交易使用

    # TODO: 交易信号生成和回测模块需改进：
    # TODO: 在交易信号生成模块不再仅仅生成+1/-1交易信号，而是同时生成交易信号
    # TODO: 和多空目标位置，这样至少可以避免以下问题：当MOQ存在时，在逐步减仓的情况下，每次加仓的交易信号强度可能
    # TODO: 都不足以买到一手股票，那么所有逐步加仓的信号都会失效。
    # TODO: 另外，将交易信号和仓位信号分开也能更好地支持交易信号型策略和仓位变化型策略

    # TODO: 需要调查：
    # TODO: 为什么在已经通过prepare_data()方法设置好了每个不同策略所需的历史数据之后，在create_signal()方法中还需要传入
    # TODO: hist_data作为参数？这个参数现在已经没什么用了，完全可以拿掉。在sel策略的generate方法中也不应该
    # TODO: 需要传入shares和dates作为参数。只需要selecting_history_data中的一部分就可以了
    def create_signal(self, hist_data: HistoryPanel):
        """ 操作信号生成方法，在输入的历史数据上分别应用选股策略、择时策略和风险控制策略，生成初步交易信号后，

        对信号进行合法性处理，最终生成合法交易信号
        input:
        :param hist_data:
            :type hist_data: HistoryPanel
            从数据仓库中导出的历史数据，包含多只股票在一定时期内特定频率的一组或多组数据
            ！！但是！！
            作为参数传入的这组历史数据并不会被直接用于交易信号的生成，用于生成交易信号的历史数据
            存储在operator对象的下面三个属性中，在生成交易信号时直接调用，避免了每次生成交易信号
            时再动态分配历史数据。
                self._selecting_history_data
                self._timing_history_data
                self._ricon_history_data

        :return=====
            list
            使用对象的策略在历史数据期间的一个子集上产生的所有合法交易信号，该信号可以输出到回测
            模块进行回测和评价分析，也可以输出到实盘操作模块触发交易操作
        """
        # 第一步，在历史数据上分别使用选股策略独立产生若干选股蒙板（sel_mask）
        # 选股策略的所有参数都通过对象属性设置，因此在这里不需要传递任何参数
        # 生成空的选股蒙板

        # 确保输入历史数据的数据格式正确；并确保择时策略和风控策略都已经关联号相应的历史数据
        assert isinstance(hist_data, HistoryPanel), \
            f'Type Error: historical data should be HistoryPanel, got {type(hist_data)}'
        assert len(self._timing_history_data) > 0, \
            f'ObjectSetupError: history data should be set before signal creation!'
        assert len(self._ricon_history_data) > 0, \
            f'ObjectSetupError: history data should be set before signal creation!'
        sel_masks = []
        shares = hist_data.shares
        date_list = hist_data.hdates
        for sel, dt in zip(self._selecting, self._selecting_history_data):  # 依次使用选股策略队列中的所有策略逐个生成选股蒙板
            # TODO: 目前选股蒙板的输入参数还比较复杂，包括shares和dates两个参数，应该消除掉这两个参数，使
            # TODO: sel.generate()函数的signature与tmg.generate()和ricon.generate()一致
            history_length = dt.shape[1]
            sel_masks.append(
                    sel.generate(hist_data=dt, shares=shares, dates=date_list[-history_length:]))
            # 生成的选股蒙板添加到选股蒙板队列中，

        sel_mask = self._selecting_blend(sel_masks)  # 根据蒙板混合前缀表达式混合所有蒙板
        # sel_mask.any(0) 生成一个行向量，每个元素对应sel_mask中的一列，如果某列全部为零，该元素为0，
        # 乘以hist_extract后，会把它对应列清零，因此不参与后续计算，降低了择时和风控计算的开销
        # TODO: 这里本意是筛选掉未中选的股票，降低择时计算的开销，使用新的数据结构后不再适用，需改进以使其适用
        # hist_selected = hist_data * selected_shares
        # 第二步，使用择时策略在历史数据上独立产生若干多空蒙板(ls_mask)
        # 生成多空蒙板时忽略在整个历史考察期内从未被选中过的股票：
        # 依次使用择时策略队列中的所有策略逐个生成多空蒙板
        ls_masks = np.array([tmg.generate(dt) for tmg, dt in zip(self._timing, self._timing_history_data)])
        ls_mask = self._ls_blend(ls_masks)  # 混合所有多空蒙板生成最终的多空蒙板
        # 第三步，风险控制交易信号矩阵生成（简称风控矩阵）
        # 依次使用风控策略队列中的所有策略生成风险控制矩阵
        ricon_mats = np.array([ricon.generate(dt) for ricon, dt in zip(self._ricon, self._ricon_history_data)])
        ricon_mat = self._ricon_blend(ricon_mats)  # 混合所有风控矩阵后得到最终的风控策略

        # 使用mask_to_signal方法将多空蒙板及选股蒙板的乘积（持仓蒙板）转化为交易信号，再加上风控交易信号矩阵，并移除所有大于1或小于-1的信号
        # 生成交易信号矩阵
        if self._ls_blender != 'none':
            # 常规情况下，所有的ls_mask会先被混合起来，然后再与sel_mask相乘后生成交易信号，与ricon_mat相加
            op_mat = (mask_to_signal(ls_mask * sel_mask) + ricon_mat).clip(-1, 1)
        else:
            # 在ls_blender为"none"的时候，代表择时多空蒙板不会进行混合，分别与sel_mask相乘后单独生成交易信号，再与ricon_mat相加
            op_mat = (mask_to_signal(ls_mask * sel_mask).sum(axis=0) + ricon_mat).clip(-1, 1)
        # 生成DataFrame，并且填充日期数据
        date_list = hist_data.hdates[-op_mat.shape[0]:]
        # TODO: 在这里似乎可以不用DataFrame，直接生成一个np.ndarray速度更快
        lst = pd.DataFrame(op_mat, index=date_list, columns=shares)
        # 定位lst中所有不全为0的行
        lst_out = lst.loc[lst.any(axis=1)]
        return lst_out

    # ================================
    # 下面是Operation模块的私有方法
    def _set_ls_blender(self, ls_blender):
        """设置多空蒙板的混合方式，包括多种混合方式, 用于产生不同的混合效果。

        input:
            :param ls_blender:
                str，合法的输入包括：
                'combo':        combo表示……
                'none':         None表示……
                'str-T':        T为浮点数，取值区间为大于零的浮点数，当所有策略的多空信号强度之和大于T时，
                                输出信号为1或-1，没有中间地带
                'pos-N-T':      N为正整数，取值区间为1到len(timing)的值，表示在N个策略为多时状态为多，
                                则为空，当策略的多空输出存在浮点数时， T表示判断某个多空值为多的阈值，
                                例如：
                                    'pos-2-0.2' 表示当至少有两个策略输出的多空值>0.2时，策略输出值
                                    为多，信号强度为所有多空信号之和。信号强度clip到(-1, 1)
                'avg_pos-N-T':  表示……
                'avg':          在每个策略发生反转时都会产生交易信号，总信号强度为所有策略信号强度的平均值

        return:=====
            None
        """
        # TODO: 使用regex判断输入的ls_blender各式是否正确
        assert isinstance(ls_blender, str), f'TypeError, expecting string but got {type(ls_blender)}'
        self._ls_blender = ls_blender

    def _set_selecting_blender(self, selecting_blender_expression):
        """ 设置选股策略的混合方式，混合方式通过选股策略混合表达式来表示

            给选股策略混合表达式赋值后，直接解析表达式，将选股策略混合表达式的前缀表达式存入选股策略混合器
        """
        if not isinstance(selecting_blender_expression, str):  # 如果输入不是str类型
            self._selecting_blender_string = '0'
            self._selecting_blender = ['0']
        else:
            self._selecting_blender_string = selecting_blender_expression
            try:
                self._selecting_blender = self._exp_to_blender
            except:
                raise ValueError(f'SimpleSelecting blender expression is not Valid: (\'{selecting_blender_expression}\')'
                                 f', all elements should be separated by blank space, for example: '
                                 f'\' 0 and ( 1 or 2 )\'')

    def _set_ricon_blender(self, ricon_blender):
        self._ricon_blender = ricon_blender

    def _ls_blend(self, ls_masks):
        """ 择时策略混合器，将各个择时策略生成的多空蒙板按规则混合成一个蒙板
            这些多空模板的混合方式由混合字符串来定义。
            混合字符串是满足以下任意一种类型的字符串：

            1，  'none'
                模式表示输入的蒙板不会被混合，所有的蒙板会被转化为一个
                三维的ndarray返回,不做任何混合，在后续计算中采用特殊计算方式
                # 分别计算每一个多空蒙板的交易信号，然后再将交易信号混合起来.

            2,  'avg':
                avg方式下，持仓取决于看多的蒙板的数量，看多蒙板越多，持仓越高，
                只有所有蒙板均看空时，最终结果才看空所有蒙板的权重相同，因此，如
                果一共五个蒙板三个看多两个看空时，持仓为60%。更简单的解释是，混
                合后的多空仓位是所有蒙版仓位的平均值.

            3,  '[pos]/[avg-pos](-N)(-T)'
                格式为满足以上正则表达式的字符串，其混合规则如下：
                在pos-N方式下，
                最终的多空信号强度取决于蒙板集合中各个蒙板的信号值，只有满足N个以
                上的蒙板信号值为多(>0)或者为空(<0)时，最终蒙板的多空信号才为多或
                为空。最终信号的强度始终为-1或1，如果希望最终信号强度为输入信号的
                平均值，应该使用avg_pos-N方式混合

                pos-N还有一种变体，即pos-N-T模式，在这种模式下，N参数仍然代表看
                多的参数个数阈值，但是并不是所有判断持仓为正的数据都会被判断为正
                只有绝对值大于T的数据才会被接受，例如，当T为0.25的时候，0.35会
                被接受为多头，但是0.15不会被接受为多头，因此尽管有两个策略在这个
                时间点判断为多头，但是实际上只有一个策略会被接受.

                avg_pos-N方式下，
                持仓同样取决于看多的蒙板的数量，只有满足N个或更多蒙板看多时，最终结果
                看多，否则看空，在看多/空情况下，最终的多空信号强度=平均多空信号强度
                。当然，avg_pos-1与avg等价，如avg_pos-2方式下，至少两个蒙板看多
                则最终看多，否则看空

                avg_pos-N还有一种变体，即avg_pos-N-T模式，在通常的模式下加
                入了一个阈值Threshold参数T，用来判断何种情况下输入的多空蒙板信号
                可以被保留，当T大于0时，只有输入信号绝对值大于T的时候才会被接受为有
                意义的信号否则就会被忽略。使用avg_pos-N-T模式，并把T设置为一个较
                小的浮点数能够过滤掉一些非常微弱的多空信号.

            4，  'str-T':
                str-T模式下，持仓只能为0或+1，只有当所有多空模版的输出的总和大于
                某一个阈值T的时候，最终结果才会是多头，否则就是空头

            5,  'combo':
                在combo模式下，所有的信号被加总合并，这样每个所有的信号都会被保留，
                虽然并不是所有的信号都有效。在这种模式下，本意是原样保存所有单个输入
                多空模板产生的交易信号，但是由于正常多空模板在生成卖出信号的时候，会
                运用"比例机制"生成卖出证券占持有份额的比例。这种比例机制会在针对
                combo模式的信号组合进行计算的过程中产生问题。
                例如：在将两组信号A和B合并到一起之后，如果A在某一天产生了一个股票
                100%卖出的信号，紧接着B在接下来的一天又产生了一次股票100%卖出的信号，
                两个信号叠加的结果，就是产生超出持有股票总数的卖出数量。将导致信号问题

        input：=====
            :type: ls_masks：object ndarray, 多空蒙板列表，包含至少一个多空蒙板
        return：=====
            :rtype: object: 一个混合后的多空蒙板
        """
        try:
            blndr = str_to_list(self._ls_blender, '-')  # 从对象的属性中读取择时混合参数
        except:
            raise TypeError(f'the timing blender converted successfully!')
        assert isinstance(blndr[0], str) and blndr[0] in self.AVAILABLE_LS_BLENDER_TYPES, \
            f'extracted blender \'{blndr[0]}\' can not be recognized, make sure ' \
            f'your input is like "str-T", "avg_pos-N-T", "pos-N-T", "combo", "none" or "avg"'
        l_m = ls_masks
        l_m_sum = np.sum(l_m, 0) # 计算所有多空模版的和
        l_count = ls_masks.shape[0]
        # 根据多空蒙板混合字符串对多空模板进行混合
        if blndr[0] == 'none':
            return l_m
        if blndr[0] == 'avg':
            return l_m_sum / l_count
        if blndr[0] == 'pos' or blndr[0] == 'avg_pos':
            l_m_sign = 0.
            n = int(blndr[1])
            if len(blndr) == 3:
                threshold = float(blndr[2])
                for msk in ls_masks:
                    l_m_sign += np.sign(np.where(np.abs(msk) < threshold, 0, msk))
            else:
                for msk in ls_masks:
                    l_m_sign += np.sign(msk)
            if blndr[0] == 'pos':
                res = np.where(np.abs(l_m_sign) >= n, l_m_sign, 0)
                return res.clip(-1, 1)
            if blndr[0] == 'avg_pos':
                res = np.where(np.abs(l_m_sign) >= n, l_m_sum, 0) / l_count
                return res.clip(-1, 1)
        if blndr[0] == 'str':
            threshold = float(blndr[1])
            return np.where(np.abs(l_m_sum) >= threshold, 1, 0) * np.sign(l_m_sum)
        if blndr[0] == 'combo':
            return l_m_sum
        raise ValueError(f'Blender text \'({blndr})\' not recognized!')

    def _selecting_blend(self, sel_masks):
        """ 选股策略混合器，将各个选股策略生成的选股蒙板按规则混合成一个蒙板

        input:
            :param sel_masks:
        :return:
            ndarray, 混合完成的选股蒙板
        """
        exp = self._selecting_blender[:]
        s = []
        while exp:  # 等同于但是更好: while exp != []
            if exp[-1].isdigit():
                s.append(sel_masks[int(exp.pop())])
            else:
                s.append(self._blend(s.pop(), s.pop(), exp.pop()))
        return unify(s[0])

    def _blend(self, n1, n2, op):
        """混合操作符函数，将两个选股、多空蒙板混合为一个

        input:
            :param n1: np.ndarray: 第一个输入矩阵
            :param n2: np.ndarray: 第二个输入矩阵
            :param op: np.ndarray: 运算符
        return:
            :return: np.ndarray

        """
        if op == 'or':
            return n1 + n2
        elif op == 'and':
            return n1 * n2
        elif op == 'orr':
            return 1 - (1 - n1) * (1 - n2)
        else:
            raise ValueError(f'ValueError, unknown operand, {op} is not an operand that can be recognized')

    def _ricon_blend(self, ricon_mats):
        if self._ricon_blender == 'add':
            return ricon_mats.sum(axis=0)
        raise NotImplementedError(f'ricon singal blender method ({self._ricon_blender}) is not supported!')

max_stars_repo_name: shepherdpp/qteasy

max_issues_repo_name: shepherdpp/qteasy

max_forks_repo_name: shepherdpp/qteasy


text:
import numpy as np
import tensorflow as tf

tf.set_random_seed(23456)  # reproducibility

class PFNNParameter:
    def __init__(self, shape, rng, phase, name):

        """shape"""
        self.control_num  = shape[0]  #4 control points
        self.weight_shape = shape
        self.bias_shape   = shape[:-1]

        """"alpha and beta"""
        self.alpha        =  tf.Variable(self.initial_alpha(rng), name=name+'alpha')
        self.beta         =  tf.Variable(self.initial_beta(), name=name+'beta')

        """index and weight for phase function"""
        self.pindex_1, self.bamount, self.wamount = self.getIndexAmount(phase)
        self.pindex_0     = (self.pindex_1-1) % self.control_num
        self.pindex_2     = (self.pindex_1+1) % self.control_num
        self.pindex_3     = (self.pindex_1+2) % self.control_num

        """weight and bias"""
        self.weight       = self.cotrol(1)
        self.bias         = self.cotrol(0)

    """initialize parameters in phase function i.e. alpha and beta"""
    def initial_alpha(self, rng):
        shape = self.weight_shape
        alpha_bound = np.sqrt(6. / np.prod(shape[-2:]))
        alpha = np.asarray(
            rng.uniform(low=-alpha_bound, high=alpha_bound, size=shape),
            dtype=np.float32)
        return tf.convert_to_tensor(alpha, dtype = tf.float32)

    def initial_beta(self):
        return tf.zeros(self.bias_shape, tf.float32)

    """ cubic function to calculate the weights and bias in nn"""

    def cotrol(self, flag):
        if flag:
            y0 = tf.nn.embedding_lookup(self.alpha, self.pindex_0)
            y1 = tf.nn.embedding_lookup(self.alpha, self.pindex_1)
            y2 = tf.nn.embedding_lookup(self.alpha, self.pindex_2)
            y3 = tf.nn.embedding_lookup(self.alpha, self.pindex_3)
            mu = self.wamount
        else:
            y0 = tf.nn.embedding_lookup(self.beta, self.pindex_0)
            y1 = tf.nn.embedding_lookup(self.beta, self.pindex_1)
            y2 = tf.nn.embedding_lookup(self.beta, self.pindex_2)
            y3 = tf.nn.embedding_lookup(self.beta, self.pindex_3)
            mu = self.bamount
        return cubic(y0, y1, y2, y3, mu)

    def getIndexAmount(self, phase):
        nslices = self.control_num                    # number of control points in phase function
        pscale = nslices * phase
        pamount = pscale % 1.0
        pindex_1 = tf.cast(pscale, 'int32') % nslices

        bamount = tf.expand_dims(pamount, 1)
        wamount = tf.expand_dims(bamount, 1)          # expand 2 dimension [n*1*1]
        return  pindex_1, bamount, wamount

def cubic(y0, y1, y2, y3, mu):
    return (
        (-0.5*y0+1.5*y1-1.5*y2+0.5*y3)*mu*mu*mu +
        (y0-2.5*y1+2.0*y2-0.5*y3)*mu*mu +
        (-0.5*y0+0.5*y2)*mu +
        (y1))

"""save 50 weights and bias for precomputation in real-time """
def save_network(alpha, beta, num_points,filename):
    nslices = 4
    for i in range(num_points):
        """calculate the index and weights in phase function """
        pscale = nslices*(float(i)/50)
        #weight
        pamount = pscale % 1.0
        #index
        pindex_1 = int(pscale) % nslices
        pindex_0 = (pindex_1-1) % nslices
        pindex_2 = (pindex_1+1) % nslices
        pindex_3 = (pindex_1+2) % nslices

        for j in range(len(alpha)):
            a = alpha[j]
            b = beta[j]
            W = cubic(a[pindex_0],a[pindex_1],a[pindex_2],a[pindex_3],pamount)
            B = cubic(b[pindex_0],b[pindex_1],b[pindex_2],b[pindex_3],pamount)

            W.tofile(filename+'/W%0i_%03i.bin' % (j,i))
            B.tofile(filename+'/b%0i_%03i.bin' % (j,i))

def regularization_penalty(alpha, gamma):
    number_alpha = len(alpha)
    penalty = 0
    for i in range(number_alpha):
        penalty += tf.reduce_mean(tf.abs(alpha[i]))
    return gamma * penalty / number_alpha

ext: py

reason: import numpy


text:
#! /usr/bin/env python

from __future__ import print_function
import sys
import os
sys.path.append(os.path.dirname(os.path.dirname(__file__)))

import tools as trtools
import logging
import numpy as np
import csv

def split_by_vessel_type(points, class_map):
    mmsis = np.unique(points['mmsi'])

    mmsi_map = {}
    for mmsi in mmsis:
        cls = class_map.get(mmsi, "Unknown")
        if cls not in mmsi_map: mmsi_map[cls] = set()
        mmsi_map[cls].add(mmsi)

    for cls, mmsis in mmsi_map.iteritems():
        filter = reduce(
            lambda a, b: a | b,
            (points['mmsi'] == mmsi for mmsi in mmsis))
        yield cls, points[filter]

def clean_format(x):
    for c in "\/ :":
        x = x.replace(c, '_')
    return x

if __name__ == "__main__":
    logging.getLogger().setLevel("INFO")
    import argparse
    parser = argparse.ArgumentParser(
        description='Split vessels as type')
    parser.add_argument(
        'sources', nargs='+', help="source (npz) files")
    parser.add_argument(
        '--classes', required=True, help="path mapping mmsi to vessel type")
    parser.add_argument(
        '--field', default="label", help="field in class file used for mapping. Defaults to `label`")
    parser.add_argument(
        '--output_template', default="./{source}_{cls}.npz", help="template for writing out files (use {cls} for class)")
    args = parser.parse_args()

    with open(args.classes) as f:
        reader =  csv.DictReader(f)
        class_map = {float(x['mmsi']): x[args.field] for x in reader if x[args.field]}

    for inpth in args.sources:
        logging.info("Loading: %s", inpth)
        points = np.load(inpth)['x']

        for cls, points_for_cls in split_by_vessel_type(points, class_map):
            logging.info("Starting conversion for class: %s", cls)
            outpth = args.output_template.format(source=os.path.splitext(os.path.basename(inpth))[0],
                                                 cls=clean_format(cls))
            sub_points = np.array(list(points_for_cls))
            if len(sub_points):
                logging.info("Writing %s rows to %s", len(sub_points), outpth)
                np.savez_compressed(outpth, x=sub_points)


text:
import collections
import os
import random
from abc import ABC
from abc import abstractmethod
from enum import Enum
from typing import Dict
from typing import Iterable
from typing import List
from typing import Optional
from typing import Union

import dgl
import numpy as np
import torch
import torchtext
import tqdm

from robustcode.analysis.graph import AstTree
from robustcode.models.modules.dgl.utransformer import GraphModel
from robustcode.models.modules.iterators import MiniBatch
from robustcode.models.modules.neural_model_base import NeuralModelBase
from robustcode.models.modules.neural_model_base import PredictionStats
from robustcode.models.robust.adversary.node_grads import each_node_grads_batch
from robustcode.models.robust.adversary.rules import NodeRenameRule
from robustcode.models.robust.adversary.rules import NodeValueIndex
from robustcode.models.robust.adversary.rules import RenameRulesForTree
from robustcode.models.robust.adversary.rules import RenameRulesIndex
from robustcode.models.robust.adversary.rules import ShuffleStrategy
from robustcode.models.robust.adversary.rules import ShuffleStrategyGradient
from robustcode.models.robust.adversary.rules import ShuffleStrategyIndividualGradient
from robustcode.models.robust.adversary.rules import ShuffleStrategyRandom
from robustcode.models.robust.adversary.tree_rules import PositionIDs
from robustcode.models.robust.adversary.tree_rules import TernaryWrapperRule
from robustcode.models.robust.dataset import Dataset
from robustcode.models.robust.dataset_util import dataset_to_trees_num
from robustcode.models.robust.visualization import FieldLabel
from robustcode.models.robust.visualization import MaskLabel
from robustcode.models.robust.visualization import TreeVisualization
from robustcode.util.misc import acc
from robustcode.util.misc import Logger
from robustcode.util.misc import Timer

class OrderedEnum(Enum):
    def __ge__(self, other):
        if self.__class__ is other.__class__:
            return self.value >= other.value
        return NotImplemented

    def __gt__(self, other):
        if self.__class__ is other.__class__:
            return self.value > other.value
        return NotImplemented

    def __le__(self, other):
        if self.__class__ is other.__class__:
            return self.value <= other.value
        return NotImplemented

    def __lt__(self, other):
        if self.__class__ is other.__class__:
            return self.value < other.value
        return NotImplemented

class AdvType(OrderedEnum):
    SOUND = 1
    SOUND_PRECISE = 2
    UNSOUND = 3

class AdversarialMode(Enum):
    RANDOM = 0
    BATCH_GRADIENT_GREEDY = 11
    BATCH_GRADIENT_ASCENT = 12
    BATCH_GRADIENT_BOOSTING = 21
    INDIVIDUAL_GRADIENT = 22
    GENETIC = 31

    def __str__(self):
        return str(self.name)

    @staticmethod
    def from_string(s):
        try:
            return AdversarialMode[s]
        except KeyError:
            raise ValueError()

class AdversaryAccuracyStats:
    """
    There are several cases we want to cover:

    Original Prediction
      Correct:
        (Sound & Precise) all the adversarial samples are correct
        (Sound)           all the adversarial samples are either correct or rejected
        (Unsound)         an incorrect adversarial sample exists
       Reject:
        (Sound & Precise) all the adversarial samples are either correct or rejected
        (Sound)           all the adversarial samples are rejected
        (Unsound)         an incorrect adversarial sample exists (i.e., rejected sample classified incorrectly)
    Incorrect:
        --
    """

    def __init__(
        self, reject_token_id, out_file=None, out_field=None, unsound_token_id=None
    ):
        self.reject_token_id = reject_token_id
        self.unsound_token_id = unsound_token_id

        self.base_correct = None
        self.base_preds = None
        self.base_reject_mask = None
        self.base_y = None

        self.correct = None
        self.reject = None
        self.incorrect = None

        self.stats = collections.Counter()

        self.batch_ids = None
        self.out_file = out_file
        self.data = {}

        self.base_stats = None
        self.adv_stats = None
        if out_field is not None:
            self.base_stats = PredictionStats(out_field)
            self.adv_stats = PredictionStats(out_field)

    def training_mask(
        self,
        base_correct: torch.Tensor,
        base_preds: torch.Tensor,
        adv_correct: torch.Tensor,
        adv_preds: torch.Tensor,
    ):
        base_reject_mask = base_preds == self.reject_token_id
        adv_reject_mask = adv_preds == self.reject_token_id

        res = {
            "correct": {
                AdvType.SOUND_PRECISE: base_correct & adv_correct,
                AdvType.SOUND: base_correct & (adv_correct | adv_reject_mask),
            },
            "reject": {
                AdvType.SOUND_PRECISE: base_reject_mask
                & (adv_correct | adv_reject_mask),
                AdvType.SOUND: base_reject_mask & adv_reject_mask,
            },
            "incorrect": {AdvType.UNSOUND: ~(adv_correct | adv_reject_mask)},
        }

        res["correct"][AdvType.UNSOUND] = base_correct & ~res["correct"][AdvType.SOUND]
        res["reject"][AdvType.UNSOUND] = (
            base_reject_mask & ~res["reject"][AdvType.SOUND_PRECISE]
        )

        unsound_mask = (
            res["correct"][AdvType.UNSOUND]
            | res["reject"][AdvType.UNSOUND]
            | res["incorrect"][AdvType.UNSOUND]
        )
        imprecise_mask = (
            ~res["correct"][AdvType.SOUND_PRECISE] & res["correct"][AdvType.SOUND]
        )
        return unsound_mask, imprecise_mask

    def set_base_predictions(
        self,
        base_correct: torch.Tensor,
        base_preds: torch.Tensor,
        base_y=None,
        batch_ids=None,
    ):
        self.base_correct = base_correct
        self.base_preds = base_preds
        self.base_reject_mask = base_preds == self.reject_token_id
        self.base_y = base_y

        self.correct = {
            AdvType.SOUND_PRECISE: base_correct.clone(),
            AdvType.SOUND: base_correct.clone(),
            # AdvType.UNSOUND: torch.zeros_like(base_correct, dtype=torch.bool),
        }

        self.reject = {
            AdvType.SOUND_PRECISE: self.base_reject_mask.clone(),
            AdvType.SOUND: self.base_reject_mask.clone(),
            # AdvType.UNSOUND: torch.zeros_like(self.base_reject_mask, dtype=torch.bool),
        }

        self.incorrect = {AdvType.UNSOUND: ~(self.base_correct | self.base_reject_mask)}
        Logger.debug(
            "correct: {}, reject: {}, incorrect: {}, total: {}".format(
                self.base_correct.sum().item(),
                self.base_reject_mask.sum().item(),
                (~(self.base_correct | self.base_reject_mask)).sum().item(),
                self.base_correct.numel(),
            )
        )
        self.batch_ids = batch_ids

    def add_adversarial_predictions(self, correct: torch.Tensor, preds: torch.Tensor):
        reject_mask = preds == self.reject_token_id
        changed = False

        def update_field(field, key, value, is_and=True):
            old = field[key]
            if is_and:
                new = old & value
            else:
                new = old | value
            if torch.all(old == new):
                return False

            field[key] = new
            return True

        changed |= update_field(self.correct, AdvType.SOUND_PRECISE, correct)
        changed |= update_field(self.correct, AdvType.SOUND, (correct | reject_mask))

        changed |= update_field(
            self.reject, AdvType.SOUND_PRECISE, (correct | reject_mask)
        )
        changed |= update_field(self.reject, AdvType.SOUND, reject_mask)
        return changed

    def update_stats(self):
        for key, value in self.correct.items():
            self.stats[("correct", key)] += torch.sum(value).item()

        for key, value in self.reject.items():
            self.stats[("reject", key)] += torch.sum(value).item()

        for key, value in self.incorrect.items():
            self.stats[("incorrect", key)] += torch.sum(value).item()

        self.stats["total"] += self.base_correct.numel()
        self.stats["correct"] += torch.sum(self.base_correct).item()
        self.stats["incorrect"] += torch.sum(
            ~(self.base_correct | self.base_reject_mask)
        ).item()
        self.stats["reject"] += torch.sum(self.base_reject_mask).item()

        self.stats[("correct", AdvType.UNSOUND)] = (
            self.stats["correct"] - self.stats[("correct", AdvType.SOUND)]
        )
        self.stats[("reject", AdvType.UNSOUND)] = (
            self.stats["reject"] - self.stats[("reject", AdvType.SOUND_PRECISE)]
        )

        self.correct[AdvType.UNSOUND] = self.base_correct & ~self.correct[AdvType.SOUND]
        self.reject[AdvType.UNSOUND] = (
            self.base_reject_mask & ~self.reject[AdvType.SOUND_PRECISE]
        )
        # correct_unsound = self.base_correct & ~self.correct[AdvType.SOUND]
        # reject_unsound = self.base_reject_mask & ~self.reject[AdvType.SOUND_PRECISE]
        unsound = (
            self.incorrect[AdvType.UNSOUND]
            | self.correct[AdvType.UNSOUND]
            | self.reject[AdvType.UNSOUND]
        )
        self.stats[("total", AdvType.UNSOUND)] += torch.sum(unsound).item()
        self.stats[("total", AdvType.SOUND)] += torch.sum(
            self.correct[AdvType.SOUND] | self.reject[AdvType.SOUND_PRECISE]
        ).item()

        self.stats["total_inputs"] += self.base_correct.shape[0]  # TODO: per document
        self.stats["correct_inputs"] += torch.sum(self.base_correct, dim=0).item()

        # self.stats['reject_inputs'] += torch.sum(self.base_reject_mask)

        if self.base_stats is not None and self.base_y is not None:
            self.base_stats.add_(self.base_y, self.base_preds, self.base_correct)

            adv_preds = self.base_preds.clone()
            adv_preds[unsound] = self.unsound_token_id
            self.adv_stats.add_(self.base_y, adv_preds, ~unsound)

        if self.out_file is not None:
            self.data[str(self.batch_ids)] = {
                "correct": self.correct,
                "reject": self.reject,
                "incorrect": self.incorrect,
                "base_correct": self.base_correct,
                "base_preds": self.base_preds,
                "base_y": self.base_y,
            }

    def save(self):
        if self.out_file is None:
            return
        assert self.out_file is not None
        Logger.debug("Saving {}...".format(self.out_file))
        torch.save(self.data, "{}.pt".format(self.out_file))

    def load(self):
        assert self.out_file is not None
        assert os.path.exists("{}.pt".format(self.out_file))
        data = torch.load("{}.pt".format(self.out_file))
        self.stats = collections.Counter()
        for batch_ids, stats in data.items():
            self.base_correct = stats["base_correct"]
            self.base_preds = stats["base_preds"]
            self.base_y = stats["base_y"]
            self.base_reject_mask = self.base_preds == self.reject_token_id

            self.correct = stats["correct"]
            self.reject = stats["reject"]
            self.incorrect = stats["incorrect"]

            self.update_stats()

        return data

    def load_fallback(self, other_out_file):
        assert self.out_file is not None
        assert os.path.exists("{}.pt".format(self.out_file))
        assert os.path.exists("{}.pt".format(other_out_file))
        data = torch.load("{}.pt".format(self.out_file))
        other_data = torch.load("{}.pt".format(other_out_file))

        def merge(A, B, reject):
            return A * ~reject + B * reject

        self.stats = collections.Counter()
        for batch_ids, stats in data.items():
            self.base_correct = stats["base_correct"]
            self.base_preds = stats["base_preds"]
            self.base_y = stats["base_y"]
            self.base_reject_mask = self.base_preds == self.reject_token_id

            self.correct = stats["correct"]
            self.reject = stats["reject"]
            self.incorrect = stats["incorrect"]

            ostats = other_data[batch_ids]
            obase_correct = ostats["base_correct"]
            obase_preds = ostats["base_preds"]

            ocorrect = ostats["correct"]
            oreject = ostats["reject"]
            oincorrect = ostats["incorrect"]

            self.base_correct = merge(
                self.base_correct, obase_correct, self.base_reject_mask
            )
            self.base_preds = merge(self.base_preds, obase_preds, self.base_reject_mask)

            for key in self.correct.keys():
                self.correct[key] = merge(
                    self.correct[key], ocorrect[key], self.base_reject_mask
                )

            for key in self.reject.keys():
                self.reject[key] = merge(
                    self.reject[key], oreject[key], self.base_reject_mask
                )

            for key in self.incorrect.keys():
                self.incorrect[key] = merge(
                    self.incorrect[key], oincorrect[key], self.base_reject_mask
                )

            self.base_reject_mask[:] = False

            self.update_stats()

    def is_sound(self):
        return (
            self.stats[("correct", AdvType.UNSOUND)]
            + self.stats[("reject", AdvType.UNSOUND)]
            + self.stats[("incorrect", AdvType.UNSOUND)]
            == 0
        )

    def to_dict(self):
        stats = {}
        shortcut = {"correct": "C", "incorrect": "I", "reject": "R", "total": "T"}
        for key in sorted(key for key in self.stats.keys() if not isinstance(key, str)):
            category, name = key
            stats["{}:{}".format(shortcut[category], name.name)] = acc(
                self.stats[key], self.stats[category]
            )
            if category not in stats:
                stats["{} ({})".format(category, shortcut[category])] = self.stats[
                    category
                ]
        return stats

    def __repr__(self):
        s = "Base Accuracy: {:6d}/{:6d} ({:.2f}%), reject: {:6d}/{:6d} ({:.2f}%)\n".format(
            self.stats["correct"],
            self.stats["total"] - self.stats["reject"],
            acc(self.stats["correct"], self.stats["total"] - self.stats["reject"]),
            self.stats["reject"],
            self.stats["total"],
            acc(self.stats["reject"], self.stats["total"]),
        )
        for key in sorted(key for key in self.stats.keys() if not isinstance(key, str)):
            category, name = key
            s += "\t{:>10s} {:>20s}: {:6d}/{:6d} ({:.2f}%)\n".format(
                category,
                name.name,
                self.stats[key],
                self.stats[category],
                acc(self.stats[key], self.stats[category]),
            )

        #     Logger.debug('Most common predictions:')
        if self.base_stats is not None:
            self.base_stats.dump_most_common("correct", 10)
            self.adv_stats.dump_most_common("unsound", 10)
        return (
            s
            + f"\nCorrect inputs: {self.stats['correct_inputs']} / {self.stats['total_inputs']}"
        )

class AdversaryBase(ABC):
    def __init__(self, dataset: Dataset):
        self.dataset = dataset

    @abstractmethod
    def attack(
        self,
        model: NeuralModelBase,
        batch: MiniBatch,
        mask=None,
        num_samples=10,
        adv_mode=AdversarialMode.RANDOM,
        shuffle: Optional[ShuffleStrategy] = None,
    ):
        raise NotImplementedError

    def compute_correct(
        self, model: NeuralModelBase, batch: MiniBatch, verbose=False, threshold=0.5
    ):
        training_mask = model.padding_mask(batch)
        mask_valid = batch.masks["mask_valid"]
        Y = batch.Y
        predictions = model.predict_with_reject(
            batch, self.dataset.reject_token_id, threshold=threshold
        )

        correct = (predictions == Y) * training_mask
        if len(Y.size()) == 2:
            """
            The adversary assumes that the order of predicted labels in the adversary and original batch are the same.
            All the modifications respect this, i.e., while code is added or removed, the order and
            number of predicted labels (via 'mask_valid') stays the same.
            However, since the sample length might change, we need to ensure that masked_select is performed
            batch_first (B x N rather than N x B) as otherwise the results are no guaranteed to preserve the order.
            """
            # TODO: check that the batch is created with batch_first=False (this is the default with torchtext)
            mask_valid = mask_valid.t()
            correct = correct.t().masked_select(mask_valid)
            predictions = predictions.t().masked_select(mask_valid)
            Y = Y.t().masked_select(mask_valid)
        else:
            correct = correct.masked_select(mask_valid)
            predictions = predictions.masked_select(mask_valid)
            Y = Y.masked_select(mask_valid)
        return correct, predictions, Y

    def visualize_adversarial(
        self, batch: MiniBatch, model: NeuralModelBase, masks, colors
    ):
        assert isinstance(model, GraphModel)
        samples = self.dataset.samples_for_batch(batch)

        training_mask = model.padding_mask(batch, mask_field="mask_valid")

        g = batch.X
        mask_data = {}
        mask_names = []
        for idx, mask in enumerate(masks):
            mask_all = torch.zeros_like(training_mask)
            mask_all[training_mask] = mask
            mask_data[str(idx)] = mask_all
            mask_names.append(str(idx))

        for name, mask in mask_data.items():
            assert name not in g.ndata
            g.ndata[name] = mask
        trees = dgl.unbatch(g)
        for name in mask_data.keys():
            del g.ndata[name]

        for tree, sample in zip(trees, samples):
            mask_labels = [
                MaskLabel.from_values(tree.ndata[name], color)
                for name, color in zip(mask_names, colors)
            ]
            TreeVisualization.visualize(
                sample,
                ["types", "values"],
                self.dataset.dtrain.fields,
                labels=mask_labels
                + [
                    FieldLabel(
                        self.dataset.dtrain.fields["types"], tree.ndata["types"]
                    ),
                    FieldLabel(
                        self.dataset.dtrain.fields["values"], tree.ndata["values"]
                    ),
                    # TODO: topk
                    # TopkLabel.from_iter([batch], sample, model, self.dataset,
                    #                     mask=MaskLabel.from_sample(sample, 'mask_valid')),
                ],
            )
            input()
            break

    def compute_gradients(
        self, batch, model, unsound_masks, imprecise_masks, threshold
    ):
        masks = [a | b for a, b in zip(unsound_masks, imprecise_masks)]

        for node_grads in each_node_grads_batch(
            batch,
            model,
            self.dataset,
            model.loss_function,
            threshold=threshold,
            masks=masks,
            max_samples=None,
        ):
            print(node_grads)

    def fit_adversarial(
        self,
        model: NeuralModelBase,
        dataset_iter: Iterable[MiniBatch],
        adversarial_iters: Union["AdversaryBatchIter", List["AdversaryBatchIter"]],
        threshold=0.5,
    ):
        if not isinstance(adversarial_iters, list):
            adversarial_iters = [adversarial_iters]
        num_refined = 0
        num_unsound = 0
        num_imprecise = 0

        def sanity_check(orig_batch: MiniBatch, adversarial_batch: MiniBatch):
            gts = orig_batch.Y
            tmasks = model.padding_mask(orig_batch, mask_field="mask_valid")

            adv_gts = adversarial_batch.Y
            adv_tmasks = model.padding_mask(adversarial_batch, mask_field="mask_valid")
            assert torch.sum(tmasks) == torch.sum(adv_tmasks), "{} vs {}".format(
                torch.sum(tmasks), torch.sum(adv_tmasks)
            )
            assert torch.all(gts[tmasks] == adv_gts[adv_tmasks])

        adv_stats = AdversaryAccuracyStats(self.dataset.reject_token_id)
        for batch in tqdm.tqdm(dataset_iter, ncols=100, leave=False):
            base_correct, base_preds, base_y = self.compute_correct(
                model, batch, verbose=False, threshold=threshold
            )

            # self.visualize_adversarial(batch, model, [base_correct], ['da_green'])
            for adversarial_iter in adversarial_iters:
                for idx, adv_batch in enumerate(adversarial_iter.iter_batch(batch)):
                    sanity_check(batch, adv_batch)

                    correct, preds, adv_y = self.compute_correct(
                        model, adv_batch, verbose=False, threshold=threshold
                    )
                    assert torch.all(base_y == adv_y)
                    unsound_mask, imprecise_mask = adv_stats.training_mask(
                        base_correct, base_preds, correct, preds
                    )
                    # self.compute_gradients(adv_batch, model, unsound_mask, imprecise_mask, threshold)

                    mask = unsound_mask | imprecise_mask
                    num_refined += torch.sum(mask).item()
                    num_unsound += torch.sum(unsound_mask).item()
                    num_imprecise += torch.sum(imprecise_mask).item()

                    # self.visualize_adversarial(adv_batch, model, [unsound_mask, imprecise_mask], ['da_red', 'da_black'])

                    model.fit_batch(
                        model.opt,
                        model.loss_function,
                        adv_batch,
                        mask_field="mask_valid",
                        training_mask=mask,
                    )

        Logger.debug(
            "Number of refined samples: {}, unsound: {}, imprecise: {}".format(
                num_refined, num_unsound, num_imprecise
            )
        )
        return num_refined

    def adversarial_accuracy(
        self,
        model: NeuralModelBase,
        base_iter: Iterable[MiniBatch],
        adversarial_iters: Union["AdversaryBatchIter", List["AdversaryBatchIter"]],
        threshold=0.5,
        out_file=None,
        verbose=False,
        approximate=False,
    ) -> AdversaryAccuracyStats:
        if not isinstance(adversarial_iters, list):
            adversarial_iters = [adversarial_iters]

        Logger.debug("out_file: {}".format(out_file))
        adv_stats = AdversaryAccuracyStats(
            self.dataset.reject_token_id,
            out_file=out_file,
            out_field=self.dataset.TARGET,
            unsound_token_id=self.dataset.unsound_token_id,
        )
        for batch_idx, batch in enumerate(base_iter):
            # if approximate and random.random() > 0.3:
            #     continue
            if approximate and batch_idx % 4 != 0:
                continue
            base_correct, base_preds, base_y = self.compute_correct(
                model, batch, verbose=False, threshold=threshold
            )
            adv_stats.set_base_predictions(
                base_correct, base_preds, base_y=base_y, batch_ids=batch.ids
            )

            "Combine results of multiple possible adversarial attacks to obtain the worst case results"
            for adversarial_iter in adversarial_iters:
                changed = False
                for idx, adv_batch in enumerate(adversarial_iter.iter_batch(batch)):
                    correct, preds, adv_y = self.compute_correct(
                        model, adv_batch, verbose=False, threshold=threshold
                    )
                    assert torch.all(
                        base_y == adv_y
                    ), "The order of predicted positions changed. The adversary will compute incorrect results!"
                    changed |= adv_stats.add_adversarial_predictions(correct, preds)

                    if idx > 0 and idx % 20 == 0:
                        if not changed:
                            Logger.debug("Stopping after {} iter".format(idx))
                            break
                        changed = False
                        Logger.debug(
                            "\titer {} {}".format(
                                idx,
                                str(
                                    [v.sum().item() for v in adv_stats.correct.values()]
                                ),
                            )
                        )

            adv_stats.update_stats()

            if verbose:
                Logger.debug(
                    "Batch: {}, Adv Accuracy\n{}".format(batch_idx, str(adv_stats))
                )
        adv_stats.save()

        Logger.debug("Adv Accuracy\n{}".format(str(adv_stats)))
        return adv_stats

class RenameAdversary(AdversaryBase):
    def __init__(self, rules_index: RenameRulesIndex, dataset: Dataset):
        super().__init__(dataset)
        self.index = rules_index

    def attack(
        self,
        model: NeuralModelBase,
        batch: MiniBatch,
        mask=None,
        num_samples=10,
        adv_mode=AdversarialMode.RANDOM,
        shuffle: Optional[ShuffleStrategy] = None,
    ):
        if num_samples == 0:
            return
        shuffle = shuffle or ShuffleStrategyRandom()
        if isinstance(batch.X, dgl.DGLGraph):
            yield from self.__attack_graph(
                model, batch, mask, num_samples, shuffle, mode=adv_mode
            )
        else:
            yield from self.__attack_seq(
                model, batch, mask, num_samples, shuffle, mode=adv_mode
            )

    def __attack_seq(
        self,
        model: NeuralModelBase,
        batch: MiniBatch,
        mask: torch.Tensor,
        num_samples,
        shuffle: ShuffleStrategy,
        mode: AdversarialMode,
    ):

        tree_sizes = batch.lengths
        if mask is not None:
            raise NotImplementedError
        else:
            tree_nodes = [None] * len(tree_sizes)

        tree_rules, tree_num_samples = self._initialize_adversarial_rules(
            num_samples, tree_nodes, model, batch, shuffle=shuffle, mode=mode
        )

        # the inputs should contain two tensors for types and values
        # assert len(batch.inputs) == 2
        # only values are being changed
        original_values = batch.X[-1]
        adversarial_masks = None

        try:
            for idx in range(max(tree_num_samples)):
                batch_values = batch.X[-1].cpu().t().numpy()
                idx_to_use = idx

                if mode == AdversarialMode.INDIVIDUAL_GRADIENT:
                    # obtain gradients w.r.t. idx-th classifiable position in each input
                    shuffle.for_next_position()
                    for rules in tree_rules:
                        for rule in rules:
                            shuffle.shuffle_candidates(rule)

                if mode == AdversarialMode.BATCH_GRADIENT_BOOSTING:
                    # mask: compute gradient only for correctly predicted positions in all previous iterations
                    batch.X[-1] = original_values
                    adversarial_masks = [
                        model.get_adversarial_mask(
                            minibatch,
                            mask_field="mask_valid",
                            previous_mask=adversarial_masks,
                        )
                        for minibatch in batch
                    ]

                if mode in [
                    AdversarialMode.BATCH_GRADIENT_ASCENT,
                    AdversarialMode.BATCH_GRADIENT_BOOSTING,
                ]:
                    shuffle.initialize(
                        model, batch, None, position_mask=adversarial_masks
                    )
                    for rules in tree_rules:
                        for rule in rules:
                            shuffle.shuffle_candidates(rule)
                    # values have been shuffled again, we can just use the argmax
                    idx_to_use = 0

                assert len(batch_values) == len(tree_rules)
                # for every example in batch ...
                for rules, values in zip(tree_rules, batch_values):
                    # one node with constant assign => one rule
                    for rule in rules:
                        rule.apply_first_valid(idx_to_use, values)

                batch.X[-1] = torch.tensor(
                    batch_values.transpose(),
                    dtype=torch.long,
                    device=original_values.device,
                )
                yield batch
        except GeneratorExit:
            return
        finally:
            batch.X[-1] = original_values

            for rules in tree_rules:
                for rule in rules:
                    rule.reset()

    def __attack_graph(
        self,
        model: NeuralModelBase,
        batch: MiniBatch,
        mask: torch.Tensor,
        num_samples,
        shuffle: ShuffleStrategy,
        mode: AdversarialMode,
    ):
        tree_sizes = batch.lengths
        offsets = (np.cumsum(tree_sizes) - np.array(tree_sizes)).tolist()

        if mask is not None:
            tree_masks = torch.split(mask, tree_sizes)
            tree_nodes = [
                np.flatnonzero(tree_mask.cpu().numpy()) for tree_mask in tree_masks
            ]
        else:
            tree_nodes = [None] * len(tree_sizes)

        g = batch.X
        assert g.number_of_nodes() == sum(tree_sizes)
        original_values = g.ndata["values"]

        tree_rules, tree_num_samples = self._initialize_adversarial_rules(
            num_samples, tree_nodes, model, batch, shuffle, mode
        )
        adversarial_mask = None

        try:
            for idx in range(max(tree_num_samples)):
                values = g.ndata["values"].cpu().numpy()

                if mode == AdversarialMode.INDIVIDUAL_GRADIENT:
                    # obtain gradients w.r.t. idx-th classifiable position in each input
                    shuffle.for_next_position()
                    for rules in tree_rules:
                        for rule in rules:
                            shuffle.shuffle_candidates(rule)

                if mode == AdversarialMode.BATCH_GRADIENT_BOOSTING:
                    # mask: compute gradient only for correctly predicted positions in all previous iterations
                    g.ndata["values"] = original_values
                    adversarial_mask = model.get_adversarial_mask(
                        batch, mask_field="mask_valid", previous_mask=adversarial_mask
                    )

                if mode in [
                    AdversarialMode.BATCH_GRADIENT_ASCENT,
                    AdversarialMode.BATCH_GRADIENT_BOOSTING,
                ]:
                    shuffle.initialize(model, batch, position_mask=adversarial_mask)
                    for rules in tree_rules:
                        for rule in rules:
                            shuffle.shuffle_candidates(rule)

                for ith_tree, (rules, offset) in enumerate(zip(tree_rules, offsets)):
                    # one node with constant assign => one rule
                    for rule in rules:
                        # apply each rule with the batch offset of its example
                        idx_to_use = idx
                        if mode in [
                            AdversarialMode.BATCH_GRADIENT_ASCENT,
                            AdversarialMode.BATCH_GRADIENT_BOOSTING,
                        ]:
                            # values have been shuffled again, we can just use the argmax
                            idx_to_use = 0
                        rule.apply_first_valid(idx_to_use, values, usage_offset=offset)

                g.ndata["values"] = torch.tensor(
                    values, dtype=torch.long, device=original_values.device
                )
                yield batch
        except GeneratorExit:
            return
        finally:
            g.ndata["values"] = original_values

    def _initialize_adversarial_rules(
        self,
        num_samples,
        tree_nodes,
        model,
        batch,
        shuffle: ShuffleStrategy,
        mode: AdversarialMode,
    ):
        ids = batch.ids
        rule_sets = [self.index.per_tree_rules.get(idx, None) for idx in ids]
        assert all(rule_set is not None for rule_set in rule_sets)

        if mode in [AdversarialMode.INDIVIDUAL_GRADIENT, AdversarialMode.RANDOM]:
            # initialize / compute gradients once
            shuffle.initialize(model, batch)

        do_shuffle = shuffle if (mode == AdversarialMode.RANDOM) else None

        tree_rules, tree_num_samples = zip(
            *[
                rule_set.init_rules(nodes, num_samples, shuffle=do_shuffle)
                for rule_set, nodes in zip(rule_sets, tree_nodes)
            ]
        )

        return tree_rules, tree_num_samples

class RenameAdversaryLimited(RenameAdversary):

    """
    A version of RenameAdversary that only renames values of new nodes that were created by other attacks.
    """

    def __init__(self, dataset: Dataset, trees: Dict[int, AstTree]):
        rules_index = RenameRulesIndex()
        self.trees = trees
        # obtain the list of all observed values in training dataset
        trees_train_num = dataset_to_trees_num(dataset.dtrain)
        self.value_index = NodeValueIndex(dataset, trees_train_num)
        super().__init__(rules_index, dataset)

    def attack(
        self,
        model,
        batch,
        mask=None,
        num_samples=10,
        adv_mode=AdversarialMode.RANDOM,
        shuffle: Optional[ShuffleStrategy] = None,
    ):
        ids = batch.ids
        if isinstance(ids, torch.Tensor):
            ids = ids.cpu().numpy()

        # the nodes provided by SubtreeAdversary will be different for each adv. example,
        # so we need to create new rules for every batch
        self.index = RenameRulesIndex()
        rule_sets = []
        for tree_id in ids:
            tree = self.trees[tree_id]
            rules = RenameRulesForTree(tree_id, tree)
            for pos, node in enumerate(tree.nodes):
                # create a rename rule for each valid newly generated node
                if node.origin == PositionIDs.ADVERSARIAL_CONSTANT:
                    rule = NodeRenameRule(
                        tree_id,
                        node.id,
                        [node.id],
                        [],
                        value="{:3d} {}".format(
                            self.dataset.VALUES.vocab.stoi[node.value], node.value
                        ),
                        candidate_values=self.value_index.values_for_type(node.type),
                        fixed_value_offset=self.dataset.fixed_value_offset,
                    )

                    rules.add(rule)
            self.index.add(rules)
        assert all(rule_set is not None for rule_set in rule_sets)

        return super().attack(model, batch, mask, num_samples, adv_mode, shuffle)

class SubtreeAdversary(AdversaryBase):
    def __init__(
        self,
        rules: List[TernaryWrapperRule],
        dataset: Dataset,
        trees: Dict[int, AstTree],
        make_iter,
    ):
        super().__init__(dataset)
        self.rules = rules
        self.trees = trees
        for key, value in dataset.ID.vocab.stoi.items():
            if key.endswith("_mod"):
                assert value not in self.trees
                idx = dataset.ID.vocab.stoi.get(key[:-4], None)
                assert idx is not None
                "if the tree was a duplicate it was removed from the dataset"
                if idx in self.trees:
                    self.trees[value] = self.trees[idx]
        self.make_iter = make_iter

        self.dataset = dataset
        self.timers = collections.defaultdict(Timer)

    def attack(
        self,
        model,
        batch,
        mask=None,
        num_samples=10,
        adv_mode=AdversarialMode.RANDOM,
        shuffle: Optional[ShuffleStrategy] = None,
    ):
        if adv_mode != AdversarialMode.RANDOM:
            raise NotImplementedError(
                "For subtree attack, only random adversarial mode is supported."
            )
            # TODO: different adv. modes for the two attacks

        if num_samples == 0:
            return

        tree_ids = batch.ids
        batch_trees = [self.trees[idx] for idx in tree_ids]

        per_node_rules = []
        for tree_id, tree in zip(tree_ids, batch_trees):
            for node in tree.nodes:
                candidate_rules = [rule for rule in self.rules if rule.matches(node)]
                if candidate_rules:
                    per_node_rules.append((tree_id, node, candidate_rules))

        vbatch_ids = [sample.id for sample in self.dataset.samples_for_batch(batch)]
        threshold = random.uniform(0.1, 0.4)
        for idx in range(num_samples):
            samples = []
            try:
                num_applied = 0
                for (tree_id, node, candidate_rules) in per_node_rules:
                    "apply the rule with prob threshold"
                    if random.random() <= threshold:
                        rule = random.choice(candidate_rules)
                        rule.apply(tree_id, node)
                        num_applied += 1
                # Logger.debug('Number of applied rules {}'.format(num_applied))

                for tree_id, tree in zip(tree_ids, batch_trees):
                    samples.append(self.dataset.tree_to_sample(tree_id, tree))

                dbatch = torchtext.data.Dataset(samples, self.dataset.fields)

                diter = self.make_iter(dbatch, cached=False)
                for adv_batch in diter:
                    vadv_batch_ids = [
                        sample.id[:-4]
                        for sample in self.dataset.samples_for_batch(adv_batch)
                    ]
                    assert vbatch_ids == vadv_batch_ids

                    yield adv_batch
            except GeneratorExit:
                return
            finally:
                for rule in self.rules:
                    rule.revert_all_changes()

                for tree in batch_trees:
                    tree.refresh()  # revert changes

class AdversaryBatchIter:
    def __init__(
        self,
        adversary: AdversaryBase,
        model: NeuralModelBase,
        dataset_iter,
        shuffle: ShuffleStrategy = None,
        num_samples=1,
        adv_mode=AdversarialMode.RANDOM,
    ):
        self.adversary = adversary
        self.model = model
        self.it = dataset_iter
        if adv_mode == AdversarialMode.RANDOM:
            self.shuffle = ShuffleStrategyRandom()
        elif adv_mode in [
            AdversarialMode.BATCH_GRADIENT_GREEDY,
            AdversarialMode.BATCH_GRADIENT_ASCENT,
            AdversarialMode.BATCH_GRADIENT_BOOSTING,
        ]:
            self.shuffle = ShuffleStrategyGradient()
        elif adv_mode == AdversarialMode.INDIVIDUAL_GRADIENT:
            self.shuffle = ShuffleStrategyIndividualGradient(num_samples)
        else:
            raise NotImplementedError(
                "Adversarial Mode " + str(adv_mode) + " is not supported"
            )

        self.num_samples = num_samples
        self.adversarial_mode = adv_mode

    def init_epoch(self):
        self.it.init_epoch()

    def __len__(self):
        return len(self.it)

    def iter_batch(self, batch):
        batch_iter = (
            [batch]
            if not isinstance(self.it, AdversaryBatchIter)
            else self.it.iter_batch(batch)
        )
        for it_batch in batch_iter:
            for adv_batch in self.adversary.attack(
                self.model,
                it_batch,
                num_samples=self.num_samples,
                adv_mode=self.adversarial_mode,
                shuffle=self.shuffle,
            ):
                yield adv_batch

    def __iter__(self):
        for batch in self.it:
            for idx, adv_batch in enumerate(
                self.adversary.attack(
                    self.model,
                    batch,
                    num_samples=self.num_samples,
                    adv_mode=self.adversarial_mode,
                    shuffle=self.shuffle,
                )
            ):
                yield adv_batch

max_stars_repo_name: david-maine/robust-code

max_issues_repo_name: david-maine/robust-code

max_forks_repo_name: eth-sri/robust-code