"""Everything needed for defining phases within an optimal control problem.

Classes:
	Phase
"""

import copy
import itertools
from typing import (Optional, Tuple)

import sympy as sym

from .bounds import PhaseBounds
from .guess import PhaseGuess
from .mesh import PhaseMesh
from .scaling import PhaseScaling
from .typing import (OptionalExprsType, OptionalSymsType, TupleSymsType)
from .utils import (check_sym_name_clash, format_as_named_tuple)

__all__ = ["Phase"]

class Phase:
    """A single continuous time phase as part of an optimal control problem.

    Attributes:
            name: The name associated with a problem. Should be something short
                    like 'A'.
            optimal_control_problem: The :obj:`OptimalControlProblem` with which
                    this phase is to be associated.
            state_variables: The continuous time state variables in this phase.
            control_variables: The continuous time control variables in this phase.
            state_equations: The dynamical state equations associated with this
                    state variables in this phase.
            integrand_functions: The integrand functions corresponding to the
                    integral variables in this phase.
            path_constraints: The continuous time path constraints associated with
                    this phase.
            bounds: The phase bounds on this phase. See :obj:PhaseBounds for more
                    details.
            scaling: The phase scaling on this phase. See :obj:PhaseScaling for
                    more details.
            guess: The initial guess at which this phase is to be solved.
            mesh: This initial mesh on which this phase is to be solved.
            _name: Protected version of :attr:`name`.
            _ocp: Protected version of :attr:`optimal_control_problem`.
            _phase_number: Protected integer number associated with this phase. If
                    not associated with any optimal control problem then defaults to
                    None until one is associated. These are ordered sequentially
                    starting at '0' in the order with which phases are added to an
                    optimal control problem.
            _phase_suffix: Protected str which is used in the naming of auto-
                    generated Pycollo variables such as the endpoint state variables.
            _y_var_user: Protected version of :attr:`state_variables`.
            _u_var_user: Protected version of :attr:`control_variables`.
            _q_var_user: Protected version of :attr:`integral_variables`.
            _t_var_user: Protected version of :attr:`time_variables`.
            _y_eqn_user: Protected version of :attr:`state_equations`.
            _c_con_user: Protected version of :attr:`path_constraints`.
            _q_fnc_user: Protected version of :attr:`integrand_functions`.
            _t0_USER: Protected version of :attr:`initial_time_variable`.
            _tF_USER: Protected version of :attr:`final_time_variable`.
            _t0: Internal Pycollo symbol for phase initial time.
            _tF: Internal Pycollo symbol for phase final time.
            _STRETCH: Convenience expression for phase time scaling stretch.
            _SHIFT: Convenience expression for phase time scaling shift.
    """

    def __init__(self,
                 name: str,
                 *,
                 optimal_control_problem: Optional["OptimalControlProblem"] = None,
                 state_variables: OptionalSymsType = None,
                 control_variables: OptionalSymsType = None,
                 state_equations: OptionalExprsType = None,
                 integrand_functions: OptionalExprsType = None,
                 path_constraints: OptionalExprsType = None,
                 bounds: Optional[PhaseBounds] = None,
                 scaling: Optional[PhaseScaling] = None,
                 guess: Optional[PhaseGuess] = None,
                 mesh: Optional[PhaseMesh] = None,
                 ):
        """Initialise the Phase object with minimum a name.

        Args:
                name: The name associated with a problem. Should be something short
                        like 'A'.
                optimal_control_problem: The :obj:`OptimalControlProblem` with
                        which this phase is to be associated. Default value is None in
                        which case the phase remain uninitialised to an optimal control
                        problem.
                state_variables: The continuous time state variables in this phase.
                        Default value is None in which case the phase has no associated
                        state variables and no phase-specific endpoint time or state
                        variables are created.
                control_variables: The continuous time control variables in this
                        phase. Default value is None in which case the phase has no
                        associated control variables.
                state_equations: The dynamical state equations associated with this
                        state variables in this phase. Default value is None in which
                        case no dynamical equations have been added to the phase yet.
                integrand_functions: The integrand functions corresponding to the
                        integral variables in this phase. Default value is None in
                        which case the phase has no integrand functions associated with
                        it and no phase-specific integral variables are created.
                path_constraints: The continuous time path constraints associated
                        with this phase. Default value is None in which case the phase
                        has no path constraints associated with it.
                bounds: The phase bounds on this phase. See :obj:PhaseBounds for
                        more details. Default value is None in which case an empty
                        :obj:`PhaseBounds` object is instantiated and associated with
                        the phase.
                scaling: The phase scaling on this phase. See :obj:PhaseScaling for
                        more details. Default value is None in which case an empty
                        :obj:`PhaseScaling` object is instantiated and associated with
                        the phase.
                guess: The initial guess at which this phase is to be solved.
                        Default value is None in which case an empty :obj:`PhaseGuess`
                        object is instantiated and associated with the phase.
                mesh: This initial mesh on which this phase is to be solved.
                        Default value is None in which case an empty :obj:`PhaseMesh`
                        object is instantiated and associated with the phase.
        """

        self._name = str(name)
        self._ocp = None
        self._phase_number = None
        self._phase_suffix = "X"

        self._y_var_user = ()
        self._u_var_user = ()
        self._q_var_user = ()
        self._t_var_user = ()

        self._y_eqn_user = ()
        self._c_con_user = ()
        self._q_fnc_user = ()

        if optimal_control_problem is not None:
            self.optimal_control_problem = optimal_control_problem

        self.state_variables = state_variables
        self.control_variables = control_variables

        self.state_equations = state_equations
        self.integrand_functions = integrand_functions
        self.path_constraints = path_constraints

        self.bounds = bounds
        self.scaling = scaling
        self.guess = guess
        self.mesh = mesh

        self.auxiliary_data = {}

    def create_new_copy(self,
                        name: str,
                        *,
                        copy_state_variables: bool = True,
                        copy_control_variables: bool = True,
                        copy_state_equations: bool = True,
                        copy_path_constraints: bool = True,
                        copy_integrand_functions: bool = True,
                        copy_state_endpoint_constraints: bool = False,
                        copy_bounds: bool = True,
                        copy_mesh: bool = True,
                        copy_scaling: bool = True,
                        copy_guess: bool = True,
                        ):

        self._check_variables_and_equations()
        new_phase = Phase(name,
                          optimal_control_problem=self.optimal_control_problem)

        if copy_state_variables:
            new_phase.state_variables = copy.deepcopy(self.state_variables)
            if copy_bounds:
                new_phase.bounds.state_variables = copy.deepcopy(self.bounds.state_variables)
            if copy_guess:
                new_phase.guess.state_variables = copy.deepcopy(self.guess.state_variables)
            if copy_scaling:
                new_phase.scaling.state_variables = copy.deepcopy(self.scaling.state_variables)

        if copy_control_variables:
            new_phase.control_variables = copy.deepcopy(self.control_variables)
            if copy_bounds:
                new_phase.bounds.control_variables = copy.deepcopy(self.bounds.control_variables)
            if copy_guess:
                new_phase.guess.control_variables = copy.deepcopy(self.guess.control_variables)
            if copy_scaling:
                new_phase.scaling.control_variables = copy.deepcopy(self.scaling.control_variables)

        if copy_state_equations:
            new_phase.state_equations = copy.deepcopy(self.state_equations)

        if copy_path_constraints:
            new_phase.path_constraints = copy.deepcopy(self.path_constraints)
            if copy_bounds:
                new_phase.bounds.path_constraints = copy.deepcopy(self.bounds.path_constraints)

        if copy_integrand_functions:
            new_phase.integrand_functions = copy.deepcopy(self.integrand_functions)
            if copy_bounds:
                new_phase.bounds.integral_variables = copy.deepcopy(self.bounds.integral_variables)

        if copy_state_endpoint_constraints and copy_bounds:
            new_phase.bounds.state_endpoint_constraints = copy.deepcopy(self.bounds.state_endpoint_constraints)

        if copy_mesh:
            new_phase.mesh = copy.deepcopy(self.mesh)

        return new_phase

    @staticmethod
    def create_new_copy_like(phase_for_copying: "Phase", name: str, **kwargs):
        """Constructor class to copy a phase."""
        return phase_for_copying.create_new_copy(name, **kwargs)

    @property
    def name(self):
        """Name of the phase."""
        return self._name

    @property
    def optimal_control_problem(self) -> Optional["OptimalControlProblem"]:
        """The optimal control problem with which this phase is associated.

        There are two allowable scenarios. In the first scenario a phase may be
        instantiated without being associated with an optimal control problem.
        If this is the case then the default values of `None` for the phase
        number and 'X' for the phase suffix remain.

        In the second scenario a phase is instantiated with an associated
        optimal control problem or is associated with an optimal control
        problem after the first type of instantiation. In this case the phase
        is appended to the protected `_phases` attribute of the
        :obj:`OptimalControlProblem`, the phase number is set according to its
        position in the order of addition to the optimal controls problem's
        phases, and its phase suffix is set as a string version of the phase
        number. Finally a replacement of any symbols that may have been used in
        supplementary information about the phase that contained the placeholder
        'X' phase suffix are renamed and substituted.

        No checking is done to see whether the phase is already associated with
        the optimal control problem in question or any other optimal control
        problem. The reason being that if the setter method for this property is
        accessed after having already been set then an `AttributeError` is
        raised (see below). The reason this class works like that is to avoid
        having to allow phases to be disassociated from an
        :obj:`OptimalControlProblem` and thus having to handled the complexities
        that would come with the phase renumbering and substitution of any
        phase-related information that has already been given to the optimal
        control problem.

        Raises:
                AttributeError: If an :obj:`OptimalControlProblem` has already been
                        associated with `self`. If a argument of any type other than
                        :obj:`OptimalControlProblem` is passed to the
                        `optimal_control_problem` property setter.
        """
        return self._ocp

    @optimal_control_problem.setter
    def optimal_control_problem(self, ocp):
        if self._ocp is not None:
            msg = ('Optimal control problem is already set for this phase and '
                   'cannot be reset.')
            raise AttributeError(msg)

        try:
            previous_phase_names = ocp._phases._fields
        except AttributeError:
            previous_phase_names = ()
        phase_names = (*previous_phase_names, self.name)
        ocp._phases = format_as_named_tuple([*ocp._phases, self],
                                            named_keys=phase_names, sympify=False)

        self._ocp = ocp
        self._phase_number = self._ocp.number_phases - 1
        self._phase_suffix = str(self.phase_number)

        self.state_variables = self.state_variables
        self.integrand_functions = self.integrand_functions

    @property
    def phase_number(self) -> Optional[int]:
        """The integer numerical identifier for the phase.

        If this phase has not yet been associated with an optimal control
        problem then None is returned.

        Corresponds to the chronological order in which it was associated with
        the optimal control problem in question.
        """
        return self._phase_number

    @property
    def initial_time_variable(self) -> sym.Symbol:
        """Symbol for the time at which this phase begins."""
        try:
            return self._t0_USER
        except AttributeError:
            msg = ("Can't access initial time until associated with an optimal "
                   "control problem.")
            raise AttributeError(msg)

    @property
    def final_time_variable(self) -> sym.Symbol:
        """Symbol for the time at which this phase begins."""
        try:
            return self._tF_USER
        except AttributeError:
            msg = ("Can't access final time until associated with an optimal "
                   "control problem.")
            raise AttributeError(msg)

    @property
    def initial_state_variables(self) -> TupleSymsType:
        """Symbols for this phase's state variables at the initial time.

        Raises:
                AttributeError: If `optimal_control_problem` property has not yet
                        been set to a not None value. See docstring for
                        `state_variables` for details about why.
        """
        try:
            return self._y_t0_user
        except AttributeError:
            msg = ("Can't access initial state until associated with an optimal "
                   "control problem.")
            raise AttributeError(msg)

    @property
    def final_state_variables(self) -> TupleSymsType:
        """Symbols for this phase's state variables at the final time.

        Raises:
                AttributeError: If `optimal_control_problem` property has not yet
                        been set to a not None value. See docstring for
                        `state_variables` for details about why.
        """
        try:
            return self._y_tF_user
        except AttributeError:
            msg = ("Can't access initial state until associated with an optimal "
                   "control problem.")
            raise AttributeError(msg)

    @property
    def state_variables(self) -> TupleSymsType:
        """Symbols for this phase's state variables in order added by user.

        The user may supply either a single symbol or an iterable of symbols.
        The supplied argument is handled by the `format_as_tuple` method from
        the `utils` module. Additional protected attributes `_y_t0_user` and
        `_y_tF_user` are set by post-appending either '_PX(t0)' or '_PX(tF)' to
        the user supplied symbols where the X is replaced by the phase suffix.
        As such if this phase has not yet been associated with an optimal
        control problem yet then `self` will not have attributes `_y_t0_user`
        and `_y_tF_user` and accessing either the `initial_state` or
        `final_state` property will raise an AttributeError.
        """
        return self._y_var_user

    @state_variables.setter
    def state_variables(self, y_vars: OptionalSymsType):

        self._y_var_user = format_as_named_tuple(y_vars)
        check_sym_name_clash(self._y_var_user)

        # Generate the state endpoint variable symbols only if phase has number
        if self.optimal_control_problem is not None:
            self._t0_USER = sym.Symbol(f't0_P{self._phase_suffix}')
            self._tF_USER = sym.Symbol(f'tF_P{self._phase_suffix}')
            self._t0 = sym.Symbol(f'_t0_P{self._phase_suffix}')
            self._tF = sym.Symbol(f'_tF_P{self._phase_suffix}')
            self._STRETCH = 0.5 * (self._tF - self._t0)
            self._SHIFT = 0.5 * (self._t0 + self._tF)
            self._t_var_user = (self._t0_USER, self._tF_USER)

            try:
                named_keys = self._y_var_user._fields
            except AttributeError:
                named_keys = ()

            self._y_t0_user = format_as_named_tuple(
                (sym.Symbol(f'{y}_P{self._phase_suffix}(t0)')
                 for y in self._y_var_user),
                named_keys=named_keys)
            self._y_tF_user = format_as_named_tuple(
                (sym.Symbol(f'{y}_P{self._phase_suffix}(tF)')
                 for y in self._y_var_user),
                named_keys=named_keys)

    @property
    def number_state_variables(self) -> int:
        """Integer number of state variables in the phase."""
        return len(self._y_var_user)

    @property
    def control_variables(self) -> TupleSymsType:
        """Symbols for this phase's control variables in order added by user.

        The user may supply either a single symbol or an iterable of symbols.
        The supplied argument is handled by the `format_as_tuple` method from
        the `utils` module.
        """
        return self._u_var_user

    @control_variables.setter
    def control_variables(self, u_vars: OptionalSymsType):
        self._u_var_user = format_as_named_tuple(u_vars)
        check_sym_name_clash(self._u_var_user)

    @property
    def number_control_variables(self) -> int:
        """Integer number of control variables in the phase."""
        return len(self._u_var_user)

    @property
    def integral_variables(self) -> TupleSymsType:
        """Symbols for this phase's integral variables.

        These symbols are auto generated as required by the user-supplied
        integrand functions.
        """
        return self._q_var_user

    @property
    def time_variables(self) -> TupleSymsType:
        """The initial and final time symbols as a pair."""
        return (self.initial_time_variable, self.final_time_variable)

    @property
    def number_integral_variables(self) -> int:
        """Integer number of integral variables in the phase."""
        return len(self._q_var_user)

    @property
    def state_equations(self) -> Tuple[sym.Expr, ...]:
        """User-supplied dynamical equations in the phase.

        These equations are the dynamical equations associated with each of the
        state variables in the phase. There should therefore be exactly one
        state equation for each dynamics symbol.

        State equations can be supplied in a compact form by the user defining additional auxiliary symbols and
        """
        return self._y_eqn_user

    @state_equations.setter
    def state_equations(self, y_eqns: OptionalExprsType):
        try:
            named_keys = self._y_var_user._fields
        except AttributeError:
            named_keys = ()
        self._y_eqn_user = format_as_named_tuple(y_eqns, use_named=True,
                                                  named_keys=named_keys)

    @property
    def number_state_equations(self) -> int:
        """Integer number of state equations in the phase.

        Should be the same as the number of state variables, i.e. there should
        be a direct mapping between the two.
        """
        return len(self._y_eqn_user)

    @property
    def path_constraints(self):
        return self._c_con_user

    @path_constraints.setter
    def path_constraints(self, c_cons):
        self._c_con_user = format_as_named_tuple(c_cons, use_named=False)

    @property
    def number_path_constraints(self):
        return len(self._c_con_user)

    @property
    def integrand_functions(self):
        return self._q_fnc_user

    @integrand_functions.setter
    def integrand_functions(self, integrands):
        self._q_fnc_user = format_as_named_tuple(integrands, use_named=False)
        self._q_var_user = tuple(sym.Symbol(f'q{i_q}_P{self._phase_suffix}')
                                  for i_q, _ in enumerate(self._q_fnc_user))

    @property
    def number_integrand_functions(self):
        return len(self._q_fnc_user)

    @property
    def bounds(self):
        return self._bounds

    @bounds.setter
    def bounds(self, bounds):
        if bounds is None:
            self._bounds = PhaseBounds(phase=self)
        else:
            self._bounds = bounds

    @property
    def scaling(self):
        return self._scaling

    @scaling.setter
    def scaling(self, scaling):
        if scaling is None:
            self._scaling = PhaseScaling(phase=self)
        else:
            self._scaling = scaling

    @property
    def mesh(self):
        return self._mesh

    @mesh.setter
    def mesh(self, mesh):
        if mesh is None:
            self._mesh = PhaseMesh(phase=self)
        else:
            self._mesh = mesh

    @property
    def guess(self):
        return self._guess

    @guess.setter
    def guess(self, guess):
        if guess is None:
            self._guess = PhaseGuess(phase=self)
        else:
            self._guess = guess

    def _check_variables_and_equations(self):
        """Check the user-supplied variables and equations for this OCP phase.

        Steps involved are:
            * Ensure that the same number of state variables and state
              equations are supplied.
            * Ensure that the symbols related to the state equations are the
              same set as the set of state variables.
            * If the two sets are not the same then issue an informative
              warning to the user describing what is wrong about the supplied
              state variables and state equations.

        Raises
        ------
        ValueError
            If the state variables and symbols associated with the state
            equations are not the same.

        """
        try:
            set_state_variables_keys = set(self.state_variables._fields)
        except AttributeError:
            set_state_variables_keys = set()

        try:
            set_state_equations_keys = set(self.state_equations._fields)
        except AttributeError:
            set_state_equations_keys = set()

        if set_state_variables_keys != set_state_equations_keys:

            intersection = set_state_variables_keys.intersection(
                set_state_equations_keys)
            spacer = "', '"
            msg_other = []

            if len(set_state_variables_keys) != len(set_state_equations_keys):
                if len(set_state_variables_keys) == 1:
                    msg_vars_len = (f"{len(set_state_variables_keys)} state "
                                    f"variable is")
                else:
                    msg_vars_len = (f"{len(set_state_variables_keys)} state "
                                    f"variables are")
                if len(set_state_equations_keys) == 1:
                    msg_eqns_len = (f"{len(set_state_equations_keys)} state "
                                    f"equation is")
                else:
                    msg_eqns_len = (f"{len(set_state_equations_keys)} state "
                                    f"equations are")
                msg_len = (f"{msg_vars_len} defined while {msg_eqns_len} "
                           f"supplied")
                msg_other.append(msg_len)

            only_in_variables = set_state_variables_keys.difference(
                intersection)
            if only_in_variables:
                immutable_only_in_variables = list(only_in_variables)
                if len(only_in_variables) == 1:
                    msg_vars = (f"the state variable "
                                f"'{immutable_only_in_variables[0]}' is defined "
                                f"without a state equation")
                else:
                    msg_vars = (f"the state variables "
                                f"'{spacer.join(immutable_only_in_variables[:-1])}' "
                                f"and '{immutable_only_in_variables[-1]}' are defined "
                                f"without state equations")
                msg_other.append(msg_vars)

            only_in_equations = set_state_equations_keys.difference(
                intersection)
            if only_in_equations:
                if len(only_in_equations) == 1:
                    immutable_only_in_equations = list(only_in_equations)
                    msg_eqns = (f"a state derivative is supplied for "
                                f"'{immutable_only_in_equations[0]}' which is not a "
                                f"state variable")
                else:
                    msg_eqns = (f"state derivatives are supplied for "
                                f"'{spacer.join(immutable_only_in_equations[:-1])}' "
                                f"and '{immutable_only_in_equations[-1]}' which are "
                                f"not defined as state variables")
                msg_other.append(msg_eqns)

            msg = ("A state equation must be supplied for each state variable "
                   f"in each phase. Currently in phase '{self.name}'")
            if len(msg_other) == 1:
                full_msg = (f"{msg}, {msg_other[0]}.")
            else:
                full_msg = (f"{msg}: {'; '.join(msg_other[:-1])}; and "
                            f"{msg_other[-1]}.")
            raise ValueError(full_msg)

    def __str__(self):
        string = (f"Phase {self.phase_number} of {self.optimal_control_problem}")
        return string

    def __repr__(self):
        string = (f"Phase({repr(self.optimal_control_problem)}, "
                  f"phase_number={self.phase_number})")
        return string


# -*- coding: utf-8 -*-

# BUG1989 is pleased to support the open source community by supporting ncnn available.
#
# Copyright (C) 2019 BUG1989. All rights reserved.
#
# Licensed under the BSD 3-Clause License (the "License"); you may not use this file except
# in compliance with the License. You may obtain a copy of the License at
#
# https://opensource.org/licenses/BSD-3-Clause
#
# Unless required by applicable law or agreed to in writing, software distributed
# under the License is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR
# CONDITIONS OF ANY KIND, either express or implied. See the License for the
# specific language governing permissions and limitations under the License.

"""
Quantization module for generating the calibration tables will be used by
quantized (INT8) models from FP32 models.with bucket split,[k, k, cin, cout]
cut into "cout" buckets.
This tool is based on Caffe Framework.
"""
from __future__ import division
from __future__ import print_function
import argparse
import numpy as np
import math, copy
import matplotlib.pyplot as plt
import sys,os
import time
import datetime
from scipy import stats

import torch
import torch.nn as nn
import torch.nn.functional as F

# np.set_printoptions(threshold='nan')
np.set_printoptions(suppress=True)

# def parse_args():
#     parser = argparse.ArgumentParser(
#         description='find the pretrained caffe models int8 quantize scale value')
#     parser.add_argument('--model', dest='model',
#                         help='path to pretrained weights', type=str)
#     parser.add_argument('--mean', dest='mean',
#                         help='value of mean', type=float, nargs=3)
#     parser.add_argument('--norm', dest='norm',
#                         help='value of normalize', type=float, nargs=1, default=1.0)
#     parser.add_argument('--images', dest='images',
#                         help='path to calibration images', type=str)
#     parser.add_argument('--output', dest='output',
#                         help='path to output calibration table file', type=str, default='calibration-dev.table')
#     parser.add_argument('--group', dest='group',
#                         help='enable the group scale', type=int, default=1)
#     parser.add_argument('--gpu', dest='gpu',
#                         help='use gpu to forward', type=int, default=0)

#     args = parser.parse_args()
#     return args, parser

# global args, parser
# args, parser = parse_args()

import torch
import torch.nn as nn
import torch.nn.functional as F

# global params
QUANTIZE_NUM = 127
QUANTIZE_WINOGRAND_NUM = 31
STATISTIC = 1
INTERVAL_NUM = 2048

# ugly global params
quantize_layer_lists = []

class QuantizeLayer:
    def __init__(self, name, blob_name, group_num):
        self.name = name
        self.blob_name = blob_name
        self.group_num = group_num
        self.weight_scale = np.zeros(group_num)
        self.blob_max = 0.0
        self.blob_distubution_interval = 0.0
        self.blob_distubution = np.zeros(INTERVAL_NUM)
        self.blob_threshold = 0
        self.blob_scale = 1.0
        self.group_zero = np.zeros(group_num)

    def quantize_weight(self, weight_data, flag):
        # spilt the weight data by cout num
        blob_group_data = np.array_split(weight_data, self.group_num)
        for i, group_data in enumerate(blob_group_data):
            max_val = np.max(group_data)
            min_val = np.min(group_data)
            threshold = max(abs(max_val), abs(min_val))
            if threshold < 0.0001:
                self.weight_scale[i] = 0
                self.group_zero[i] = 1
            else:
                if(flag == True):
                    self.weight_scale[i] = QUANTIZE_WINOGRAND_NUM / threshold
                else:
                    self.weight_scale[i] = QUANTIZE_NUM / threshold
            print("%-20s group : %-5d max_val : %-10f scale_val : %-10f" % (self.name + "_param0", i, threshold, self.weight_scale[i]))

    def initial_blob_max(self, blob_data):
        # get the max value of blob
        max_val = np.max(blob_data)
        min_val = np.min(blob_data)
        self.blob_max = max(self.blob_max, max(abs(max_val), abs(min_val)))

    def initial_blob_distubution_interval(self):
        self.blob_distubution_interval = STATISTIC * self.blob_max / INTERVAL_NUM
        print("%-20s max_val : %-10.8f distribution_intervals : %-10.8f" % (self.name, self.blob_max, self.blob_distubution_interval))

    def initial_histograms(self, blob_data):
        # collect histogram of every group channel blob
        th = self.blob_max
        hist, hist_edge = np.histogram(blob_data, bins=INTERVAL_NUM, range=(0, th))
        self.blob_distubution += hist

    def quantize_blob(self):
        # calculate threshold
        distribution = np.array(self.blob_distubution)
        # pick threshold which minimizes KL divergence
        threshold_bin = threshold_distribution(distribution)
        self.blob_threshold = threshold_bin
        threshold = (threshold_bin + 0.5) * self.blob_distubution_interval
        # get the activation calibration value
        self.blob_scale = QUANTIZE_NUM / threshold
        print("%-20s bin : %-8d threshold : %-10f interval : %-10f scale : %-10f" % (self.name, threshold_bin, threshold, self.blob_distubution_interval, self.blob_scale))

def _smooth_distribution(p, eps=0.0001):
    """Given a discrete distribution (may have not been normalized to 1),
    smooth it by replacing zeros with eps multiplied by a scaling factor and taking the
    corresponding amount off the non-zero values.
    Ref: http://web.engr.illinois.edu/~hanj/cs412/bk3/KL-divergence.pdf
    """
    is_zeros = (p == 0).astype(np.float32)
    is_nonzeros = (p != 0).astype(np.float32)
    n_zeros = is_zeros.sum()
    n_nonzeros = p.size - n_zeros
    if not n_nonzeros:
        raise ValueError('The discrete probability distribution is malformed. All entries are 0.')
    eps1 = eps * float(n_zeros) / float(n_nonzeros)
    assert eps1 < 1.0, 'n_zeros=%d, n_nonzeros=%d, eps1=%f' % (n_zeros, n_nonzeros, eps1)
    hist = p.astype(np.float32)
    hist += eps * is_zeros + (-eps1) * is_nonzeros
    assert (hist <= 0).sum() == 0
    return hist

def threshold_distribution(distribution, target_bin=128):
    """
    Return the best threshold value.
    Ref: https://github.com//apache/incubator-mxnet/blob/master/python/mxnet/contrib/quantization.py
    Args:
        distribution: list, activations has been processed by histogram and normalize,size is 2048
        target_bin: int, the num of bin that is used by quantize, Int8 default value is 128
    Returns:
        target_threshold: int, num of bin with the minimum KL
    """
    distribution = distribution[1:]
    length = distribution.size
    threshold_sum = sum(distribution[target_bin:])
    kl_divergence = np.zeros(length - target_bin)

    for threshold in range(target_bin, length):
        sliced_nd_hist = copy.deepcopy(distribution[:threshold])

        # generate reference distribution p
        p = sliced_nd_hist.copy()
        p[threshold-1] += threshold_sum
        threshold_sum = threshold_sum - distribution[threshold]

        # is_nonzeros[k] indicates whether hist[k] is nonzero
        is_nonzeros = (p != 0).astype(np.int64)
        #
        quantized_bins = np.zeros(target_bin, dtype=np.int64)
        # calculate how many bins should be merged to generate quantized distribution q
        num_merged_bins = sliced_nd_hist.size // target_bin

        # merge hist into num_quantized_bins bins
        for j in range(target_bin):
            start = j * num_merged_bins
            stop = start + num_merged_bins
            quantized_bins[j] = sliced_nd_hist[start:stop].sum()
        quantized_bins[-1] += sliced_nd_hist[target_bin * num_merged_bins:].sum()

        # expand quantized_bins into p.size bins
        q = np.zeros(sliced_nd_hist.size, dtype=np.float64)
        for j in range(target_bin):
            start = j * num_merged_bins
            if j == target_bin - 1:
                stop = -1
            else:
                stop = start + num_merged_bins
            norm = is_nonzeros[start:stop].sum()
            if norm != 0:
                q[start:stop] = float(quantized_bins[j]) / float(norm)
        q[p == 0] = 0
        # p = _smooth_distribution(p) # with some bugs, need to fix
        # q = _smooth_distribution(q)
        p[p == 0] = 0.0001
        q[q == 0] = 0.0001

        # calculate kl_divergence between q and p
        kl_divergence[threshold - target_bin] = stats.entropy(p, q)

    min_kl_divergence = np.argmin(kl_divergence)
    threshold_value = min_kl_divergence + target_bin

    return threshold_value

def net_forward(net, image):
    """
    network inference and statistics the cost time
    Args:
        net: the instance of PyTorch inference
        image: a image need to be inference
    Returns:
        none
    """
    # load image

    # transformer.preprocess the image

    # net forward
    net(image)

def file_name(file_dir):
    """
    Find the all file path with the directory
    Args:
        file_dir: The source file directory
    Returns:
        files_path: all the file path into a list
    """
    files_path = []

    for root, dir, files in os.walk(file_dir):
        for name in files:
            file_path = root + "/" + name
            print(file_path)
            files_path.append(file_path)

    return files_path

def network_prepare(net, mean, norm):
    """
    instance the prepare process param of PyTorch network inference
    Args:
        net: the instance of PyTorch inference
        mean: the value of mean
        norm: the value of normalize
    Returns:
        none
    """
    print("Network initial")

    # img_mean = np.array(mean)

    # # initial transformer
    # transformer = caffe.io.Transformer({'data': net.blobs['data'].data.shape})
    # # convert hwc to cwh
    # transformer.set_transpose('data', (2,0,1))
    # # load meanfile
    # transformer.set_mean('data', img_mean)
    # # resize image data from [0,1] to [0,255]
    # transformer.set_raw_scale('data', 255)
    # # convert RGB -> BGR
    # transformer.set_channel_swap('data', (2,1,0))
    # # normalize
    # transformer.set_input_scale('data', norm)

    # return transformer

def weight_quantize(net):
    """
    PyTorch convolution weight blob Int8 quantize
    Args:
        net: the instance of PyTorch inference
    Returns:
        none
    """
    print("\nQuantize the kernel weight:")

    for name, layer in net.named_modules():
        # find the convolution layers to get out the weight_scale
        if isinstance(layer, nn.Conv2d):
            weight_blob = layer.weight.detach().numpy()
            quanitze_layer = QuantizeLayer(name, weight_blob, layer.out_channels)
            # quantize the weight value using 6bit for conv3x3s1 layer to winograd F(4,3)
            if layer.kernel_size[0] == 3 and layer.stride[0] == 1 and layer.groups != layer.out_channels:
                quanitze_layer.quantize_weight(weight_blob, True)
            # quantize the weight value using 8bit for another conv layers
            else:
                quanitze_layer.quantize_weight(weight_blob, False)
            # add the quantize_layer into the save list
            quantize_layer_lists.append(quanitze_layer)

    return None

class Hook_struct:
    def __init__(self, name, hook):
        self.name = name
        self.hook = hook

hook_list = []
input_list = []

def get_feature(modules, input):
    input_list.append(input)

def activation_quantize(net, images_files):
    """
    Activation Int8 quantize, optimaize threshold selection with KL divergence,
    given a dataset, find the optimal threshold for quantizing it.
    Ref: http://on-demand.gputechconf.com/gtc/2017/presentation/s7310-8-bit-inference-with-tensorrt.pdf
    Args:
        net: the instance of Caffe inference
        transformer:
        images_files: calibration dataset
    Returns:
        none
    """
    print("\nQuantize the Activation:")
    print("image num:%d" % len(images_files))

    # register the hook1
    for name, layer in net.named_modules():
        if isinstance(layer, nn.Conv2d):
            hook = layer.register_forward_pre_hook(get_feature)
            hook_status = Hook_struct(name, hook)
            hook_list.append(hook_status)

    # run float32 inference on calibration dataset to find the activations range
    for i, image in enumerate(images_files):
        # inference
        # just to test
        image = torch.rand(1, 3, 32, 32)
        net(image)
        # find max threshold
        for j, layer in enumerate(quantize_layer_lists):
            blob = input_list[j][0].detach().numpy().flatten()
            layer.initial_blob_max(blob)
        # if i % 100 == 0:
            print("loop stage 1 : %d/%d" % (i, len(images_files)))

    # calculate statistic blob scope and interval distribution
    for layer in quantize_layer_lists:
        layer.initial_blob_distubution_interval()

    input_list.clear()

    # for each layers
    # collect histograms of activations
    print("\nCollect histograms of activations:")
    for i, image in enumerate(images_files):
        image = torch.rand(1, 3, 32, 32)
        net(image)
        for j, layer in enumerate(quantize_layer_lists):
            blob = input_list[j][0].detach().numpy().flatten()
            layer.initial_histograms(blob)
        # if i % 100 == 0:
            print("loop stage 2 : %d/%d" % (i, len(images_files)))

    # calculate threshold with KL divergence
    for layer in quantize_layer_lists:
        layer.quantize_blob()

    # remove the hook
    for hook in hook_list:
        hook.hook.remove()

    return None

def save_calibration_file(calibration_path):
    calibration_file = open(calibration_path, 'w')
    # save temp
    save_temp = []
    # save weight scale
    for layer in quantize_layer_lists:
        save_string = layer.name + "_param_0"
        for i in range(layer.group_num):
            save_string = save_string + " " + str(layer.weight_scale[i])
        save_temp.append(save_string)

    # save bottom blob scales
    for layer in quantize_layer_lists:
        save_string = layer.name + " " + str(layer.blob_scale)
        save_temp.append(save_string)

    # save into txt file
    for data in save_temp:
        calibration_file.write(data + "\n")

    calibration_file.close()

    # save calibration logs
    save_temp_log = []
    calibration_file_log = open(calibration_path + ".log", 'w')
    for layer in quantize_layer_lists:
        save_string = layer.name + ": value range 0 - " + str(layer.blob_max) \
                                 + ", interval " + str(layer.blob_distubution_interval) \
                                 + ", interval num " + str(INTERVAL_NUM) \
                                 + ", threshold num " + str(layer.blob_threshold) + "\n" \
                                 + str(layer.blob_distubution.astype(dtype=np.int64))
        save_temp_log.append(save_string)

    # save into txt file
    for data in save_temp_log:
        calibration_file_log.write(data + "\n")

def usage_info():
    """
    usage info
    """
    print("Input params is illegal...╮(╯3╰)╭")
    print("try it again:\n python pytorch-int8-convert-tool.py -h")

class Net(nn.Module):
    def __init__(self):
        super(Net, self).__init__()
        self.conv1 = nn.Conv2d(3, 6, 5)
        self.pool = nn.MaxPool2d(2, 2)
        self.conv2 = nn.Conv2d(6, 16, 5)
        self.fc1 = nn.Linear(16 * 5 * 5, 120)
        self.fc2 = nn.Linear(120, 84)
        self.fc3 = nn.Linear(84, 10)

    def forward(self, x):
        x = self.pool(F.relu(self.conv1(x)))
        x = self.pool(F.relu(self.conv2(x)))
        x = x.view(-1, 16 * 5 * 5)
        x = F.relu(self.fc1(x))
        x = F.relu(self.fc2(x))
        x = self.fc3(x)
        return x

def main():
    """
    main function
    """

    # time start
    time_start = datetime.datetime.now()

    # print(args)

    # if args.proto == None or args.model == None or args.mean == None or args.images == None:
    #     usage_info()
    #     return None

    # # trained pytorch path
    # pytorch_model = args.model

    # # mean value
    # mean = args.mean

    # # norm value
    # norm = 1.0
    # if args.norm != 1.0:
    #     norm = args.norm[0]

    # # calibration dataset
    # images_path = args.images
    images_path = './Images'

    # # the output calibration file
    # calibration_path = args.output

    # # enable the group scale
    # group_on = args.group

    # # default use CPU to forwark
    # if args.gpu != 0:
    #     print("gpu status: %d" % (args.gpu))

    # initial caffe net and the forword model(GPU or CPU)
    net = torch.load('./cifar10.pkl')
    print(net)

    # prepare the cnn network
    # transformer = network_prepare(net, mean, norm)

    # get the calibration datasets images files path
    images_files = file_name(images_path)

    # quanitze kernel weight of the caffemodel to find it's calibration table
    weight_quantize(net)

    # quantize activation value of the caffemodel to find it's calibration table
    activation_quantize(net, images_files)

    # save the calibration tables,best wish for your INT8 inference have low accuracy loss :)
    save_calibration_file('./cifar10.table')

    # time end
    time_end = datetime.datetime.now()

    print("\nPyTorch Int8 Calibration table is done, it's cost %s, best wish for your INT8 inference has a low accuracy loss...\(^▽^)/..." % (time_end - time_start))

if __name__ == "__main__":
    main()


"""
desitarget.cuts
===============

An old copy of the Main Survey cuts (../cuts.py) that were used for commissioning (cmx).

.. _`the Gaia data model`: https://gea.esac.esa.int/archive/documentation/GDR2/Gaia_archive/chap_datamodel/sec_dm_main_tables/ssec_dm_gaia_source.html
.. _`the Legacy Surveys`: http://www.legacysurvey.org/
.. _`the wiki`: https://desi.lbl.gov/trac/wiki/TargetSelectionWG/TargetSelection
.. _`Legacy Surveys mask`: http://www.legacysurvey.org/dr8/bitmasks/
"""

import numpy as np
from pkg_resources import resource_filename
from desitarget.geomask import imaging_mask

# ADM set up the DESI default logger
from desiutil.log import get_logger
log = get_logger()

def shift_photo_north(gflux=None, rflux=None, zflux=None):
    """Convert fluxes in the northern (BASS/MzLS) to the southern (DECaLS) system.

    Parameters
    ----------
    gflux, rflux, zflux : :class:`array_like` or `float`
        The flux in nano-maggies of g, r, z bands.

    Returns
    -------
    The equivalent fluxes shifted to the southern system.

    Notes
    -----
    - see also https://desi.lbl.gov/DocDB/cgi-bin/private/RetrieveFile?docid=3390;filename=Raichoor_DESI_05Dec2017.pdf;version=1
    - Update for DR9 https://desi.lbl.gov/trac/attachment/wiki/TargetSelectionWG/TargetSelection/North_vs_South_dr9.png
    """
    # ADM if floats were sent, treat them like arrays.
    flt = False
    if _is_row(gflux):
        flt = True
        gflux = np.atleast_1d(gflux)
        rflux = np.atleast_1d(rflux)
        zflux = np.atleast_1d(zflux)

    # ADM only use the g-band color shift when r and g are non-zero
    gshift = gflux * 10**(-0.4*0.004)
    w = np.where((gflux != 0) & (rflux != 0))
    gshift[w] = (gflux[w] * 10**(-0.4*0.004) * (gflux[w]/rflux[w])**complex(-0.059)).real

    # ADM only use the r-band color shift when r and z are non-zero
    # ADM and only use the z-band color shift when r and z are non-zero
    w = np.where((rflux != 0) & (zflux != 0))
    rshift = rflux * 10**(0.4*0.003)
    zshift = zflux * 10**(0.4*0.013)

    rshift[w] = (rflux[w] * 10**(0.4*0.003) * (rflux[w]/zflux[w])**complex(-0.024)).real
    zshift[w] = (zflux[w] * 10**(0.4*0.013) * (rflux[w]/zflux[w])**complex(+0.015)).real

    if flt:
        return gshift[0], rshift[0], zshift[0]

    return gshift, rshift, zshift

def isLRG(gflux=None, rflux=None, zflux=None, w1flux=None, w2flux=None,
          zfiberflux=None, rfluxivar=None, zfluxivar=None, w1fluxivar=None,
          gnobs=None, rnobs=None, znobs=None, maskbits=None, primary=None,
          south=True):
    """
    Parameters
    ----------
    south: boolean, defaults to ``True``
        Use cuts appropriate to the Northern imaging surveys (BASS/MzLS)
        if ``south=False``, otherwise use cuts appropriate to the
        Southern imaging survey (DECaLS).

    Returns
    -------
    :class:`array_like`
        ``True`` if and only if the object is an LRG target.

    Notes
    -----
    - Current version (12/07/2020) is version 232 on `the wiki`_.
    - See :func:`~desitarget.cuts.set_target_bits` for other parameters.
    """
    # ADM LRG targets.
    if primary is None:
        primary = np.ones_like(rflux, dtype='?')
    lrg = primary.copy()

    # ADM basic quality cuts.
    lrg &= notinLRG_mask(
        primary=primary, rflux=rflux, zflux=zflux, w1flux=w1flux,
        zfiberflux=zfiberflux, gnobs=gnobs, rnobs=rnobs, znobs=znobs,
        rfluxivar=rfluxivar, zfluxivar=zfluxivar, w1fluxivar=w1fluxivar,
        maskbits=maskbits
    )

    # ADM color-based selection of LRGs.
    lrg &= isLRG_colors(
        gflux=gflux, rflux=rflux, zflux=zflux, w1flux=w1flux,
        zfiberflux=zfiberflux, south=south, primary=primary
    )

    return lrg

def notinLRG_mask(primary=None, rflux=None, zflux=None, w1flux=None,
                  zfiberflux=None, gnobs=None, rnobs=None, znobs=None,
                  rfluxivar=None, zfluxivar=None, w1fluxivar=None,
                  maskbits=None):
    """See :func:`~desitarget.cuts.isLRG` for details.

    Returns
    -------
    :class:`array_like`
        ``True`` if and only if the object is NOT masked for poor quality.
    """
    if primary is None:
        primary = np.ones_like(rflux, dtype='?')
    lrg = primary.copy()

    # ADM to maintain backwards-compatibility with mocks.
    if zfiberflux is None:
        log.warning('Setting zfiberflux to zflux!!!')
        zfiberflux = zflux.copy()

    lrg &= (rfluxivar > 0) & (rflux > 0)   # ADM quality in r.
    lrg &= (zfluxivar > 0) & (zflux > 0) & (zfiberflux > 0)   # ADM quality in z.
    lrg &= (w1fluxivar > 0) & (w1flux > 0)  # ADM quality in W1.

    # ADM observed in every band.
    lrg &= (gnobs > 0) & (rnobs > 0) & (znobs > 0)

    # ADM default mask bits from the Legacy Surveys not set.
    lrg &= imaging_mask(maskbits)

    return lrg

def isLRG_colors(gflux=None, rflux=None, zflux=None, w1flux=None,
                 zfiberflux=None, ggood=None,
                 w2flux=None, primary=None, south=True):
    """(see, e.g., :func:`~desitarget.cuts.isLRG`).

    Notes:
        - the `ggood` and `w2flux` inputs are an attempt to maintain
          backwards-compatibility with the mocks.
    """
    if primary is None:
        primary = np.ones_like(rflux, dtype='?')
    lrg = primary.copy()

    # ADM to maintain backwards-compatibility with mocks.
    if zfiberflux is None:
        log.warning('Setting zfiberflux to zflux!!!')
        zfiberflux = zflux.copy()

    gmag = 22.5 - 2.5 * np.log10(gflux.clip(1e-7))
    # ADM safe as these fluxes are set to > 0 in notinLRG_mask.
    rmag = 22.5 - 2.5 * np.log10(rflux.clip(1e-7))
    zmag = 22.5 - 2.5 * np.log10(zflux.clip(1e-7))
    w1mag = 22.5 - 2.5 * np.log10(w1flux.clip(1e-7))
    zfibermag = 22.5 - 2.5 * np.log10(zfiberflux.clip(1e-7))

    if south:
        lrg &= zmag - w1mag > 0.8 * (rmag-zmag) - 0.6    # non-stellar cut.
        lrg &= (
            ((gmag - w1mag > 2.6) & (gmag - rmag > 1.4))
            | (rmag - w1mag > 1.8)                       # low-z cut.
        )
        lrg &= rmag - zmag > (zmag - 16.83) * 0.45       # double sliding cut 1.
        lrg &= rmag - zmag > (zmag - 13.80) * 0.19       # double sliding cut 2.
    else:
        lrg &= zmag - w1mag > 0.8 * (rmag-zmag) - 0.6   # non-stellar cut.
        lrg &= (
            ((gmag - w1mag > 2.67) & (gmag - rmag > 1.45))
            | (rmag - w1mag > 1.85)                      # low-z cut.
        )
        lrg &= rmag - zmag > (zmag - 16.79) * 0.45       # double sliding cut 1.
        lrg &= rmag - zmag > (zmag - 13.76) * 0.19       # double sliding cut 2.

    lrg &= zfibermag < 21.5    # faint limit.

    return lrg

def isELG(gflux=None, rflux=None, zflux=None, w1flux=None, w2flux=None,
          gsnr=None, rsnr=None, zsnr=None, gnobs=None, rnobs=None, znobs=None,
          maskbits=None, south=True, primary=None):
    """Definition of ELG target classes. Returns a boolean array.
    (see :func:`~desitarget.cuts.set_target_bits` for parameters).

    Notes:
    - Current version (12/09/20) is version 233 on `the wiki`_.
    """
    if primary is None:
        primary = np.ones_like(rflux, dtype='?')
    elg = primary.copy()

    elg &= notinELG_mask(maskbits=maskbits, gsnr=gsnr, rsnr=rsnr, zsnr=zsnr,
                         gnobs=gnobs, rnobs=rnobs, znobs=znobs, primary=primary)

    elg &= isELG_colors(gflux=gflux, rflux=rflux, zflux=zflux, w1flux=w1flux,
                        w2flux=w2flux, south=south, primary=primary)

    return elg

def notinELG_mask(maskbits=None, gsnr=None, rsnr=None, zsnr=None,
                  gnobs=None, rnobs=None, znobs=None, primary=None):
    """Standard set of masking cuts used by all ELG target selection classes.
    (see :func:`~desitarget.cuts.set_target_bits` for parameters).
    """
    if primary is None:
        primary = np.ones_like(maskbits, dtype='?')
    elg = primary.copy()

    # ADM good signal-to-noise in all bands.
    elg &= (gsnr > 0) & (rsnr > 0) & (zsnr > 0)

    # ADM observed in every band.
    elg &= (gnobs > 0) & (rnobs > 0) & (znobs > 0)

    # ADM default mask bits from the Legacy Surveys not set.
    elg &= imaging_mask(maskbits)

    return elg

def isELG_colors(gflux=None, rflux=None, zflux=None, w1flux=None,
                 w2flux=None, south=True, primary=None):
    """Color cuts for ELG target selection classes
    (see, e.g., :func:`~desitarget.cuts.set_target_bits` for parameters).
    """
    if primary is None:
        primary = np.ones_like(rflux, dtype='?')
    elg = primary.copy()

    # ADM work in magnitudes instead of fluxes. NOTE THIS IS ONLY OK AS
    # ADM the snr masking in ALL OF g, r AND z ENSURES positive fluxes.
    g = 22.5 - 2.5*np.log10(gflux.clip(1e-16))
    r = 22.5 - 2.5*np.log10(rflux.clip(1e-16))
    z = 22.5 - 2.5*np.log10(zflux.clip(1e-16))

    # ADM cuts shared by the northern and southern selections.
    elg &= g > 20                       # bright cut.
    elg &= r - z > 0.3                  # blue cut.
    elg &= r - z < 1.6                  # red cut.
    elg &= g - r < -1.2*(r - z) + 1.6   # OII flux cut.

    # ADM cuts that are unique to the north or south.
    if south:
        elg &= g < 23.4  # faint cut.
        # ADM south has the FDR cut to remove stars and low-z galaxies.
        elg &= g - r < 1.15*(r - z) - 0.15
    else:
        elg &= g < 23.5  # faint cut.
        elg &= g - r < 1.15*(r - z) - 0.20  # remove stars and low-z galaxies.

    return elg

def _check_BGS_targtype(targtype):
    """Fail if `targtype` is not one of the strings 'bright', 'faint' or 'wise'.
    """
    targposs = ['faint', 'bright', 'wise']

    if targtype not in targposs:
        msg = 'targtype must be one of {} not {}'.format(targposs, targtype)
        log.critical(msg)
        raise ValueError(msg)

def isBGS(rfiberflux=None, gflux=None, rflux=None, zflux=None, w1flux=None, w2flux=None,
          gnobs=None, rnobs=None, znobs=None, gfracmasked=None, rfracmasked=None, zfracmasked=None,
          gfracflux=None, rfracflux=None, zfracflux=None, gfracin=None, rfracin=None, zfracin=None,
          gfluxivar=None, rfluxivar=None, zfluxivar=None, maskbits=None, Grr=None, refcat=None,
          w1snr=None, gaiagmag=None, objtype=None, primary=None, south=True, targtype=None):
    """Definition of BGS target classes. Returns a boolean array.

    Args
    ----
    targtype: str, optional, defaults to ``faint``
        Pass ``bright`` to use colors appropriate to the ``BGS_BRIGHT`` selection
        or ``faint`` to use colors appropriate to the ``BGS_FAINT`` selection
        or ``wise`` to use colors appropriate to the ``BGS_WISE`` selection.

    Returns
    -------
    :class:`array_like`
        ``True`` if and only if the object is a BGS target of type ``targtype``.

    Notes
    -----
    - Current version (10/24/18) is version 143 on `the wiki`_.
    - See :func:`~desitarget.cuts.set_target_bits` for other parameters.
    """
    _check_BGS_targtype(targtype)

    # ------ Bright Galaxy Survey
    if primary is None:
        primary = np.ones_like(rflux, dtype='?')
    bgs = primary.copy()

    bgs &= notinBGS_mask(gnobs=gnobs, rnobs=rnobs, znobs=znobs, primary=primary,
                         gfracmasked=gfracmasked, rfracmasked=rfracmasked, zfracmasked=zfracmasked,
                         gfracflux=gfracflux, rfracflux=rfracflux, zfracflux=zfracflux,
                         gfracin=gfracin, rfracin=rfracin, zfracin=zfracin, w1snr=w1snr,
                         gfluxivar=gfluxivar, rfluxivar=rfluxivar, zfluxivar=zfluxivar, Grr=Grr,
                         gaiagmag=gaiagmag, maskbits=maskbits, targtype=targtype)

    bgs &= isBGS_colors(rfiberflux=rfiberflux, gflux=gflux, rflux=rflux, zflux=zflux, w1flux=w1flux,
                        w2flux=w2flux, south=south, targtype=targtype, primary=primary)

    bgs |= isBGS_lslga(gflux=gflux, rflux=rflux, zflux=zflux, w1flux=w1flux, refcat=refcat,
                       maskbits=maskbits, south=south, targtype=targtype)

    return bgs

def notinBGS_mask(gnobs=None, rnobs=None, znobs=None, primary=None,
                  gfracmasked=None, rfracmasked=None, zfracmasked=None,
                  gfracflux=None, rfracflux=None, zfracflux=None,
                  gfracin=None, rfracin=None, zfracin=None, w1snr=None,
                  gfluxivar=None, rfluxivar=None, zfluxivar=None, Grr=None,
                  gaiagmag=None, maskbits=None, targtype=None):
    """Standard set of masking cuts used by all BGS target selection classes
    (see, e.g., :func:`~desitarget.cuts.isBGS` for parameters).
    """
    _check_BGS_targtype(targtype)

    if primary is None:
        primary = np.ones_like(gnobs, dtype='?')
    bgs = primary.copy()

    bgs &= (gnobs >= 1) & (rnobs >= 1) & (znobs >= 1)
    bgs &= (gfracmasked < 0.4) & (rfracmasked < 0.4) & (zfracmasked < 0.4)
    bgs &= (gfracflux < 5.0) & (rfracflux < 5.0) & (zfracflux < 5.0)
    bgs &= (gfracin > 0.3) & (rfracin > 0.3) & (zfracin > 0.3)
    bgs &= (gfluxivar > 0) & (rfluxivar > 0) & (zfluxivar > 0)

    # ADM geometric masking cuts from the Legacy Surveys.
    bgs &= imaging_mask(maskbits)

    if targtype == 'bright':
        bgs &= ((Grr > 0.6) | (gaiagmag == 0))
    elif targtype == 'faint':
        bgs &= ((Grr > 0.6) | (gaiagmag == 0))
    elif targtype == 'wise':
        bgs &= Grr < 0.4
        bgs &= Grr > -1
        bgs &= w1snr > 5

    return bgs

def isBGS_colors(rfiberflux=None, gflux=None, rflux=None, zflux=None, w1flux=None,
                 w2flux=None, south=True, targtype=None, primary=None):
    """Standard set of color-based cuts used by all BGS target selection classes
    (see, e.g., :func:`~desitarget.cuts.isBGS` for parameters).
    """
    _check_BGS_targtype(targtype)

    # ADM to maintain backwards-compatibility with mocks.
    if rfiberflux is None:
        log.warning('Setting rfiberflux to rflux!!!')
        rfiberflux = rflux.copy()

    if primary is None:
        primary = np.ones_like(rflux, dtype='?')
    bgs = primary.copy()
    fmc = np.zeros_like(rflux, dtype='?')

    if south:
        bgs &= rflux > gflux * 10**(-1.0/2.5)
        bgs &= rflux < gflux * 10**(4.0/2.5)
        bgs &= zflux > rflux * 10**(-1.0/2.5)
        bgs &= zflux < rflux * 10**(4.0/2.5)
    else:
        bgs &= rflux > gflux * 10**(-1.0/2.5)
        bgs &= rflux < gflux * 10**(4.0/2.5)
        bgs &= zflux > rflux * 10**(-1.0/2.5)
        bgs &= zflux < rflux * 10**(4.0/2.5)

    g = 22.5 - 2.5*np.log10(gflux.clip(1e-16))
    r = 22.5 - 2.5*np.log10(rflux.clip(1e-16))
    z = 22.5 - 2.5*np.log10(zflux.clip(1e-16))
    rfib = 22.5 - 2.5*np.log10(rfiberflux.clip(1e-16))

    # Fibre Magnitude Cut (FMC) -- This is a low surface brightness cut
    # with the aim of increase the redshift success rate.
    fmc |= ((rfib < (2.9 + 1.2 + 1.0) + r) & (r < 17.8))
    fmc |= ((rfib < 22.9) & (r < 20.0) & (r > 17.8))
    fmc |= ((rfib < 2.9 + r) & (r > 20))

    bgs &= fmc

    if targtype == 'bright':
        bgs &= rflux > 10**((22.5-19.5)/2.5)
    elif targtype == 'faint':
        bgs &= rflux > 10**((22.5-20.0)/2.5)
        bgs &= rflux <= 10**((22.5-19.5)/2.5)
    elif targtype == 'wise':
        bgs &= rflux > 10**((22.5-20.0)/2.5)
        bgs &= w1flux*gflux > (zflux*rflux)*10**(-0.2)

    return bgs

def isBGS_lslga(gflux=None, rflux=None, zflux=None, w1flux=None, refcat=None,
                maskbits=None, south=True, targtype=None):
    """Module to recover the LSLGA objects in all BGS target selection classes
    (see, e.g., :func:`~desitarget.cuts.isBGS` for parameters).
    """
    _check_BGS_targtype(targtype)

    bgs = np.zeros_like(rflux, dtype='?')

    # the LSLGA galaxies.
    LX = bgs.copy()
    # ADM Could check on "L2" for DR8, need to check on "LX" post-DR8.
    if refcat is not None:
        rc1d = np.atleast_1d(refcat)
        if isinstance(rc1d[0], str):
            LX = [(rc[0] == "L") if len(rc) > 0 else False for rc in rc1d]
        else:
            LX = [(rc.decode()[0] == "L") if len(rc) > 0 else False for rc in rc1d]
        if np.ndim(refcat) == 0:
            LX = np.array(LX[0], dtype=bool)
        else:
            LX = np.array(LX, dtype=bool)

    bgs |= LX
    # ADM geometric masking cuts from the Legacy Surveys.
    bgs &= imaging_mask(maskbits, bgsmask=True)

    if targtype == 'bright':
        bgs &= rflux > 10**((22.5-19.5)/2.5)
    elif targtype == 'faint':
        bgs &= rflux > 10**((22.5-20.0)/2.5)
        bgs &= rflux <= 10**((22.5-19.5)/2.5)
    elif targtype == 'wise':
        bgs &= rflux > 10**((22.5-20.0)/2.5)
        bgs &= w1flux*gflux > (zflux*rflux)*10**(-0.2)

    return bgs

def isQSO_randomforest(gflux=None, rflux=None, zflux=None, maskbits=None,
                       w1flux=None, w2flux=None, objtype=None, release=None,
                       gnobs=None, rnobs=None, znobs=None, deltaChi2=None,
                       primary=None, ra=None, dec=None, south=True, return_probs=False):
    """Define QSO targets from a Random Forest. Returns a boolean array.

    Parameters
    ----------
    south : :class:`boolean`, defaults to ``True``
        If ``False``, shift photometry to the Northern (BASS/MzLS)
        imaging system.
    return_probs : :class:`boolean`, defaults to ``False``
        If ``True``, return the QSO/high-z QSO probabilities in addition
        to the QSO target booleans. Only coded up for DR8 or later of the
        Legacy Surveys. Will return arrays of zeros for earlier DRs.

    Returns
    -------
    :class:`array_like`
        ``True`` for objects that are Random Forest quasar targets.
    :class:`array_like`
        ``True`` for objects that are high-z RF quasar targets.
    :class:`array_like`
        The (float) probability that a target is a quasar. Only returned
        if `return_probs` is ``True``.
    :class:`array_like`
        The (float) probability that a target is a high-z quasar. Only
        returned if `return_probs` is ``True``.

    Notes
    -----
    - Current version (20/11/20) is version 173 on `the wiki`_.
    - See :func:`~desitarget.cuts.set_target_bits` for other parameters.
    """
    # ADM Primary (True for anything to initially consider as a possible target).
    if primary is None:
        primary = np.ones_like(gflux, dtype=bool)

    # RELEASE
    # ADM default to RELEASE of 5000 if nothing is passed.
    if release is None:
        release = np.zeros_like(gflux, dtype='?') + 5000
    release = np.atleast_1d(release)

    # Build variables for random forest
    nFeatures = 11   # Number of attributes describing each object to be classified by the rf
    nbEntries = rflux.size
    if not south:
        gflux, rflux, zflux = shift_photo_north(gflux, rflux, zflux)

    colors, r, photOK = _getColors(nbEntries, nFeatures, gflux, rflux, zflux, w1flux, w2flux)
    r = np.atleast_1d(r)

    # Preselection to speed up the process
    rMax = 22.7   # r < 22.7
    rMin = 17.5   # r > 17.5
    preSelection = (r < rMax) & (r > rMin) & photOK & primary

    # ADM targets have to be observed in every band.
    preSelection &= (gnobs > 0) & (rnobs > 0) & (znobs > 0)

    if objtype is not None:
        preSelection &= _psflike(objtype)
    if deltaChi2 is not None:
        deltaChi2 = np.atleast_1d(deltaChi2)
        preSelection[release < 5000] &= deltaChi2[release < 5000] > 30.
    # ADM Reject objects in masks.
    # ADM BRIGHT BAILOUT GALAXY CLUSTER (1, 10, 12, 13) bits not set.
    # ALLMASK_G	| ALLMASK_R | ALLMASK_Z (5, 6, 7) bits not set.
    # Now only 1, 12, 13
    if maskbits is not None:
        # ADM default mask bits from the Legacy Surveys not set.
        preSelection &= imaging_mask(maskbits)

    # "qso" mask initialized to "preSelection" mask.
    qso = np.copy(preSelection)
    # ADM to specifically store the selection from the "HighZ" RF.
    qsohiz = np.copy(preSelection)

    # ADM these store the probabilities, should they need returned.
    pqso = np.zeros_like(qso, dtype='>f4')
    pqsohiz = np.zeros_like(qso, dtype='>f4')

    if np.any(preSelection):

        from desitarget.myRF import myRF

        # Data reduction to preselected objects
        colorsReduced = colors[preSelection]
        releaseReduced = release[preSelection]
        r_Reduced = r[preSelection]
        colorsIndex = np.arange(0, nbEntries, dtype=np.int64)
        colorsReducedIndex = colorsIndex[preSelection]

        # Path to random forest files
        pathToRF = resource_filename('desitarget', 'data')
        # rf filenames
        rf_DR3_fileName = pathToRF + '/rf_model_dr3.npz'
        rf_DR7_fileName = pathToRF + '/rf_model_dr7.npz'
        rf_DR7_HighZ_fileName = pathToRF + '/rf_model_dr7_HighZ.npz'
        rf_DR8_fileName = pathToRF + '/rf_model_dr8.npz'
        rf_DR8_HighZ_fileName = pathToRF + '/rf_model_dr8_HighZ.npz'
        rf_DR9_fileName = pathToRF + '/rf_model_dr9.npz'
        rf_DR9_HighZ_fileName = pathToRF + '/rf_model_dr9_HighZ.npz'

        tmpReleaseOK = releaseReduced < 5000
        if np.any(tmpReleaseOK):
            # rf initialization - colors data duplicated within "myRF"
            rf_DR3 = myRF(colorsReduced[tmpReleaseOK], pathToRF,
                          numberOfTrees=200, version=1)
            # rf loading
            rf_DR3.loadForest(rf_DR3_fileName)
            # Compute rf probabilities
            tmp_rf_proba = rf_DR3.predict_proba()
            tmp_r_Reduced = r_Reduced[tmpReleaseOK]
            # Compute optimized proba cut
            pcut = np.where(tmp_r_Reduced > 20.0,
                            0.95 - (tmp_r_Reduced - 20.0) * 0.08, 0.95)
            # Add rf proba test result to "qso" mask
            qso[colorsReducedIndex[tmpReleaseOK]] = tmp_rf_proba >= pcut
            # ADM no high-z selection for DR3.
            qsohiz &= False

        tmpReleaseOK = (releaseReduced >= 5000) & (releaseReduced < 8000)
        if np.any(tmpReleaseOK):
            # rf initialization - colors data duplicated within "myRF"
            rf = myRF(colorsReduced[tmpReleaseOK], pathToRF,
                      numberOfTrees=500, version=2)
            rf_HighZ = myRF(colorsReduced[tmpReleaseOK], pathToRF,
                            numberOfTrees=500, version=2)
            # rf loading
            rf.loadForest(rf_DR7_fileName)
            rf_HighZ.loadForest(rf_DR7_HighZ_fileName)
            # Compute rf probabilities
            tmp_rf_proba = rf.predict_proba()
            tmp_rf_HighZ_proba = rf_HighZ.predict_proba()
            # Compute optimized proba cut
            tmp_r_Reduced = r_Reduced[tmpReleaseOK]
            pcut = np.where(tmp_r_Reduced > 20.8,
                            0.83 - (tmp_r_Reduced - 20.8) * 0.025, 0.83)
            pcut[tmp_r_Reduced > 21.5] = 0.8125 - 0.15 * (tmp_r_Reduced[tmp_r_Reduced > 21.5] - 21.5)
            pcut[tmp_r_Reduced > 22.3] = 0.6925 - 0.70 * (tmp_r_Reduced[tmp_r_Reduced > 22.3] - 22.3)
            pcut_HighZ = np.where(tmp_r_Reduced > 20.5,
                                  0.55 - (tmp_r_Reduced - 20.5) * 0.025, 0.55)

            # Add rf proba test result to "qso" mask
            qso[colorsReducedIndex[tmpReleaseOK]] = \
                (tmp_rf_proba >= pcut) | (tmp_rf_HighZ_proba >= pcut_HighZ)
            # ADM populate a mask specific to the "HighZ" selection.
            qsohiz[colorsReducedIndex[tmpReleaseOK]] = \
                (tmp_rf_HighZ_proba >= pcut_HighZ)

        tmpReleaseOK = (releaseReduced >= 8000) & (releaseReduced < 9000)
        if np.any(tmpReleaseOK):
            # rf initialization - colors data duplicated within "myRF"
            rf = myRF(colorsReduced[tmpReleaseOK], pathToRF,
                      numberOfTrees=500, version=2)
            rf_HighZ = myRF(colorsReduced[tmpReleaseOK], pathToRF,
                            numberOfTrees=500, version=2)
            # rf loading
            rf.loadForest(rf_DR8_fileName)
            rf_HighZ.loadForest(rf_DR8_HighZ_fileName)
            # Compute rf probabilities
            tmp_rf_proba = rf.predict_proba()
            tmp_rf_HighZ_proba = rf_HighZ.predict_proba()
            # Compute optimized proba cut
            tmp_r_Reduced = r_Reduced[tmpReleaseOK]
            pcut = 0.88 - 0.03*np.tanh(tmp_r_Reduced - 20.5)
            pcut_HighZ = 0.55

            # Add rf proba test result to "qso" mask
            qso[colorsReducedIndex[tmpReleaseOK]] = \
                (tmp_rf_proba >= pcut) | (tmp_rf_HighZ_proba >= pcut_HighZ)
            # ADM populate a mask specific to the "HighZ" selection.
            qsohiz[colorsReducedIndex[tmpReleaseOK]] = \
                (tmp_rf_HighZ_proba >= pcut_HighZ)
            # ADM store the probabilities in case they need returned.
            pqso[colorsReducedIndex[tmpReleaseOK]] = tmp_rf_proba
            # ADM populate a mask specific to the "HighZ" selection.
            pqsohiz[colorsReducedIndex[tmpReleaseOK]] = tmp_rf_HighZ_proba

        tmpReleaseOK = releaseReduced >= 9000
        if np.any(tmpReleaseOK):
            # rf initialization - colors data duplicated within "myRF"
            rf = myRF(colorsReduced[tmpReleaseOK], pathToRF,
                      numberOfTrees=500, version=2)
            rf_HighZ = myRF(colorsReduced[tmpReleaseOK], pathToRF,
                            numberOfTrees=500, version=2)
            # rf loading
            rf.loadForest(rf_DR9_fileName)
            rf_HighZ.loadForest(rf_DR9_HighZ_fileName)
            # Compute rf probabilities
            tmp_rf_proba = rf.predict_proba()
            tmp_rf_HighZ_proba = rf_HighZ.predict_proba()
            # Compute optimized proba cut
            tmp_r_Reduced = r_Reduced[tmpReleaseOK]
            if not south:
                # threshold selection for North footprint.
                pcut = 0.857 - 0.03*np.tanh(tmp_r_Reduced - 20.5)
                pcut_HighZ = 0.7
            else:
                pcut = np.ones(tmp_rf_proba.size)
                pcut_HighZ = np.ones(tmp_rf_HighZ_proba.size)
                is_des = (gnobs[preSelection][tmpReleaseOK] > 4) &\
                         (rnobs[preSelection][tmpReleaseOK] > 4) &\
                         (znobs[preSelection][tmpReleaseOK] > 4) &\
                         ((ra[preSelection][tmpReleaseOK] >= 320) | (ra[preSelection][tmpReleaseOK] <= 100)) &\
                         (dec[preSelection][tmpReleaseOK] <= 10)
                # threshold selection for DES footprint.
                pcut[is_des] = 0.75 - 0.05*np.tanh(tmp_r_Reduced[is_des] - 20.5)
                pcut_HighZ[is_des] = 0.50
                # threshold selection for South footprint.
                pcut[~is_des] = 0.85 - 0.04*np.tanh(tmp_r_Reduced[~is_des] - 20.5)
                pcut_HighZ[~is_des] = 0.65

            # Add rf proba test result to "qso" mask
            qso[colorsReducedIndex[tmpReleaseOK]] = \
                (tmp_rf_proba >= pcut) | (tmp_rf_HighZ_proba >= pcut_HighZ)
            # ADM populate a mask specific to the "HighZ" selection.
            qsohiz[colorsReducedIndex[tmpReleaseOK]] = \
                (tmp_rf_HighZ_proba >= pcut_HighZ)
            # ADM store the probabilities in case they need returned.
            pqso[colorsReducedIndex[tmpReleaseOK]] = tmp_rf_proba
            # ADM populate a mask specific to the "HighZ" selection.
            pqsohiz[colorsReducedIndex[tmpReleaseOK]] = tmp_rf_HighZ_proba

    # In case of call for a single object passed to the function with
    # scalar arguments. Return "numpy.bool_" instead of "~numpy.ndarray".
    if nbEntries == 1:
        qso = qso[0]
        qsohiz = qsohiz[0]
        pqso = pqso[0]
        pqsohiz = pqsohiz[0]

    # ADM if requested, return the probabilities as well.
    if return_probs:
        return qso, qsohiz, pqso, pqsohiz
    return qso, qsohiz

def _psflike(psftype):
    """ If the object is PSF """
    # ADM explicitly checking for NoneType. I can't see why we'd ever want to
    # ADM run this test on empty information. In the past we have had bugs where
    # ADM we forgot to pass objtype=objtype in, e.g., isSTD
    if psftype is None:
        raise ValueError("NoneType submitted to _psfflike function")

    psftype = np.asarray(psftype)
    # ADM in Python3 these string literals become byte-like
    # ADM so to retain Python2 compatibility we need to check
    # ADM against both bytes and unicode
    # ADM, also 'PSF' for astropy.io.fits; 'PSF ' for fitsio (sigh)
    psflike = ((psftype == 'PSF') | (psftype == b'PSF') |
               (psftype == 'PSF ') | (psftype == b'PSF '))
    return psflike

def _getColors(nbEntries, nfeatures, gflux, rflux, zflux, w1flux, w2flux):

    limitInf = 1.e-04
    gflux = gflux.clip(limitInf)
    rflux = rflux.clip(limitInf)
    zflux = zflux.clip(limitInf)
    w1flux = w1flux.clip(limitInf)
    w2flux = w2flux.clip(limitInf)

    g = np.where(gflux > limitInf, 22.5-2.5*np.log10(gflux), 0.)
    r = np.where(rflux > limitInf, 22.5-2.5*np.log10(rflux), 0.)
    z = np.where(zflux > limitInf, 22.5-2.5*np.log10(zflux), 0.)
    W1 = np.where(w1flux > limitInf, 22.5-2.5*np.log10(w1flux), 0.)
    W2 = np.where(w2flux > limitInf, 22.5-2.5*np.log10(w2flux), 0.)

    photOK = (g > 0.) & (r > 0.) & (z > 0.) & (W1 > 0.) & (W2 > 0.)

    colors = np.zeros((nbEntries, nfeatures))
    colors[:, 0] = g-r
    colors[:, 1] = r-z
    colors[:, 2] = g-z
    colors[:, 3] = g-W1
    colors[:, 4] = r-W1
    colors[:, 5] = z-W1
    colors[:, 6] = g-W2
    colors[:, 7] = r-W2
    colors[:, 8] = z-W2
    colors[:, 9] = W1-W2
    colors[:, 10] = r

    return colors, r, photOK

def _is_row(table):
    """Return True/False if this is a row of a table instead of a full table.

    supports numpy.ndarray, astropy.io.fits.FITS_rec, and astropy.table.Table
    """
    import astropy.io.fits.fitsrec
    import astropy.table.row
    if isinstance(table, (astropy.io.fits.fitsrec.FITS_record, astropy.table.row.Row)) or \
       np.isscalar(table):
        return True
    else:
        return False